diff --git a/routes/chat_routes.py b/routes/chat_routes.py index 764c7e0..8df0af0 100644 --- a/routes/chat_routes.py +++ b/routes/chat_routes.py @@ -436,10 +436,11 @@ def setup_chat_routes( else: logger.warning(f"[doc-inject] NOT FOUND by ID {active_doc_id}") if not active_doc: - active_doc = _doc_db.query(DBDocument).filter( + _session_doc_q = _doc_db.query(DBDocument).filter( DBDocument.session_id == session, DBDocument.is_active == True - ).order_by(DBDocument.updated_at.desc()).first() + ) + active_doc = _owner_session_filter(_session_doc_q, ctx.user).order_by(DBDocument.updated_at.desc()).first() if active_doc: logger.info(f"[doc-inject] found by session fallback: title={active_doc.title!r}") # Last resort: the document the agent itself just created/edited @@ -453,7 +454,8 @@ def setup_chat_routes( from src.tool_implementations import get_active_document _mem_id = get_active_document() if _mem_id: - cand = _doc_db.query(DBDocument).filter(DBDocument.id == _mem_id).first() + _mem_q = _doc_db.query(DBDocument).filter(DBDocument.id == _mem_id) + cand = _owner_session_filter(_mem_q, ctx.user).first() if cand and (not cand.session_id or cand.session_id == session): active_doc = cand logger.info(f"[doc-inject] found by in-memory active id: title={active_doc.title!r} (session_id={cand.session_id!r})") diff --git a/tests/test_security_regressions.py b/tests/test_security_regressions.py index 488c38f..6b3a691 100644 --- a/tests/test_security_regressions.py +++ b/tests/test_security_regressions.py @@ -947,14 +947,18 @@ def test_chat_active_document_lookup_is_owner_scoped(): """The explicit `active_doc_id` path in /api/chat_stream must scope the document lookup to the caller. Resolving by id alone let any user inject another user's document into their own chat context (the session and - in-memory fallbacks were already owner/session-bound; this branch wasn't).""" + in-memory fallbacks also need the same owner gate because active document + state is process-global).""" import re src = Path(__file__).resolve().parents[1] / "routes" / "chat_routes.py" text = src.read_text() # The frontend-supplied id is resolved through the shared owner filter. assert "_owner_session_filter(_doc_q, ctx.user)" in text + assert "_owner_session_filter(_session_doc_q, ctx.user)" in text + assert "_owner_session_filter(_mem_q, ctx.user)" in text # And never by id alone (the previous IDOR shape, whitespace-insensitive). flat = re.sub(r"\s+", " ", text) assert "filter( DBDocument.id == active_doc_id, ).first()" not in flat assert "filter(DBDocument.id == active_doc_id).first()" not in flat + assert "filter(DBDocument.id == _mem_id).first()" not in flat