320 lines
9.2 KiB
Go
320 lines
9.2 KiB
Go
package ws
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"homestead/homestead_gateway/util/cache"
|
|
"homestead/homestead_gateway/util/config"
|
|
"homestead/homestead_gateway/util/queue"
|
|
"log/slog"
|
|
"net"
|
|
"net/http"
|
|
"os/signal"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
func NewWebsocketGateway(cfg config.GatewayConfig, logger *slog.Logger, closefn func() error) *WebsocketGateway {
|
|
return &WebsocketGateway{
|
|
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,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
return true // local by default; change for production
|
|
},
|
|
},
|
|
bodySizeBytes: int64(cfg.BodySize) * 1024 * 1024,
|
|
}
|
|
}
|
|
|
|
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("/sync", wsg.handleSync)
|
|
mux.HandleFunc("/health", wsg.handleHealth)
|
|
|
|
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(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, _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 {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
wsg.sendWebsocketPing(conn)
|
|
}
|
|
}
|
|
}()
|
|
|
|
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
|
|
}
|
|
|
|
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})
|
|
}
|
|
}
|