From cd336945854669e69b2d62b4604d212ca508ee00 Mon Sep 17 00:00:00 2001 From: shuwenn <47200617+alphabetc1@users.noreply.github.com> Date: Wed, 14 Jan 2026 12:21:55 +0800 Subject: [PATCH] feat: add --admin-api-key for finer-grained endpoint auth (#15908) Co-authored-by: Simo Lin --- docs/advanced_features/server_arguments.md | 1 + python/sglang/srt/entrypoints/http_server.py | 51 ++- python/sglang/srt/server_args.py | 10 + python/sglang/srt/utils/auth.py | 185 +++++++++++ python/sglang/srt/utils/common.py | 15 - python/sglang/test/test_http_server_auth.py | 319 +++++++++++++++++++ 6 files changed, 563 insertions(+), 18 deletions(-) create mode 100644 python/sglang/srt/utils/auth.py create mode 100644 python/sglang/test/test_http_server_auth.py diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 38ad385cf..6f419ed62 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -190,6 +190,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | Argument | Description | Defaults | Options | | --- | --- | --- | --- | | `--api-key` | Set API key of the server. It is also used in the OpenAI API compatible server. | `None` | Type: str | +| `--admin-api-key` | Set **admin API key** for administrative/control endpoints (e.g., weights update, cache flush, `/get_server_info`). Endpoints marked as admin-only require `Authorization: Bearer ` when this is set. | `None` | Type: str | | `--served-model-name` | Override the model name returned by the v1/models endpoint in OpenAI API server. | `None` | Type: str | | `--weight-version` | Version identifier for the model weights. Defaults to 'default' if not specified. | `default` | Type: str | | `--chat-template` | The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server. | `None` | Type: str | diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 0d9e4e433..1d44d66bd 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -141,7 +141,6 @@ from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.tracing.trace import process_tracing_init, trace_set_thread_info from sglang.srt.utils import ( - add_api_key_middleware, add_prometheus_middleware, add_prometheus_track_response_middleware, delete_directory, @@ -149,6 +148,7 @@ from sglang.srt.utils import ( kill_process_tree, set_uvicorn_logging_configs, ) +from sglang.srt.utils.auth import AuthLevel, app_has_admin_force_endpoints, auth_level from sglang.utils import get_exception_traceback from sglang.version import __version__ @@ -607,6 +607,7 @@ async def get_load(): # example usage: # curl -s -X POST http://localhost:30000/set_internal_state -H "Content-Type: application/json" -d '{"server_args": {"pp_max_micro_batch_size": 8}}' @app.api_route("/set_internal_state", methods=["POST", "PUT"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def set_internal_state(obj: SetInternalStateReq, request: Request): res = await _global_state.tokenizer_manager.set_internal_state(obj) return res @@ -699,6 +700,7 @@ async def classify_request(obj: EmbeddingReqInput, request: Request): @app.api_route("/flush_cache", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def flush_cache(): """Flush the radix cache.""" ret = await _global_state.tokenizer_manager.flush_cache() @@ -710,6 +712,7 @@ async def flush_cache(): @app.api_route("/clear_hicache_storage_backend", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def clear_hicache_storage_backend(): """Clear the hierarchical cache storage backend.""" ret = await _global_state.tokenizer_manager.clear_hicache_storage() @@ -720,6 +723,7 @@ async def clear_hicache_storage_backend(): @app.api_route("/start_profile", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def start_profile_async(obj: Optional[ProfileReqInput] = None): """Start profiling.""" if obj is None: @@ -744,6 +748,7 @@ async def start_profile_async(obj: Optional[ProfileReqInput] = None): @app.api_route("/stop_profile", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def stop_profile_async(): """Stop profiling.""" await _global_state.tokenizer_manager.stop_profile() @@ -754,6 +759,7 @@ async def stop_profile_async(): @app.api_route("/freeze_gc", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def freeze_gc_async(): """ See engine.freeze_gc for more details. @@ -766,6 +772,7 @@ async def freeze_gc_async(): @app.api_route("/start_expert_distribution_record", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def start_expert_distribution_record_async(): """Start recording the expert distribution. Clear the previous record if any.""" await _global_state.tokenizer_manager.start_expert_distribution_record() @@ -776,6 +783,7 @@ async def start_expert_distribution_record_async(): @app.api_route("/stop_expert_distribution_record", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def stop_expert_distribution_record_async(): """Stop recording the expert distribution.""" await _global_state.tokenizer_manager.stop_expert_distribution_record() @@ -786,6 +794,7 @@ async def stop_expert_distribution_record_async(): @app.api_route("/dump_expert_distribution_record", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def dump_expert_distribution_record_async(): """Dump expert distribution record.""" await _global_state.tokenizer_manager.dump_expert_distribution_record() @@ -796,6 +805,7 @@ async def dump_expert_distribution_record_async(): @app.post("/update_weights_from_disk") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: Request): """Update the weights from disk inplace without re-launching the server.""" success, message, num_paused_requests = ( @@ -820,6 +830,7 @@ async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: R @app.post("/init_weights_send_group_for_remote_instance") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_send_group_for_remote_instance( obj: InitWeightsSendGroupForRemoteInstanceReqInput, request: Request ): @@ -836,6 +847,7 @@ async def init_weights_send_group_for_remote_instance( @app.post("/send_weights_to_remote_instance") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def send_weights_to_remote_instance( obj: SendWeightsToRemoteInstanceReqInput, request: Request ): @@ -852,6 +864,7 @@ async def send_weights_to_remote_instance( @app.get("/get_remote_instance_transfer_engine_info") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def get_remote_instance_transfer_engine_info(rank: int = None): if rank is None or rank < 0: return Response(status_code=HTTPStatus.BAD_REQUEST) @@ -876,6 +889,7 @@ async def get_remote_instance_transfer_engine_info(rank: int = None): @app.post("/init_weights_update_group") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_update_group( obj: InitWeightsUpdateGroupReqInput, request: Request ): @@ -891,6 +905,7 @@ async def init_weights_update_group( @app.post("/destroy_weights_update_group") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def destroy_weights_update_group( obj: DestroyWeightsUpdateGroupReqInput, request: Request ): @@ -905,6 +920,7 @@ async def destroy_weights_update_group( @app.post("/update_weights_from_tensor") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_tensor( obj: UpdateWeightsFromTensorReqInput, request: Request ): @@ -926,6 +942,7 @@ async def update_weights_from_tensor( @app.post("/update_weights_from_distributed") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_distributed( obj: UpdateWeightsFromDistributedReqInput, request: Request ): @@ -944,6 +961,7 @@ async def update_weights_from_distributed( @app.post("/update_weights_from_ipc") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Request): """Update the weights from IPC (Inter-Process Communication) for checkpoint-engine integration.""" success, message = await _global_state.tokenizer_manager.update_weights_from_ipc( @@ -960,6 +978,7 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re @app.post("/update_weight_version") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Request): """Update the weight version. This operation requires no active requests.""" if obj.abort_all_requests: @@ -990,6 +1009,7 @@ async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Reque @app.api_route("/get_weights_by_name", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request): """Get model parameter by name.""" try: @@ -1003,6 +1023,7 @@ async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request): @app.api_route("/release_memory_occupation", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def release_memory_occupation( obj: ReleaseMemoryOccupationReqInput, request: Request ): @@ -1014,6 +1035,7 @@ async def release_memory_occupation( @app.api_route("/resume_memory_occupation", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def resume_memory_occupation( obj: ResumeMemoryOccupationReqInput, request: Request ): @@ -1025,6 +1047,7 @@ async def resume_memory_occupation( @app.post("/weights_checker") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def check_weights(obj: CheckWeightsReqInput, request: Request): success, message = await _global_state.tokenizer_manager.check_weights(obj, request) return ORJSONResponse( @@ -1034,6 +1057,7 @@ async def check_weights(obj: CheckWeightsReqInput, request: Request): @app.api_route("/slow_down", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def slow_down(obj: SlowDownReqInput, request: Request): """Slow down the system deliberately. Only for testing. Example scenario: when we want to test performance of D in large-scale PD disaggregation and have no enough nodes for P, @@ -1047,6 +1071,7 @@ async def slow_down(obj: SlowDownReqInput, request: Request): @app.api_route("/load_lora_adapter", methods=["POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def load_lora_adapter(obj: LoadLoRAAdapterReqInput, request: Request): """Load a new LoRA adapter without re-launching the server.""" result = await _global_state.tokenizer_manager.load_lora_adapter(obj, request) @@ -1079,6 +1104,7 @@ async def load_lora_adapter_from_tensors( @app.api_route("/unload_lora_adapter", methods=["POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def unload_lora_adapter(obj: UnloadLoRAAdapterReqInput, request: Request): """Load a new LoRA adapter without re-launching the server.""" result = await _global_state.tokenizer_manager.unload_lora_adapter(obj, request) @@ -1120,6 +1146,7 @@ async def close_session(obj: CloseSessionReqInput, request: Request): @app.api_route("/configure_logging", methods=["GET", "POST"]) +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def configure_logging(obj: ConfigureLoggingReq, request: Request): """Configure the request logging options.""" _global_state.tokenizer_manager.configure_logging(obj) @@ -1127,6 +1154,7 @@ async def configure_logging(obj: ConfigureLoggingReq, request: Request): @app.post("/abort_request") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def abort_request(obj: AbortReq, request: Request): """Abort a request.""" try: @@ -1181,6 +1209,7 @@ async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Re @app.post("/pause_generation") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def pause_generation(obj: PauseGenerationReqInput, request: Request): """Pause generation.""" await _global_state.tokenizer_manager.pause_generation(obj) @@ -1191,6 +1220,7 @@ async def pause_generation(obj: PauseGenerationReqInput, request: Request): @app.post("/continue_generation") +@auth_level(AuthLevel.ADMIN_OPTIONAL) async def continue_generation(obj: ContinueGenerationReqInput, request: Request): """Continue generation.""" await _global_state.tokenizer_manager.continue_generation(obj) @@ -1750,8 +1780,23 @@ def launch_server( # Add api key authorization # This is only supported in single tokenizer mode. - if server_args.api_key: - add_api_key_middleware(app, server_args.api_key) + # + # Backward compatibility: + # - api_key only: behavior matches legacy (all endpoints require api_key) + # - no keys: legacy had no restriction; ADMIN_FORCE endpoints must still be rejected when + # admin_api_key is not configured. + if ( + server_args.api_key + or server_args.admin_api_key + or app_has_admin_force_endpoints(app) + ): + from sglang.srt.utils.auth import add_api_key_middleware + + add_api_key_middleware( + app, + api_key=server_args.api_key, + admin_api_key=server_args.admin_api_key, + ) else: # If it is multi-tokenizer mode, we need to write the arguments to shared memory # for other worker processes to read. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index fffed38fa..e30e53ee0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -370,6 +370,7 @@ class ServerArgs: # API related api_key: Optional[str] = None + admin_api_key: Optional[str] = None served_model_name: Optional[str] = None weight_version: str = "default" chat_template: Optional[str] = None @@ -3250,6 +3251,15 @@ class ServerArgs: default=ServerArgs.api_key, help="Set API key of the server. It is also used in the OpenAI API compatible server.", ) + parser.add_argument( + "--admin-api-key", + type=str, + default=ServerArgs.admin_api_key, + help=( + "Set admin API key for sensitive management endpoints (e.g. /clear_hicache_storage_backend). " + "When set, admin endpoints require this key and do NOT accept --api-key." + ), + ) parser.add_argument( "--served-model-name", type=str, diff --git a/python/sglang/srt/utils/auth.py b/python/sglang/srt/utils/auth.py new file mode 100644 index 000000000..0536d5a35 --- /dev/null +++ b/python/sglang/srt/utils/auth.py @@ -0,0 +1,185 @@ +"""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 + + @app.middleware("http") + async def authentication(request, call_next): + path = request.url.path + authz = request.headers.get("Authorization") + level = _get_auth_level_from_app_and_scope(request.app, request.scope) + decision = decide_request_auth( + method=request.method, + path=path, + authorization_header=authz, + api_key=api_key, + admin_api_key=admin_api_key, + auth_level=level, + ) + + if not decision.allowed: + return ORJSONResponse( + content={ + "error": ( + "Unauthorized" + if decision.error_status_code == 401 + else "Forbidden" + ) + }, + status_code=decision.error_status_code, + ) + + return await call_next(request) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index f95bcf126..6b2a4703c 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -82,7 +82,6 @@ import torch.distributed import torch.distributed as dist import triton import zmq -from fastapi.responses import ORJSONResponse from packaging import version as pkg_version from PIL import Image from starlette.routing import Mount @@ -1131,20 +1130,6 @@ def rank0_log(msg: str): logger.info(msg) -def add_api_key_middleware(app, api_key: str): - @app.middleware("http") - async def authentication(request, call_next): - if request.method == "OPTIONS": - return await call_next(request) - if request.url.path.startswith("/health") or request.url.path.startswith( - "/metrics" - ): - return await call_next(request) - if request.headers.get("Authorization") != "Bearer " + api_key: - return ORJSONResponse(content={"error": "Unauthorized"}, status_code=401) - return await call_next(request) - - def configure_logger(server_args, prefix: str = ""): if SGLANG_LOGGING_CONFIG_PATH := os.getenv("SGLANG_LOGGING_CONFIG_PATH"): if not os.path.exists(SGLANG_LOGGING_CONFIG_PATH): diff --git a/python/sglang/test/test_http_server_auth.py b/python/sglang/test/test_http_server_auth.py new file mode 100644 index 000000000..37e04eecc --- /dev/null +++ b/python/sglang/test/test_http_server_auth.py @@ -0,0 +1,319 @@ +""" +Unit tests for HTTP server admin auth. + +Usage: + python3 -m pytest test/test_http_server_auth.py -v +""" + +import importlib.util +import os +import sys +import unittest + + +def _load_auth_module(): + """Load auth.py directly, avoiding importing the full sglang package. + + This keeps the test importable even if optional runtime deps (e.g. orjson/httpx) + are not installed in the unit test environment. + """ + this_dir = os.path.dirname(__file__) + python_dir = os.path.abspath(os.path.join(this_dir, "..", "..")) + auth_path = os.path.join(python_dir, "sglang", "srt", "utils", "auth.py") + + module_name = "_sglang_srt_utils_auth_for_test" + spec = importlib.util.spec_from_file_location(module_name, auth_path) + assert spec is not None and spec.loader is not None + m = importlib.util.module_from_spec(spec) + # dataclasses (py3.12) may consult sys.modules during class processing + sys.modules[module_name] = m + spec.loader.exec_module(m) + return m + + +_auth = _load_auth_module() +decide_request_auth = _auth.decide_request_auth +AuthLevel = _auth.AuthLevel + + +class TestHttpServerAdminAuth(unittest.TestCase): + def _decide( + self, + *, + method: str, + path: str, + authorization_header: str | None, + api_key: str | None, + admin_api_key: str | None, + auth_level: AuthLevel, + ): + return decide_request_auth( + method=method, + path=path, + authorization_header=authorization_header, + api_key=api_key, + admin_api_key=admin_api_key, + auth_level=auth_level, + ) + + def test_no_keys_configured(self): + # No keys configured -> NORMAL + ADMIN_OPTIONAL are open (legacy), + # but ADMIN_FORCE must be rejected (403) explicitly. + self.assertTrue( + self._decide( + method="GET", + path="/v1/models", + authorization_header=None, + api_key=None, + admin_api_key=None, + auth_level=AuthLevel.NORMAL, + ).allowed + ) + self.assertTrue( + self._decide( + method="POST", + path="/admin_optional_demo", + authorization_header=None, + api_key=None, + admin_api_key=None, + auth_level=AuthLevel.ADMIN_OPTIONAL, + ).allowed + ) + + d = self._decide( + method="POST", + path="/admin_force_demo", + authorization_header=None, + api_key=None, + admin_api_key=None, + auth_level=AuthLevel.ADMIN_FORCE, + ) + self.assertFalse(d.allowed) + self.assertEqual(d.error_status_code, 403) + + def test_api_key_only(self): + # api_key configured -> NORMAL requires api_key (legacy). + self.assertFalse( + self._decide( + method="GET", + path="/v1/models", + authorization_header=None, + api_key="user", + admin_api_key=None, + auth_level=AuthLevel.NORMAL, + ).allowed + ) + self.assertTrue( + self._decide( + method="GET", + path="/v1/models", + authorization_header="Bearer user", + api_key="user", + admin_api_key=None, + auth_level=AuthLevel.NORMAL, + ).allowed + ) + + # ADMIN_OPTIONAL requires api_key when only api_key is configured. + self.assertFalse( + self._decide( + method="POST", + path="/admin_optional_demo", + authorization_header="Bearer wrong", + api_key="user", + admin_api_key=None, + auth_level=AuthLevel.ADMIN_OPTIONAL, + ).allowed + ) + self.assertTrue( + self._decide( + method="POST", + path="/admin_optional_demo", + authorization_header="Bearer user", + api_key="user", + admin_api_key=None, + auth_level=AuthLevel.ADMIN_OPTIONAL, + ).allowed + ) + + # ADMIN_FORCE must be rejected even if api_key is configured (403). + d = self._decide( + method="POST", + path="/admin_force_demo", + authorization_header="Bearer user", + api_key="user", + admin_api_key=None, + auth_level=AuthLevel.ADMIN_FORCE, + ) + self.assertFalse(d.allowed) + self.assertEqual(d.error_status_code, 403) + + def test_admin_api_key_only(self): + # admin_api_key only: + # - normal endpoints open + # - optional/force endpoints require admin_api_key + self.assertTrue( + self._decide( + method="GET", + path="/v1/models", + authorization_header="Bearer user", + api_key=None, + admin_api_key="admin", + auth_level=AuthLevel.NORMAL, + ).allowed + ) + self.assertTrue( + self._decide( + method="GET", + path="/v1/models", + authorization_header=None, + api_key=None, + admin_api_key="admin", + auth_level=AuthLevel.NORMAL, + ).allowed + ) + + # Optional endpoints require admin_api_key when admin_api_key is configured. + self.assertTrue( + self._decide( + method="POST", + path="/admin_optional_demo", + authorization_header="Bearer admin", + api_key=None, + admin_api_key="admin", + auth_level=AuthLevel.ADMIN_OPTIONAL, + ).allowed + ) + self.assertFalse( + self._decide( + method="POST", + path="/admin_optional_demo", + authorization_header="Bearer user", + api_key=None, + admin_api_key="admin", + auth_level=AuthLevel.ADMIN_OPTIONAL, + ).allowed + ) + + d = self._decide( + method="POST", + path="/admin_force_demo", + authorization_header="Bearer admin", + api_key=None, + admin_api_key="admin", + auth_level=AuthLevel.ADMIN_FORCE, + ) + self.assertTrue(d.allowed) + + def test_with_both_api_keys(self): + # both api_key and admin_api_key configured: + # - normal endpoints require api_key + # - optional endpoints require admin_api_key (api_key is NOT accepted) + # - force endpoints require admin_api_key + self.assertTrue( + self._decide( + method="GET", + path="/v1/models", + authorization_header="Bearer user", + api_key="user", + admin_api_key="admin", + auth_level=AuthLevel.NORMAL, + ).allowed + ) + self.assertFalse( + self._decide( + method="GET", + path="/v1/models", + authorization_header="Bearer admin", + api_key="user", + admin_api_key="admin", + auth_level=AuthLevel.NORMAL, + ).allowed + ) + # Optional endpoints must require admin_api_key when both keys are configured. + self.assertFalse( + self._decide( + method="POST", + path="/admin_optional_demo", + authorization_header="Bearer user", + api_key="user", + admin_api_key="admin", + auth_level=AuthLevel.ADMIN_OPTIONAL, + ).allowed + ) + self.assertTrue( + self._decide( + method="POST", + path="/admin_optional_demo", + authorization_header="Bearer admin", + api_key="user", + admin_api_key="admin", + auth_level=AuthLevel.ADMIN_OPTIONAL, + ).allowed + ) + self.assertFalse( + self._decide( + method="POST", + path="/admin_force_demo", + authorization_header="Bearer user", + api_key="user", + admin_api_key="admin", + auth_level=AuthLevel.ADMIN_FORCE, + ).allowed + ) + self.assertTrue( + self._decide( + method="POST", + path="/admin_force_demo", + authorization_header="Bearer admin", + api_key="user", + admin_api_key="admin", + auth_level=AuthLevel.ADMIN_FORCE, + ).allowed + ) + + def test_options_is_always_allowed(self): + # CORS preflight should never be blocked. + self.assertTrue( + self._decide( + method="OPTIONS", + path="/v1/models", + authorization_header=None, + api_key="user", + admin_api_key="admin", + auth_level=AuthLevel.ADMIN_FORCE, + ).allowed + ) + + def test_health_and_metrics_are_always_allowed(self): + # Health/metrics endpoints are always public by design, regardless of auth level / keys. + combos = [ + dict(api_key=None, admin_api_key=None), + dict(api_key="user", admin_api_key=None), + dict(api_key=None, admin_api_key="admin"), + dict(api_key="user", admin_api_key="admin"), + ] + paths_allowed = [ + "/health", + "/health_generate", + "/metrics", + "/metrics/", + "/metrics/prometheus", + ] + for keys in combos: + for path in paths_allowed: + self.assertTrue( + self._decide( + method="GET", + path=path, + authorization_header=None, + api_key=keys["api_key"], + admin_api_key=keys["admin_api_key"], + auth_level=AuthLevel.ADMIN_FORCE, + ).allowed, + msg=f"expected allowed for {path=} with {keys=}", + ) + + +if __name__ == "__main__": + unittest.main()