package ws import ( "context" "encoding/json" "fmt" "homestead/homestead_gateway/util/config" "log/slog" "net" "net/http" "os/signal" "syscall" "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]connectionMetaData), bodySizeBytes: int64(cfg.BodySize) * 1024 * 1024, port: cfg.HttpPort, modHandler: modH, botHandler: botH, } } 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("/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(wsg.logger, mux), BaseContext: func(l net.Listener) context.Context { return ctx }, } errCh := make(chan error, 1) go wsg.listen(srv, listenAddr, errCh) select { case <-ctx.Done(): wsg.deafen(srv) return nil case err := <-errCh: return err } } // func (wsg *WebsocketGateway) modReadLoop(c *websocket.Conn, meta connectionMetaData) { defer func() { wsg.unregisterConn(c) _ = c.Close() wsg.logger.Info("mod disconnected", "server_id", meta.id, "remote", c.RemoteAddr().String()) }() pingTicker := time.NewTicker(30 * time.Second) defer pingTicker.Stop() // Send pings in a separate goroutine go func() { for range pingTicker.C { _ = c.SetWriteDeadline(time.Now().Add(5 * time.Second)) if err := c.WriteMessage(websocket.PingMessage, nil); err != nil { wsg.logger.Debug("write ping failed", "server_id", meta.id, "err", err) return } wsg.logger.Debug("sent ping to mod", "server_id", meta.id) } }() for { typ, data, err := c.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) { wsg.logger.Warn("unexpected mod close", "server_id", meta.id, "err", err) } else { wsg.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()}) 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()}) 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 := wsg.modHandler.Handle(context.Background(), msg); err != nil { _ = writeJSONSafe(c, map[string]string{"error": "handler error: " + err.Error()}) wsg.logger.Error("mod handler error", "server_id", meta.id, "err", err) continue } _ = writeJSONSafe(c, map[string]string{"status": "ok"}) } } func (wsg *WebsocketGateway) botReadLoop(c *websocket.Conn, meta connectionMetaData) { defer func() { wsg.unregisterConn(c) _ = c.Close() wsg.logger.Info("bot disconnected", "bot_id", meta.id, "remote", c.RemoteAddr().String()) }() pingTicker := time.NewTicker(30 * time.Second) defer pingTicker.Stop() // Send pings in a separate goroutine go func() { for range pingTicker.C { _ = c.SetWriteDeadline(time.Now().Add(5 * time.Second)) if err := c.WriteMessage(websocket.PingMessage, nil); err != nil { wsg.logger.Debug("write ping failed", "bot_id", meta.id, "err", err) return } wsg.logger.Debug("sent ping to bot", "bot_id", meta.id) } }() for { typ, data, err := c.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) { wsg.logger.Warn("unexpected bot close", "bot_id", meta.id, "err", err) } else { wsg.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()}) 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()}) 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 := wsg.botHandler.Handle(context.Background(), msg); err != nil { _ = writeJSONSafe(c, map[string]string{"error": "handler error: " + err.Error()}) wsg.logger.Error("bot handler error", "bot_id", meta.id, "err", err) continue } _ = writeJSONSafe(c, map[string]string{"status": "ok"}) } }