293 lines
8.6 KiB
Go
293 lines
8.6 KiB
Go
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) handlePush(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := wsg.validateAndUpgradeConnection(w, r)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
if wsg.bodySizeBytes > 0 {
|
|
conn.SetReadLimit(wsg.bodySizeBytes)
|
|
} else {
|
|
conn.SetReadLimit(1 << 20) // sensible default 1MiB
|
|
}
|
|
|
|
_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
|
conn.SetPongHandler(func(appData string) error {
|
|
_ = conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
|
return nil
|
|
})
|
|
|
|
typ, data, err := conn.ReadMessage()
|
|
if err != nil {
|
|
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 {
|
|
wsg.sendWebsocketError(conn, "First message must be a handshake.", 400)
|
|
wsg.logger.Warn("Invalid handshake message type.", "remote", conn.RemoteAddr().String())
|
|
return
|
|
}
|
|
|
|
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 := connectionMetaData{connectionType: handshake.Type}
|
|
|
|
switch handshake.Type {
|
|
case "mod":
|
|
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 = 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 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 = 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:
|
|
wsg.sendWebsocketError(conn, "Unknown handshake.", 400)
|
|
wsg.logger.Warn("Unknown connection type.", "remote", conn.RemoteAddr().String(), "type", handshake.Type)
|
|
return
|
|
}
|
|
}
|
|
|
|
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() {
|
|
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()
|
|
|
|
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"})
|
|
|
|
// Handle pings
|
|
select {
|
|
case <-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
|
|
}
|
|
default:
|
|
}
|
|
}
|
|
}
|
|
|
|
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()
|
|
|
|
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"})
|
|
|
|
// Handle pings
|
|
select {
|
|
case <-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
|
|
}
|
|
default:
|
|
}
|
|
}
|
|
}
|