Security: fixes for v3 beta (#12838)

src/documents/admin.py

@@ -198,6 +198,7 @@ class ShareLinksAdmin(GuardedModelAdmin):
 class ShareLinkBundleAdmin(GuardedModelAdmin):
     list_display = ("created", "status", "expiration", "owner", "slug")
     list_filter = ("status", "created", "expiration", "owner")
+    readonly_fields = ("file_path",)
     search_fields = ("slug",)
 
     def get_queryset(self, request):  # pragma: no cover

src/documents/models.py

@@ -1019,7 +1019,17 @@ class ShareLinkBundle(models.Model):
     def absolute_file_path(self) -> Path | None:
         if not self.file_path:
             return None
-        return (settings.SHARE_LINK_BUNDLE_DIR / Path(self.file_path)).resolve()
+        relative_path = Path(self.file_path)
+        if relative_path.is_absolute():
+            return None
+
+        bundle_dir = settings.SHARE_LINK_BUNDLE_DIR.resolve()
+        absolute_path = (bundle_dir / relative_path).resolve()
+        try:
+            absolute_path.relative_to(bundle_dir)
+        except ValueError:
+            return None
+        return absolute_path
 
     def remove_file(self) -> None:
         if self.absolute_file_path is not None and self.absolute_file_path.exists():

src/documents/tests/test_share_link_bundles.py

@@ -124,7 +124,7 @@ class ShareLinkBundleAPITests(DirectoriesMixin, APITestCase):
         self.assertIn("document_ids", response.data)
 
     def test_download_ready_bundle_streams_file(self) -> None:
-        bundle_file = Path(self.dirs.media_dir) / "bundles" / "ready.zip"
+        bundle_file = settings.SHARE_LINK_BUNDLE_DIR / "bundles" / "ready.zip"
         bundle_file.parent.mkdir(parents=True, exist_ok=True)
         bundle_file.write_bytes(b"binary-zip-content")
 
@@ -132,7 +132,7 @@ class ShareLinkBundleAPITests(DirectoriesMixin, APITestCase):
             slug="readyslug",
             file_version=ShareLink.FileVersion.ARCHIVE,
             status=ShareLinkBundle.Status.READY,
-            file_path=str(bundle_file),
+            file_path=str(bundle_file.relative_to(settings.SHARE_LINK_BUNDLE_DIR)),
         )
         bundle.documents.set([self.document])
 
@@ -199,11 +199,11 @@ class ShareLinkBundleTaskTests(DirectoriesMixin, APITestCase):
         self.document = DocumentFactory.create()
 
     def test_cleanup_expired_share_link_bundles(self) -> None:
-        expired_path = Path(self.dirs.media_dir) / "expired.zip"
+        expired_path = settings.SHARE_LINK_BUNDLE_DIR / "expired.zip"
         expired_path.parent.mkdir(parents=True, exist_ok=True)
         expired_path.write_bytes(b"expired")
 
-        active_path = Path(self.dirs.media_dir) / "active.zip"
+        active_path = settings.SHARE_LINK_BUNDLE_DIR / "active.zip"
         active_path.write_bytes(b"active")
 
         expired_bundle = ShareLinkBundle.objects.create(
@@ -211,7 +211,7 @@ class ShareLinkBundleTaskTests(DirectoriesMixin, APITestCase):
             file_version=ShareLink.FileVersion.ARCHIVE,
             status=ShareLinkBundle.Status.READY,
             expiration=timezone.now() - timedelta(days=1),
-            file_path=str(expired_path),
+            file_path=expired_path.name,
         )
         expired_bundle.documents.set([self.document])
 
@@ -220,7 +220,7 @@ class ShareLinkBundleTaskTests(DirectoriesMixin, APITestCase):
             file_version=ShareLink.FileVersion.ARCHIVE,
             status=ShareLinkBundle.Status.READY,
             expiration=timezone.now() + timedelta(days=1),
