temp push

ws/websocket.go

@@ -8,29 +8,12 @@ import (
 	"log/slog"
 	"net"
 	"net/http"
-	"strings"
 	"time"
 
 	"github.com/gorilla/websocket"
 )
 
-func (m *GatewayMessageIn) Validate() error {
-	if strings.TrimSpace(m.MsgID) == "" {
-		return errors.New("msg_id missing")
-	}
-	if strings.TrimSpace(m.Server) == "" {
-		return errors.New("server missing")
-	}
-	if strings.TrimSpace(m.User.ID) == "" {
-		return errors.New("user.mod_uid missing")
-	}
-	if strings.TrimSpace(m.Content) == "" {
-		return errors.New("content missing")
-	}
-	return nil
-}
-
-func NewWsGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func() error) *WebsocketGateway {
+func NewWsGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func() error, modH ModHandler, botH BotHandler) *WebsocketGateway {
 	return &WebsocketGateway{
 		logger:  logger,
 		closefn: closefn,
@@ -42,14 +25,14 @@ func NewWsGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func()
 				return true // local by default; change for production
 			},
 		},
-		outgoingCh:    make(chan GatewayMessageOut, cfg.QueueSize),
-		modConns:      make(map[*websocket.Conn]struct{}),
+		conns:         make(map[*websocket.Conn]connMetadata),
 		bodySizeBytes: int64(cfg.BodySize) * 1024 * 1024,
 		port:          cfg.HttpPort,
+		modHandler:    modH,
+		botHandler:    botH,
 	}
 }
 
-// Serve starts the HTTP server and /ws endpoint and blocks until ctx cancelled or server fails.
 func (g *WebsocketGateway) Serve(ctx context.Context, listenAddr string) error {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/ws", g.handleWS)
@@ -67,7 +50,7 @@ func (g *WebsocketGateway) Serve(ctx context.Context, listenAddr string) error {
 	errCh := make(chan error, 1)
 	go func() {
 		g.logger.Info("ws gateway listening", "addr", listenAddr)
-		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
 			errCh <- err
 		}
 		close(errCh)
@@ -99,9 +82,6 @@ func (g *WebsocketGateway) handleWS(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	g.registerConn(conn)
-	g.logger.Info("ws connected", "remote", conn.RemoteAddr().String())
-
 	// Configure read limits & pong handler
 	if g.bodySizeBytes > 0 {
 		conn.SetReadLimit(g.bodySizeBytes)
@@ -115,69 +95,182 @@ func (g *WebsocketGateway) handleWS(w http.ResponseWriter, r *http.Request) {
 		return nil
 	})
 
-	go g.readLoop(conn)
+	// 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) readLoop(c *websocket.Conn) {
+func (g *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connMetadata) {
 	defer func() {
 		g.unregisterConn(c)
 		_ = c.Close()
-		g.logger.Info("ws disconnected", "remote", c.RemoteAddr().String())
+		g.logger.Info("mod disconnected", "server_id", meta.id, "remote", c.RemoteAddr().String())
 	}()
 
 	pingTicker := time.NewTicker(30 * time.Second)
 	defer pingTicker.Stop()
 
 	for {
-		// Read one message (blocks until message arrives)
 		typ, data, err := c.ReadMessage()
 		if err != nil {
 			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
-				g.logger.Warn("unexpected ws close", "err", err)
+				g.logger.Warn("unexpected mod close", "server_id", meta.id, "err", err)
 			} else {
-				g.logger.Debug("ws read error", "err", err)
+				g.logger.Debug("mod read error", "server_id", meta.id, "err", err)
 			}
 			return
 		}
+
 		if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
 			continue
 		}
 
-		var in GatewayMessageIn
-		if err := json.Unmarshal(data, &in); err != nil {
+		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 client", "remote", c.RemoteAddr().String(), "err", err)
+			g.logger.Warn("invalid json from mod", "server_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
 			continue
 		}
-		in.ReceivedAt = time.Now().UTC()
-		if err := in.Validate(); err != nil {
+
+		msg.ReceivedAt = time.Now().UTC()
+		if err := msg.Validate(); err != nil {
 			_ = writeJSONSafe(c, map[string]string{"error": err.Error()})
-			g.logger.Warn("message validation failed", "remote", c.RemoteAddr().String(), "err", err)
+			g.logger.Warn("mod message validation failed", "server_id", meta.id, "remote", c.RemoteAddr().String(), "err", err)
 			continue
 		}
 
-		out := GatewayMessageOut{
-			Type:        "message",
-			Payload:     in,
-			ForwardedAt: time.Now().UTC(),
+		// 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
 		}
 
-		// Non-blocking enqueue with backpressure
-		select {
-		case g.outgoingCh <- out:
-			_ = writeJSONSafe(c, map[string]string{"status": "queued"})
-			g.logger.Debug("enqueued message", "msg_id", in.MsgID, "server", in.Server)
-		default:
-			_ = writeJSONSafe(c, map[string]string{"error": "gateway busy"})
-			g.logger.Warn("outgoing queue full", "msg_id", in.MsgID, "remote", c.RemoteAddr().String())
-		}
+		_ = writeJSONSafe(c, map[string]string{"status": "ok"})
 
-		// also handle pings periodically (so client sees ping frequently)
+		// 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", "err", err)
+				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: