"""Regression tests for the polish-pass security hardening. Covers the review findings fixed after the initial implementation: - Host-header parser edge cases (case-folding, bracketed-IPv6 trailing garbage) - error-envelope secret scrubbing + generic message for unexpected bugs - pre-buffer upload size limit (Content-Length rejected before the body is read) - 204 responses carry no body - the pending registry is bounded """ from __future__ import annotations import os import re from typing import Any import pytest pytest.importorskip("fastapi") from notebooklm.server import app as app_module # noqa: E402 from notebooklm.server._auth import _host_is_loopback # noqa: E402 from notebooklm.server._errors import _redact # noqa: E402 from notebooklm.server._pending import _MAX_ENTRIES, PendingRegistry # noqa: E402 from .conftest import TEST_TOKEN # noqa: E402 class TestHostLoopbackGuard: @pytest.mark.parametrize( "host", [ "127.0.0.1", "localhost", "LOCALHOST", "Localhost", "::1", "[::1]", "[::1]:8000", "127.0.0.1:9000", "localhost:8000", ], ) def test_accepts_loopback(self, host: str) -> None: assert _host_is_loopback(host) is True @pytest.mark.parametrize( "host", [ "[::1]evil.com", "[::1]@evil.com", "[::1]:bad", "evil.com", "0.0.0.0", "127.0.0.1.evil.com", "", "[::1", "2130706433", ], ) def test_rejects_non_loopback(self, host: str) -> None: assert _host_is_loopback(host) is False class TestErrorScrubbing: def test_redact_masks_credential_shaped_text(self) -> None: out = _redact("failed Authorization: Bearer abcSECRET123 while decoding") assert "abcSECRET123" not in out assert "***" in out def test_redact_caps_length(self) -> None: out = _redact("x" * 5000) assert len(out) <= 301 # 300 + the ellipsis def test_unexpected_bug_message_is_generic(self) -> None: from notebooklm.server._errors import error_response # A non-library exception (a bug) must never echo its str() — which could # carry anything — to the client. resp = error_response(RuntimeError("Bearer leakedTOKEN in a stray bug")) body = resp.body.decode() assert "leakedTOKEN" not in body assert "Internal server error" in body assert resp.status_code == 500 class TestUploadPreBufferLimit: def test_oversized_content_length_is_413_before_handler( self, app: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: from fastapi.testclient import TestClient # Spy: the upload handler must never run when the declared length is over cap. called = {"add_file": False} orig = fake_client.sources.add_file async def spy(*args: Any, **kwargs: Any) -> Any: called["add_file"] = True return await orig(*args, **kwargs) monkeypatch.setattr(fake_client.sources, "add_file", spy) monkeypatch.setattr(app_module, "MAX_UPLOAD_BYTES", 8) headers = {"Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1"} with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post( "/v1/notebooks/nb-1/sources/file", files={"file": ("big.txt", b"x" * 4096, "text/plain")}, ) assert resp.status_code == 413 assert called["add_file"] is False assert resp.json()["error"]["category"] == "validation" def test_oversized_urlencoded_upload_form_is_413_before_parse( self, app: Any, monkeypatch: pytest.MonkeyPatch ) -> None: from fastapi.testclient import TestClient monkeypatch.setattr(app_module, "MAX_UPLOAD_BYTES", 8) headers = {"Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1"} with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post( "/v1/notebooks/nb-1/sources/file", data={"file": "x" * 64}, ) assert resp.status_code == 413 assert resp.json()["error"]["category"] == "validation" def test_oversized_non_form_upload_route_body_is_413_before_parse( self, app: Any, monkeypatch: pytest.MonkeyPatch ) -> None: from fastapi.testclient import TestClient monkeypatch.setattr(app_module, "MAX_UPLOAD_BYTES", 8) headers = { "Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1", "Content-Type": "text/plain", } with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post( "/v1/notebooks/nb-1/sources/file", content=b"x" * 64, ) assert resp.status_code == 413 assert resp.json()["error"]["category"] == "validation" class TestJsonBodyLimits: def test_oversized_json_content_length_is_413_before_handler( self, app: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: from fastapi.testclient import TestClient called = {"ask": False} orig = fake_client.chat.ask async def spy(*args: Any, **kwargs: Any) -> Any: called["ask"] = True return await orig(*args, **kwargs) monkeypatch.setattr(fake_client.chat, "ask", spy) monkeypatch.setattr( app_module, "JSON_BODY_LIMITS", ( app_module._BodyLimit( "POST", re.compile(r"^/v1/notebooks/[^/]+/chat$"), 32, "chat ask", ), ), ) headers = { "Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1", "Content-Type": "application/json", } with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post( "/v1/notebooks/nb-1/chat", content=b'{"question":"' + (b"x" * 64) + b'"}', ) assert resp.status_code == 413 assert called["ask"] is False assert resp.json()["error"]["category"] == "validation" def test_chunked_json_to_capped_route_is_411_before_handler( self, app: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: from fastapi.testclient import TestClient called = {"ask": False} orig = fake_client.chat.ask async def spy(*args: Any, **kwargs: Any) -> Any: called["ask"] = True return await orig(*args, **kwargs) def chunks() -> Any: yield b'{"question":"' yield b"hello" yield b'"}' monkeypatch.setattr(fake_client.chat, "ask", spy) headers = { "Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1", "Content-Type": "application/json", } with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post("/v1/notebooks/nb-1/chat", content=chunks()) assert resp.status_code == 411 assert called["ask"] is False assert resp.json()["error"]["category"] == "validation" def test_multipart_upload_ignores_json_caps( self, app: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: from fastapi.testclient import TestClient monkeypatch.setattr(app_module, "DEFAULT_JSON_BODY_BYTES", 8) headers = {"Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1"} with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post( "/v1/notebooks/nb-1/sources/file", files={"file": ("ok.txt", b"x" * 4096, "text/plain")}, ) assert resp.status_code == 201 assert len(fake_client.uploaded_paths) == 1 def test_multipart_to_json_route_uses_route_cap( self, app: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: from fastapi.testclient import TestClient called = {"ask": False} orig = fake_client.chat.ask async def spy(*args: Any, **kwargs: Any) -> Any: called["ask"] = True return await orig(*args, **kwargs) monkeypatch.setattr(fake_client.chat, "ask", spy) monkeypatch.setattr( app_module, "JSON_BODY_LIMITS", ( app_module._BodyLimit( "POST", re.compile(r"^/v1/notebooks/[^/]+/chat$"), 32, "chat ask", ), ), ) headers = {"Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1"} with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post( "/v1/notebooks/nb-1/chat", files={"file": ("not-json.txt", b"x" * 4096, "text/plain")}, ) assert resp.status_code == 413 assert called["ask"] is False assert len(fake_client.uploaded_paths) == 0 class TestChunkedMultipartRejected: def test_multipart_without_content_length_is_411(self, app: Any) -> None: from fastapi.testclient import TestClient # A generator body makes httpx use chunked transfer-encoding (no # Content-Length), the exact bypass a running-cap-after-spool leaves open. def _chunks() -> Any: yield b"--b\r\nContent-Disposition: form-data; name=file; filename=x\r\n\r\n" yield b"data\r\n--b--\r\n" headers = { "Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1", "Content-Type": "multipart/form-data; boundary=b", } with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post("/v1/notebooks/nb-1/sources/file", content=_chunks()) assert resp.status_code == 411 assert resp.json()["error"]["category"] == "validation" def test_urlencoded_upload_form_without_content_length_is_411(self, app: Any) -> None: from fastapi.testclient import TestClient def _chunks() -> Any: yield b"file=" yield b"x" * 32 headers = { "Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1", "Content-Type": "application/x-www-form-urlencoded", } with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post("/v1/notebooks/nb-1/sources/file", content=_chunks()) assert resp.status_code == 411 assert resp.json()["error"]["category"] == "validation" def test_non_form_upload_route_body_without_content_length_is_411(self, app: Any) -> None: from fastapi.testclient import TestClient def _chunks() -> Any: yield b"x" * 32 headers = { "Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1", "Content-Type": "text/plain", } with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.post("/v1/notebooks/nb-1/sources/file", content=_chunks()) assert resp.status_code == 411 assert resp.json()["error"]["category"] == "validation" class TestPeerLoopbackGuard: def test_non_loopback_peer_rejected_even_with_spoofed_host(self, app: Any) -> None: from fastapi.testclient import TestClient # Off-loopback PEER address + a forged loopback Host header: the # unspoofable peer-address check rejects it (the Host spoof does not help). headers = {"Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1"} with TestClient( app, headers=headers, client=("8.8.8.8", 5555), raise_server_exceptions=False ) as c: resp = c.get("/v1/notebooks") assert resp.status_code == 403 assert resp.json()["error"]["category"] == "auth" @pytest.mark.parametrize( ("addr", "expected"), [ ("127.0.0.1", True), ("::1", True), # IPv4-mapped IPv6 loopback: ``ipaddress.is_loopback`` only resolves the # mapped IPv4 in newer CPython patch levels, so this must be pinned # version-independently (it regressed on some macOS 3.10/3.11 runners). ("::ffff:127.0.0.1", True), ("8.8.8.8", False), ("::ffff:8.8.8.8", False), # mapped *external* — must NOT read as loopback ("0.0.0.0", False), ("not-an-ip", False), ], ) def test_addr_is_loopback_version_independent(self, addr: str, expected: bool) -> None: from notebooklm.server._auth import _addr_is_loopback assert _addr_is_loopback(addr) is expected @pytest.mark.parametrize("peer_host", ["::1", "::ffff:127.0.0.1"]) def test_ipv6_loopback_peer_allowed(self, app: Any, peer_host: str) -> None: from fastapi.testclient import TestClient # An IPv6 loopback peer (``::1``) and the IPv4-mapped loopback # (``::ffff:127.0.0.1``) both satisfy the unspoofable peer-loopback guard. headers = {"Authorization": f"Bearer {TEST_TOKEN}", "Host": "127.0.0.1"} with TestClient( app, headers=headers, client=(peer_host, 5555), raise_server_exceptions=False ) as c: resp = c.get("/v1/notebooks") assert resp.status_code == 200 def test_external_bind_optin_allows_non_loopback_peer( self, app: Any, monkeypatch: pytest.MonkeyPatch ) -> None: from fastapi.testclient import TestClient from notebooklm.server._auth import ALLOW_EXTERNAL_BIND_ENV monkeypatch.setenv(ALLOW_EXTERNAL_BIND_ENV, "1") headers = {"Authorization": f"Bearer {TEST_TOKEN}", "Host": "example.com"} with TestClient( app, headers=headers, client=("8.8.8.8", 5555), raise_server_exceptions=False ) as c: resp = c.get("/v1/notebooks") assert resp.status_code == 200 class TestNoContentResponses: def test_delete_notebook_has_empty_body(self, authed_client: Any) -> None: resp = authed_client.delete("/v1/notebooks/nb-1") assert resp.status_code == 204 assert resp.content == b"" def test_delete_source_has_empty_body(self, authed_client: Any) -> None: resp = authed_client.delete("/v1/notebooks/nb-1/sources/src-1") assert resp.status_code == 204 assert resp.content == b"" class TestUploadPathSafety: """The upload spools into a private ``mkdtemp`` dir named after the caller's *basename* — directory components are stripped (traversal guard) and the file is isolated, so a malicious filename can neither escape nor reach a real path. """ def test_malicious_filename_is_basenamed_and_isolated( self, authed_client: Any, fake_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: captured: dict[str, Any] = {} orig = fake_client.sources.add_file async def spy(notebook_id: str, path: str, *args: Any, **kwargs: Any) -> Any: captured["path"] = path return await orig(notebook_id, path, *args, **kwargs) monkeypatch.setattr(fake_client.sources, "add_file", spy) evil = "../../etc/passwd" resp = authed_client.post( "/v1/notebooks/nb-1/sources/file", files={"file": (evil, b"data", "application/pdf")}, ) assert resp.status_code == 201 path = captured["path"] # The traversal is stripped: the file lives directly under a unique # server-owned ``nblm-upload-`` dir, named only by the basename. assert os.path.basename(path) == "passwd" assert os.path.basename(os.path.dirname(path)).startswith("nblm-upload-") # Absolute, canonical, no traversal component escapes our temp dir. assert os.path.isabs(path) and ".." not in path def test_safe_upload_name(self) -> None: # Now the shared neutral helper (_app.source_add.safe_upload_name): control # chars stripped, . / .. rejected, extension preserved on truncation. from notebooklm.server.routes.sources import _safe_upload_name assert _safe_upload_name("report.pdf") == "report.pdf" assert _safe_upload_name("../../etc/passwd") == "passwd" # traversal stripped assert _safe_upload_name("a/b/c.txt") == "c.txt" assert _safe_upload_name("a\x00b.pdf") == "ab.pdf" # control chars stripped assert _safe_upload_name("a\x7fb.pdf") == "ab.pdf" # DEL stripped assert _safe_upload_name("a\x85b.pdf") == "ab.pdf" # C1 control stripped assert _safe_upload_name("..") == "upload.bin" # directory cursor rejected assert _safe_upload_name("") == "upload.bin" # empty fallback assert _safe_upload_name(None) == "upload.bin" # Length-bounded AND extension-preserving (stem truncated, ".pdf" kept). long = _safe_upload_name("x" * 500 + ".pdf") assert len(long) <= 255 and long.endswith(".pdf") # Multibyte names are bounded by BYTE length (255), not char count. emoji = _safe_upload_name("😀" * 300 + ".pdf") assert len(emoji.encode("utf-8")) <= 255 and emoji.endswith(".pdf") class TestGenerateSourceDefaulting: """A bare generate (no ``source_ids``) scopes to ALL sources, like the CLI: ``passthrough_source_ids`` resolves an empty selection to ``None`` (the client's all-sources sentinel), not an empty tuple (which the API rejects as "… generation is unavailable").""" async def test_empty_resolves_to_none(self) -> None: from notebooklm.server.routes._passthrough import passthrough_source_ids assert await passthrough_source_ids(None, "nb", ()) is None assert await passthrough_source_ids(None, "nb", []) is None async def test_nonempty_passes_through(self) -> None: from notebooklm.server.routes._passthrough import passthrough_source_ids assert await passthrough_source_ids(None, "nb", ("s1", "s2")) == ("s1", "s2") class TestErrorEnvelopeShape: """Hand-raised ``HTTPException``s use the same ``{error:{category,message}}`` envelope as classified library errors (R9 single-shape contract), not FastAPI's default ``{"detail": ...}``.""" def test_loopback_guard_403_uses_envelope(self, raw_client: Any) -> None: # raw_client sends Host: testserver (not loopback) → 403 from the guard. resp = raw_client.get("/v1/notebooks") assert resp.status_code == 403 body = resp.json() assert "detail" not in body assert body["error"]["category"] == "auth" assert isinstance(body["error"]["message"], str) def test_missing_token_401_uses_envelope(self, app: Any) -> None: from fastapi.testclient import TestClient # Loopback peer + Host clear the rebinding guard; the wrong token trips 401. headers = {"Authorization": "Bearer wrong-token", "Host": "127.0.0.1"} with TestClient( app, headers=headers, client=("127.0.0.1", 5555), raise_server_exceptions=False ) as c: resp = c.get("/v1/notebooks") assert resp.status_code == 401 body = resp.json() assert "detail" not in body assert body["error"]["category"] == "auth" def test_route_404_uses_envelope(self, authed_client: Any) -> None: # An unknown source is an in-route HTTPException(404) — not a classified # library error — and must still render the envelope. resp = authed_client.get("/v1/notebooks/nb-1/sources/missing") assert resp.status_code == 404 body = resp.json() assert "detail" not in body assert body["error"]["category"] == "not_found" class TestPendingRegistryBounded: def test_eviction_past_cap(self) -> None: reg = PendingRegistry() # Record cap + 5; the 5 oldest are evicted (their later poll → 404). for i in range(_MAX_ENTRIES + 5): reg.record("nb", f"id-{i}") assert reg.knows("nb", "id-0") is False assert reg.knows("nb", "id-4") is False assert reg.knows("nb", f"id-{_MAX_ENTRIES + 4}") is True def test_record_is_idempotent(self) -> None: reg = PendingRegistry() reg.record("nb", "x") reg.record("nb", "x") reg.drop("nb", "x") assert reg.knows("nb", "x") is False def test_drop_then_rerecord_survives_eviction_churn(self) -> None: # #1872: drop() must remove the stale FIFO tuple. Fill to the cap, then # drop + re-record "x" (moving it to newest) and push one more entry so a # REAL eviction fires. The fix evicts the genuinely-oldest entry ("y-0") # and keeps the live, re-recorded "x". With the bug, drop() leaves x's # stale tuple in _order, the re-record duplicates it, and that stale copy # sits at the front — so it is evicted first, dropping "x" from _ids and # making knows() wrongly False. reg = PendingRegistry() reg.record("nb", "x") for i in range(_MAX_ENTRIES - 1): reg.record("nb", f"y-{i}") # now at the cap: _order = [x, y-0 .. y-(MAX-2)] reg.drop("nb", "x") reg.record("nb", "x") # re-recorded as the newest entry reg.record("nb", "z") # one over the cap -> exactly one real eviction assert reg.knows("nb", "x") is True # the live re-recorded id survives assert reg.knows("nb", "z") is True assert reg.knows("nb", "y-0") is False # the genuinely-oldest entry was evicted def test_drop_removes_fifo_tuple_no_duplicate_on_rerecord(self) -> None: reg = PendingRegistry() reg.record("nb", "x") reg.drop("nb", "x") assert ("nb", "x") not in reg._order reg.record("nb", "x") assert list(reg._order).count(("nb", "x")) == 1