-            file_path=str(active_path),
+            file_path=active_path.name,
         )
         active_bundle.documents.set([self.document])
 
@@ -424,7 +424,7 @@ class ShareLinkBundleFilterSetTests(DirectoriesMixin, APITestCase):
 
 
 class ShareLinkBundleModelTests(DirectoriesMixin, APITestCase):
-    def test_absolute_file_path_handles_relative_and_absolute(self) -> None:
+    def test_absolute_file_path_handles_relative_path(self) -> None:
         relative_path = Path("relative.zip")
         bundle = ShareLinkBundle.objects.create(
             slug="relative-bundle",
@@ -437,10 +437,23 @@ class ShareLinkBundleModelTests(DirectoriesMixin, APITestCase):
             (settings.SHARE_LINK_BUNDLE_DIR / relative_path).resolve(),
         )
 
-        absolute_path = Path(self.dirs.media_dir) / "absolute.zip"
-        bundle.file_path = str(absolute_path)
+    def test_absolute_file_path_rejects_absolute_path(self) -> None:
+        bundle = ShareLinkBundle.objects.create(
+            slug="absolute-bundle",
+            file_version=ShareLink.FileVersion.ORIGINAL,
+            file_path=str(Path(self.dirs.media_dir) / "absolute.zip"),
+        )
 
-        self.assertEqual(bundle.absolute_file_path.resolve(), absolute_path.resolve())
+        self.assertIsNone(bundle.absolute_file_path)
+
+    def test_absolute_file_path_rejects_traversal_outside_bundle_dir(self) -> None:
+        bundle = ShareLinkBundle.objects.create(
+            slug="traversal-bundle",
+            file_version=ShareLink.FileVersion.ORIGINAL,
+            file_path="../escaped.zip",
+        )
+
+        self.assertIsNone(bundle.absolute_file_path)
 
     def test_str_returns_translated_slug(self) -> None:
         bundle = ShareLinkBundle.objects.create(

src/documents/tests/test_views.py

@@ -437,8 +437,14 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase):
         )
         super().setUp()
 
+    def grant_view_document_permission(self) -> None:
+        self.user.user_permissions.add(
+            *Permission.objects.filter(codename="view_document"),
+        )
+
     @override_settings(AI_ENABLED=False)
     def test_post_ai_disabled(self) -> None:
+        self.grant_view_document_permission()
         response = self.client.post(
             self.ENDPOINT,
             data='{"q": "question"}',
@@ -451,6 +457,7 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase):
     @patch("documents.views.get_objects_for_user_owner_aware")
     @override_settings(AI_ENABLED=True)
     def test_post_no_document_id(self, mock_get_objects, mock_stream_chat) -> None:
+        self.grant_view_document_permission()
         mock_get_objects.return_value = [self.document]
         mock_stream_chat.return_value = iter([b"data"])
         response = self.client.post(
@@ -464,6 +471,7 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase):
     @patch("documents.views.stream_chat_with_documents")
     @override_settings(AI_ENABLED=True)
     def test_post_with_document_id(self, mock_stream_chat) -> None:
+        self.grant_view_document_permission()
         mock_stream_chat.return_value = iter([b"data"])
         response = self.client.post(
             self.ENDPOINT,
@@ -475,6 +483,7 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase):
 
     @override_settings(AI_ENABLED=True)
     def test_post_with_invalid_document_id(self) -> None:
+        self.grant_view_document_permission()
         response = self.client.post(
             self.ENDPOINT,
             data='{"q": "question", "document_id": 999999}',
@@ -486,6 +495,7 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase):
     @patch("documents.views.has_perms_owner_aware")
     @override_settings(AI_ENABLED=True)
     def test_post_with_document_id_no_permission(self, mock_has_perms) -> None:
