brain-of-reese/tests/unit/test_security_headers.py

"""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"


# ---------------------------------------------------------------------------
# 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]