from __future__ import annotations from datetime import datetime, timezone from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException from sqlalchemy import create_engine, select from sqlalchemy.orm import Session from app.database import Base from app.main import _check_rate_limit, _is_trusted_proxy from app.models import SubmissionAttempt def test_rate_limit_is_persistent_and_does_not_store_raw_client() -> None: engine = create_engine("sqlite://") Base.metadata.create_all(engine) with Session(engine) as db: for _ in range(5): mock_request = MagicMock() mock_request.client.host = "203.0.113.42" mock_request.headers.get.return_value = None _check_rate_limit(mock_request, db) with pytest.raises(HTTPException) as blocked: mock_request = MagicMock() mock_request.client.host = "203.0.113.42" mock_request.headers.get.return_value = None _check_rate_limit(mock_request, db) assert blocked.value.status_code == 429 attempts = list(db.scalars(select(SubmissionAttempt))) assert len(attempts) == 5 assert all(item.client_hash != "203.0.113.42" and len(item.client_hash) == 64 for item in attempts) assert all(item.created_at.replace(tzinfo=timezone.utc) <= datetime.now(timezone.utc) for item in attempts) def test_rate_limit_uses_forwarded_for_header_from_trusted_proxy() -> None: """X-Forwarded-For should be used when client is trusted proxy.""" engine = create_engine("sqlite://") Base.metadata.create_all(engine) with Session(engine) as db: with patch("app.main.settings") as mock_settings: mock_settings.rate_limit_secret = "test-secret-for-testing" mock_settings.trusted_proxy_cidrs = ["127.0.0.1/32"] mock_proxy = MagicMock() mock_proxy.client.host = "127.0.0.1" mock_proxy.headers.get.return_value = "198.51.100.10" for _ in range(5): _check_rate_limit(mock_proxy, db) with pytest.raises(HTTPException) as blocked: mock_other = MagicMock() mock_other.client.host = "127.0.0.1" mock_other.headers.get.return_value = "198.51.100.10" _check_rate_limit(mock_other, db) assert blocked.value.status_code == 429 def test_trusted_proxy_checks_cidrs() -> None: assert _is_trusted_proxy("127.0.0.1", ["127.0.0.1/32"]) is True assert _is_trusted_proxy("10.0.0.1", ["10.0.0.0/8"]) is True assert _is_trusted_proxy("192.168.1.1", ["192.168.1.0/24"]) is True assert _is_trusted_proxy("::1", ["::1/128"]) is True assert _is_trusted_proxy("203.0.113.5", ["127.0.0.1/32"]) is False assert _is_trusted_proxy("invalid", []) is False def test_rate_limit_ignores_forwarded_for_from_untrusted_client() -> None: """X-Forwarded-For should be ignored when client is not in trusted CIDRs.""" engine = create_engine("sqlite://") Base.metadata.create_all(engine) with Session(engine) as db: # Client 203.0.113.42 is NOT trusted by default mock_untrusted = MagicMock() mock_untrusted.client.host = "203.0.113.42" mock_untrusted.headers.get.return_value = "10.0.0.99" for i in range(3): _check_rate_limit(mock_untrusted, db) # Should use real client 203.0.113.42, not forwarded 10.0.0.99 # So 3 attempts from 203.0.113.42 should be allowed (limit is 5) mock_different = MagicMock() mock_different.client.host = "203.0.113.42" mock_different.headers.get.return_value = "10.0.0.88" _check_rate_limit(mock_different, db) # Should succeed, not blocked def test_rate_limit_uses_forwarded_for_from_trusted_proxy() -> None: """X-Forwarded-For should be used when client IS in trusted CIDRs.""" engine = create_engine("sqlite://") Base.metadata.create_all(engine) with Session(engine) as db: # 127.0.0.1 IS trusted by default with patch("app.main.settings") as mock_settings: mock_settings.rate_limit_secret = "test-secret-for-testing" mock_settings.trusted_proxy_cidrs = ["127.0.0.1/32", "::1/128"] mock_proxy = MagicMock() mock_proxy.client.host = "127.0.0.1" mock_proxy.headers.get.return_value = "203.0.113.100" for _ in range(5): _check_rate_limit(mock_proxy, db) # Should use forwarded IP 203.0.113.100, so a different forwarded IP should be allowed mock_other_forwarded = MagicMock() mock_other_forwarded.client.host = "127.0.0.1" mock_other_forwarded.headers.get.return_value = "198.51.100.50" _check_rate_limit(mock_other_forwarded, db) # Should succeed def test_rate_limit_independent_limits_for_two_clients_through_proxy() -> None: """Two clients behind trusted proxy should have independent rate limits.""" engine = create_engine("sqlite://") Base.metadata.create_all(engine) with Session(engine) as db: with patch("app.main.settings") as mock_settings: mock_settings.rate_limit_secret = "test-secret-for-testing" # Trusted proxy is the Astro container IP mock_settings.trusted_proxy_cidrs = ["172.17.0.0/16"] # Docker network # Client 1: 198.51.100.10 mock_client1 = MagicMock() mock_client1.client.host = "172.17.0.3" # Astro container mock_client1.headers.get.return_value = "198.51.100.10" # Client 2: 198.51.100.20 mock_client2 = MagicMock() mock_client2.client.host = "172.17.0.3" # Same Astro container mock_client2.headers.get.return_value = "198.51.100.20" # Client 1 makes 5 requests for _ in range(5): _check_rate_limit(mock_client1, db) # Client 1 should be blocked with pytest.raises(HTTPException) as blocked: _check_rate_limit(mock_client1, db) assert blocked.value.status_code == 429 # Client 2 should still be allowed (independent limit) _check_rate_limit(mock_client2, db) # Should succeed def test_forged_xff_rejected_on_untrusted_port() -> None: """XFF should be rejected when connection is not from trusted proxy.""" engine = create_engine("sqlite://") Base.metadata.create_all(engine) with Session(engine) as db: with patch("app.main.settings") as mock_settings: mock_settings.rate_limit_secret = "test-secret-for-testing" # Only trust Docker network, NOT direct connections mock_settings.trusted_proxy_cidrs = ["172.17.0.0/16"] # Direct connection with forged XFF mock_direct = MagicMock() mock_direct.client.host = "203.0.113.50" # Not in trusted CIDR mock_direct.headers.get.return_value = "10.0.0.1" # Forged XFF # Should use real client 203.0.113.50, not forged 10.0.0.1 for i in range(3): _check_rate_limit(mock_direct, db) # Another request from same real client should count mock_direct2 = MagicMock() mock_direct2.client.host = "203.0.113.50" mock_direct2.headers.get.return_value = "10.0.0.2" # Different forged XFF _check_rate_limit(mock_direct2, db) # Should succeed (4th request from 203.0.113.50) def test_direct_access_without_xff_header() -> None: """Direct access without X-Forwarded-For should use real client IP.""" engine = create_engine("sqlite://") Base.metadata.create_all(engine) with Session(engine) as db: with patch("app.main.settings") as mock_settings: mock_settings.rate_limit_secret = "test-secret-for-testing" mock_settings.trusted_proxy_cidrs = ["127.0.0.1/32"] # Direct connection without XFF mock_direct = MagicMock() mock_direct.client.host = "192.168.1.100" mock_direct.headers.get.return_value = None # No XFF # Should use real client 192.168.1.100 _check_rate_limit(mock_direct, db) attempts = list(db.scalars(select(SubmissionAttempt))) assert len(attempts) == 1 # Hash should be of the real IP, not empty assert len(attempts[0].client_hash) == 64