AegisGitea-MCP/tests/security/test_oauth_input_bounds.py

"""Input bounds of the unauthenticated OAuth routes (SEC-01, SEC-09)."""

from __future__ import annotations

import threading
from typing import Any

import pytest
from mcp.server.auth.provider import RegistrationError
from starlette.applications import Starlette
from starlette.testclient import TestClient

from aegis_gitea_mcp.oauth import provider as provider_module
from aegis_gitea_mcp.oauth.callback import MAX_CONSENT_BODY_BYTES
from aegis_gitea_mcp.oauth.provider import MAX_PENDING_PER_SOURCE, MAX_REGISTRATIONS_PER_SOURCE
from aegis_gitea_mcp.request_context import set_client_address
from tests.security.helpers import asgi_request
from tests.test_oauth_helpers import (
    BASE,
    REDIRECT,
    build_oauth_app,
    oauth_settings,
    pkce,
    register,
)
from tests.test_oauth_provider import client_info

CHUNK = 64 * 1024


@pytest.fixture(autouse=True)
def _env(mock_env: None) -> None:
    """The audit logger and the rate limiter read the process-wide settings."""


def chunks_of(total: int) -> list[bytes]:
    return [b"a" * CHUNK for _ in range(total // CHUNK)]


async def test_consent_post_over_the_cap_is_413_without_reading_the_body() -> None:
    app = build_oauth_app().app
    result = await asgi_request(
        app,
        "POST",
        "/oauth/consent",
        headers={
            "content-type": "application/x-www-form-urlencoded",
            "content-length": str(1024 * 1024),
            "origin": BASE,
        },
        chunks=chunks_of(1024 * 1024),
    )
    assert result.status == 413
    assert result.bytes_read == 0


async def test_consent_post_without_declared_length_stops_reading_at_the_cap() -> None:
    app = build_oauth_app().app
    result = await asgi_request(
        app,
        "POST",
        "/oauth/consent",
        headers={"content-type": "application/x-www-form-urlencoded", "origin": BASE},
        chunks=chunks_of(1024 * 1024),
    )
    assert result.status == 413
    assert result.bytes_read <= MAX_CONSENT_BODY_BYTES + CHUNK


async def test_consent_multipart_is_415_without_reading_any_body() -> None:
    app = build_oauth_app().app
    result = await asgi_request(
        app,
        "POST",
        "/oauth/consent",
        headers={"content-type": "multipart/form-data; boundary=x", "origin": BASE},
        chunks=chunks_of(512 * 1024),
    )
    assert result.status == 415
    assert result.bytes_read == 0


async def test_consent_json_body_is_415() -> None:
    app = build_oauth_app().app
    result = await asgi_request(
        app,
        "POST",
        "/oauth/consent",
        headers={"content-type": "application/json", "origin": BASE},
        chunks=[b"{}"],
    )
    assert result.status == 415


async def test_small_urlencoded_consent_post_still_reaches_the_handler() -> None:
    app = build_oauth_app().app
    result = await asgi_request(
        app,
        "POST",
        "/oauth/consent",
        headers={"content-type": "application/x-www-form-urlencoded", "origin": BASE},
        chunks=[b"request=unknown&csrf=x&decision=approve"],
    )
    # Reaches the handler, which answers its fixed 400 page for an unknown request.
    assert result.status == 400
    assert b"invalid or has expired" in result.body


# SEC-09: pre-sign-in state cannot be flushed by one address


def client_at(app: Starlette, address: str) -> TestClient:
    """A browser-like client whose requests come from ``address``."""
    return TestClient(app, base_url=BASE, client=(address, 50000), follow_redirects=False)


def authorize(client: TestClient, registration: dict[str, str]) -> Any:
    return client.get(
        "/authorize",
        params={
            "client_id": registration["client_id"],
            "redirect_uri": REDIRECT,
            "response_type": "code",
            "code_challenge": pkce()[1],
            "code_challenge_method": "S256",
        },
    )


def test_a_flood_of_authorize_calls_from_one_address_does_not_evict_another_users_request() -> None:
    app = build_oauth_app(oauth_settings(rate_limit_per_minute=1_000_000)).app
    victim = client_at(app, "198.51.100.20")
    consent_url = authorize(victim, register(victim).json()).headers["location"]
    assert victim.get(consent_url).status_code == 200

    attacker = client_at(app, "203.0.113.66")
    registration = register(attacker).json()
    accepted = refused = 0
    for _ in range(1100):
        location = authorize(attacker, registration).headers["location"]
        if "/oauth/consent" in location:
            accepted += 1
        else:
            refused += 1
            assert "error=temporarily_unavailable" in location

    assert accepted == MAX_PENDING_PER_SOURCE
    assert refused == 1100 - MAX_PENDING_PER_SOURCE
    assert victim.get(consent_url).status_code == 200  # still pending


def test_registrations_per_address_are_limited() -> None:
    app = build_oauth_app(oauth_settings(rate_limit_per_minute=1_000_000)).app
    flooder = client_at(app, "203.0.113.66")

    statuses = [register(flooder).status_code for _ in range(MAX_REGISTRATIONS_PER_SOURCE + 5)]

    assert statuses[:MAX_REGISTRATIONS_PER_SOURCE] == [201] * MAX_REGISTRATIONS_PER_SOURCE
    assert statuses[MAX_REGISTRATIONS_PER_SOURCE:] == [400] * 5
    # Another address is unaffected.
    assert register(client_at(app, "198.51.100.20")).status_code == 201


async def test_a_young_client_survives_a_registration_flood(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(provider_module, "MAX_CLIENTS", 10)
    oauth = build_oauth_app()
    provider = oauth.provider
    set_client_address("198.51.100.20")
    await provider.register_client(client_info("legit"))
    oauth.clock.advance(60)

    refused = 0
    for number in range(40):
        set_client_address(f"203.0.113.{number}")
        try:
            await provider.register_client(client_info(f"flood-{number}"))
        except RegistrationError:
            refused += 1

    assert "legit" in provider._store.clients
    assert refused > 0 and len(provider._store.clients) == 10

    # Once clients are older than the pending window, old ones can be evicted again.
    oauth.clock.advance(provider_module.PENDING_TTL + 1)
    set_client_address("198.51.100.99")
    await provider.register_client(client_info("later"))
    assert "later" in provider._store.clients


async def test_the_store_write_runs_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None:
    store = build_oauth_app().store
    threads: list[str] = []
    original = store._write

    def recording(data: bytes) -> None:
        threads.append(threading.current_thread().name)
        original(data)

    monkeypatch.setattr(store, "_write", recording)

    await store.save()

    assert threads and threads[0] != threading.main_thread().name