diff --git a/litellm/proxy/middleware/guard_middleware.py b/litellm/proxy/middleware/guard_middleware.py new file mode 100644 index 00000000000..3012cdf9168 --- /dev/null +++ b/litellm/proxy/middleware/guard_middleware.py @@ -0,0 +1,210 @@ +"""Optional fastapi-guard security middleware for the proxy. + +fastapi-guard (https://github.com/Guard-Core/fastapi-guard) provides IP +block/allow lists, rate limiting with auto-ban, user-agent blocking, +penetration-attempt detection, passive mode, optional Redis-backed shared +state, and optional IPInfo geo/cloud-provider lookups. + +Positioning: this is edge protection at the ASGI layer, complementary to +litellm's existing controls. It does not replace the virtual-key RPM/TPM +limiter (identity-level), the Enterprise ``allowed_ips`` auth-layer check, +or the login throttle. It runs before auth so screened requests never +reach key validation. + +Design constraints honored here: + +- Pure ASGI (``__call__(scope, receive, send)``), never BaseHTTPMiddleware, + matching the middleware house style and its streaming-preservation test. +- Opt-in via ``LITELLM_GUARD_ENABLED``; when unset, the wrapper delegates + straight through with no guard import and no config work. +- Fail-closed when enabled: the proxy does not depend on fastapi-guard (it + is deliberately absent from pyproject), so the import is lazy. If the + flag is set and the package is missing, every screened request is + rejected with a 503 in litellm's OpenAI-error envelope and an error is + logged once: an operator who explicitly enabled the guard must not get + unscreened passthrough. (Fork PRs may not modify dependency files per + the guard-fork-dependencies CI policy, so the dependency addition is + left to the maintainers.) +- Zero-trust client identity by default: ``X-Forwarded-For`` is honored + only from proxies explicitly listed in ``LITELLM_GUARD_TRUSTED_PROXIES`` + (guard ignores forwarded headers from unlisted peers), so a caller + cannot spoof its way past allowlists, limits, or bans by setting the + header. Deployments behind a load balancer must list their ingress + addresses there. +- Rejections guard itself builds use guard's response shapes; the one + rejection this wrapper builds (guard unavailable) uses litellm's + OpenAI-error JSON shape (the ``AdmissionControlMiddleware`` model), so + clients see a familiar envelope. +""" + +from __future__ import annotations + +import importlib +import os +from collections.abc import Callable +from typing import Final + +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send + +from litellm._logging import verbose_proxy_logger + +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") + + +def _env_str(name: str, default: str = "") -> str: + return os.getenv(name, default) + + +def _env_int(name: str, default: int) -> int: + raw: Final = os.getenv(name) + if raw is None or raw == "": + return default + return int(raw) + + +def _env_bool(name: str) -> bool: + return os.getenv(name, "").strip().lower() in ("1", "true", "yes") + + +def _csv(raw: str | None) -> tuple[str, ...]: + if not raw: + return () + return tuple(item.strip() for item in raw.split(",") if item.strip()) + + +def guard_enabled() -> bool: + """Read lazily per request so YAML ``environment_variables`` entries are + seen (the ``LITELLM_ENABLE_HSTS`` pattern).""" + return _env_bool(GUARD_ENABLED_ENV) + + +def _build_guard_config() -> object: + """Build the guard config from the LITELLM_GUARD_* environment variables. + + Returns guard's ``SecurityConfig``. The ``object`` return and the + dynamic import are the seam: fastapi-guard is an optional dependency + that is deliberately absent from the environments litellm type-checks, + so the class can only be reached through a runtime lookup. + """ + guard_module: Final = importlib.import_module("guard") + # fastapi-guard is an optional dependency that is deliberately absent + # from the environments litellm type-checks, so attributes read off the + # dynamically imported module are Any; the Callable annotation re-types + # the seam so no Any escapes this line. + security_config_cls: Final[Callable[..., object]] = guard_module.SecurityConfig # pyright: ignore[reportAny] # optional dep absent from the typecheck venv; seam re-typed below + + redis_url_env: Final = _env_str("LITELLM_GUARD_REDIS_URL") + ipinfo_token_env: Final = _env_str("LITELLM_GUARD_IPINFO_TOKEN") + # One literal, no conditional inserts: optional list-valued features are + # passed as their empty/None form, which SecurityConfig treats the same + # as unset. In-memory state unless a Redis URL is configured: never + # implicitly depend on a Redis server being reachable. The proxy + # defaults to a single uvicorn worker; multi-worker deployments should + # set LITELLM_GUARD_REDIS_URL so counters and bans are shared. + kwargs: Final = { + "enable_rate_limiting": True, + "rate_limit": _env_int("LITELLM_GUARD_RATE_LIMIT", 100), + "rate_limit_window": _env_int("LITELLM_GUARD_RATE_LIMIT_WINDOW", 60), + "enable_ip_banning": True, + "auto_ban_threshold": _env_int("LITELLM_GUARD_AUTO_BAN_THRESHOLD", 10), + "auto_ban_duration": _env_int("LITELLM_GUARD_AUTO_BAN_DURATION", 300), + "enable_penetration_detection": True, + "enable_redis": bool(redis_url_env), + "redis_url": redis_url_env or None, + "redis_prefix": "litellm_guard:", + "passive_mode": _env_bool("LITELLM_GUARD_PASSIVE_MODE"), + "blacklist": _csv(_env_str("LITELLM_GUARD_BLOCKED_IPS")), + "whitelist": _csv(_env_str("LITELLM_GUARD_ALLOWED_IPS")) or None, + "blocked_user_agents": list(_csv(_env_str("LITELLM_GUARD_BLOCKED_USER_AGENTS"))), + "blocked_countries": frozenset(_csv(_env_str("LITELLM_GUARD_BLOCKED_COUNTRIES"))), + "whitelist_countries": frozenset(_csv(_env_str("LITELLM_GUARD_ALLOWED_COUNTRIES"))), + "block_cloud_providers": frozenset(_csv(_env_str("LITELLM_GUARD_BLOCK_CLOUD_PROVIDERS"))) or None, + # Zero-trust default: no proxies are trusted, so forwarded headers + # are ignored and the connecting peer is the client. List real + # ingress addresses to enable X-Forwarded-For resolution. + "trusted_proxies": _csv(_env_str("LITELLM_GUARD_TRUSTED_PROXIES")), + "trusted_proxy_depth": _env_int("LITELLM_GUARD_TRUSTED_PROXY_DEPTH", 1), + "exclude_paths": list(_csv(_env_str(GUARD_EXCLUDED_PATHS_ENV)) or DEFAULT_EXCLUDED_PATHS), + "ipinfo_token": ipinfo_token_env or None, + } + return security_config_cls(**kwargs) + + +def _guard_unavailable_response() -> JSONResponse: + """Fail-closed rejection for when the guard is enabled but cannot be + initialized (litellm's OpenAI-error envelope, the + ``AdmissionControlMiddleware`` model).""" + return JSONResponse( + status_code=503, + content={ + "error": { + "message": ( + "LITELLM_GUARD_ENABLED is set but the fastapi-guard package " + "is not installed. Install it with: pip install fastapi-guard. " + "Requests are rejected until then." + ), + "type": "guard_unavailable_error", + "code": "503", + } + }, + ) + + +class GuardMiddleware: + """Pure-ASGI fastapi-guard wrapper. + + The inner fastapi-guard middleware is instantiated on the first request + after the flag is seen (so YAML-configured environment variables are + honored), and the downstream app is wrapped by it; everything the flag + admits flows through guard's pipeline and then the normal proxy stack. + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + self._guard_asgi: ASGIApp | None = None + self._guard_failed = False + + def _ensure_guard(self) -> ASGIApp | None: + """Build and return the guard ASGI app, or None when the package is + missing (logged loudly once; the caller then fails closed).""" + if self._guard_asgi is not None: + return self._guard_asgi + if self._guard_failed: + return None + try: + guard_module: Final = importlib.import_module("guard") + # Same optional-dep seam as _build_guard_config: the attribute + # is Any; the Callable annotation keeps it from escaping. + middleware_cls: Final[Callable[..., ASGIApp]] = guard_module.SecurityMiddleware # pyright: ignore[reportAny] # optional dep absent from the typecheck venv; seam re-typed below + except ImportError as exc: + self._guard_failed = True + # Log loudly once; security misconfiguration must be visible. + verbose_proxy_logger.error( + "LITELLM_GUARD_ENABLED is set but fastapi-guard is not " + "installed: %s. Install it with: pip install fastapi-guard. " + "Requests are REJECTED until then.", + exc, + ) + return None + + self._guard_asgi = middleware_cls(self.app, config=_build_guard_config()) + return self._guard_asgi + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http" or not guard_enabled(): + await self.app(scope, receive, send) + return + + guard_asgi: Final = self._ensure_guard() + if guard_asgi is None: + # Flag on, package missing: fail closed with 503 rather than + # forward unscreened. The operator explicitly enabled the + # guard; silently honoring the opposite of that intent would + # make the flag a no-op. The misconfiguration is logged once. + await _guard_unavailable_response()(scope, receive, send) + return + + await guard_asgi(scope, receive, send) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 089647c6297..775591e6c71 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2563,6 +2563,18 @@ current_dir = os.path.dirname(os.path.abspath(__file__)) # app.mount("/ui", StaticFiles(directory=ui_path, html=True), name="ui") +# Optional fastapi-guard security middleware (off unless LITELLM_GUARD_ENABLED +# is set). Added before CORSMiddleware and SecurityHeadersMiddleware so it +# nests *inside* both (Starlette makes the last-added middleware outermost) +# and guard's own 403/429 rejections still leave with CORS and security +# headers. It stays outside the auth dependencies, so screened-and-blocked +# requests never reach key validation. See middleware/guard_middleware.py. +from litellm.proxy.middleware.guard_middleware import ( # noqa: E402 # mid-file attach: the middleware stack freezes once the lifespan runs + GuardMiddleware, +) + +app.add_middleware(GuardMiddleware) + app.add_middleware( CORSMiddleware, allow_origins=origins, diff --git a/tests/unit/proxy/middleware/test_guard_middleware.py b/tests/unit/proxy/middleware/test_guard_middleware.py new file mode 100644 index 00000000000..c0f7fd7bac7 --- /dev/null +++ b/tests/unit/proxy/middleware/test_guard_middleware.py @@ -0,0 +1,257 @@ +""" +Unit tests for the optional fastapi-guard middleware wiring. + +Mirrors test_security_headers_middleware.py: minimal app, TestClient, +monkeypatched env, default-off assertion, and the pure-ASGI house-style +assertion. fastapi-guard is an optional dependency (not added to +pyproject by this change), so behavior tests skip when it is absent; the +fail-closed path is exercised unconditionally via a poisoned sys.modules +entry. +""" + +import itertools +import sys +from typing import Final + +import pytest +from starlette.testclient import TestClient + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.middleware.guard_middleware import ( + GuardMiddleware, + guard_enabled, +) + +_IP: Final = itertools.count(1) + +try: + import guard # noqa: F401 # availability probe for skip logic + + _GUARD_AVAILABLE: Final = True +except ImportError: + _GUARD_AVAILABLE: Final = False + + +def _unique_ip() -> str: + n: Final = next(_IP) + return f"198.51.{n // 250}.{(n % 250) + 1}" + + +def test_is_pure_asgi_not_base_http_middleware(): + """BaseHTTPMiddleware degrades streaming; this must be pure ASGI.""" + from starlette.middleware.base import BaseHTTPMiddleware + + assert not issubclass(GuardMiddleware, BaseHTTPMiddleware) + assert "__call__" in GuardMiddleware.__dict__ + + +def test_guard_enabled_flag(monkeypatch): + monkeypatch.delenv("LITELLM_GUARD_ENABLED", raising=False) + assert guard_enabled() is False + monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") + assert guard_enabled() is True + + +def _build_app(): + from starlette.applications import Starlette + from starlette.responses import JSONResponse + from starlette.routing import Route + + async def ping(request): + return JSONResponse({"ok": True}) + + app: Final = Starlette(routes=[Route("/ping", ping, methods=["GET"])]) + app.add_middleware(GuardMiddleware) + return app + + +def _build_stacked_app(): + """The proxy's nesting: guard innermost, then CORS, then security + headers outermost (Starlette makes the last-added middleware + outermost), so guard rejections leave with their headers.""" + from starlette.applications import Starlette + from starlette.middleware.cors import CORSMiddleware + from starlette.responses import JSONResponse + from starlette.routing import Route + + from litellm.proxy.middleware.security_headers_middleware import ( + SecurityHeadersMiddleware, + ) + + async def ping(request): + return JSONResponse({"ok": True}) + + app: Final = Starlette(routes=[Route("/ping", ping, methods=["GET"])]) + app.add_middleware(GuardMiddleware) + app.add_middleware( + CORSMiddleware, + allow_origins=["https://ok.example"], + allow_methods=["GET"], + allow_headers=["*"], + ) + app.add_middleware(SecurityHeadersMiddleware) + return app + + +class _ClientIPInjector: + """TestClient hardcodes scope["client"]; the middleware stack must see + the test IP, so mutate the scope before it reaches the app.""" + + def __init__(self, asgi_app, client_ip): + self._asgi_app: Final = asgi_app + self._client_ip: Final = client_ip + + async def __call__(self, 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, asgi_app=None): + return TestClient(_ClientIPInjector(asgi_app or _build_app(), client_ip)) + + +def test_disabled_by_default(monkeypatch): + monkeypatch.delenv("LITELLM_GUARD_ENABLED", raising=False) + client: Final = _client(_unique_ip()) + assert client.get("/ping").status_code == 200 + + +def test_enabled_without_package_fails_closed(monkeypatch): + """Flag on + package missing: every request is rejected with 503 and + the misconfiguration is logged exactly once. An operator who enabled + the guard must not get unscreened passthrough.""" + monkeypatch.setitem(sys.modules, "guard", None) + monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") + error_messages: Final = [] + monkeypatch.setattr(verbose_proxy_logger, "error", lambda msg, *args: error_messages.append(msg)) + client: Final = _client(_unique_ip()) + first: Final = client.get("/ping") + second: Final = client.get("/ping") + assert first.status_code == 503 + assert second.status_code == 503 + assert first.json()["error"]["type"] == "guard_unavailable_error" + assert len(error_messages) == 1 + + +class _StubSecurityConfig: + """Captures the kwargs the middleware assembles from the environment.""" + + def __init__(self, **kwargs): # kwargs-ok: mirrors the SecurityConfig keyword surface the wrapper calls + self.kwargs: Final = kwargs + + +class _StubSecurityMiddleware: + """Pass-through stand-in recording how the wrapper wired it.""" + + instances: Final = [] + + def __init__(self, app, config=None): + self.app: Final = app + self.config: Final = config + _StubSecurityMiddleware.instances.append(self) + + async def __call__(self, scope, receive, send): + await self.app(scope, receive, send) + + +def test_enabled_wires_guard_with_env_config(monkeypatch): + """Flag on + importable guard module: the wrapper builds the config + from the environment and delegates through the guard middleware. Uses + a stub module, so the wiring is exercised in CI where the real + optional package is absent.""" + import types + + stub_module: Final = types.ModuleType("guard") + stub_module.SecurityConfig = _StubSecurityConfig + stub_module.SecurityMiddleware = _StubSecurityMiddleware + monkeypatch.setitem(sys.modules, "guard", stub_module) + _StubSecurityMiddleware.instances.clear() + monkeypatch.delenv("LITELLM_GUARD_REDIS_URL", raising=False) + monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") + monkeypatch.setenv("LITELLM_GUARD_RATE_LIMIT", "33") + client: Final = _client(_unique_ip()) + assert client.get("/ping").status_code == 200 + wired: Final = _StubSecurityMiddleware.instances[-1] + assert isinstance(wired.config, _StubSecurityConfig) + assert wired.config.kwargs["rate_limit"] == 33 + assert wired.config.kwargs["enable_redis"] is False + assert wired.config.kwargs["trusted_proxies"] == () + assert wired.config.kwargs["exclude_paths"] == [ + "/health", + "/health/liveliness", + "/health/readiness", + "/metrics", + ] + + +@pytest.mark.skipif(not _GUARD_AVAILABLE, reason="fastapi-guard not installed") +def test_blocked_ip_is_rejected(monkeypatch): + blocked: Final = _unique_ip() + monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") + monkeypatch.setenv("LITELLM_GUARD_BLOCKED_IPS", blocked) + with _client(blocked) as client: + assert client.get("/ping").status_code == 403 + + +@pytest.mark.skipif(not _GUARD_AVAILABLE, reason="fastapi-guard not installed") +def test_untrusted_forwarded_header_is_ignored(monkeypatch): + """No trusted proxies configured: X-Forwarded-For must not override the + connecting address, or a private-network caller could spoof its way + past allowlists, limits, and bans.""" + blocked: Final = "10.4.5.6" + monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") + monkeypatch.setenv("LITELLM_GUARD_BLOCKED_IPS", blocked) + with _client(blocked) as client: + rejected: Final = client.get("/ping", headers={"x-forwarded-for": "198.51.100.9"}) + assert rejected.status_code == 403 + + +@pytest.mark.skipif(not _GUARD_AVAILABLE, reason="fastapi-guard not installed") +def test_rejected_response_keeps_cors_and_security_headers(monkeypatch): + """Guard is nested inside the CORS and security-headers wrappers, so a + guard rejection carries their headers (readable cross-origin errors).""" + blocked: Final = _unique_ip() + monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") + monkeypatch.setenv("LITELLM_GUARD_BLOCKED_IPS", blocked) + with _client(blocked, _build_stacked_app()) as client: + rejected: Final = client.get("/ping", headers={"Origin": "https://ok.example"}) + assert rejected.status_code == 403 + # Both wrappers ran on the guard rejection. guard sets its own + # anti-framing header value and litellm's setdefault defers to it, so + # presence (not a specific value) is the property pinned here. + assert rejected.headers["x-frame-options"] in ("DENY", "SAMEORIGIN") + assert rejected.headers["content-security-policy"] == "frame-ancestors 'none'" + assert rejected.headers["access-control-allow-origin"] == "https://ok.example" + + +@pytest.mark.skipif(not _GUARD_AVAILABLE, reason="fastapi-guard not installed") +def test_rate_limit_returns_429(monkeypatch): + 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: 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: Final = _unique_ip() + monkeypatch.setenv("LITELLM_GUARD_ENABLED", "true") + monkeypatch.setenv("LITELLM_GUARD_PASSIVE_MODE", "true") + monkeypatch.setenv("LITELLM_GUARD_BLOCKED_IPS", blocked) + with _client(blocked) as client: + assert client.get("/ping").status_code == 200 + + +@pytest.mark.skipif(not _GUARD_AVAILABLE, reason="fastapi-guard not installed") +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: Final = _build_guard_config() + assert config.enable_redis is False + assert config.redis_url is None