diff --git a/litellm/proxy/middleware/guard_middleware.py b/litellm/proxy/middleware/guard_middleware.py index 45b1f67f08e..fd618f02c67 100644 --- a/litellm/proxy/middleware/guard_middleware.py +++ b/litellm/proxy/middleware/guard_middleware.py @@ -30,7 +30,7 @@ Design constraints honored here: from __future__ import annotations import os -from typing import TYPE_CHECKING +from typing import Final, TYPE_CHECKING from starlette.types import ASGIApp, Receive, Scope, Send @@ -39,10 +39,10 @@ from litellm._logging import verbose_proxy_logger if TYPE_CHECKING: from guard import SecurityConfig -GUARD_ENABLED_ENV = "LITELLM_GUARD_ENABLED" -GUARD_EXCLUDED_PATHS_ENV = "LITELLM_GUARD_EXCLUDED_PATHS" -DEFAULT_EXCLUDED_PATHS = "/health,/health/liveliness,/health/readiness,/metrics" -DEFAULT_TRUSTED_PROXIES = "10.0.0.0/8,172.16.0.0/12,192.168.0.0/16" +GUARD_ENABLED_ENV: Final = "LITELLM_GUARD_ENABLED" +GUARD_EXCLUDED_PATHS_ENV: Final = "LITELLM_GUARD_EXCLUDED_PATHS" +DEFAULT_EXCLUDED_PATHS: Final = "/health,/health/liveliness,/health/readiness,/metrics" +DEFAULT_TRUSTED_PROXIES: Final = "10.0.0.0/8,172.16.0.0/12,192.168.0.0/16" def _env_str(name: str, default: str = "") -> str: @@ -50,7 +50,7 @@ def _env_str(name: str, default: str = "") -> str: def _env_int(name: str, default: int) -> int: - raw = os.getenv(name) + raw: Final = os.getenv(name) if raw is None or raw == "": return default return int(raw) @@ -76,7 +76,7 @@ def _build_guard_config() -> SecurityConfig: """Build the SecurityConfig from the LITELLM_GUARD_* environment variables.""" from guard import SecurityConfig - kwargs: dict[str, object] = { # mutable-ok: assembled conditionally, then passed once to SecurityConfig + kwargs: Final[dict[str, object]] = { # mutable-ok: assembled conditionally, then passed once to SecurityConfig "enable_rate_limiting": True, "rate_limit": _env_int("LITELLM_GUARD_RATE_LIMIT", 100), "rate_limit_window": _env_int("LITELLM_GUARD_RATE_LIMIT_WINDOW", 60), diff --git a/tests/unit/proxy/middleware/test_guard_middleware.py b/tests/unit/proxy/middleware/test_guard_middleware.py index 51a94f86f79..aa7d47040f3 100644 --- a/tests/unit/proxy/middleware/test_guard_middleware.py +++ b/tests/unit/proxy/middleware/test_guard_middleware.py @@ -10,6 +10,7 @@ pyproject by this change), so behavior tests skip when it is absent. import asyncio import itertools import os +from typing import Final import pytest from starlette.testclient import TestClient @@ -19,18 +20,18 @@ from litellm.proxy.middleware.guard_middleware import ( guard_enabled, ) -_IP = itertools.count(1) +_IP: Final = itertools.count(1) try: import guard # noqa: F401 # availability probe for skip logic - _GUARD_AVAILABLE = True + _GUARD_AVAILABLE: Final = True except ImportError: - _GUARD_AVAILABLE = False + _GUARD_AVAILABLE: Final = False def _unique_ip() -> str: - n = next(_IP) + n: Final = next(_IP) return f"198.51.{n // 250}.{(n % 250) + 1}" @@ -57,7 +58,7 @@ def _build_app(): async def ping(request): return JSONResponse({"ok": True}) - app = Starlette(routes=[Route("/ping", ping, methods=["GET"])]) + app: Final = Starlette(routes=[Route("/ping", ping, methods=["GET"])]) app.add_middleware(GuardMiddleware) return app @@ -67,13 +68,14 @@ class _ClientIPInjector: the test IP, so mutate the scope before it reaches the app.""" def __init__(self, asgi_app, client_ip): - self._asgi_app = asgi_app - self._client_ip = client_ip + self._asgi_app: Final = asgi_app + self._client_ip: Final = client_ip async def __call__(self, scope, receive, send): - if scope["type"] == "http": - scope["client"] = (self._client_ip, 50000) - await self._asgi_app(scope, receive, send) + # TestClient hardcodes scope["client"]; the middleware stack must see + # the TEST-NET address, so build a new scope rather than mutate. + screened: Final = {**scope, "client": (self._client_ip, 50000)} if scope["type"] == "http" else scope + await self._asgi_app(screened, receive, send) def _client(client_ip): @@ -82,7 +84,7 @@ def _client(client_ip): def test_disabled_by_default(monkeypatch): monkeypatch.delenv("LITELLM_GUARD_ENABLED", raising=False) - client = _client(_unique_ip()) + client: Final = _client(_unique_ip()) assert client.get("/ping").status_code == 200 @@ -90,7 +92,7 @@ def test_fail_loud_without_package(monkeypatch): if _GUARD_AVAILABLE: pytest.skip("fastapi-guard is installed") monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") - client = _client(_unique_ip()) + client: Final = _client(_unique_ip()) # Flag on, package missing: request still passes (logged loudly), but # the guard attribute proves the wiring attempted the import. assert client.get("/ping").status_code == 200 @@ -98,7 +100,7 @@ def test_fail_loud_without_package(monkeypatch): @pytest.mark.skipif(not _GUARD_AVAILABLE, reason="fastapi-guard not installed") def test_blocked_ip_is_rejected(monkeypatch): - blocked = _unique_ip() + blocked: Final = _unique_ip() monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") monkeypatch.setenv("LITELLM_GUARD_BLOCKED_IPS", blocked) with _client(blocked) as client: @@ -107,18 +109,18 @@ def test_blocked_ip_is_rejected(monkeypatch): @pytest.mark.skipif(not _GUARD_AVAILABLE, reason="fastapi-guard not installed") def test_rate_limit_returns_429(monkeypatch): - client_ip = _unique_ip() + client_ip: Final = _unique_ip() monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") monkeypatch.setenv("LITELLM_GUARD_RATE_LIMIT", "2") monkeypatch.setenv("LITELLM_GUARD_RATE_LIMIT_WINDOW", "60") with _client(client_ip) as client: - codes = [client.get("/ping").status_code for _ in range(3)] + codes: Final = [client.get("/ping").status_code for _ in range(3)] assert codes == [200, 200, 429] @pytest.mark.skipif(not _GUARD_AVAILABLE, reason="fastapi-guard not installed") def test_passive_mode_never_blocks(monkeypatch): - blocked = _unique_ip() + blocked: Final = _unique_ip() monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") monkeypatch.setenv("LITELLM_GUARD_PASSIVE_MODE", "true") monkeypatch.setenv("LITELLM_GUARD_BLOCKED_IPS", blocked) @@ -131,5 +133,5 @@ def test_redis_stays_off_without_config(monkeypatch): monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") from litellm.proxy.middleware.guard_middleware import _build_guard_config - config = _build_guard_config() + config: Final = _build_guard_config() assert config.enable_redis is False