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:
Claude 2026-04-14 21:16:39 +00:00
parent 8289ac7de3
commit ff48c99d6c
No known key found for this signature in database

View file

@ -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.