AegisGitea-MCP/tests/security/test_rate_limits.py

"""Rate limiting covers the OAuth endpoints, IPv6 subscribers and untrusted proxies (SEC-08)."""

from __future__ import annotations

import logging
from typing import Any

import pytest
from starlette.types import Message, Receive, Scope, Send

from aegis_gitea_mcp.rate_limit import RequestRateLimiter, client_bucket
from aegis_gitea_mcp.transports import ForwardedHeadersMiddleware
from tests.test_executor_helpers import make_settings
from tests.test_oauth_helpers import browser, build_oauth_app, oauth_settings

LIMIT = 3


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


@pytest.mark.parametrize(
    ("method", "path"),
    [
        ("POST", "/token"),
        ("POST", "/revoke"),
        ("POST", "/register"),
        ("GET", "/authorize"),
        ("GET", "/oauth/callback"),
        ("POST", "/oauth/consent"),
        ("POST", "/mcp"),
    ],
)
def test_every_unauthenticated_endpoint_is_rate_limited(method: str, path: str) -> None:
    client = browser(build_oauth_app(oauth_settings(rate_limit_per_minute=LIMIT)).app)

    statuses = [client.request(method, path).status_code for _ in range(LIMIT + 1)]

    assert statuses[LIMIT] == 429
    assert 429 not in statuses[:LIMIT]


def test_the_429_answer_has_retry_after_and_no_detail() -> None:
    client = browser(build_oauth_app(oauth_settings(rate_limit_per_minute=1)).app)
    client.post("/token")

    response = client.post("/token")

    assert response.status_code == 429
    assert response.headers["retry-after"] == "60"
    assert response.json() == {"error": {"kind": "rate_limited", "message": "Too many requests"}}


def test_ipv6_addresses_of_one_slash_64_share_a_bucket() -> None:
    limiter = RequestRateLimiter(make_settings(rate_limit_per_minute=2))

    first = limiter.check("2001:db8:1:2::1", None)
    second = limiter.check("2001:db8:1:2:ffff:ffff:ffff:ffff", None)
    third = limiter.check("2001:db8:1:2::abcd", None)
    other_network = limiter.check("2001:db8:1:3::1", None)

    assert first.allowed and second.allowed
    assert not third.allowed
    assert other_network.allowed


def test_ipv4_clients_and_mapped_addresses_keep_their_own_bucket() -> None:
    assert client_bucket("192.0.2.10") == "192.0.2.10"
    assert client_bucket("::ffff:192.0.2.10") == "192.0.2.10"
    assert client_bucket("2001:db8:1:2::1") == "2001:db8:1:2::/64"
    assert client_bucket("fe80::1%eth0") == "fe80::/64"
    assert client_bucket("testclient") == "testclient"
    limiter = RequestRateLimiter(make_settings(rate_limit_per_minute=1))
    assert limiter.check("192.0.2.10", None).allowed
    assert limiter.check("192.0.2.11", None).allowed
    assert not limiter.check("192.0.2.10", None).allowed


# Forwarded headers


class Recorder:
    """Downstream ASGI app that remembers the client address it saw."""

    def __init__(self) -> None:
        self.clients: list[Any] = []

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        self.clients.append(scope.get("client"))


def http_scope(peer: str, forwarded: str | None) -> Scope:
    headers = [(b"host", b"mcp.example.com")]
    if forwarded is not None:
        headers.append((b"x-forwarded-for", forwarded.encode()))
    return {
        "type": "http",
        "method": "GET",
        "path": "/",
        "scheme": "http",
        "headers": headers,
        "client": (peer, 5000),
    }


async def noop_receive() -> Message:
    return {"type": "http.disconnect"}


async def noop_send(message: Message) -> None:
    return None


async def test_forwarded_for_from_an_untrusted_peer_is_ignored_and_warned_about_once(
    caplog: pytest.LogCaptureFixture,
) -> None:
    app = Recorder()
    middleware = ForwardedHeadersMiddleware(app, "127.0.0.1")

    with caplog.at_level(logging.WARNING, logger="aegis_gitea_mcp.transports"):
        for _ in range(3):
            await middleware(http_scope("172.18.0.5", "203.0.113.9"), noop_receive, noop_send)

    # Spoofed header: the peer address stays the client address.
    assert app.clients == [("172.18.0.5", 5000)] * 3
    warnings = [r for r in caplog.records if "FORWARDED_ALLOW_IPS" in r.getMessage()]
    assert len(warnings) == 1
    assert "203.0.113.9" not in warnings[0].getMessage()
    assert "172.18.0.5" not in warnings[0].getMessage()


async def test_forwarded_for_from_a_trusted_peer_is_applied_without_a_warning(
    caplog: pytest.LogCaptureFixture,
) -> None:
    app = Recorder()
    middleware = ForwardedHeadersMiddleware(app, "172.18.0.0/16, 127.0.0.1")

    with caplog.at_level(logging.WARNING, logger="aegis_gitea_mcp.transports"):
        await middleware(http_scope("172.18.0.5", "203.0.113.9"), noop_receive, noop_send)
        await middleware(http_scope("127.0.0.1", "203.0.113.10"), noop_receive, noop_send)

    assert [client[0] for client in app.clients] == ["203.0.113.9", "203.0.113.10"]
    assert caplog.records == []


async def test_requests_without_forwarded_for_never_warn(caplog: pytest.LogCaptureFixture) -> None:
    app = Recorder()
    middleware = ForwardedHeadersMiddleware(app, "127.0.0.1")

    with caplog.at_level(logging.WARNING, logger="aegis_gitea_mcp.transports"):
        await middleware(http_scope("198.51.100.7", None), noop_receive, noop_send)

    assert caplog.records == []
    assert app.clients == [("198.51.100.7", 5000)]


async def test_wildcard_trusts_every_peer_without_a_warning(
    caplog: pytest.LogCaptureFixture,
) -> None:
    app = Recorder()
    middleware = ForwardedHeadersMiddleware(app, "*")

    with caplog.at_level(logging.WARNING, logger="aegis_gitea_mcp.transports"):
        await middleware(http_scope("198.51.100.7", "203.0.113.9"), noop_receive, noop_send)

    assert app.clients[0][0] == "203.0.113.9"
    assert caplog.records == []