almost done, queue fix next

ws/websocket.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"homestead/homestead_gateway/util/cache"
 	"homestead/homestead_gateway/util/config"
+	"homestead/homestead_gateway/util/queue"
 	"log/slog"
 	"net"
 	"net/http"
@@ -16,16 +17,15 @@ import (
 	"github.com/gorilla/websocket"
 )
 
-func NewWebsocketGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func() error, modH ModHandler, botH BotHandler) *WebsocketGateway {
+func NewWebsocketGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func() error) *WebsocketGateway {
 	return &WebsocketGateway{
-		logger:     logger,
-		closeFn:    closefn,
-		modHandler: modH,
-		botHandler: botH,
-		port:       cfg.HttpPort,
-		apiKey:     cfg.Websocket,
-		cache:      cache.NewCache(),
-		conns:      cache.NewConnectionCache(),
+		logger:  logger,
+		closeFn: closefn,
+		port:    cfg.HttpPort,
+		apiKey:  cfg.Websocket,
+		cache:   cache.NewCache(),
+		queue:   queue.NewQueue[GatewayMessageOut](32),
+		conns:   cache.NewConnectionCache(),
 		upgrader: websocket.Upgrader{
 			ReadBufferSize:  1024,
 			WriteBufferSize: 1024,
@@ -69,139 +69,251 @@ func (wsg *WebsocketGateway) Serve(ctx context.Context, listenAddr string) error
 
 //
 
-func (wsg *WebsocketGateway) modReadLoop(conn *websocket.Conn, meta cache.ConnectionMetaData) {
+//func (wsg *WebsocketGateway) modReadLoop(conn *websocket.Conn, meta cache.ConnectionMetaData) {
+//	defer func() {
+//		wsg.unregisterConn(conn, meta, "mod")
+//		wsg.logger.Info("Client disconnected.", "remote", conn.RemoteAddr().String(), "server_id", meta.ID)
+//	}()
+//
+//	ticker := time.NewTicker(30 * time.Second)
+//	defer ticker.Stop()
+//
+//	go func() {
+//		for range ticker.C {
+//			wsg.sendWebsocketPing(conn)
+//		}
+//	}()
+//
+//	for {
+//		typ, data, err := conn.ReadMessage()
+//
+//		if err != nil {
+//			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+//				wsg.logger.Warn("Mod-Client unexpectedly closed the connection.", "err", err)
+//			}
+//			return
+//		}
+//
+//		if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
+//			continue
+//		}
+//
+//		var msg GatewayModMessageIn
+//		if err := json.Unmarshal(data, &msg); err != nil {
+//			_ = writeJSONSafe(conn, map[string]string{"error": "invalid json: " + err.Error()})
+//			wsg.logger.Warn("invalid json from mod", "server_id", meta.ID, "remote", conn.RemoteAddr().String(), "err", err)
+//			continue
+//		}
+//
+//		msg.ReceivedAt = time.Now().UTC()
+//		if err := msg.Validate(); err != nil {
+//			_ = writeJSONSafe(conn, map[string]string{"error": err.Error()})
+//			wsg.logger.Warn("mod message validation failed", "server_id", meta.ID, "remote", conn.RemoteAddr().String(), "err", err)
+//			continue
+//		}
+//
+//		// Handle the message (forward to bot, enrich, etc.)
+//		if err := wsg.modHandler.Handle(conn, msg); err != nil {
+//			_ = writeJSONSafe(conn, map[string]string{"error": "handler error: " + err.Error()})
+//			wsg.logger.Error("mod handler error", "server_id", meta.ID, "err", err)
+//			continue
+//		}
+//
+//		_ = writeJSONSafe(conn, map[string]string{"status": "completed"}) // or "queued"
+//	}
+//}
+//
+//func (wsg *WebsocketGateway) botReadLoop(conn *websocket.Conn, meta cache.ConnectionMetaData) {
+//	defer func() {
+//		wsg.unregisterConn(conn, meta, "bot")
+//		wsg.logger.Info("bot disconnected", "bot_id", meta.ID, "remote", conn.RemoteAddr().String())
+//	}()
+//
+//	pingTicker := time.NewTicker(30 * time.Second)
+//	defer pingTicker.Stop()
+//
+//	// Send pings in a separate goroutine
+//	go func() {
+//		for range pingTicker.C {
+//			_ = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
+//			if err := conn.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 := conn.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(conn, map[string]string{"error": "invalid json: " + err.Error()})
+//			wsg.logger.Warn("invalid json from bot", "bot_id", meta.ID, "remote", conn.RemoteAddr().String(), "err", err)
+//			continue
+//		}
+//
+//		msg.ReceivedAt = time.Now().UTC()
+//		if err := msg.Validate(); err != nil {
+//			_ = writeJSONSafe(conn, map[string]string{"error": err.Error()})
+//			wsg.logger.Warn("bot message validation failed", "bot_id", meta.ID, "remote", conn.RemoteAddr().String(), "err", err)
+//			continue
+//		}
+//
+//		// Handle the message (forward to mod, enrich, etc.)
+//		if err := wsg.botHandler.Handle(conn, msg); err != nil {
+//			_ = writeJSONSafe(conn, map[string]string{"error": "handler error: " + err.Error()})
+//			wsg.logger.Error("bot handler error", "bot_id", meta.ID, "err", err)
+//			continue
+//		}
+//
+//		_ = writeJSONSafe(conn, map[string]string{"status": "ok"})
+//	}
+//}
+
+func (wsg *WebsocketGateway) read(conn *websocket.Conn, meta cache.ConnectionMetaData, _type string) {
 	defer func() {
-		wsg.unregisterConn(conn, meta, "mod")
-		wsg.logger.Info("Client disconnected.", "remote", conn.RemoteAddr().String(), "server_id", meta.ID)
-	}()
-
-	ticker := time.NewTicker(30 * time.Second)
-	defer ticker.Stop()
-
-	go func() {
-		for range ticker.C {
-			wsg.sendWebsocketPing(conn)
-		}
-	}()
-
-	for {
-		typ, data, err := conn.ReadMessage()
-
-		if err != nil {
-			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
-				wsg.logger.Warn("Mod-Client unexpectedly closed the connection.", "err", err)
-			}
-			return
-		}
-
-		if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
-			continue
-		}
-
-		var msg GatewayModMessageIn
-		if err := json.Unmarshal(data, &msg); err != nil {
-			_ = writeJSONSafe(conn, map[string]string{"error": "invalid json: " + err.Error()})
-			wsg.logger.Warn("invalid json from mod", "server_id", meta.ID, "remote", conn.RemoteAddr().String(), "err", err)
-			continue
-		}
-
-		msg.ReceivedAt = time.Now().UTC()
-		if err := msg.Validate(); err != nil {
-			_ = writeJSONSafe(conn, map[string]string{"error": err.Error()})
-			wsg.logger.Warn("mod message validation failed", "server_id", meta.ID, "remote", conn.RemoteAddr().String(), "err", err)
-			continue
-		}
-
-		// Handle the message (forward to bot, enrich, etc.)
-		if err := wsg.modHandler.Handle(conn, msg); err != nil {
-			_ = writeJSONSafe(conn, map[string]string{"error": "handler error: " + err.Error()})
-			wsg.logger.Error("mod handler error", "server_id", meta.ID, "err", err)
-			continue
-		}
-
-		_ = writeJSONSafe(conn, map[string]string{"status": "completed"}) // or "queued"
-	}
-}
-
-func (wsg *WebsocketGateway) botReadLoop(conn *websocket.Conn, meta cache.ConnectionMetaData) {
-	defer func() {
-		wsg.unregisterConn(conn, meta, "bot")
-		wsg.logger.Info("bot disconnected", "bot_id", meta.ID, "remote", conn.RemoteAddr().String())
-	}()
-
-	pingTicker := time.NewTicker(30 * time.Second)
-	defer pingTicker.Stop()
-
-	// Send pings in a separate goroutine
-	go func() {
-		for range pingTicker.C {
-			_ = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
-			if err := conn.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 := conn.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(conn, map[string]string{"error": "invalid json: " + err.Error()})
-			wsg.logger.Warn("invalid json from bot", "bot_id", meta.ID, "remote", conn.RemoteAddr().String(), "err", err)
-			continue
-		}
-
-		msg.ReceivedAt = time.Now().UTC()
-		if err := msg.Validate(); err != nil {
-			_ = writeJSONSafe(conn, map[string]string{"error": err.Error()})
-			wsg.logger.Warn("bot message validation failed", "bot_id", meta.ID, "remote", conn.RemoteAddr().String(), "err", err)
-			continue
-		}
-
-		// Handle the message (forward to mod, enrich, etc.)
-		if err := wsg.botHandler.Handle(conn, msg); err != nil {
-			_ = writeJSONSafe(conn, map[string]string{"error": "handler error: " + err.Error()})
-			wsg.logger.Error("bot handler error", "bot_id", meta.ID, "err", err)
-			continue
-		}
-
-		_ = writeJSONSafe(conn, map[string]string{"status": "ok"})
-	}
-}
-
-func (wsg *WebsocketGateway) read(conn *websocket.Conn, meta cache.ConnectionMetaData, typ string) {
-	defer func() {
-		wsg.unregisterConn(conn, meta, typ)
+		wsg.unregisterConn(conn, meta, _type)
 		wsg.logger.Info("Client disconnected.", "remote", conn.RemoteAddr().String())
 	}()
 
 	ticker := time.NewTicker(30 * time.Second)
+	ctx, cancel := context.WithCancel(context.Background())
 	defer ticker.Stop()
+	defer cancel()
 
 	go func() {
-		for range ticker.C {
-			wsg.sendWebsocketPing(conn)
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				wsg.sendWebsocketPing(conn)
+			}
 		}
 	}()
 
-	switch typ {
-	case "mod":
+	for {
+		typ, data, err := conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
+				wsg.logger.Error("Client unexpectedly closed the connection.", "err", err)
+			}
+			return
+		}
 
-	case "bot":
+		if typ != websocket.TextMessage && typ != websocket.BinaryMessage {
+			continue
+		}
+
+		ts := time.Now().UTC()
+		var message GatewayMessageIn
+		if err := json.Unmarshal(data, &message); err != nil {
+			wsg.sendWebsocketError(conn, "Malformed message.", 400, false)
+			wsg.logger.Warn("Received malformed message json from client.", "remote", conn.RemoteAddr().String(), "err", err)
+			continue
+		}
+
+		message.Type = _type
+		message.ReceivedAt = ts
+		if err := message.Validate(); err != nil {
+			wsg.sendWebsocketError(conn, "Malformed message.", 400, false)
+			wsg.logger.Warn("Received malformed message json from client; validation failed.", "remote", conn.RemoteAddr().String(), "err", err)
+			continue
+		}
+
+		var ok bool
+		var destConn *websocket.Conn
+
+		switch message.Type {
+		case "mod":
+			if wsg.bot != nil {
+				ok = true
+				destConn = wsg.bot
+			} else {
+				ok = false
+			}
+		case "bot":
+			var id string
+			id, ok = wsg.cache.GetByChannelId(message.ID)
+			if ok {
+				var dest *websocket.Conn
+				dest, _, ok = wsg.conns.GetById(id)
+				if ok {
+					destConn = dest
+				} else {
+					wsg.sendWebsocketError(conn, "Internal Server Error", 500, true)
+					wsg.logger.Error("Invalid cache structure.", "remote", conn.RemoteAddr().String(), "id", id)
+					return
+				}
+			}
+		default:
+			panic("invalid message type")
+		}
+
+		if ok {
+			if destConn == nil {
+				wsg.sendWebsocketError(conn, "Internal Server Error", 500, true)
+				wsg.logger.Error("Destination connection unavailable.", "remote", conn.RemoteAddr().String())
+				return
+			}
+
+			err = wsg.sendWebsocketResponse(destConn, GatewayMessageOut{
+				Type:        message.Type,
+				ID:          message.Destination.ID,
+				Author:      message.Author,
+				Content:     message.Content,
+				Meta:        message.Meta,
+				Ts:          message.Ts,
+				ReceivedAt:  message.ReceivedAt,
+				ForwardedAt: time.Now().UTC(),
+			})
+
+			if err != nil {
+				_ = wsg.sendWebsocketResponse(conn, GatewayAck{Status: "failed", Type: message.Type})
+				wsg.logger.Error("Failed to forward message.", "remote", conn.RemoteAddr().String(), "err", err)
+				continue
+			}
+
+			_ = wsg.sendWebsocketResponse(conn, GatewayAck{Status: "completed", Type: message.Type})
+			continue
+		}
+
+		if message.Type == "mod" {
+			wsg.cache.Set(message.ID, message.Destination.ID)
+		}
+
+		out := GatewayMessageOut{
+			Type:        message.Type,
+			ID:          message.Destination.ID,
+			Author:      message.Author,
+			Content:     message.Content,
+			Meta:        message.Meta,
+			Ts:          message.Ts,
+			ReceivedAt:  message.ReceivedAt,
+			ForwardedAt: time.Now().UTC(),
+		}
+
+		queued := wsg.queue.Enqueue(out)
+		if !queued {
+			_ = wsg.sendWebsocketResponse(conn, GatewayAck{Status: "failed", Type: message.Type})
+			wsg.logger.Warn("Failed to queue message.", "remote", conn.RemoteAddr().String())
+			continue
+		}
+
+		_ = wsg.sendWebsocketResponse(conn, GatewayAck{Status: "queued", Type: message.Type})
 	}
 }