280 lines
8.4 KiB
Go
280 lines
8.4 KiB
Go
package ws
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"homestead/homestead_gateway/util/config"
|
|
"log/slog"
|
|
"net"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
func NewWsGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func() error, modH ModHandler, botH BotHandler) *WebsocketGateway {
|
|
return &WebsocketGateway{
|
|
logger: logger,
|
|
closefn: closefn,
|
|
apiKey: cfg.Websocket,
|
|
upgrader: websocket.Upgrader{
|
|
ReadBufferSize: 1024,
|
|
WriteBufferSize: 1024,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
return true // local by default; change for production
|
|
},
|
|
},
|
|
conns: make(map[*websocket.Conn]connMetadata),
|
|
bodySizeBytes: int64(cfg.BodySize) * 1024 * 1024,
|
|
port: cfg.HttpPort,
|
|
modHandler: modH,
|
|
botHandler: botH,
|
|
}
|
|
}
|
|
|
|
func (g *WebsocketGateway) Serve(ctx context.Context, listenAddr string) error {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/ws", g.handleWS)
|
|
mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(200)
|
|
_, _ = w.Write([]byte("ok"))
|
|
})
|
|
|
|
srv := &http.Server{
|
|
Addr: listenAddr,
|
|
Handler: loggingMiddleware(g.logger, mux),
|
|
BaseContext: func(l net.Listener) context.Context { return ctx },
|
|
}
|
|
|
|
errCh := make(chan error, 1)
|
|
go func() {
|
|
g.logger.Info("ws gateway listening", "addr", listenAddr)
|
|
if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
|
errCh <- err
|
|
}
|
|
close(errCh)
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
g.logger.Info("shutting down http server")
|
|
_ = srv.Shutdown(shutdownCtx)
|
|
g.closeAll()
|
|
return nil
|
|
case err := <-errCh:
|
|
return err
|
|
}
|
|
}
|
|
|
|
func (g *WebsocketGateway) handleWS(w http.ResponseWriter, r *http.Request) {
|
|
if !g.validateApiKey(r) {
|
|
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
|
g.logger.Warn("ws auth failed", "remote", r.RemoteAddr)
|
|
return
|
|
}
|
|
|
|
conn, err := g.upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
g.logger.Error("ws upgrade error", err, "remote", r.RemoteAddr)
|
|
return
|
|
}
|
|
|
|
// Configure read limits & pong handler
|
|
if g.bodySizeBytes > 0 {
|
|
conn.SetReadLimit(g.bodySizeBytes)
|
|
} else {
|
|
conn.SetReadLimit(1 << 20) // sensible default 1MiB
|
|
}
|
|
|
|
_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
|
conn.SetPongHandler(func(appData string) error {
|
|
_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
|
return nil
|
|
})
|
|
|
|
// First message must be a handshake identifying the connection type
|
|
typ, data, err := conn.ReadMessage()
|
|
if err != nil {
|
|
g.logger.Error("failed to read handshake", "err", err, "remote", conn.RemoteAddr().String())
|
|
_ = conn.Close()
|
|
return
|
|
}
|
|
|
|
if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
|
|
g.logger.Warn("invalid handshake message type", "remote", conn.RemoteAddr().String())
|
|
_ = writeJSONSafe(conn, map[string]string{"error": "first message must be handshake"})
|
|
_ = conn.Close()
|
|
return
|
|
}
|
|
|
|
var envelope MessageEnvelope
|
|
if err := json.Unmarshal(data, &envelope); err != nil {
|
|
g.logger.Warn("invalid handshake json", "err", err, "remote", conn.RemoteAddr().String())
|
|
_ = writeJSONSafe(conn, map[string]string{"error": "invalid handshake: " + err.Error()})
|
|
_ = conn.Close()
|
|
return
|
|
}
|
|
|
|
meta := connMetadata{connType: envelope.Type}
|
|
|
|
// Validate handshake based on type
|
|
switch envelope.Type {
|
|
case "mod":
|
|
var hs ModHandshake
|
|
if err := json.Unmarshal(envelope.Data, &hs); err != nil {
|
|
g.logger.Warn("invalid mod handshake", "err", err, "remote", conn.RemoteAddr().String())
|
|
_ = writeJSONSafe(conn, map[string]string{"error": "invalid mod handshake: " + err.Error()})
|
|
_ = conn.Close()
|
|
return
|
|
}
|
|
meta.id = hs.ServerID
|
|
g.registerConn(conn, meta)
|
|
g.logger.Info("mod connected", "server_id", hs.ServerID, "remote", conn.RemoteAddr().String())
|
|
go g.modReadLoop(conn, meta)
|
|
|
|
case "bot":
|
|
var hs BotHandshake
|
|
if err := json.Unmarshal(envelope.Data, &hs); err != nil {
|
|
g.logger.Warn("invalid bot handshake", "err", err, "remote", conn.RemoteAddr().String())
|
|
_ = writeJSONSafe(conn, map[string]string{"error": "invalid bot handshake: " + err.Error()})
|
|
_ = conn.Close()
|
|
return
|
|
}
|
|
meta.id = hs.BotID
|
|
g.registerConn(conn, meta)
|
|
g.logger.Info("bot connected", "bot_id", hs.BotID, "remote", conn.RemoteAddr().String())
|
|
go g.botReadLoop(conn, meta)
|
|
|
|
default:
|
|
g.logger.Warn("unknown connection type", "type", envelope.Type, "remote", conn.RemoteAddr().String())
|
|
_ = writeJSONSafe(conn, map[string]string{"error": "unknown connection type: " + envelope.Type})
|
|
_ = conn.Close()
|
|
return
|
|
}
|
|
}
|
|
|
|
func (g *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connMetadata) {
|
|
defer func() {
|
|
g.unregisterConn(c)
|
|
_ = c.Close()
|
|
g.logger.Info("mod disconnected", "server_id", meta.id, "remote", c.RemoteAddr().String())
|
|
}()
|
|
|
|
pingTicker := time.NewTicker(30 * time.Second)
|
|
defer pingTicker.Stop()
|
|
|
|
for {
|
|
typ, data, err := c.ReadMessage()
|
|
if err != nil {
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
|
|
g.logger.Warn("unexpected mod close", "server_id", meta.id, "err", err)
|
|
} else {
|
|
g.logger.Debug("mod read error", "server_id", meta.id, "err", err)
|
|
}
|
|
return
|
|
}
|
|
|
|
if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
|
|
continue
|
|
}
|
|
|
|
var msg GatewayModMessageIn
|
|
if err := json.Unmarshal(data, &msg); err != nil {
|
|
_ = writeJSONSafe(c, map[string]string{"error": "invalid json: " + err.Error()})
|
|
g.logger.Warn("invalid json from mod", "server_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
|
|
continue
|
|
}
|
|
|
|
msg.ReceivedAt = time.Now().UTC()
|
|
if err := msg.Validate(); err != nil {
|
|
_ = writeJSONSafe(c, map[string]string{"error": err.Error()})
|
|
g.logger.Warn("mod message validation failed", "server_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
|
|
continue
|
|
}
|
|
|
|
// Handle the message (forward to bot, enrich, etc.)
|
|
if err := g.modHandler.Handle(context.Background(), msg); err != nil {
|
|
_ = writeJSONSafe(c, map[string]string{"error": "handler error: " + err.Error()})
|
|
g.logger.Error("mod handler error", "server_id", meta.id, "err", err)
|
|
continue
|
|
}
|
|
|
|
_ = writeJSONSafe(c, map[string]string{"status": "ok"})
|
|
|
|
// Handle pings
|
|
select {
|
|
case <-pingTicker.C:
|
|
_ = c.SetWriteDeadline(time.Now().Add(5 * time.Second))
|
|
if err := c.WriteMessage(websocket.PingMessage, nil); err != nil {
|
|
g.logger.Debug("write ping failed", "server_id", meta.id, "err", err)
|
|
return
|
|
}
|
|
default:
|
|
}
|
|
}
|
|
}
|
|
|
|
func (g *WebsocketGateway) botReadLoop(c *websocket.Conn, meta connMetadata) {
|
|
defer func() {
|
|
g.unregisterConn(c)
|
|
_ = c.Close()
|
|
g.logger.Info("bot disconnected", "bot_id", meta.id, "remote", c.RemoteAddr().String())
|
|
}()
|
|
|
|
pingTicker := time.NewTicker(30 * time.Second)
|
|
defer pingTicker.Stop()
|
|
|
|
for {
|
|
typ, data, err := c.ReadMessage()
|
|
if err != nil {
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
|
|
g.logger.Warn("unexpected bot close", "bot_id", meta.id, "err", err)
|
|
} else {
|
|
g.logger.Debug("bot read error", "bot_id", meta.id, "err", err)
|
|
}
|
|
return
|
|
}
|
|
|
|
if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
|
|
continue
|
|
}
|
|
|
|
var msg GatewayBotMessageIn
|
|
if err := json.Unmarshal(data, &msg); err != nil {
|
|
_ = writeJSONSafe(c, map[string]string{"error": "invalid json: " + err.Error()})
|
|
g.logger.Warn("invalid json from bot", "bot_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
|
|
continue
|
|
}
|
|
|
|
msg.ReceivedAt = time.Now().UTC()
|
|
if err := msg.Validate(); err != nil {
|
|
_ = writeJSONSafe(c, map[string]string{"error": err.Error()})
|
|
g.logger.Warn("bot message validation failed", "bot_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
|
|
continue
|
|
}
|
|
|
|
// Handle the message (forward to mod, enrich, etc.)
|
|
if err := g.botHandler.Handle(context.Background(), msg); err != nil {
|
|
_ = writeJSONSafe(c, map[string]string{"error": "handler error: " + err.Error()})
|
|
g.logger.Error("bot handler error", "bot_id", meta.id, "err", err)
|
|
continue
|
|
}
|
|
|
|
_ = writeJSONSafe(c, map[string]string{"status": "ok"})
|
|
|
|
// Handle pings
|
|
select {
|
|
case <-pingTicker.C:
|
|
_ = c.SetWriteDeadline(time.Now().Add(5 * time.Second))
|
|
if err := c.WriteMessage(websocket.PingMessage, nil); err != nil {
|
|
g.logger.Debug("write ping failed", "bot_id", meta.id, "err", err)
|
|
return
|
|
}
|
|
default:
|
|
}
|
|
}
|
|
}
|