diff --git a/src/documents/tests/test_permission_filtering_security.py b/src/documents/tests/test_permission_filtering_security.py index 0db6a692a..a49bb9f34 100644 --- a/src/documents/tests/test_permission_filtering_security.py +++ b/src/documents/tests/test_permission_filtering_security.py @@ -1,10 +1,15 @@ from __future__ import annotations +from unittest.mock import patch + import pytest from django.contrib.auth.models import AnonymousUser from django.contrib.auth.models import Group +from django.contrib.auth.models import Permission from django.contrib.auth.models import User +from django.test import override_settings from guardian.shortcuts import assign_perm +from rest_framework.test import APIClient from documents.permissions import permitted_document_ids from documents.tests.factories import DocumentFactory @@ -151,3 +156,46 @@ class TestPermittedDocumentIdsIncludeDeleted: expected_visible=[], expected_hidden=[doc.pk], ) + + +@pytest.mark.django_db +class TestAiChatAllDocumentsPermissionBoundary: + """ + Regression test pinning the "ask across all documents" AI chat behavior + (ChatStreamingView.post, no document_id) to the same owner/permission + boundary enforced by permitted_document_ids(). This call site was + migrated from get_objects_for_user_owner_aware() to + permitted_document_ids() in Task 5; this test must stay green across + that swap. + """ + + ENDPOINT = "/api/documents/chat/" + + @override_settings(AI_ENABLED=True) + @patch("documents.views.stream_chat_with_documents") + def test_chat_all_documents_excludes_unshared_document(self, mock_stream_chat): + mock_stream_chat.return_value = iter([b"data"]) + + owner = User.objects.create_user(username="owner") + asker = User.objects.create_user(username="asker") + asker.user_permissions.add( + *Permission.objects.filter(codename="view_document"), + ) + shared = DocumentFactory(owner=owner) + not_shared = DocumentFactory(owner=owner) + assign_perm("view_document", asker, shared) + + client = APIClient() + client.force_authenticate(user=asker) + response = client.post( + self.ENDPOINT, + data={"q": "question"}, + format="json", + ) + + assert response.status_code == 200 + mock_stream_chat.assert_called_once() + _, kwargs = mock_stream_chat.call_args + visible_ids = {doc.pk for doc in kwargs["documents"]} + assert shared.pk in visible_ids + assert not_shared.pk not in visible_ids diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 0cf1bcdae..c385e9b57 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -648,11 +648,11 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): self.assertIn(b"AI is required for this feature", response.content) @patch("documents.views.stream_chat_with_documents") - @patch("documents.views.get_objects_for_user_owner_aware") + @patch("documents.views.permitted_document_ids") @override_settings(AI_ENABLED=True) - def test_post_no_document_id(self, mock_get_objects, mock_stream_chat) -> None: + def test_post_no_document_id(self, mock_permitted_ids, mock_stream_chat) -> None: self.grant_view_document_permission() - mock_get_objects.return_value = [self.document] + mock_permitted_ids.return_value = [self.document.pk] mock_stream_chat.return_value = iter([b"data"]) response = self.client.post( self.ENDPOINT, @@ -661,23 +661,23 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): ) self.assertEqual(response.status_code, 200) self.assertEqual(response["Content-Type"], "text/event-stream") - mock_stream_chat.assert_called_once_with( - query_str="question", - documents=[self.document], - output_language=None, - ) + mock_stream_chat.assert_called_once() + call_kwargs = mock_stream_chat.call_args.kwargs + self.assertEqual(call_kwargs["query_str"], "question") + self.assertEqual(list(call_kwargs["documents"]), [self.document]) + self.assertIsNone(call_kwargs["output_language"]) @patch("documents.views.stream_chat_with_documents") - @patch("documents.views.get_objects_for_user_owner_aware") + @patch("documents.views.permitted_document_ids") @override_settings(AI_ENABLED=True) def test_post_uses_user_display_language( self, - mock_get_objects, + mock_permitted_ids, mock_stream_chat, ) -> None: UiSettings.objects.create(user=self.user, settings={"language": "de-de"}) self.grant_view_document_permission() - mock_get_objects.return_value = [self.document] + mock_permitted_ids.return_value = [self.document.pk] mock_stream_chat.return_value = iter([b"data"]) response = self.client.post( @@ -687,11 +687,11 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase): ) self.assertEqual(response.status_code, 200) - mock_stream_chat.assert_called_once_with( - query_str="question", - documents=[self.document], - output_language="de-de", - ) + mock_stream_chat.assert_called_once() + call_kwargs = mock_stream_chat.call_args.kwargs + self.assertEqual(call_kwargs["query_str"], "question") + self.assertEqual(list(call_kwargs["documents"]), [self.document]) + self.assertEqual(call_kwargs["output_language"], "de-de") @patch("documents.views.stream_chat_with_documents") @override_settings(AI_ENABLED=True) diff --git a/src/documents/views.py b/src/documents/views.py index 405913f95..4f1b2ba50 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -177,6 +177,7 @@ from documents.permissions import get_objects_for_user_owner_aware from documents.permissions import has_global_statistics_permission from documents.permissions import has_perms_owner_aware from documents.permissions import has_system_status_permission +from documents.permissions import permitted_document_ids from documents.permissions import set_permissions_for_object from documents.plugins.date_parsing import get_date_parser from documents.schema import generate_object_with_permissions_schema @@ -2270,10 +2271,8 @@ class ChatStreamingView(GenericAPIView[Any]): documents = [document] else: - documents = get_objects_for_user_owner_aware( - request.user, - "view_document", - Document, + documents = Document.objects.filter( + id__in=permitted_document_ids(request.user), ) output_language = _get_llm_output_language(ai_config=ai_config, request=request) @@ -2741,10 +2740,8 @@ class DocumentSelectionMixin: for key, value in filters.items() if key not in _TANTIVY_SEARCH_PARAM_NAMES } - permitted_documents = get_objects_for_user_owner_aware( - user, - permission_codename, - Document, + permitted_documents = Document.objects.filter( + id__in=permitted_document_ids(user), ) # orm-filtered docs filtered_documents = DocumentFilterSet( @@ -3352,10 +3349,8 @@ class SelectionDataView(GenericAPIView[Any]): serializer.is_valid(raise_exception=True) ids = serializer.validated_data.get("documents") - permitted_documents = get_objects_for_user_owner_aware( - request.user, - "documents.view_document", - Document, + permitted_documents = Document.objects.filter( + id__in=permitted_document_ids(request.user), ) if permitted_documents.filter(pk__in=ids).count() != len(ids): return HttpResponseForbidden("Insufficient permissions") @@ -3527,10 +3522,8 @@ class GlobalSearchView(PassUserMixin): OBJECT_LIMIT = 3 docs = [] if request.user.has_perm("documents.view_document"): - all_docs = get_objects_for_user_owner_aware( - request.user, - "view_document", - Document, + all_docs = Document.objects.filter( + id__in=permitted_document_ids(request.user), ) if db_only: docs = all_docs.filter(title__icontains=query)[:OBJECT_LIMIT] @@ -3734,11 +3727,7 @@ class StatisticsView(GenericAPIView[Any]): documents = ( Document.objects.all() if can_view_global_stats - else get_objects_for_user_owner_aware( - user, - "documents.view_document", - Document, - ) + else Document.objects.filter(id__in=permitted_document_ids(user)) ).filter(root_document__isnull=True) tags = ( Tag.objects.all()