diff --git a/src/documents/tests/test_permission_filtering_security.py b/src/documents/tests/test_permission_filtering_security.py index ec31ec22a..7d84bcb86 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 @@ -145,3 +150,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 f64b6ea04..8c2a84996 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -612,11 +612,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, @@ -625,23 +625,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( @@ -651,11 +651,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 528d6ced5..10987f448 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 @@ -2265,10 +2266,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) @@ -2736,10 +2735,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( @@ -3347,10 +3344,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") @@ -3522,10 +3517,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] @@ -3729,11 +3722,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()