package ws import ( "context" "fmt" "log/slog" "net/http" "os/signal" "syscall" "time" "github.com/gorilla/websocket" ) func (g *WebsocketGateway) Start() error { ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) defer stop() return g.Serve(ctx, fmt.Sprintf(":%d", g.port)) } // util func (g *WebsocketGateway) validateApiKey(r *http.Request) bool { apiKey := r.URL.Query().Get("api_key") if apiKey == "" { apiKey = r.Header.Get("X-API-Key") } return !(apiKey == "" || apiKey != g.apiKey) } func writeJSONSafe(c *websocket.Conn, v interface{}) error { _ = c.SetWriteDeadline(time.Now().Add(5 * time.Second)) if err := c.WriteJSON(v); err != nil { // caller handles logging return err } return nil } func loggingMiddleware(logger *slog.Logger, next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { start := time.Now() next.ServeHTTP(w, r) logger.Info("http request", "remote", r.RemoteAddr, "method", r.Method, "path", r.URL.Path, "duration", time.Since(start)) }) } // connections func (g *WebsocketGateway) registerConn(c *websocket.Conn, meta connMetadata) { g.connsMu.Lock() g.conns[c] = meta g.connsMu.Unlock() } func (g *WebsocketGateway) unregisterConn(c *websocket.Conn) { g.connsMu.Lock() delete(g.conns, c) g.connsMu.Unlock() } func (g *WebsocketGateway) closeAll() { g.connsMu.Lock() defer g.connsMu.Unlock() g.logger.Info("closing websocket connections") for c := range g.conns { _ = c.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "shutting down"), time.Now().Add(time.Second)) _ = c.Close() } }