mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-16 23:43:03 +00:00
fix(stream): address code review findings on resume-stream
- 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.
This commit is contained in:
parent
8289ac7de3
commit
ff48c99d6c
1 changed files with 45 additions and 23 deletions
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue