Merge cc605f89ba into c2fc628d36

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)

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,

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