"""Unit tests: the pure-ASGI security-headers middleware (phase 82, task 01). Pins the middleware against the phase design: * the ``CSP`` constant is the exact owner-approved A1 string; * a plain 200 response carries all three headers with the exact values; * a 404-shaped response carries them too (the static catch-all's 404s get the headers as soon as the middleware is registered in ``app/main.py``, task 02); * a **streaming** SSE response passes through byte-identical — both chunks, in order, chunk boundaries preserved (no buffering, no coalescing) — AND carries the headers (the SSE passthrough pin); * the HTTP path wraps ``send`` (and leaves ``receive`` alone); * a non-HTTP scope (websocket / lifespan) passes through untouched — the app is called with the ORIGINAL ``send``/``receive`` (no wrapping, no header injection attempt, no crash). """ from __future__ import annotations import asyncio from collections.abc import AsyncIterator from typing import Any import pytest from starlette.responses import StreamingResponse from starlette.types import ASGIApp, Message, Scope from app.core.security_headers import CSP, SecurityHeadersMiddleware #: The exact expected header set (decision A1 for the CSP, A4 for the #: other two). EXPECTED_HEADERS = { "content-security-policy": "default-src 'self'; base-uri 'none'; frame-ancestors 'none'", "x-frame-options": "DENY", "x-content-type-options": "nosniff", } # --------------------------------------------------------------------------- # ASGI drivers / test apps (the repo's unit suite has no pytest-asyncio — # ``asyncio.run`` matches the existing style, see test_caching.py) # --------------------------------------------------------------------------- def _drive( app: ASGIApp, scope: Scope, receive_messages: list[Message] | None = None ) -> list[Message]: """Run ``app`` for one connection; return every ``send`` message in order.""" queued: list[Message] = list( receive_messages if receive_messages is not None else [{"type": "http.request", "body": b"", "more_body": False}] ) sent: list[Message] = [] async def receive() -> Message: return queued.pop(0) if queued else {"type": "http.disconnect"} async def send(message: Message) -> None: sent.append(message) async def _run() -> None: await app(scope, receive, send) asyncio.run(_run()) return sent def _http_scope() -> Scope: return { "type": "http", "method": "GET", "path": "/", "query_string": b"", "headers": [], "raw_path": b"/", "root_path": "", "scheme": "http", "server": ("testserver", 80), } def _header(message: Message, name: str) -> str | None: for key, value in message.get("headers", []): if key.lower() == name.lower().encode("ascii"): return value.decode("latin-1") return None def _plain_app(status: int, body: bytes, headers: list | None = None) -> ASGIApp: """A raw ASGI app answering with exactly one start + one body message.""" async def app(scope: Scope, receive: Any, send: Any) -> None: await send( {"type": "http.response.start", "status": status, "headers": list(headers or [])} ) await send({"type": "http.response.body", "body": body, "more_body": False}) return app def _assert_security_headers(start: Message, expected_extra: dict[str, str] | None = None) -> None: for name, value in EXPECTED_HEADERS.items(): assert _header(start, name) == value, f"expected header {name!r} == {value!r}" for name, value in (expected_extra or {}).items(): assert _header(start, name) == value, f"expected pre-existing header {name!r} kept" # --------------------------------------------------------------------------- # The policy constant # --------------------------------------------------------------------------- def test_csp_constant_is_the_exact_a1_policy() -> None: """The owner-approved A1 string, verbatim: same-origin default, no base-tag hijack, no framing — no 'unsafe-inline', no report sink.""" assert CSP == "default-src 'self'; base-uri 'none'; frame-ancestors 'none'" assert "unsafe-inline" not in CSP # --------------------------------------------------------------------------- # Plain responses # --------------------------------------------------------------------------- def test_plain_200_response_carries_all_three_headers_with_exact_values() -> None: wrapped = SecurityHeadersMiddleware(_plain_app(200, b"hello", headers=[[b"content-type", b"text/plain"]])) sent = _drive(wrapped, _http_scope()) assert [m["type"] for m in sent] == ["http.response.start", "http.response.body"] start, body_msg = sent assert start["status"] == 200 _assert_security_headers(start, expected_extra={"content-type": "text/plain"}) # The body passes through byte-identical. assert body_msg["body"] == b"hello" assert body_msg["more_body"] is False def test_404_shaped_response_carries_all_three_headers() -> None: """A 404 (e.g. the static catch-all's "file not found" JSON) gets the headers too — with an EMPTY downstream header list (no pre-existing headers to preserve).""" wrapped = SecurityHeadersMiddleware(_plain_app(404, b"not found")) sent = _drive(wrapped, _http_scope()) start, body_msg = sent assert start["status"] == 404 _assert_security_headers(start) assert body_msg["body"] == b"not found" def test_pre_existing_csp_from_an_inner_layer_is_preserved() -> None: """Phase 91 (task 05): the caching middleware publishes, on themed HTML pages only, the A1 string EXTENDED with a ``style-src`` sha256 hash for the inline theme tag (the A1 policy would block the tag in every real browser). A CSP an inner layer has already set is that layer's deliberate one and must survive the outer middleware — while the other two headers are still added.""" themed = ( "default-src 'self'; base-uri 'none'; frame-ancestors 'none'; " "style-src 'self' 'sha256-2rm3wPcQfXmE8q1s9vBzK7hN4tY5uJ6gW3oR0cAeDfH='" ) wrapped = SecurityHeadersMiddleware( _plain_app( 200, b"", headers=[[b"content-security-policy", themed.encode("ascii")]], ) ) sent = _drive(wrapped, _http_scope()) start = sent[0] assert _header(start, "content-security-policy") == themed # not clobbered assert _header(start, "x-frame-options") == "DENY" assert _header(start, "x-content-type-options") == "nosniff" # --------------------------------------------------------------------------- # The SSE streaming passthrough pin # --------------------------------------------------------------------------- def test_sse_stream_passes_through_byte_identical_with_headers() -> None: """A ``text/event-stream`` response driven through a REAL starlette ``StreamingResponse`` (async generator yielding two SSE frames): the body arrives as exactly the two chunks, in order, boundaries intact — no buffering, no coalescing — and the three headers are on the start message (the SSE passthrough pin; a body-touching regression would break this and the chat E2E suites).""" async def frames() -> AsyncIterator[bytes]: yield b"data: a\n\n" yield b"data: b\n\n" response = StreamingResponse(frames(), media_type="text/event-stream") sent = _drive(SecurityHeadersMiddleware(response), _http_scope()) start = sent[0] assert start["type"] == "http.response.start" assert start["status"] == 200 _assert_security_headers(start) # starlette's own StreamingResponse appends ``; charset=utf-8`` — the # media type itself must be SSE (the middleware never rewrites it). content_type = _header(start, "content-type") assert content_type is not None assert content_type.split(";", 1)[0].strip() == "text/event-stream" body_msgs = [m for m in sent[1:] if m["type"] == "http.response.body"] # starlette's StreamingResponse closes the generator with a final # EMPTY trailer frame (no bytes) after the last chunk — the two data # frames must still arrive untouched, in order, boundaries intact. assert [m["body"] for m in body_msgs] == [b"data: a\n\n", b"data: b\n\n", b""] assert b"".join(m["body"] for m in body_msgs) == b"data: a\n\ndata: b\n\n" assert [m["more_body"] for m in body_msgs] == [True, True, False] def test_sse_frames_are_forwarded_in_order_without_delayed_buffering() -> None: """Message-level pin: every ``http.response.body`` message the app sends is forwarded as-is (same bytes, same ``more_body``) — the wrapper mutates ONLY the ``http.response.start`` message.""" original_frames: list[Message] = [] async def recording_app(scope: Scope, receive: Any, send: Any) -> None: # Same response shape as the SSE route: start, then two frames. await send( { "type": "http.response.start", "status": 200, "headers": [[b"content-type", b"text/event-stream"]], } ) for frame in (b"data: a\n\n", b"data: b\n\n"): original_frames.append( {"type": "http.response.body", "body": frame, "more_body": frame != b"data: b\n\n"} ) await send(original_frames[-1]) sent = _drive(SecurityHeadersMiddleware(recording_app), _http_scope()) forwarded = [m for m in sent[1:] if m["type"] == "http.response.body"] # Identical objects: forwarded byte-for-byte, in order, untouched. assert forwarded == original_frames # --------------------------------------------------------------------------- # Scope handling: http wrapping vs non-http passthrough # --------------------------------------------------------------------------- def test_http_scope_wraps_send_but_not_receive() -> None: seen: dict[str, Any] = {} async def app(scope: Scope, receive: Any, send: Any) -> None: seen["receive"] = receive seen["send"] = send await send({"type": "http.response.start", "status": 200, "headers": []}) await send({"type": "http.response.body", "body": b"", "more_body": False}) queued: list[Message] = [{"type": "http.request", "body": b"", "more_body": False}] sent: list[Message] = [] async def receive() -> Message: return queued.pop(0) if queued else {"type": "http.disconnect"} async def send(message: Message) -> None: sent.append(message) asyncio.run(SecurityHeadersMiddleware(app)(_http_scope(), receive, send)) assert seen["send"] is not send # the wrapped send did the header writes assert seen["receive"] is receive # receive is never touched assert sent[0]["type"] == "http.response.start" _assert_security_headers(sent[0]) @pytest.mark.parametrize("scope_type", ["websocket", "lifespan"], ids=["websocket", "lifespan"]) def test_non_http_scope_passes_through_untouched(scope_type: str) -> None: """Non-HTTP connections (websocket, lifespan) reach the app with the ORIGINAL ``send``/``receive`` — no wrapping, no header injection attempt, no crash.""" seen: dict[str, Any] = {} first, second = { "websocket": ( {"type": "websocket.accept"}, {"type": "websocket.send", "text": "hi"}, ), "lifespan": ( {"type": "lifespan.startup.complete"}, {"type": "lifespan.shutdown.complete"}, ), }[scope_type] async def app(scope: Scope, receive: Any, send: Any) -> None: seen["scope"] = scope seen["receive"] = receive seen["send"] = send await send(first) await send(second) queued: list[Message] = [ {"type": f"{scope_type}.connect"} if scope_type == "websocket" else {"type": "lifespan.startup"} ] sent: list[Message] = [] async def receive() -> Message: return queued.pop(0) if queued else {"type": "lifespan.shutdown"} async def send(message: Message) -> None: sent.append(message) scope = {"type": scope_type, "path": "/", "raw_path": b"/"} asyncio.run(SecurityHeadersMiddleware(app)(scope, receive, send)) assert seen["scope"] is scope assert seen["send"] is send # the ORIGINAL send — the fast path assert seen["receive"] is receive # Both messages pass through untouched — no headers injected anywhere. assert sent == [first, second]