+        self.grant_view_document_permission()
         mock_has_perms.return_value = False
         response = self.client.post(
             self.ENDPOINT,
@@ -494,3 +504,31 @@ class TestAIChatStreamingView(DirectoriesMixin, TestCase):
         )
         self.assertEqual(response.status_code, 403)
         self.assertIn(b"Insufficient permissions", response.content)
+
+    @patch("documents.views.stream_chat_with_documents")
+    @override_settings(AI_ENABLED=True)
+    def test_post_no_document_id_requires_view_document_permission(
+        self,
+        mock_stream_chat,
+    ) -> None:
+        response = self.client.post(
+            self.ENDPOINT,
+            data='{"q": "question"}',
+            content_type="application/json",
+        )
+        self.assertEqual(response.status_code, 403)
+        mock_stream_chat.assert_not_called()
+
+    @patch("documents.views.stream_chat_with_documents")
+    @override_settings(AI_ENABLED=True)
+    def test_post_with_document_id_requires_view_document_permission(
+        self,
+        mock_stream_chat,
+    ) -> None:
+        response = self.client.post(
+            self.ENDPOINT,
+            data=f'{{"q": "question", "document_id": {self.document.pk}}}',
+            content_type="application/json",
+        )
+        self.assertEqual(response.status_code, 403)
+        mock_stream_chat.assert_not_called()

src/documents/views.py

@@ -2150,7 +2150,7 @@ class ChatStreamingSerializer(serializers.Serializer[dict[str, Any]]):
     name="dispatch",
 )
 class ChatStreamingView(GenericAPIView[Any]):
-    permission_classes = (IsAuthenticated,)
+    permission_classes = (IsAuthenticated, ViewDocumentsPermissions)
     serializer_class = ChatStreamingSerializer
 
     def post(self, request, *args, **kwargs):

src/documents/workflows/webhooks.py

@@ -4,70 +4,12 @@ import httpx
 from celery import shared_task
 from django.conf import settings
 
-from paperless.network import format_host_for_url
-from paperless.network import is_public_ip
-from paperless.network import resolve_hostname_ips
+from paperless.network import PinnedHostHTTPTransport
 from paperless.network import validate_outbound_http_url
 
 logger = logging.getLogger("paperless.workflows.webhooks")
 
 
