re-formatted, better structure

2025-12-01 18:09:10 +01:00
parent 4a668493c4
commit 46aef47e21
9 changed files with 328 additions and 235 deletions
--- a/ws/websocket.go
+++ b/ws/websocket.go
@@ -3,11 +3,13 @@ package ws
 import (
 	"context"
 	"encoding/json"
-	"errors"
+	"fmt"
 	"homestead/homestead_gateway/util/config"
 	"log/slog"
 	"net"
 	"net/http"
+	"os/signal"
+	"syscall"
 	"time"

 	"github.com/gorilla/websocket"
@@ -16,7 +18,7 @@ import (
 func NewWsGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func() error, modH ModHandler, botH BotHandler) *WebsocketGateway {
 	return &WebsocketGateway{
 		logger:  logger,
-		closefn: closefn,
+		closeFn: closefn,
 		apiKey:  cfg.Websocket,
 		upgrader: websocket.Upgrader{
 			ReadBufferSize:  1024,
@@ -25,7 +27,7 @@ func NewWsGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func()
 				return true // local by default; change for production
 			},
 		},
-		conns:         make(map[*websocket.Conn]connMetadata),
+		conns:         make(map[*websocket.Conn]connectionMetaData),
 		bodySizeBytes: int64(cfg.BodySize) * 1024 * 1024,
 		port:          cfg.HttpPort,
 		modHandler:    modH,
@@ -33,58 +35,48 @@ func NewWsGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func()
 	}
 }

-func (g *WebsocketGateway) Serve(ctx context.Context, listenAddr string) error {
+func (wsg *WebsocketGateway) Start() error {
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer stop()
+
+	return wsg.Serve(ctx, fmt.Sprintf(":%d", wsg.port))
+}
+
+func (wsg *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"))
-	})
+	mux.HandleFunc("/push", wsg.handlePush)
+	mux.HandleFunc("/ready", wsg.handleReady)
+	mux.HandleFunc("/health", wsg.handleHealth)
+	mux.HandleFunc("/register", wsg.handleRegister)

 	srv := &http.Server{
 		Addr:        listenAddr,
-		Handler:     loggingMiddleware(g.logger, mux),
+		Handler:     loggingMiddleware(wsg.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)
-	}()
+
+	go wsg.listen(srv, listenAddr, 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()
+		wsg.deafen(srv)
 		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)
+func (wsg *WebsocketGateway) handlePush(w http.ResponseWriter, r *http.Request) {
+	conn, err := wsg.validateAndUpgradeConnection(w, r)
 	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)
+	if wsg.bodySizeBytes > 0 {
+		conn.SetReadLimit(wsg.bodySizeBytes)
 	} else {
 		conn.SetReadLimit(1 << 20) // sensible default 1MiB
 	}
@@ -95,72 +87,93 @@ func (g *WebsocketGateway) handleWS(w http.ResponseWriter, r *http.Request) {
 		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()
+		wsg.sendWebsocketError(conn, "Internal Server Error", 500)
+		wsg.logger.Error("Failed to read handshake.", "remote", conn.RemoteAddr().String(), "err", err)
 		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()
+		wsg.sendWebsocketError(conn, "First message must be a handshake.", 400)
+		wsg.logger.Warn("Invalid handshake message type.", "remote", conn.RemoteAddr().String())
 		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()
+	var handshake Handshake
+	if err := json.Unmarshal(data, &handshake); err != nil {
+		wsg.sendWebsocketError(conn, "Malformed handshake.", 400)
+		wsg.logger.Warn("Malformed handshake.", "remote", conn.RemoteAddr().String(), "err", err)
 		return
 	}

-	meta := connMetadata{connType: envelope.Type}
+	meta := connectionMetaData{connectionType: handshake.Type}

-	// Validate handshake based on type
-	switch envelope.Type {
+	switch handshake.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()
+		var mhs ModHandshake
+
+		if err := json.Unmarshal(handshake.Data, &mhs); err != nil {
+			wsg.sendWebsocketError(conn, "Malformed mod handshake.", 400)
+			wsg.logger.Warn("Malformed mod handshake.", "remote", conn.RemoteAddr().String(), "err", err)
 			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)
+
+		meta.id = mhs.ServerID
+
+		if err = wsg.sendWebsocketResponse(conn, GatewayAck{Status: "connected", Type: "mod"}); err != nil {
+			return
+		}
+
+		wsg.registerConn(conn, meta)
+		wsg.logger.Info("Mod connected via Websocket.", "remote", conn.RemoteAddr().String(), "server_id", mhs.ServerID)
+
+		go wsg.modReadLoop(conn, meta) // replace with external handler mayhaps

 	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()
+		var bhs BotHandshake
+
+		if err := json.Unmarshal(handshake.Data, &bhs); err != nil {
+			wsg.sendWebsocketError(conn, "Malformed bot handshake.", 400)
+			wsg.logger.Warn("Malformed bot handshake.", "remote", conn.RemoteAddr().String(), "err", err)
 			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)
+
+		meta.id = bhs.BotID
+
+		if err = wsg.sendWebsocketResponse(conn, GatewayAck{Status: "connected", Type: "bot"}); err != nil {
+			return
+		}
+
+		wsg.registerConn(conn, meta)
+		wsg.logger.Info("Bot connected via Websocket.", "remote", conn.RemoteAddr().String(), "bot_id", bhs.BotID)
+
+		go wsg.botReadLoop(conn, meta) // replace with external handler mayhaps

 	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()
+		wsg.sendWebsocketError(conn, "Unknown handshake.", 400)
+		wsg.logger.Warn("Unknown connection type.", "remote", conn.RemoteAddr().String(), "type", handshake.Type)
 		return
 	}
 }

