test_hardening.py

"""FIX-10: app-level deploy-grade server hardening.

Covers the guards that must hold even without a reverse proxy in front of the app —
security headers, the request-size cap, per-IP rate limiting, and the /api/identify
concurrency bound — plus the ``readyz``/``health`` fix that drops the private ``_store``
access. See ``docs/audits/asvs-l2-delta.md`` for the checklist this closes.
"""

from __future__ import annotations

import base64

from fastapi.testclient import TestClient

from sprout.answer import Assistant
from sprout.config import Config
from sprout.hardening import SECURITY_HEADERS, ConcurrencyLimiter, _TokenBucket
from sprout.server import create_app


def _client(assistant: Assistant, config: Config) -> TestClient:
    return TestClient(create_app(config, assistant=assistant))


def test_security_headers_present_on_every_response(assistant: Assistant, config: Config) -> None:
    c = _client(assistant, config)
    r = c.get("/livez")
    for name, value in SECURITY_HEADERS.items():
        assert r.headers.get(name) == value


def test_security_headers_present_on_rejected_responses(
    assistant: Assistant, config: Config
) -> None:
    """Headers must land even on responses generated by the size/rate-limit middleware
    itself, not just ones the route handlers produce."""
    small_server = config.server.model_copy(update={"max_body_bytes": 10})
    cfg = config.model_copy(update={"server": small_server})
    c = _client(assistant, cfg)
    r = c.post("/api/chat", json={"question": "why are my monstera leaves yellowing?"})
    assert r.status_code == 413
    assert r.headers.get("Content-Security-Policy")


def test_request_size_cap_rejects_oversized_body(assistant: Assistant, config: Config) -> None:
    small_server = config.server.model_copy(update={"max_body_bytes": 100})
    cfg = config.model_copy(update={"server": small_server})
    c = _client(assistant, cfg)
    big_image = base64.b64encode(b"x" * 500).decode("ascii")
    r = c.post("/api/identify", json={"image_b64": big_image})
    assert r.status_code == 413
    assert "error" in r.json()


def test_request_size_cap_allows_normal_body(assistant: Assistant, config: Config) -> None:
    c = _client(assistant, config)
    r = c.post("/api/chat", json={"question": "why are my monstera leaves yellowing?"})
    assert r.status_code == 200


def test_rate_limit_returns_429_once_exhausted(assistant: Assistant, config: Config) -> None:
    cfg = config.model_copy(
        update={
            "server": config.server.model_copy(
                update={"rate_limit_requests": 2, "rate_limit_window_s": 60.0}
            )
        }
    )
    c = _client(assistant, cfg)
    codes = [c.get("/api/disclosure").status_code for _ in range(4)]
    assert codes[:2] == [200, 200]
    assert 429 in codes[2:]


def test_identify_rate_limit_is_independent_of_general_limit(
    assistant: Assistant, config: Config
) -> None:
    cfg = config.model_copy(
        update={
            "server": config.server.model_copy(
                update={
                    "rate_limit_requests": 1000,
                    "identify_rate_limit_requests": 1,
                    "identify_rate_limit_window_s": 60.0,
                }
            )
        }
    )
    c = _client(assistant, cfg)
    img = base64.b64encode(b"jpeg-bytes").decode("ascii")
    first = c.post("/api/identify", json={"image_b64": img})
    second = c.post("/api/identify", json={"image_b64": img})
    assert first.status_code != 429
    assert second.status_code == 429
    # The general endpoint is unaffected by the identify-specific bucket.
    assert c.get("/api/disclosure").status_code == 200


def test_identify_concurrency_limit_returns_503_when_saturated(
    assistant: Assistant, config: Config
) -> None:
    cfg = config.model_copy(
        update={"server": config.server.model_copy(update={"identify_max_concurrency": 1})}
    )
    c = _client(assistant, cfg)
    img = base64.b64encode(b"jpeg-bytes").decode("ascii")

    limiter = ConcurrencyLimiter(1)
    assert limiter.try_acquire() is True
    assert limiter.try_acquire() is False  # already saturated
    limiter.release()
    assert limiter.try_acquire() is True

    # And through the real route: a single slot is still enough for sequential requests.
    r = c.post("/api/identify", json={"image_b64": img})
    assert r.status_code == 200


def test_token_bucket_refills_over_time() -> None:
    bucket = _TokenBucket(capacity=1, window_s=60.0)
    assert bucket.allow() is True
    assert bucket.allow() is False
    bucket.updated -= 61.0  # simulate a minute elapsing without a real sleep
    assert bucket.allow() is True