-class WebhookTransport(httpx.HTTPTransport):
-    """
-    Transport that resolves/validates hostnames and rewrites to a vetted IP
-    while keeping Host/SNI as the original hostname.
-    """
-
-    def __init__(
-        self,
-        hostname: str,
-        *args,
-        allow_internal: bool = False,
-        **kwargs,
-    ) -> None:
-        super().__init__(*args, **kwargs)
-        self.hostname = hostname
-        self.allow_internal = allow_internal
-
-    def handle_request(self, request: httpx.Request) -> httpx.Response:
-        hostname = request.url.host
-
-        if not hostname:
-            raise httpx.ConnectError("No hostname in request URL")
-
-        try:
-            ips = resolve_hostname_ips(hostname)
-        except ValueError as e:
-            raise httpx.ConnectError(str(e)) from e
-
-        if not self.allow_internal:
-            for ip_str in ips:
-                if not is_public_ip(ip_str):
-                    raise httpx.ConnectError(
-                        f"Connection blocked: {hostname} resolves to a non-public address",
-                    )
-
-        ip_str = ips[0]
-        formatted_ip = format_host_for_url(ip_str)
-
-        new_headers = httpx.Headers(request.headers)
-        if "host" in new_headers:
-            del new_headers["host"]
-        new_headers["Host"] = hostname
-        new_url = request.url.copy_with(host=formatted_ip)
-
-        request = httpx.Request(
-            method=request.method,
-            url=new_url,
-            headers=new_headers,
-            content=request.stream,
-            extensions=request.extensions,
-        )
-        request.extensions["sni_hostname"] = hostname
-
-        return super().handle_request(request)
-
-
 @shared_task(
     retry_backoff=True,
     autoretry_for=(httpx.HTTPStatusError,),
@@ -83,7 +25,7 @@ def send_webhook(
     as_json: bool = False,
 ):
     try:
-        parsed = validate_outbound_http_url(
+        validate_outbound_http_url(
             url,
             allowed_schemes=settings.WEBHOOKS_ALLOWED_SCHEMES,
             allowed_ports=settings.WEBHOOKS_ALLOWED_PORTS,
@@ -94,12 +36,7 @@ def send_webhook(
         logger.warning("Webhook blocked: %s", e)
         raise
 
-    hostname = parsed.hostname
-    if hostname is None:  # pragma: no cover
-        raise ValueError("Invalid URL scheme or hostname.")
-
-    transport = WebhookTransport(
-        hostname=hostname,
+    transport = PinnedHostHTTPTransport(
         allow_internal=settings.WEBHOOKS_ALLOW_INTERNAL_REQUESTS,
     )
 

src/paperless/network.py

@@ -4,6 +4,8 @@ from collections.abc import Collection
 from urllib.parse import ParseResult
 from urllib.parse import urlparse
 
+import httpx
+
 
 def is_public_ip(ip: str | int) -> bool:
     try:
@@ -74,3 +76,121 @@ def validate_outbound_http_url(
                 )
 
     return parsed
+
+
+def _rewrite_request_to_pinned_ip(
+    request: httpx.Request,
+    *,
+    allow_internal: bool,
+) -> httpx.Request:
+    hostname = request.url.host
+
+    if not hostname:
+        raise httpx.ConnectError("No hostname in request URL")
+
+    try:
+        ips = resolve_hostname_ips(hostname)
+    except ValueError as e:
+        raise httpx.ConnectError(str(e)) from e
+
+    if not allow_internal:
+        for ip_str in ips:
+            if not is_public_ip(ip_str):
+                raise httpx.ConnectError(
+                    f"Connection blocked: {hostname} resolves to a non-public address",
+                )
+
+    ip_str = ips[0]
+    formatted_ip = format_host_for_url(ip_str)
+
+    new_headers = httpx.Headers(request.headers)
+    if "host" in new_headers:
+        del new_headers["host"]
+    host_header = format_host_for_url(hostname)
+    default_port = 443 if request.url.scheme == "https" else 80
+    if request.url.port and request.url.port != default_port:
+        host_header = f"{host_header}:{request.url.port}"
+    new_headers["Host"] = host_header
+    new_url = request.url.copy_with(host=formatted_ip)
+
+    rewritten_request = httpx.Request(
+        method=request.method,
+        url=new_url,
+        headers=new_headers,
+        content=request.stream,
+        extensions=request.extensions,
+    )
+    rewritten_request.extensions["sni_hostname"] = hostname
+
+    return rewritten_request
+
+
+class PinnedHostHTTPTransport(httpx.HTTPTransport):
+    """
+    HTTP transport that resolves/validates hostnames per request and connects to
+    a vetted IP while preserving the original Host header and TLS SNI hostname.
+    """
+
+    def __init__(
+        self,
+        *args,
+        allow_internal: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.allow_internal = allow_internal
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        request = _rewrite_request_to_pinned_ip(
+            request,
+            allow_internal=self.allow_internal,
+        )
+        return super().handle_request(request)
+
+
+class PinnedHostAsyncHTTPTransport(httpx.AsyncHTTPTransport):
+    """
+    Async variant of PinnedHostHTTPTransport.
+    """
+
+    def __init__(
+        self,
+        *args,
+        allow_internal: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.allow_internal = allow_internal
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        request = _rewrite_request_to_pinned_ip(
+            request,
+            allow_internal=self.allow_internal,
+        )
+        return await super().handle_async_request(request)
+
+
+def create_pinned_httpx_client(
+    url: str,
+    *,
+    allow_internal: bool = False,
+    **kwargs,
+) -> httpx.Client:
+    validate_outbound_http_url(url, allow_internal=allow_internal)
+    return httpx.Client(
+        transport=PinnedHostHTTPTransport(allow_internal=allow_internal),
+        **kwargs,
+    )
+
+
+def create_pinned_async_httpx_client(
+    url: str,
+    *,
+    allow_internal: bool = False,
+    **kwargs,
+) -> httpx.AsyncClient:
+    validate_outbound_http_url(url, allow_internal=allow_internal)
+    return httpx.AsyncClient(
+        transport=PinnedHostAsyncHTTPTransport(allow_internal=allow_internal),
+        **kwargs,
+    )

src/paperless/tests/test_network.py

@@ -0,0 +1,50 @@
+from unittest import mock
+
+import httpx
+import pytest
+
+from paperless.network import PinnedHostHTTPTransport
+
+
+def test_pinned_host_transport_blocks_internal_rebinding():
+    transport = PinnedHostHTTPTransport(allow_internal=False)
+    request = httpx.Request("GET", "http://example.com/test")
+
+    with (
+        mock.patch(
+            "paperless.network.resolve_hostname_ips",
+            return_value=["127.0.0.1"],
+        ),
+        pytest.raises(httpx.ConnectError, match="non-public address"),
+    ):
+        transport.handle_request(request)
+
+
+def test_pinned_host_transport_rewrites_to_vetted_ip():
+    transport = PinnedHostHTTPTransport(allow_internal=False)
+    request = httpx.Request("GET", "https://example.com:8443/test")
+
+    def assert_rewritten_request(
+        self,
+        rewritten_request,
+    ):
+        assert str(rewritten_request.url) == "https://93.184.216.34:8443/test"
+        assert rewritten_request.headers["Host"] == "example.com:8443"
+        assert rewritten_request.extensions["sni_hostname"] == "example.com"
+        return httpx.Response(200, request=rewritten_request)
+
+    with (
+        mock.patch(
+            "paperless.network.resolve_hostname_ips",
+            return_value=["93.184.216.34"],
+        ),
+        mock.patch.object(
+            httpx.HTTPTransport,
+            "handle_request",
+            autospec=True,
+            side_effect=assert_rewritten_request,
+        ),
+    ):
+        response = transport.handle_request(request)
+
+    assert response.status_code == 200

src/paperless_ai/ai_classifier.py

@@ -60,11 +60,14 @@ def get_context_for_document(
         if user
         else None
     )
+    visible_document_ids = (
+        list(visible_documents.values_list("pk", flat=True))
+        if visible_documents is not None
+        else None
+    )
     similar_docs = query_similar_documents(
         document=doc,
-        document_ids=[document.pk for document in visible_documents]
-        if visible_documents
-        else None,
+        document_ids=visible_document_ids,
     )[:max_docs]
     context_blocks = []
     for similar in similar_docs:

src/paperless_ai/client.py

@@ -9,6 +9,10 @@ if TYPE_CHECKING:
     from llama_index.llms.openai_like import OpenAILike
 
 from paperless.config import AIConfig
+from paperless.network import PinnedHostAsyncHTTPTransport
+from paperless.network import PinnedHostHTTPTransport
+from paperless.network import create_pinned_async_httpx_client
+from paperless.network import create_pinned_httpx_client
 from paperless.network import validate_outbound_http_url
 from paperless_ai.base_model import DocumentClassifierSchema
 
@@ -27,23 +31,47 @@ class AIClient:
     def get_llm(self) -> "Ollama | OpenAILike":
         if self.settings.llm_backend == LLMBackend.OLLAMA:
             from llama_index.llms.ollama import Ollama
+            from ollama import AsyncClient
+            from ollama import Client
 
             endpoint = self.settings.llm_endpoint or "http://localhost:11434"
             validate_outbound_http_url(
                 endpoint,
                 allow_internal=self.settings.llm_allow_internal_endpoints,
             )
+            transport = PinnedHostHTTPTransport(
+                allow_internal=self.settings.llm_allow_internal_endpoints,
+            )
+            async_transport = PinnedHostAsyncHTTPTransport(
+                allow_internal=self.settings.llm_allow_internal_endpoints,
+            )
             return Ollama(
                 model=self.settings.llm_model or "llama3.1",
                 base_url=endpoint,
                 request_timeout=120,
+                client=Client(
+                    host=endpoint,
+                    timeout=120,
+                    transport=transport,
+                ),
+                async_client=AsyncClient(
+                    host=endpoint,
+                    timeout=120,
+                    transport=async_transport,
+                ),
             )
         elif self.settings.llm_backend == LLMBackend.OPENAI_LIKE:
             from llama_index.llms.openai_like import OpenAILike
 
             endpoint = self.settings.llm_endpoint or None
+            http_client = None
+            async_http_client = None
             if endpoint:
-                validate_outbound_http_url(
+                http_client = create_pinned_httpx_client(
+                    endpoint,
+                    allow_internal=self.settings.llm_allow_internal_endpoints,
+                )
+                async_http_client = create_pinned_async_httpx_client(
                     endpoint,
                     allow_internal=self.settings.llm_allow_internal_endpoints,
                 )
@@ -53,6 +81,8 @@ class AIClient:
                 api_key=self.settings.llm_api_key,
                 is_chat_model=True,
                 is_function_calling_model=True,
+                http_client=http_client,
+                async_http_client=async_http_client,
             )
         else:
             raise ValueError(f"Unsupported LLM backend: {self.settings.llm_backend}")

src/paperless_ai/embedding.py

@@ -13,6 +13,10 @@ from documents.models import Document
 from documents.models import Note
 from paperless.config import AIConfig
 from paperless.models import LLMEmbeddingBackend
+from paperless.network import PinnedHostAsyncHTTPTransport
+from paperless.network import PinnedHostHTTPTransport
+from paperless.network import create_pinned_async_httpx_client
+from paperless.network import create_pinned_httpx_client
 from paperless.network import validate_outbound_http_url
 
 OCR_LEADER_REGEX = re.compile(r"[._\-\u00b7]{4,}")
@@ -27,8 +31,14 @@ def get_embedding_model() -> "BaseEmbedding":
             from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 
             endpoint = config.llm_embedding_endpoint or config.llm_endpoint or None
+            http_client = None
+            async_http_client = None
             if endpoint:
-                validate_outbound_http_url(
+                http_client = create_pinned_httpx_client(
+                    endpoint,
+                    allow_internal=config.llm_allow_internal_endpoints,
+                )
+                async_http_client = create_pinned_async_httpx_client(
                     endpoint,
                     allow_internal=config.llm_allow_internal_endpoints,
                 )
@@ -36,6 +46,8 @@ def get_embedding_model() -> "BaseEmbedding":
                 model_name=config.llm_embedding_model or "text-embedding-3-small",
                 api_key=config.llm_api_key,
                 api_base=endpoint,
+                http_client=http_client,
+                async_http_client=async_http_client,
             )
         case LLMEmbeddingBackend.HUGGINGFACE:
             from llama_index.embeddings.huggingface import HuggingFaceEmbedding
@@ -47,6 +59,8 @@ def get_embedding_model() -> "BaseEmbedding":
             )
         case LLMEmbeddingBackend.OLLAMA:
             from llama_index.embeddings.ollama import OllamaEmbedding
+            from ollama import AsyncClient
+            from ollama import Client
 
             endpoint = (
                 config.llm_embedding_endpoint
@@ -57,10 +71,23 @@ def get_embedding_model() -> "BaseEmbedding":
                 endpoint,
                 allow_internal=config.llm_allow_internal_endpoints,
             )
-            return OllamaEmbedding(
+            embedding = OllamaEmbedding(
                 model_name=config.llm_embedding_model or "embeddinggemma",
                 base_url=endpoint,
             )
+            embedding._client = Client(
+                host=endpoint,
+                transport=PinnedHostHTTPTransport(
+                    allow_internal=config.llm_allow_internal_endpoints,
+                ),
+            )
+            embedding._async_client = AsyncClient(
+                host=endpoint,
+                transport=PinnedHostAsyncHTTPTransport(
+                    allow_internal=config.llm_allow_internal_endpoints,
+                ),
+            )
+            return embedding
         case _:
             raise ValueError(
                 f"Unsupported embedding backend: {config.llm_embedding_backend}",

src/paperless_ai/indexing.py

@@ -1,5 +1,6 @@
 import logging
 import shutil
+from collections.abc import Iterable
 from datetime import timedelta
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -327,14 +328,24 @@ def truncate_content(content: str) -> str:
     return " ".join(truncated_chunks)
 
 
+def normalize_document_ids(document_ids: Iterable[int | str] | None) -> set[str] | None:
+    if document_ids is None:
+        return None
+    return {str(document_id) for document_id in document_ids}
+
+
 def query_similar_documents(
     document: Document,
     top_k: int = 5,
-    document_ids: list[int] | None = None,
+    document_ids: Iterable[int | str] | None = None,
 ) -> list[Document]:
     """
     Runs a similarity query and returns top-k similar Document objects.
     """
+    allowed_document_ids = normalize_document_ids(document_ids)
+    if allowed_document_ids is not None and not allowed_document_ids:
+        return []
+
     if not vector_store_file_exists():
         queue_llm_index_update_if_needed(
             rebuild=False,
@@ -349,11 +360,13 @@ def query_similar_documents(
         [
             node.node_id
             for node in index.docstore.docs.values()
-            if node.metadata.get("document_id") in document_ids
+            if node.metadata.get("document_id") in allowed_document_ids
         ]
-        if document_ids
+        if allowed_document_ids is not None
         else None
     )
+    if doc_node_ids is not None and not doc_node_ids:
+        return []
 
     from llama_index.core.retrievers import VectorIndexRetriever
 
@@ -368,10 +381,23 @@ def query_similar_documents(
     )
     results = retriever.retrieve(query_text)
 
-    document_ids = [
-        int(node.metadata["document_id"])
-        for node in results
-        if "document_id" in node.metadata
-    ]
+    retrieved_document_ids: list[int] = []
+    for node in results:
+        document_id = node.metadata.get("document_id")
+        if document_id is None:
+            continue
+        normalized_document_id = str(document_id)
+        if (
+            allowed_document_ids is not None
+            and normalized_document_id not in allowed_document_ids
+        ):
+            continue
+        try:
+            retrieved_document_ids.append(int(normalized_document_id))
+        except ValueError:
+            logger.warning(
+                "Skipping LLM index result with invalid document_id %r.",
+                document_id,
+            )
 
-    return list(Document.objects.filter(pk__in=document_ids))
+    return list(Document.objects.filter(pk__in=retrieved_document_ids))

src/paperless_ai/tests/test_ai_indexing.py

@@ -3,6 +3,7 @@ from unittest.mock import MagicMock
 from unittest.mock import patch
 
 import pytest
+from django.contrib.auth.models import User
 from django.test import override_settings
 from django.utils import timezone
 from llama_index.core.base.embeddings.base import BaseEmbedding
@@ -428,3 +429,78 @@ def test_query_similar_documents_triggers_update_when_index_missing(
     )
     mock_load.assert_not_called()
     assert result == []
+
+
+@pytest.mark.django_db
+def test_query_similar_documents_normalizes_and_post_filters_allowed_ids(
+    real_document,
+) -> None:
+    real_document.owner = User.objects.create_user(username="rag-owner")
+    real_document.save()
+    private_owner = User.objects.create_user(username="rag-private-owner")
+    private_document = Document.objects.create(
+        title="Private similar document",
+        content="Similar private content that must not reach RAG.",
+        owner=private_owner,
+        added=timezone.now(),
+    )
+
+    with (
+        patch(
+            "paperless_ai.indexing.vector_store_file_exists",
+            return_value=True,
+        ),
+        patch("paperless_ai.indexing.load_or_build_index") as mock_load_or_build_index,
+        patch("llama_index.core.retrievers.VectorIndexRetriever") as mock_retriever_cls,
+    ):
+        allowed_node = MagicMock()
+        allowed_node.node_id = "allowed-node"
+        allowed_node.metadata = {"document_id": str(real_document.pk)}
+        private_node = MagicMock()
+        private_node.node_id = "private-node"
+        private_node.metadata = {"document_id": str(private_document.pk)}
+
+        mock_index = MagicMock()
+        mock_index.docstore.docs.values.return_value = [allowed_node, private_node]
+        mock_load_or_build_index.return_value = mock_index
+
+        mock_retriever = MagicMock()
+        mock_retriever.retrieve.return_value = [private_node, allowed_node]
+        mock_retriever_cls.return_value = mock_retriever
+
+        result = indexing.query_similar_documents(
+            real_document,
+            top_k=2,
+            document_ids=[real_document.pk],
+        )
+
+    mock_retriever_cls.assert_called_once_with(
+        index=mock_index,
+        similarity_top_k=2,
+        doc_ids=["allowed-node"],
+    )
+    assert result == [real_document]
+    assert private_document not in result
+
+
+@pytest.mark.django_db
+def test_query_similar_documents_empty_allow_list_fails_closed(
+    real_document,
+) -> None:
+    with (
+        patch(
+            "paperless_ai.indexing.vector_store_file_exists",
+            return_value=True,
+        ) as mock_vector_store_exists,
+        patch("paperless_ai.indexing.load_or_build_index") as mock_load_or_build_index,
+        patch("llama_index.core.retrievers.VectorIndexRetriever") as mock_retriever_cls,
+    ):
+        result = indexing.query_similar_documents(
+            real_document,
+            document_ids=[],
+        )
+
+    assert result == []
+    mock_vector_store_exists.assert_not_called()
+    mock_load_or_build_index.assert_not_called()
+    mock_retriever_cls.assert_not_called()

src/paperless_ai/tests/test_client.py

@@ -1,3 +1,4 @@
+from unittest.mock import ANY
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
@@ -40,6 +41,8 @@ def test_get_llm_ollama(mock_ai_config, mock_ollama_llm):
         model="test_model",
         base_url="http://test-url",
         request_timeout=120,
+        client=ANY,
+        async_client=ANY,
     )
     assert client.llm == mock_ollama_llm.return_value
 
@@ -58,6 +61,8 @@ def test_get_llm_openai(mock_ai_config, mock_openai_llm):
         api_key="test_api_key",
         is_chat_model=True,
         is_function_calling_model=True,
+        http_client=ANY,
+        async_http_client=ANY,
     )
     assert client.llm == mock_openai_llm.return_value
 

src/paperless_ai/tests/test_embedding.py

@@ -1,4 +1,5 @@
 import json
+from unittest.mock import ANY
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
@@ -70,6 +71,8 @@ def test_get_embedding_model_openai(mock_ai_config):
             model_name="text-embedding-3-small",
             api_key="test_api_key",
             api_base="http://test-url",
+            http_client=ANY,
+            async_http_client=ANY,
         )
         assert model == MockOpenAIEmbedding.return_value
 
@@ -89,6 +92,8 @@ def test_get_embedding_model_openai_prefers_embedding_endpoint(mock_ai_config):
             model_name="text-embedding-3-small",
             api_key="test_api_key",
             api_base="http://embedding-url",
+            http_client=ANY,
+            async_http_client=ANY,
         )
         assert model == MockOpenAIEmbedding.return_value