From ff48c99d6c88e50e5655750c2dfdbe6bc6f58e1b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 14 Apr 2026 21:16:39 +0000 Subject: [PATCH] fix(stream): address code review findings on resume-stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Critical (auth): the resume-stream handler verified chat ownership but not that the requested message_id actually lives in that chat. Because the Redis log is keyed by message_id alone, an attacker who learned a victim's message_id could pass one of their own chat_ids to satisfy the ownership check and replay the victim's stream. Added an explicit message-to-chat binding check via Chats.get_message_by_id_and_message_id after the ownership check. - Safety: reject non-dict payloads (`if not isinstance(data, dict)`) so stray client input — string/list/null — doesn't raise AttributeError on `data.get(...)`. - Perf: the per-token log append was doing two Redis round-trips (XADD + EXPIRE) on the streaming hot path. Collapse them into a single pipeline execute (one RTT), and refresh the TTL only every 64 appends instead of every append. With a 1h TTL that still leaves comfortable headroom for even pathologically long responses without EXPIRE ever risking mid-stream expiry. - Protocol cleanup: removed the resume-stream:ack emission. The frontend doesn't consume it and YAGNI — the seq idempotency guard in chatEventHandler already delivers the observability (you can see last_seq advance as replays arrive). Can be added back when a concrete client-side use case appears. --- backend/open_webui/socket/main.py | 68 ++++++++++++++++++++----------- 1 file changed, 45 insertions(+), 23 deletions(-) diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index ed3b821f92..38a08db402 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -196,19 +196,37 @@ def _stream_key(message_id: str) -> str: return f'{REDIS_KEY_PREFIX}:stream:{message_id}' +# Refresh the resume-log TTL only every N writes instead of every write. +# XADD resets the key's idle time but not its absolute TTL, so we must +# EXPIRE occasionally. Doing it once per N appends amortizes that extra +# round-trip away from the per-token hot path. With N=64 and the default +# 1-hour TTL, even a pathologically long 128k-token response only +# triggers ~2000 EXPIRE calls total and never risks TTL expiry mid-stream. +RESUME_STREAM_TTL_REFRESH_EVERY = 64 + + async def _stream_log_append(message_id: str, envelope: dict, seq: int) -> None: - """Append an outbound WS envelope to the resume log.""" + """Append an outbound WS envelope to the resume log. + + Uses a Redis pipeline to collapse XADD + (occasional) EXPIRE into a + single network round-trip per call, keeping the added latency on the + streaming hot path bounded to one Redis RTT. + """ if REDIS is None or not message_id: return try: - await REDIS.xadd( - _stream_key(message_id), + refresh_ttl = (seq == 1) or (seq % RESUME_STREAM_TTL_REFRESH_EVERY == 0) + key = _stream_key(message_id) + pipe = REDIS.pipeline(transaction=False) + pipe.xadd( + key, {'seq': str(seq), 'payload': json.dumps(envelope)}, maxlen=RESUME_STREAM_MAXLEN, approximate=True, ) - # Refresh TTL on every append so active streams don't expire mid-run. - await REDIS.expire(_stream_key(message_id), RESUME_STREAM_TTL_SEC) + if refresh_ttl: + pipe.expire(key, RESUME_STREAM_TTL_SEC) + await pipe.execute() except Exception as e: log.debug(f'stream resume log append failed for {message_id}: {e}') @@ -620,14 +638,18 @@ async def resume_stream(sid, data): Client payload: `{chat_id, message_id, last_seq}`. Flow: - 1. Authenticate the session and verify the user owns the chat. + 1. Authenticate the session and verify the user owns the chat AND + that the requested message_id actually belongs to that chat. + Both checks are required because the Redis stream log is keyed + by message_id alone — without the message-to-chat binding check, + an attacker who obtained a victim's message_id could satisfy the + chat-ownership check with any chat they own and then read the + victim's stream. 2. Read entries from the Redis resume log for this message_id whose `seq` is greater than `last_seq`. 3. Emit each entry as a normal `events` event to THIS session only (via `to=sid`). Other sessions for the same user keep receiving the live stream unchanged. - 4. Send a final `resume-stream:ack` so the client can resolve its - pending resume state. No-op when Redis is not configured — in that deployment mode, refresh during streaming falls back to the existing behavior (wait for the @@ -636,6 +658,11 @@ async def resume_stream(sid, data): if REDIS is None: return + # Reject malformed payloads early so `data.get(...)` never throws on a + # non-object client input (string/list/null/etc). + if not isinstance(data, dict): + return + user = SESSION_POOL.get(sid) if not user: return @@ -651,13 +678,20 @@ async def resume_stream(sid, data): if not user_id or not chat_id or not message_id: return - # Auth: the stream log is keyed by message_id only, so a chat-ownership - # check is essential to prevent a client from reading another user's - # stream by guessing a message_id. + # Step 1a: user owns the chat. chat = await Chats.get_chat_by_id_and_user_id(chat_id, user_id) if not chat: return + # Step 1b: message_id actually lives in this chat. Without this, a + # caller who knows a victim's message_id could pass one of their OWN + # chat_ids (satisfying 1a) and read the victim's stream. + # `get_message_by_id_and_message_id` returns {} when the chat exists + # but the message_id isn't present, so check for a real id field. + message = await Chats.get_message_by_id_and_message_id(chat_id, message_id) + if not message or not message.get('id'): + return + envelopes = await _stream_log_read(message_id, last_seq) for envelope in envelopes: # Replay to the requesting session only. Live listeners in @@ -665,18 +699,6 @@ async def resume_stream(sid, data): # emit path and must not see these duplicates. await sio.emit('events', envelope, to=sid) - last_replayed_seq = envelopes[-1].get('seq') if envelopes else last_seq - await sio.emit( - 'resume-stream:ack', - { - 'chat_id': chat_id, - 'message_id': message_id, - 'replayed': len(envelopes), - 'last_seq': last_replayed_seq, - }, - to=sid, - ) - def normalize_document_id(document_id: str) -> str: """Canonicalize document IDs to prevent auth bypass via prefix variants.