229 lines
10 KiB
Python
229 lines
10 KiB
Python
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.admin_security import verify_admin
|
|
from app.main import _check_rate_limit, _is_trusted_proxy
|
|
from app.models import AdminAuthAttempt, SubmissionAttempt
|
|
|
|
|
|
def _admin_config() -> MagicMock:
|
|
config = MagicMock()
|
|
config.admin_token = "correct-token"
|
|
config.admin_auth_attempt_limit = 3
|
|
config.admin_auth_window_seconds = 600
|
|
config.rate_limit_secret = "test-secret-for-testing"
|
|
config.trusted_proxy_cidrs = ["127.0.0.1/32"]
|
|
return config
|
|
|
|
|
|
def test_admin_auth_failures_are_hashed_and_limited() -> None:
|
|
engine = create_engine("sqlite://")
|
|
Base.metadata.create_all(engine)
|
|
request = MagicMock()
|
|
request.client.host = "203.0.113.42"
|
|
request.headers.get.return_value = None
|
|
with Session(engine) as db:
|
|
for _ in range(3):
|
|
with pytest.raises(HTTPException) as denied:
|
|
verify_admin(request, db, "Bearer wrong", _admin_config())
|
|
assert denied.value.status_code == 401
|
|
with pytest.raises(HTTPException) as limited:
|
|
verify_admin(request, db, "Bearer correct-token", _admin_config())
|
|
assert limited.value.status_code == 429
|
|
attempts = list(db.scalars(select(AdminAuthAttempt)))
|
|
assert len(attempts) == 3
|
|
assert all(item.client_hash != "203.0.113.42" and len(item.client_hash) == 64 for item in attempts)
|
|
|
|
|
|
def test_successful_admin_auth_clears_previous_failures() -> None:
|
|
engine = create_engine("sqlite://")
|
|
Base.metadata.create_all(engine)
|
|
request = MagicMock()
|
|
request.client.host = "203.0.113.43"
|
|
request.headers.get.return_value = None
|
|
with Session(engine) as db:
|
|
with pytest.raises(HTTPException):
|
|
verify_admin(request, db, "Bearer wrong", _admin_config())
|
|
assert verify_admin(request, db, "Bearer correct-token", _admin_config()) == "admin"
|
|
assert list(db.scalars(select(AdminAuthAttempt))) == []
|
|
|
|
|
|
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
|