73 lines
1.7 KiB
Go
73 lines
1.7 KiB
Go
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()
|
|
}
|
|
}
|