chore: import upstream snapshot with attribution

python/sglang/srt/utils/auth.py

@@ -0,0 +1,208 @@
+"""Auth utilities for HTTP servers.
+
+This module is intentionally lightweight (no torch import) so it can be used in unit tests.
+"""
+
+from __future__ import annotations
+
+import secrets
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional
+
+
+@dataclass(frozen=True)
+class AuthDecision:
+    allowed: bool
+    error_status_code: int = 401  # Only meaningful when allowed=False
+
+
+class AuthLevel(str, Enum):
+    """Per-endpoint auth level (attached to endpoint function via `@auth_level`)."""
+
+    NORMAL = "normal"
+    ADMIN_OPTIONAL = "admin_optional"
+    ADMIN_FORCE = "admin_force"
+
+
+def auth_level(level: AuthLevel):
+    """Mark endpoint with auth level (stored in endpoint metadata)."""
+
+    def decorator(func):
+        func._auth_level = level
+        return func
+
+    return decorator
+
+
+def _get_auth_level_from_app_and_scope(app: Any, scope: dict) -> AuthLevel:
+    """Best-effort resolve auth level by matching the request to a route."""
+    # Import lazily to keep this module unit-test friendly (FastAPI/Starlette are not
+    # required unless you actually use the middleware / route matching).
+    from starlette.routing import Match
+
+    # Prefer app.router.routes when available; fall back to app.routes.
+    routes = getattr(getattr(app, "router", None), "routes", None) or getattr(
+        app, "routes", []
+    )
+
+    for route in routes:
+        try:
+            match, child_scope = route.matches(scope)
+        except Exception:
+            continue
+        if match == Match.FULL:
+            endpoint = child_scope.get("endpoint") or getattr(route, "endpoint", None)
+            level = getattr(endpoint, "_auth_level", None)
+            return level if isinstance(level, AuthLevel) else AuthLevel.NORMAL
+
+    return AuthLevel.NORMAL
+
+
+def app_has_admin_force_endpoints(app: Any) -> bool:
+    """Return True if any route endpoint is marked as ADMIN_FORCE."""
+    routes = getattr(getattr(app, "router", None), "routes", None) or getattr(
+        app, "routes", []
+    )
+    for route in routes:
+        endpoint = getattr(route, "endpoint", None)
+        if getattr(endpoint, "_auth_level", None) == AuthLevel.ADMIN_FORCE:
+            return True
+    return False
+
+
+def decide_request_auth(
+    *,
+    method: str,
+    path: str,
+    authorization_header: Optional[str],
+    api_key: Optional[str],
+    admin_api_key: Optional[str],
+    auth_level: AuthLevel,
+) -> AuthDecision:
+    """Pure auth decision function (easy to unit test).
+
+    Auth levels:
+    - NORMAL: legacy behavior (api_key protects all endpoints when configured)
+    - ADMIN_OPTIONAL: can be accessed without any key (if no keys configured),
+      or with api_key/admin_api_key depending on server config.
+    - ADMIN_FORCE: requires admin_api_key; if admin_api_key is NOT configured,
+      it must be rejected (403) even if api_key is provided.
+
+    NOTE :
+    - Health/metrics endpoints are always allowed (even when api_key/admin_api_key is set),
+      to support k8s/liveness/readiness and Prometheus scraping without embedding secrets.
+    - We match them by prefix to cover common variants like /health_generate.
+    """
+    if method == "OPTIONS":
+        return AuthDecision(allowed=True)
+
+    if path.startswith("/health") or path.startswith("/metrics"):
+        return AuthDecision(allowed=True)
+
+    def _check_bearer_token(
+        authorization_header: Optional[str], expected_token: str
+    ) -> bool:
+        """Check bearer token with constant-time comparison."""
+        if not authorization_header:
+            return False
+        parts = authorization_header.split(" ", 1)
+        if len(parts) != 2 or parts[0].lower() != "bearer":
+            return False
+        return secrets.compare_digest(parts[1], expected_token)
+
+    # Force-auth endpoints: only admin_api_key can unlock them; if admin_api_key is unset,
+    # reject them unconditionally (explicitly "not allowed").
+    if auth_level == AuthLevel.ADMIN_FORCE:
+        if not admin_api_key:
+            return AuthDecision(allowed=False, error_status_code=403)
+        if not _check_bearer_token(authorization_header, admin_api_key):
+            return AuthDecision(allowed=False)
+        return AuthDecision(allowed=True)
+
+    # Optional-auth endpoints:
+    # - no keys configured: allow
+    # - only api_key: require api_key
+    # - only admin_api_key: require admin_api_key
+    # - both: require admin_api_key (api_key is NOT accepted)
+    if auth_level == AuthLevel.ADMIN_OPTIONAL:
+        if admin_api_key:
+            return AuthDecision(
+                allowed=_check_bearer_token(authorization_header, admin_api_key)
+            )
+        elif api_key:
+            return AuthDecision(
+                allowed=_check_bearer_token(authorization_header, api_key)
+            )
+        else:
+            return AuthDecision(allowed=True)
+
+    # Normal endpoints:
+    # - if api_key is configured, require api_key (even if admin_api_key is also configured)
+    # - otherwise allow (including the "admin_api_key only" case)
+    if api_key:
+        return AuthDecision(allowed=_check_bearer_token(authorization_header, api_key))
+
+    return AuthDecision(allowed=True)
+
+
+def add_api_key_middleware(
+    app,
+    *,
+    api_key: Optional[str],
+    admin_api_key: Optional[str],
+):
+    """Add middleware for three endpoint auth levels: normal/admin_optional/admin_force."""
+    # Import lazily so `decide_request_auth()` can be unit-tested without FastAPI installed.
+    from fastapi.responses import ORJSONResponse
+    from starlette.requests import Request
+
+    class _ApiKeyASGIMiddleware:
+        """ASGI-native middleware to preserve client disconnect events."""
+
+        def __init__(self, app, *, api_key, admin_api_key, fastapi_app):
+            self.app = app
+            self.api_key = api_key
+            self.admin_api_key = admin_api_key
+            self.fastapi_app = fastapi_app
+
+        async def __call__(self, scope, receive, send):
+            if scope["type"] != "http":
+                await self.app(scope, receive, send)
+                return
+
+            request = Request(scope, receive=receive)
+            path = request.url.path
+            authz = request.headers.get("Authorization")
+            level = _get_auth_level_from_app_and_scope(self.fastapi_app, scope)
+            decision = decide_request_auth(
+                method=request.method,
+                path=path,
+                authorization_header=authz,
+                api_key=self.api_key,
+                admin_api_key=self.admin_api_key,
+                auth_level=level,
+            )
+
+            if not decision.allowed:
+                response = ORJSONResponse(
+                    content={
+                        "error": (
+                            "Unauthorized"
+                            if decision.error_status_code == 401
+                            else "Forbidden"
+                        )
+                    },
+                    status_code=decision.error_status_code,
+                )
+                await response(scope, receive, send)
+                return
+
+            await self.app(scope, receive, send)
+
+    app.add_middleware(
+        _ApiKeyASGIMiddleware,
+        api_key=api_key,
+        admin_api_key=admin_api_key,
+        fastapi_app=app,
+    )