diff --git a/app/routes/document_routes.py b/app/routes/document_routes.py index 15f5b8d0..0d207b1c 100644 --- a/app/routes/document_routes.py +++ b/app/routes/document_routes.py @@ -1066,6 +1066,7 @@ async def embed_file_upload( @router.post("/query_multiple") async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody): + user_authorized = get_user_id(request) try: # Get the embedding of the query text embedding = get_cached_query_embedding(body.query) @@ -1085,13 +1086,30 @@ async def query_embeddings_by_file_ids(request: Request, body: QueryMultipleBody documents = _apply_distance_threshold(documents) + # Authorization: return only chunks the requester owns (or unowned ones), + # mirroring the per-document check in /query. The filter above scopes by + # file_id only, so without this any caller could read another tenant's + # documents by supplying their file_id. + authorized_documents = [ + (doc, score) + for doc, score in documents + if doc.metadata.get("user_id") in (None, user_authorized) + ] + withheld = len(documents) - len(authorized_documents) + if withheld: + logger.warning( + "query_multiple: withheld %d chunk(s) not owned by user %s", + withheld, + user_authorized, + ) + # Ensure documents list is not empty - if not documents: + if not authorized_documents: raise HTTPException( status_code=404, detail="No documents found for the given query" ) - return documents + return authorized_documents except HTTPException as http_exc: logger.error( "HTTP Exception in query_embeddings_by_file_ids | Status: %d | Detail: %s", diff --git a/tests/test_main.py b/tests/test_main.py index a7fce89a..baa2becf 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -279,3 +279,43 @@ def test_extract_text_from_file(tmp_path, auth_headers): assert json_data["file_id"] == "test_text_123" assert json_data["filename"] == "test_text_extraction.txt" assert json_data["known_type"] is True # text files are known types + + +def test_query_multiple_filters_unauthorized(auth_headers, monkeypatch): + """/query_multiple must withhold chunks owned by another user. + + Regression: it filtered by file_id only (no user_id check), so any caller + could read another tenant's documents by supplying their file_id. + """ + from app.services.vector_store.async_pg_vector import AsyncPgVector + + async def cross_tenant_results(self, embedding, k, filter=None, executor=None): + return [ + ( + Document( + page_content="mine", + metadata={"file_id": "f1", "user_id": "testuser"}, + ), + 0.1, + ), + ( + Document( + page_content="theirs", + metadata={"file_id": "f2", "user_id": "someoneelse"}, + ), + 0.2, + ), + ] + + monkeypatch.setattr( + AsyncPgVector, + "asimilarity_search_with_score_by_vector", + cross_tenant_results, + ) + + data = {"query": "q", "file_ids": ["f1", "f2"], "k": 4} + response = client.post("/query_multiple", json=data, headers=auth_headers) + assert response.status_code == 200, response.text + contents = [row[0]["page_content"] for row in response.json()] + assert "mine" in contents + assert "theirs" not in contents