-func (g *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connMetadata) {
+func (wsg *WebsocketGateway) handleReady(w http.ResponseWriter, r *http.Request) {}
+
+func (wsg *WebsocketGateway) handleHealth(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(200)
+	_ = json.NewEncoder(w).Encode(map[string]interface{}{"status": "healthy"})
+}
+
+func (wsg *WebsocketGateway) handleRegister(w http.ResponseWriter, r *http.Request) {}
+
+//
+
+func (wsg *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connectionMetaData) {
 	defer func() {
-		g.unregisterConn(c)
+		wsg.unregisterConn(c)
 		_ = c.Close()
-		g.logger.Info("mod disconnected", "server_id", meta.id, "remote", c.RemoteAddr().String())
+		wsg.logger.Info("mod disconnected", "server_id", meta.id, "remote", c.RemoteAddr().String())
 	}()

 	pingTicker := time.NewTicker(30 * time.Second)
@@ -170,9 +183,9 @@ func (g *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connMetadata) {
 		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)
+				wsg.logger.Warn("unexpected mod close", "server_id", meta.id, "err", err)
 			} else {
-				g.logger.Debug("mod read error", "server_id", meta.id, "err", err)
+				wsg.logger.Debug("mod read error", "server_id", meta.id, "err", err)
 			}
 			return
 		}
@@ -184,21 +197,21 @@ func (g *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connMetadata) {
 		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)
+			wsg.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)
+			wsg.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 {
+		if err := wsg.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)
+			wsg.logger.Error("mod handler error", "server_id", meta.id, "err", err)
 			continue
 		}

@@ -209,7 +222,7 @@ func (g *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connMetadata) {
 		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)
+				wsg.logger.Debug("write ping failed", "server_id", meta.id, "err", err)
 				return
 			}
 		default:
@@ -217,11 +230,11 @@ func (g *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connMetadata) {
 	}
 }

-func (g *WebsocketGateway) botReadLoop(c *websocket.Conn, meta connMetadata) {
+func (wsg *WebsocketGateway) botReadLoop(c *websocket.Conn, meta connectionMetaData) {
 	defer func() {
-		g.unregisterConn(c)
+		wsg.unregisterConn(c)
 		_ = c.Close()
-		g.logger.Info("bot disconnected", "bot_id", meta.id, "remote", c.RemoteAddr().String())
+		wsg.logger.Info("bot disconnected", "bot_id", meta.id, "remote", c.RemoteAddr().String())
 	}()

 	pingTicker := time.NewTicker(30 * time.Second)
@@ -231,9 +244,9 @@ func (g *WebsocketGateway) botReadLoop(c *websocket.Conn, meta connMetadata) {
 		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)
+				wsg.logger.Warn("unexpected bot close", "bot_id", meta.id, "err", err)
 			} else {
-				g.logger.Debug("bot read error", "bot_id", meta.id, "err", err)
+				wsg.logger.Debug("bot read error", "bot_id", meta.id, "err", err)
 			}
 			return
 		}
@@ -245,21 +258,21 @@ func (g *WebsocketGateway) botReadLoop(c *websocket.Conn, meta connMetadata) {
 		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)
+			wsg.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)
+			wsg.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 {
+		if err := wsg.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)
+			wsg.logger.Error("bot handler error", "bot_id", meta.id, "err", err)
 			continue
 		}

@@ -270,7 +283,7 @@ func (g *WebsocketGateway) botReadLoop(c *websocket.Conn, meta connMetadata) {
 		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)
+				wsg.logger.Debug("write ping failed", "bot_id", meta.id, "err", err)
 				return
 			}
 		default: