From 445f76c4926e88e44ac99b3151146269f7e217ad Mon Sep 17 00:00:00 2001 From: Ferran Llamas Date: Mon, 1 Jun 2026 12:12:35 +0200 Subject: [PATCH 1/2] Done --- nucliadb/src/nucliadb/middleware/__init__.py | 40 +++++++++++++++++++- nucliadb/src/nucliadb/writer/app.py | 3 +- 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/nucliadb/src/nucliadb/middleware/__init__.py b/nucliadb/src/nucliadb/middleware/__init__.py index 6ab122a3fb..3bd2f7a213 100644 --- a/nucliadb/src/nucliadb/middleware/__init__.py +++ b/nucliadb/src/nucliadb/middleware/__init__.py @@ -18,17 +18,26 @@ # along with this program. If not, see . 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_PATH_RE = re.compile(r"/kb/([^/]+)") +_UUID_RE = re.compile(r"^[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 +140,32 @@ 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: + match = _KB_PATH_RE.search(request.url.path) + if match: + kbid = match.group(1) + if _UUID_RE.match(kbid): + exists = await self._kb_exists(kbid) + if not exists: + return JSONResponse( + status_code=404, + content={"detail": f"Knowledge Box '{kbid}' not found"}, + ) + return await call_next(request) + + 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 diff --git a/nucliadb/src/nucliadb/writer/app.py b/nucliadb/src/nucliadb/writer/app.py index 88f5351e43..69642435a6 100644 --- a/nucliadb/src/nucliadb/writer/app.py +++ b/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), ] ) From 6b050935f16071dc85e1fdce03b06739d92240e2 Mon Sep 17 00:00:00 2001 From: Ferran Llamas Date: Mon, 1 Jun 2026 14:23:51 +0200 Subject: [PATCH 2/2] Simpler --- nucliadb/src/nucliadb/middleware/__init__.py | 27 ++--- .../unit/middleware/test_middleware.py | 102 +++++++++++++++++- 2 files changed, 116 insertions(+), 13 deletions(-) diff --git a/nucliadb/src/nucliadb/middleware/__init__.py b/nucliadb/src/nucliadb/middleware/__init__.py index 3bd2f7a213..b2080c12c9 100644 --- a/nucliadb/src/nucliadb/middleware/__init__.py +++ b/nucliadb/src/nucliadb/middleware/__init__.py @@ -33,8 +33,9 @@ PROCESS_TIME_HEADER = "X-PROCESS-TIME" ACCESS_CONTROL_EXPOSE_HEADER = "Access-Control-Expose-Headers" -_KB_PATH_RE = re.compile(r"/kb/([^/]+)") -_UUID_RE = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.IGNORECASE) +_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) @@ -150,18 +151,20 @@ class KBExistsMiddleware(BaseHTTPMiddleware): """ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: - match = _KB_PATH_RE.search(request.url.path) - if match: - kbid = match.group(1) - if _UUID_RE.match(kbid): - exists = await self._kb_exists(kbid) - if not exists: - return JSONResponse( - status_code=404, - content={"detail": f"Knowledge Box '{kbid}' not found"}, - ) + 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: diff --git a/nucliadb/tests/nucliadb/unit/middleware/test_middleware.py b/nucliadb/tests/nucliadb/unit/middleware/test_middleware.py index 3a9ea65403..61f1776957 100644 --- a/nucliadb/tests/nucliadb/unit/middleware/test_middleware.py +++ b/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/ 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")