diff --git a/server/internal/metrics/realtime.go b/server/internal/metrics/realtime.go index 45fa7fe7f9..bf807626dd 100644 --- a/server/internal/metrics/realtime.go +++ b/server/internal/metrics/realtime.go @@ -15,6 +15,7 @@ type RealtimeCollector struct { slowEvictionsTotal *prometheus.Desc messagesSentTotal *prometheus.Desc messagesDropped *prometheus.Desc + inboundTooLarge *prometheus.Desc redisConnected *prometheus.Desc redisXAddTotal *prometheus.Desc redisXAddErrors *prometheus.Desc @@ -35,6 +36,7 @@ func NewRealtimeCollector(m *realtime.Metrics) *RealtimeCollector { slowEvictionsTotal: newRealtimeDesc("slow_evictions_total", "Total realtime clients evicted for slow consumption."), messagesSentTotal: newRealtimeDesc("messages_sent_total", "Total realtime messages sent."), messagesDropped: newRealtimeDesc("messages_dropped_total", "Total realtime messages dropped."), + inboundTooLarge: newRealtimeDesc("inbound_too_large_total", "Total realtime connections closed for exceeding the inbound message size limit."), redisConnected: newRealtimeDesc("redis_connected", "Whether the realtime Redis relay is connected."), redisXAddTotal: newRealtimeDesc("redis_xadd_total", "Total Redis XADD operations by the realtime relay."), redisXAddErrors: newRealtimeDesc("redis_xadd_errors_total", "Total Redis XADD errors by the realtime relay."), @@ -58,6 +60,7 @@ func (c *RealtimeCollector) Describe(ch chan<- *prometheus.Desc) { c.slowEvictionsTotal, c.messagesSentTotal, c.messagesDropped, + c.inboundTooLarge, c.redisConnected, c.redisXAddTotal, c.redisXAddErrors, @@ -82,6 +85,7 @@ func (c *RealtimeCollector) Collect(ch chan<- prometheus.Metric) { ch <- prometheus.MustNewConstMetric(c.slowEvictionsTotal, prometheus.CounterValue, float64(m.SlowEvictionsTotal.Load())) ch <- prometheus.MustNewConstMetric(c.messagesSentTotal, prometheus.CounterValue, float64(m.MessagesSentTotal.Load())) ch <- prometheus.MustNewConstMetric(c.messagesDropped, prometheus.CounterValue, float64(m.MessagesDroppedTotal.Load())) + ch <- prometheus.MustNewConstMetric(c.inboundTooLarge, prometheus.CounterValue, float64(m.InboundTooLargeTotal.Load())) ch <- prometheus.MustNewConstMetric(c.redisConnected, prometheus.GaugeValue, boolFloat(m.RedisConnected.Load())) ch <- prometheus.MustNewConstMetric(c.redisXAddTotal, prometheus.CounterValue, float64(m.RedisXAddTotal.Load())) ch <- prometheus.MustNewConstMetric(c.redisXAddErrors, prometheus.CounterValue, float64(m.RedisXAddErrors.Load())) diff --git a/server/internal/metrics/realtime_test.go b/server/internal/metrics/realtime_test.go index 7799e27e06..e4c490b2d7 100644 --- a/server/internal/metrics/realtime_test.go +++ b/server/internal/metrics/realtime_test.go @@ -13,6 +13,7 @@ func TestRealtimeCollectorExposesCounters(t *testing.T) { m := &realtime.Metrics{} m.ActiveConnections.Store(3) m.MessagesSentTotal.Store(11) + m.InboundTooLargeTotal.Store(7) m.RedisConnected.Store(true) m.RedisMirrorPrimaryErrors.Store(2) m.RedisMirrorSecondaryErrors.Store(5) @@ -25,6 +26,7 @@ func TestRealtimeCollectorExposesCounters(t *testing.T) { for _, want := range []string{ "multica_realtime_active_connections 3", "multica_realtime_messages_sent_total 11", + "multica_realtime_inbound_too_large_total 7", "multica_realtime_redis_connected 1", `multica_realtime_redis_mirror_errors_total{target="primary"} 2`, `multica_realtime_redis_mirror_errors_total{target="secondary"} 5`, diff --git a/server/internal/realtime/hub.go b/server/internal/realtime/hub.go index 2cba7c30e8..1b78721bc1 100644 --- a/server/internal/realtime/hub.go +++ b/server/internal/realtime/hub.go @@ -3,6 +3,7 @@ package realtime import ( "context" "encoding/json" + "errors" "log/slog" "net" "net/http" @@ -194,6 +195,15 @@ const ( writeWait = 10 * time.Second pongWait = 60 * time.Second pingPeriod = (pongWait * 9) / 10 + + // inboundReadLimit caps a single inbound message. Every frame a client + // legitimately sends is tiny — the largest is the token auth frame, well + // under 1 KiB — but gorilla buffers a whole message in memory before + // handing it over, and a fragmented message keeps the read deadline alive + // through interleaved pongs. Without a limit one connection can therefore + // grow that buffer without bound and OOM the process. Matches the daemon + // hub limit so both WebSocket surfaces answer this question the same way. + inboundReadLimit = 64 * 1024 ) var upgrader = websocket.Upgrader{ @@ -701,13 +711,25 @@ func authenticateToken(tokenStr string, pr PATResolver, ctx context.Context) (st } // firstMessageAuth reads the first WebSocket message expecting an auth payload. -func firstMessageAuth(conn *websocket.Conn) (string, string) { +// A non-empty errMsg is for the caller to write back before closing the +// connection. closed=true means the connection is already torn down and the +// caller must return without writing anything further. +func firstMessageAuth(conn *websocket.Conn) (token, errMsg string, closed bool) { conn.SetReadDeadline(time.Now().Add(10 * time.Second)) defer conn.SetReadDeadline(time.Time{}) _, raw, err := conn.ReadMessage() if err != nil { - return "", `{"error":"auth timeout or read error"}` + if errors.Is(err, websocket.ErrReadLimit) { + // gorilla has already replied CloseMessageTooBig (1009), so an + // auth_error frame here would be data sent after a close frame. + // Counted separately to keep the breach out of ordinary churn. + M.InboundTooLargeTotal.Add(1) + slog.Warn("ws: pre-auth frame exceeded read limit", "limit_bytes", inboundReadLimit) + conn.Close() + return "", "", true + } + return "", `{"error":"auth timeout or read error"}`, false } var msg struct { @@ -717,10 +739,10 @@ func firstMessageAuth(conn *websocket.Conn) (string, string) { } `json:"payload"` } if err := json.Unmarshal(raw, &msg); err != nil || msg.Type != "auth" || msg.Payload.Token == "" { - return "", `{"error":"expected auth message as first frame"}` + return "", `{"error":"expected auth message as first frame"}`, false } - return msg.Payload.Token, "" + return msg.Payload.Token, "", false } type wsMessageWriter interface { @@ -780,8 +802,16 @@ func HandleWebSocket(hub *Hub, mc MembershipChecker, pr PATResolver, resolveSlug return } + // Bound inbound messages here rather than in readPump: the token auth + // path below reads its first frame before the caller is authenticated, so + // a limit installed any later leaves that read unbounded. + conn.SetReadLimit(inboundReadLimit) + if userID == "" { - tokenStr, errMsg := firstMessageAuth(conn) + tokenStr, errMsg, closed := firstMessageAuth(conn) + if closed { + return + } if errMsg != "" { writeWSAuthErrorAndClose(conn, []byte(errMsg), "workspace_id", workspaceID) return @@ -869,7 +899,17 @@ func (c *Client) readPump() { for { _, raw, err := c.conn.ReadMessage() if err != nil { - if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) { + switch { + case errors.Is(err, websocket.ErrReadLimit): + // Counted separately so an over-limit close stays visible + // instead of blending into ordinary connection churn. + M.InboundTooLargeTotal.Add(1) + slog.Warn("ws: inbound frame exceeded read limit", + "limit_bytes", inboundReadLimit, + "user_id", c.userID, + "workspace_id", c.workspaceID, + ) + case websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure): slog.Debug("websocket read error", "error", err, "user_id", c.userID, "workspace_id", c.workspaceID) } break diff --git a/server/internal/realtime/hub_test.go b/server/internal/realtime/hub_test.go index 8e268cc7d2..2ee1c2b118 100644 --- a/server/internal/realtime/hub_test.go +++ b/server/internal/realtime/hub_test.go @@ -463,3 +463,113 @@ func TestCheckOrigin(t *testing.T) { }) } } + +// waitFor polls cond until it holds or the deadline expires. Hub registration +// and the metrics bump both happen off the connection's own goroutine, so +// asserting on them right after a read would race. +func waitFor(t *testing.T, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("timed out waiting for %s", what) +} + +// The token-auth path reads its first frame before the caller has presented +// any credential, so the read limit has to be in place by then (#6210). +func TestHandleWebSocket_RejectsOversizedFrameBeforeAuth(t *testing.T) { + hub, server := newTestHub(t) + defer server.Close() + + before := M.InboundTooLargeTotal.Load() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws?workspace_id=" + testWorkspaceID + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("failed to connect WebSocket: %v", err) + } + defer conn.Close() + + // Write errors are expected here: the server rejects the frame from its + // declared length and closes before the payload is drained. + _ = conn.WriteMessage(websocket.TextMessage, bytes.Repeat([]byte("a"), inboundReadLimit+1)) + + conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + if _, _, err := conn.ReadMessage(); !websocket.IsCloseError(err, websocket.CloseMessageTooBig) { + t.Fatalf("read after oversized pre-auth frame = %v, want close code %d", err, websocket.CloseMessageTooBig) + } + + waitFor(t, "inbound_too_large_total to increment", func() bool { + return M.InboundTooLargeTotal.Load()-before == 1 + }) + if n := totalClients(hub); n != 0 { + t.Fatalf("oversized pre-auth frame registered %d clients, want 0", n) + } +} + +func TestReadPump_RejectsOversizedFrameAfterAuth(t *testing.T) { + hub, server := newTestHub(t) + defer server.Close() + + conn := connectWS(t, server) + defer conn.Close() + + waitFor(t, "client registration", func() bool { return totalClients(hub) == 1 }) + before := M.InboundTooLargeTotal.Load() + + oversized, err := json.Marshal(map[string]any{ + "type": "ping", + "pad": string(bytes.Repeat([]byte("a"), inboundReadLimit)), + }) + if err != nil { + t.Fatalf("marshal oversized frame: %v", err) + } + _ = conn.WriteMessage(websocket.TextMessage, oversized) + + conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + if _, _, err := conn.ReadMessage(); !websocket.IsCloseError(err, websocket.CloseMessageTooBig) { + t.Fatalf("read after oversized frame = %v, want close code %d", err, websocket.CloseMessageTooBig) + } + + waitFor(t, "inbound_too_large_total to increment", func() bool { + return M.InboundTooLargeTotal.Load()-before == 1 + }) + waitFor(t, "client to be unregistered", func() bool { return totalClients(hub) == 0 }) +} + +// The limit must not clip legitimate traffic: real frames are ~1 KiB, so +// anything comfortably below the cap has to keep working. +func TestReadPump_AcceptsFrameUnderReadLimit(t *testing.T) { + _, server := newTestHub(t) + defer server.Close() + + conn := connectWS(t, server) + defer conn.Close() + + frame, err := json.Marshal(map[string]any{ + "type": "ping", + "pad": string(bytes.Repeat([]byte("a"), inboundReadLimit/2)), + }) + if err != nil { + t.Fatalf("marshal ping frame: %v", err) + } + if len(frame) >= inboundReadLimit { + t.Fatalf("test frame is %d bytes, not under the %d byte limit", len(frame), inboundReadLimit) + } + if err := conn.WriteMessage(websocket.TextMessage, frame); err != nil { + t.Fatalf("write ping: %v", err) + } + + conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + _, raw, err := conn.ReadMessage() + if err != nil { + t.Fatalf("read pong: %v", err) + } + if !strings.Contains(string(raw), "pong") { + t.Fatalf("got %s, want a pong frame", raw) + } +} diff --git a/server/internal/realtime/metrics.go b/server/internal/realtime/metrics.go index 2031e63f5a..d3f9c2c642 100644 --- a/server/internal/realtime/metrics.go +++ b/server/internal/realtime/metrics.go @@ -19,6 +19,11 @@ type Metrics struct { MessagesSentTotal atomic.Int64 MessagesDroppedTotal atomic.Int64 + // InboundTooLargeTotal counts connections closed because a peer sent a + // message over inboundReadLimit, on either the pre-auth or the + // post-auth read path. + InboundTooLargeTotal atomic.Int64 + // Per-event-type send counters keyed by event type string. // Value is *atomic.Int64. eventSent sync.Map @@ -132,17 +137,18 @@ func (m *Metrics) Snapshot() map[string]any { nodeID, _ = v.(string) } return map[string]any{ - "connects_total": m.ConnectsTotal.Load(), - "disconnects_total": m.DisconnectsTotal.Load(), - "active_connections": m.ActiveConnections.Load(), - "slow_evictions_total": m.SlowEvictionsTotal.Load(), - "messages_sent_total": m.MessagesSentTotal.Load(), - "messages_dropped_total": m.MessagesDroppedTotal.Load(), - "events_sent_by_type": snapshotCounters(&m.eventSent), - "subscribes_total": snapshotCounters(&m.subscribeTotal), - "unsubscribes_total": snapshotCounters(&m.unsubscribeTotal), - "subscribe_denied_total": snapshotCounters(&m.subscribeDeniedTotal), - "active_scope_rooms": snapshotCounters(&m.scopeRooms), + "connects_total": m.ConnectsTotal.Load(), + "disconnects_total": m.DisconnectsTotal.Load(), + "active_connections": m.ActiveConnections.Load(), + "slow_evictions_total": m.SlowEvictionsTotal.Load(), + "messages_sent_total": m.MessagesSentTotal.Load(), + "messages_dropped_total": m.MessagesDroppedTotal.Load(), + "inbound_too_large_total": m.InboundTooLargeTotal.Load(), + "events_sent_by_type": snapshotCounters(&m.eventSent), + "subscribes_total": snapshotCounters(&m.subscribeTotal), + "unsubscribes_total": snapshotCounters(&m.unsubscribeTotal), + "subscribe_denied_total": snapshotCounters(&m.subscribeDeniedTotal), + "active_scope_rooms": snapshotCounters(&m.scopeRooms), "redis": map[string]any{ "connected": m.RedisConnected.Load(), "node_id": nodeID, @@ -168,6 +174,7 @@ func (m *Metrics) Reset() { m.SlowEvictionsTotal.Store(0) m.MessagesSentTotal.Store(0) m.MessagesDroppedTotal.Store(0) + m.InboundTooLargeTotal.Store(0) m.eventSent.Range(func(k, _ any) bool { m.eventSent.Delete(k); return true }) m.subscribeTotal.Range(func(k, _ any) bool { m.subscribeTotal.Delete(k); return true }) m.unsubscribeTotal.Range(func(k, _ any) bool { m.unsubscribeTotal.Delete(k); return true })