Kb exists middleware

nucliadb/src/nucliadb/middleware/__init__.py

@@ -18,17 +18,27 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import logging
+import re
 import time
 from collections import deque
 from typing import ClassVar
 
+from cachetools import TTLCache
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from starlette.requests import Request
-from starlette.responses import Response
+from starlette.responses import JSONResponse, Response
+
+from nucliadb.common import datamanagers
 
 PROCESS_TIME_HEADER = "X-PROCESS-TIME"
 ACCESS_CONTROL_EXPOSE_HEADER = "Access-Control-Expose-Headers"
 
+_KB_UUID_PATH_RE = re.compile(
+    r"/kb/([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})(?:/|$)", re.IGNORECASE
+)
+
+_kb_exists_cache: TTLCache[str, bool] = TTLCache(maxsize=4096, ttl=5 * 60)
+
 
 logger = logging.getLogger("nucliadb.middleware")
 
@@ -131,3 +141,34 @@ def get_count(self) -> int:
 class HourlyLogCounter(EventCounter):
     def __init__(self):
         super().__init__(window_seconds=3600)
+
+
+class KBExistsMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware that checks whether the kbid found in the request path exists.
+    If the kbid is not found in the path the request is passed through unchanged.
+    Results are cached in memory for 5 minutes.
+    """
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        kbid = self.get_kbid_from_path(request.url.path)
+        if kbid and not await self._kb_exists(kbid):
+            return JSONResponse(
+                status_code=404,
+                content={"detail": f"Knowledge Box '{kbid}' not found"},
+            )
+        return await call_next(request)
+
+    def get_kbid_from_path(self, path: str) -> str | None:
+        match = _KB_UUID_PATH_RE.search(path)
+        if match:
+            return match.group(1)
+        return None
+
+    async def _kb_exists(self, kbid: str) -> bool:
+        cached = _kb_exists_cache.get(kbid)
+        if cached is not None:
+            return cached
+        exists = await datamanagers.atomic.kb.exists_kb(kbid=kbid)
+        _kb_exists_cache[kbid] = exists
+        return exists

nucliadb/src/nucliadb/writer/app.py

@@ -27,7 +27,7 @@
 from starlette.requests import ClientDisconnect
 from starlette.responses import HTMLResponse
 
-from nucliadb.middleware import ClientErrorPayloadLoggerMiddleware
+from nucliadb.middleware import ClientErrorPayloadLoggerMiddleware, KBExistsMiddleware
 from nucliadb.writer import API_PREFIX
 from nucliadb.writer.api.v1.router import api as api_v1
 from nucliadb.writer.lifecycle import lifespan
@@ -46,6 +46,7 @@
 middleware.extend(
     [
         Middleware(AuthenticationMiddleware, backend=NucliaCloudAuthenticationBackend()),
+        Middleware(KBExistsMiddleware),
         Middleware(ClientErrorPayloadLoggerMiddleware),
     ]
 )

nucliadb/tests/nucliadb/unit/middleware/test_middleware.py

@@ -19,6 +19,7 @@
 #
 
 import time
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from starlette.applications import Starlette
@@ -27,7 +28,12 @@
 from starlette.routing import Route
 from starlette.testclient import TestClient
 
-from nucliadb.middleware import EventCounter, ProcessTimeHeaderMiddleware
+from nucliadb.middleware import (
+    _KB_UUID_PATH_RE,
+    EventCounter,
+    KBExistsMiddleware,
+    ProcessTimeHeaderMiddleware,
+)
 
 
 class TestCaseProcessTimeHeaderMiddleware:
@@ -66,3 +72,97 @@ def test_event_counter():
     assert counter.get_count() == 200
     time.sleep(2.1)
     assert counter.get_count() == 0
+
+
+KBID = "4b9a1e20-df3c-4e97-a9b6-1c2d3e4f5a6b"
+
+# (path, should_match)
+KB_UUID_PATH_CASES = [
+    # --- writer resource endpoints ---
+    (f"/api/v1/kb/{KBID}/resources", True),
+    (f"/api/v1/kb/{KBID}/resource/somerid", True),
+    (f"/api/v1/kb/{KBID}/slug/my-slug", True),
+    (f"/api/v1/kb/{KBID}/resource/somerid/reprocess", True),
+    (f"/api/v1/kb/{KBID}/resource/somerid/reindex", True),
+    (f"/api/v1/kb/{KBID}/upload", True),
+    (f"/api/v1/kb/{KBID}/tusupload/some-upload-id", True),
+    # --- writer service endpoints ---
+    (f"/api/v1/kb/{KBID}/labelset/my-labelset", True),
+    (f"/api/v1/kb/{KBID}/export", True),
+    (f"/api/v1/kb/{KBID}/configuration", True),
+    # --- path without a kbid (no match expected) ---
+    ("/api/v1/kbs", False),
+    ("/api/v1/kb/not-a-uuid/resources", False),
+    ("/api/v1/kb/my-slug/resources", False),
+    # --- bare /kb/<uuid> with no trailing slash/segment ---
+    (f"/api/v1/kb/{KBID}", True),
+]
+
+
+@pytest.mark.parametrize("path,should_match", KB_UUID_PATH_CASES)
+def test_kb_uuid_path_regex(path, should_match):
+    match = _KB_UUID_PATH_RE.search(path)
+    if should_match:
+        assert match is not None, f"Expected match for path: {path}"
+        assert match.group(1).lower() == KBID.lower()
+    else:
+        assert match is None, f"Expected no match for path: {path}"
+
+
+class TestKBExistsMiddleware:
+    @pytest.fixture()
+    def app(self):
+        def endpoint(request):
+            return PlainTextResponse("ok")
+
+        return Starlette(
+            routes=[
+                Route(f"/api/v1/kb/{{kbid}}/resources", endpoint),
+                Route("/api/v1/kbs", endpoint),
+            ],
+            middleware=[Middleware(KBExistsMiddleware)],
+        )
+
+    @pytest.fixture()
+    def client(self, app):
+        return TestClient(app, raise_server_exceptions=False)
+
+    @pytest.fixture(autouse=True)
+    def clear_cache(self):
+        # Clear the _kb_exists_cache before each test to avoid interference between tests
+        from nucliadb.middleware import _kb_exists_cache
+
+        _kb_exists_cache.clear()
+
+    def test_kb_exists_passes_through(self, client):
+        with patch(
+            "nucliadb.middleware.datamanagers.atomic.kb.exists_kb",
+            new=AsyncMock(return_value=True),
+        ):
+            response = client.get(f"/api/v1/kb/{KBID}/resources")
+        assert response.status_code == 200
+
+    def test_kb_not_found_returns_404(self, client):
+        with patch(
+            "nucliadb.middleware.datamanagers.atomic.kb.exists_kb",
+            new=AsyncMock(return_value=False),
+        ):
+            response = client.get(f"/api/v1/kb/{KBID}/resources")
+        assert response.status_code == 404
+        assert KBID in response.json()["detail"]
+
+    def test_non_uuid_path_skips_check(self, client):
+        with patch(
+            "nucliadb.middleware.datamanagers.atomic.kb.exists_kb",
+            new=AsyncMock(side_effect=AssertionError("should not be called")),
+        ):
+            # /kbs has no kbid in path — middleware must not call exists_kb
+            response = client.get("/api/v1/kbs")
+        assert response.status_code == 200
+
+    def test_slug_path_skips_check(self, client):
+        with patch(
+            "nucliadb.middleware.datamanagers.atomic.kb.exists_kb",
+            new=AsyncMock(side_effect=AssertionError("should not be called")),
+        ):
+            client.get("/api/v1/kb/my-slug/resources")