diff --git a/.github/workflows/test-unit-proxy-infra.yml b/.github/workflows/test-unit-proxy-infra.yml index 2436c0498e1..2681400e8b7 100644 --- a/.github/workflows/test-unit-proxy-infra.yml +++ b/.github/workflows/test-unit-proxy-infra.yml @@ -26,7 +26,6 @@ jobs: tests/test_litellm/proxy/middleware tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/pass_through_endpoints - tests/test_litellm/proxy/rust_control_plane tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/experimental tests/test_litellm/proxy/common_utils diff --git a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py index a96a5431294..3c8cfb8f89b 100644 --- a/litellm/proxy/logging_endpoints/callback_logs_endpoints.py +++ b/litellm/proxy/logging_endpoints/callback_logs_endpoints.py @@ -14,15 +14,19 @@ The endpoint is generic: realtime is the first producer, but the contract is the self-describing `StandardLoggingPayload`, so completions/responses can use it too. """ +import hmac +import json +import os import uuid from datetime import datetime, timezone -from typing import Any +from typing import Any, Optional -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Request +from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.proxy.callback_logs_endpoints import ( CallbackLogFailure, @@ -32,12 +36,118 @@ from litellm.types.proxy.callback_logs_endpoints import ( ) # Routes the Python proxy exposes for the Rust data-plane gateway to call into -# (logging today; auth/budgets later). Namespaced under /v1/rust_control_plane so -# they're clearly distinct from the proxy's own control-plane/management routes. +# (auth + logging today; budgets later). Namespaced under /v1/rust_control_plane +# so they're clearly distinct from the proxy's own control-plane/management routes. rust_control_plane_router = APIRouter( prefix="/v1/rust_control_plane", tags=["rust control plane"] ) +DATA_PLANE_KEY_ENV_VAR = "LITELLM_DATA_PLANE_KEY" +DATA_PLANE_KEY_HEADER = "X-LiteLLM-Data-Plane-Key" + + +def require_data_plane_key(request: Request) -> None: + """ + Authenticate requests from the Rust data plane with a dedicated secret. + + This intentionally uses ``LITELLM_DATA_PLANE_KEY`` instead of the proxy + master key: the data plane is a separate trust boundary and must not get + admin privileges. + """ + expected_key: Optional[str] = os.getenv(DATA_PLANE_KEY_ENV_VAR) + if not expected_key: + raise HTTPException(status_code=500, detail="data-plane auth not configured") + + provided_key: Optional[str] = request.headers.get(DATA_PLANE_KEY_HEADER) + if not provided_key or not hmac.compare_digest(provided_key, expected_key): + raise HTTPException(status_code=401, detail="invalid data-plane key") + + +class VerifyKeyRequest(BaseModel): + api_key: str + # The actual route the gateway is serving this key on (e.g. "/v1/realtime"). + # REQUIRED and not defaulted: the gateway always sends its own request path, + # so validation runs route/model restrictions against the real route. + route: str + # Forwarded so user_api_key_auth's model access checks run for key, team, + # and access-group restrictions. + model: Optional[str] = None + + +def _synthetic_request( + route: str, authorization_header: str, model: Optional[str] +) -> Request: + """ + Build a minimal ASGI request for user_api_key_auth to validate the key + against the data-plane route and model instead of this internal endpoint. + """ + body = json.dumps({"model": model} if model is not None else {}).encode() + + async def receive() -> dict[str, Any]: + return {"type": "http.request", "body": body, "more_body": False} + + scope = { + "type": "http", + "method": "POST", + "path": route, + "raw_path": route.encode(), + "headers": [ + (b"authorization", authorization_header.encode()), + (b"content-type", b"application/json"), + ], + "query_string": b"", + "scheme": "http", + "client": ("127.0.0.1", 0), + "server": ("127.0.0.1", 4000), + } + request = Request(scope, receive) + # Admission checks only. Realtime spend is reconciled later through callback + # logs, so optimistic reservation here would have no request lifecycle to + # release it. + request.state.skip_budget_reservation = True + return request + + +@rust_control_plane_router.post( + "/authentication", + dependencies=[Depends(require_data_plane_key)], + include_in_schema=False, +) +async def verify_key(body: VerifyKeyRequest) -> dict[str, Any]: + """ + Verify a virtual key on behalf of the Rust ai-gateway (data plane). + + The Rust gateway forwards the client's virtual key here; this delegates to + the proxy's existing ``user_api_key_auth`` validation and returns the + resolved ``UserAPIKeyAuth`` as JSON so the data plane can enforce the same + limits the control plane would. + + Gated by ``require_data_plane_key`` (the dedicated data-plane secret, not + the master key). + + On any auth failure the response is a 401 with a minimal body so internals + are not leaked to the caller. + """ + bearer_key = ( + body.api_key if body.api_key.startswith("Bearer ") else f"Bearer {body.api_key}" + ) + synthetic_request = _synthetic_request( + route=body.route, authorization_header=bearer_key, model=body.model + ) + try: + auth = await user_api_key_auth(request=synthetic_request, api_key=bearer_key) + except (ProxyException, HTTPException) as exc: + status_code = getattr(exc, "status_code", None) or getattr(exc, "code", None) + try: + is_server_error = status_code is not None and int(status_code) >= 500 + except (TypeError, ValueError): + is_server_error = False + if is_server_error: + raise + raise HTTPException(status_code=401, detail="invalid api key") + + return auth.model_dump(exclude_none=True, mode="json") + class CallbackLogsReplayer: """ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b134414f65e..50db86c59ef 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -349,9 +349,6 @@ from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( rust_control_plane_router, ) -from litellm.proxy.rust_control_plane.auth_endpoints import ( - router as rust_control_plane_auth_router, -) from litellm.proxy.management_endpoints.budget_management_endpoints import ( router as budget_management_router, ) @@ -16644,7 +16641,6 @@ app.include_router(caching_router) app.include_router(analytics_router) app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) -app.include_router(rust_control_plane_auth_router) app.include_router(rust_control_plane_router) app.include_router(ui_crud_endpoints_router) app.include_router(openai_files_router) diff --git a/litellm/proxy/rust_control_plane/__init__.py b/litellm/proxy/rust_control_plane/__init__.py deleted file mode 100644 index 10f215aebc1..00000000000 --- a/litellm/proxy/rust_control_plane/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Internal routes exposed by the Python control plane for Rust data planes.""" diff --git a/litellm/proxy/rust_control_plane/auth_endpoints.py b/litellm/proxy/rust_control_plane/auth_endpoints.py deleted file mode 100644 index b8aad7e7ecb..00000000000 --- a/litellm/proxy/rust_control_plane/auth_endpoints.py +++ /dev/null @@ -1,162 +0,0 @@ -""" -Rust control-plane auth endpoints. - -This module exposes the authentication seam used by the Rust ai-gateway -(the data plane). The Rust gateway terminates client connections and needs to -validate the virtual keys it receives WITHOUT reimplementing LiteLLM's key -validation logic. Instead, it calls back into this Python control plane to -verify a key and get the resolved ``UserAPIKeyAuth`` object. - -Security model: - - These endpoints are gated by a DEDICATED data-plane secret - (``LITELLM_DATA_PLANE_KEY``), NOT the proxy master key. The data plane is a - distinct trust boundary from proxy admins, so it gets its own credential - that can be rotated independently and never grants admin access. - - The data-plane key is compared in constant time to avoid leaking the secret - via timing side-channels. -""" - -import hmac -import json -import os -from typing import Any, Optional - -from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import BaseModel - -DATA_PLANE_KEY_ENV_VAR = "LITELLM_DATA_PLANE_KEY" -DATA_PLANE_KEY_HEADER = "X-LiteLLM-Data-Plane-Key" - - -def require_data_plane_key(request: Request) -> None: - """ - FastAPI dependency that authenticates a request from the Rust data plane. - - Reads the ``X-LiteLLM-Data-Plane-Key`` header and compares it, in constant - time, against the ``LITELLM_DATA_PLANE_KEY`` environment variable. - - This is intentionally a SEPARATE secret from ``LITELLM_MASTER_KEY`` — the - data plane is its own trust boundary and must not be granted master-key - privileges. - - Raises: - HTTPException 500: if ``LITELLM_DATA_PLANE_KEY`` is unset/empty - (misconfiguration — fail closed rather than allow unauthenticated - access). - HTTPException 401: if the header is missing or does not match. - """ - expected_key: Optional[str] = os.getenv(DATA_PLANE_KEY_ENV_VAR) - if not expected_key: - raise HTTPException(status_code=500, detail="data-plane auth not configured") - - provided_key: Optional[str] = request.headers.get(DATA_PLANE_KEY_HEADER) - if not provided_key or not hmac.compare_digest(provided_key, expected_key): - raise HTTPException(status_code=401, detail="invalid data-plane key") - - -class VerifyKeyRequest(BaseModel): - api_key: str - # The actual route the gateway is serving this key on (e.g. "/v1/realtime"). - # REQUIRED and not defaulted: the gateway always sends its own request path, - # so validation runs route/model restrictions against the real route — never a - # hardcoded guess, and never this internal endpoint's path (which a normal - # virtual key isn't allowed to call). - route: str - # The model being requested, if any. Forwarded so user_api_key_auth's model - # access checks (key + team + access-group, via can_key_call_model) run — this - # is what stops a valid key from reaching a model it isn't allowed to call. - model: Optional[str] = None - - -def _synthetic_request( - route: str, authorization_header: str, model: Optional[str] -) -> Request: - """ - Build a minimal ASGI request standing in for the client's real call, so - ``user_api_key_auth`` evaluates the key against the intended data-plane - ``route`` and ``model`` instead of this internal endpoint's path. The model - goes in the JSON body, where the proxy reads it for model-access enforcement. - """ - body = json.dumps({"model": model} if model is not None else {}).encode() - - async def receive() -> dict[str, Any]: - return {"type": "http.request", "body": body, "more_body": False} - - scope = { - "type": "http", - "method": "POST", - "path": route, - "raw_path": route.encode(), - "headers": [ - (b"authorization", authorization_header.encode()), - (b"content-type", b"application/json"), - ], - "query_string": b"", - "scheme": "http", - "client": ("127.0.0.1", 0), - "server": ("127.0.0.1", 4000), - } - request = Request(scope, receive) - # This endpoint verifies admission only. The realtime gateway records actual - # session spend through callback logs, so a pre-call optimistic reservation - # here would have no matching request lifecycle to reconcile. - request.state.skip_budget_reservation = True - return request - - -router = APIRouter(prefix="/v1/rust_control_plane", tags=["rust control plane"]) - - -@router.post( - "/authentication", - dependencies=[Depends(require_data_plane_key)], - # Internal data-plane route: keep it out of the public OpenAPI spec / docs - # (and the generated UI schema.d.ts). It's not a client- or UI-facing API. - include_in_schema=False, -) -async def verify_key(body: VerifyKeyRequest) -> dict[str, Any]: - """ - Verify a virtual key on behalf of the Rust ai-gateway (data plane). - - The Rust gateway forwards the client's virtual key here; this delegates to - the proxy's existing ``user_api_key_auth`` validation and returns the - resolved ``UserAPIKeyAuth`` as JSON so the data plane can enforce the same - limits the control plane would. - - Gated by ``require_data_plane_key`` (the dedicated data-plane secret, not - the master key). - - On any auth failure the response is a 401 with a minimal body so internals - are not leaked to the caller. - """ - from litellm.proxy._types import ProxyException - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - - # user_api_key_auth expects the value exactly as the Authorization header - # arrives — i.e. WITH the "Bearer " prefix (it strips it itself). The gateway - # sends the bare key, so normalize: add the prefix unless already present. - bearer_key = ( - body.api_key if body.api_key.startswith("Bearer ") else f"Bearer {body.api_key}" - ) - synthetic_request = _synthetic_request( - route=body.route, authorization_header=bearer_key, model=body.model - ) - try: - auth = await user_api_key_auth(request=synthetic_request, api_key=bearer_key) - except (ProxyException, HTTPException) as exc: - # Expected auth failures (invalid / expired / over-budget / blocked) → 401 - # with a minimal body so internals aren't leaked. A 5xx (e.g. a DB outage - # surfaced as a 500, on either HTTPException.status_code or - # ProxyException.code) is NOT masked as "invalid key" — it propagates so - # operators see the real error. The data plane still fails closed: it - # rejects any non-200 from this endpoint. - status_code = getattr(exc, "status_code", None) or getattr(exc, "code", None) - try: - is_server_error = status_code is not None and int(status_code) >= 500 - except (TypeError, ValueError): - is_server_error = False - if is_server_error: - raise - raise HTTPException(status_code=401, detail="invalid api key") - - return auth.model_dump(exclude_none=True, mode="json") diff --git a/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py b/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py index 40e89329b8d..c65b1e220e3 100644 --- a/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py +++ b/tests/test_litellm/proxy/logging_endpoints/test_callback_logs_endpoints.py @@ -3,13 +3,20 @@ import time import pytest -from fastapi import HTTPException +from fastapi import HTTPException, Request from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( CallbackLogsReplayer, + DATA_PLANE_KEY_ENV_VAR, + DATA_PLANE_KEY_HEADER, + VerifyKeyRequest, + _synthetic_request, ingest_callback_logs, + require_data_plane_key, + rust_control_plane_router, + verify_key, ) from litellm.types.proxy.callback_logs_endpoints import ( CallbackLogRecord, @@ -19,6 +26,20 @@ from litellm.types.proxy.callback_logs_endpoints import ( REQ_ID = "cb-logs-unit-test-1" +def _make_request(headers: dict) -> Request: + """Build a minimal ASGI Request with the given headers.""" + raw_headers = [ + (k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in headers.items() + ] + scope = { + "type": "http", + "method": "POST", + "path": "/v1/rust_control_plane/authentication", + "headers": raw_headers, + } + return Request(scope) + + def _sample_payload(**overrides): payload = { "id": REQ_ID, @@ -44,6 +65,225 @@ def _sample_payload(**overrides): return payload +def test_require_data_plane_key_500_when_env_unset(monkeypatch): + monkeypatch.delenv(DATA_PLANE_KEY_ENV_VAR, raising=False) + request = _make_request({DATA_PLANE_KEY_HEADER: "anything"}) + with pytest.raises(HTTPException) as exc_info: + require_data_plane_key(request) + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "data-plane auth not configured" + + +def test_require_data_plane_key_500_when_env_empty(monkeypatch): + monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "") + request = _make_request({DATA_PLANE_KEY_HEADER: "anything"}) + with pytest.raises(HTTPException) as exc_info: + require_data_plane_key(request) + assert exc_info.value.status_code == 500 + + +def test_require_data_plane_key_401_when_header_missing(monkeypatch): + monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "secret-dp-key") + request = _make_request({}) + with pytest.raises(HTTPException) as exc_info: + require_data_plane_key(request) + assert exc_info.value.status_code == 401 + + +def test_require_data_plane_key_401_when_header_wrong(monkeypatch): + monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "secret-dp-key") + request = _make_request({DATA_PLANE_KEY_HEADER: "wrong-key"}) + with pytest.raises(HTTPException) as exc_info: + require_data_plane_key(request) + assert exc_info.value.status_code == 401 + + +def test_require_data_plane_key_does_not_accept_master_key(monkeypatch): + """The data-plane key must be a dedicated secret, not the master key.""" + monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "secret-dp-key") + monkeypatch.setenv("LITELLM_MASTER_KEY", "sk-master-1234") + request = _make_request({DATA_PLANE_KEY_HEADER: "sk-master-1234"}) + with pytest.raises(HTTPException) as exc_info: + require_data_plane_key(request) + assert exc_info.value.status_code == 401 + + +def test_require_data_plane_key_passes_when_correct(monkeypatch): + monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "secret-dp-key") + request = _make_request({DATA_PLANE_KEY_HEADER: "secret-dp-key"}) + # Should not raise. + assert require_data_plane_key(request) is None + + +def test_router_mounts_auth_verify_under_rust_control_plane(): + assert any( + getattr(route, "path", None) == "/v1/rust_control_plane/authentication" + for route in rust_control_plane_router.routes + ) + + +@pytest.mark.asyncio +async def test_synthetic_request_skips_budget_reservation(): + request = _synthetic_request( + route="/v1/realtime", + authorization_header="Bearer sk-test-key", + model="gpt-realtime", + ) + + assert request.url.path == "/v1/realtime" + assert request.state.skip_budget_reservation is True + assert (await request.json()) == {"model": "gpt-realtime"} + + +@pytest.mark.asyncio +async def test_verify_key_returns_model_dump(monkeypatch): + expected_auth = UserAPIKeyAuth( + api_key="hashed-key", user_id="user-123", max_budget=100.0 + ) + + captured = {} + + async def fake_user_api_key_auth(request, api_key): + captured["api_key"] = api_key + captured["request"] = request + return expected_auth + + monkeypatch.setattr( + "litellm.proxy.logging_endpoints.callback_logs_endpoints.user_api_key_auth", + fake_user_api_key_auth, + ) + + body = VerifyKeyRequest( + api_key="sk-test-key", route="/v1/realtime", model="gpt-realtime" + ) + result = await verify_key(body=body) + + # The key is forwarded WITH the Bearer prefix (user_api_key_auth strips it). + assert captured["api_key"] == "Bearer sk-test-key" + # Validation runs against a synthetic request carrying the gateway's route... + assert captured["request"].url.path == "/v1/realtime" + assert captured["request"].headers["authorization"] == "Bearer sk-test-key" + # ...and the requested model in the body, so model-access checks enforce it. + assert (await captured["request"].json())["model"] == "gpt-realtime" + assert result == expected_auth.model_dump(exclude_none=True, mode="json") + assert result["user_id"] == "user-123" + + +@pytest.mark.asyncio +async def test_verify_key_omits_model_when_absent(monkeypatch): + captured = {} + + async def fake_user_api_key_auth(request, api_key): + captured["request"] = request + return UserAPIKeyAuth(api_key="hashed-key") + + monkeypatch.setattr( + "litellm.proxy.logging_endpoints.callback_logs_endpoints.user_api_key_auth", + fake_user_api_key_auth, + ) + + body = VerifyKeyRequest(api_key="sk-test-key", route="/v1/realtime") + await verify_key(body=body) + # No model requested -> empty body, not {"model": null}. + assert (await captured["request"].json()) == {} + + +@pytest.mark.asyncio +async def test_verify_key_does_not_double_prefix_existing_bearer(monkeypatch): + captured = {} + + async def fake_user_api_key_auth(request, api_key): + captured["api_key"] = api_key + captured["request"] = request + return UserAPIKeyAuth(api_key="hashed-key") + + monkeypatch.setattr( + "litellm.proxy.logging_endpoints.callback_logs_endpoints.user_api_key_auth", + fake_user_api_key_auth, + ) + + body = VerifyKeyRequest(api_key="Bearer sk-test-key", route="/v1/realtime") + await verify_key(body=body) + + assert captured["api_key"] == "Bearer sk-test-key" + assert captured["request"].headers["authorization"] == "Bearer sk-test-key" + + +@pytest.mark.asyncio +async def test_verify_key_401_on_proxy_exception(monkeypatch): + async def fake_user_api_key_auth(request, api_key): + raise ProxyException( + message="bad key", + type="auth_error", + param=None, + code="401", + ) + + monkeypatch.setattr( + "litellm.proxy.logging_endpoints.callback_logs_endpoints.user_api_key_auth", + fake_user_api_key_auth, + ) + + body = VerifyKeyRequest(api_key="sk-bad-key", route="/v1/realtime") + with pytest.raises(HTTPException) as exc_info: + await verify_key(body=body) + assert exc_info.value.status_code == 401 + assert exc_info.value.detail == "invalid api key" + + +@pytest.mark.asyncio +async def test_verify_key_401_on_http_exception(monkeypatch): + async def fake_user_api_key_auth(request, api_key): + raise HTTPException(status_code=403, detail="forbidden internals") + + monkeypatch.setattr( + "litellm.proxy.logging_endpoints.callback_logs_endpoints.user_api_key_auth", + fake_user_api_key_auth, + ) + + body = VerifyKeyRequest(api_key="sk-bad-key", route="/v1/realtime") + with pytest.raises(HTTPException) as exc_info: + await verify_key(body=body) + assert exc_info.value.status_code == 401 + # Internals must not leak. + assert exc_info.value.detail == "invalid api key" + + +@pytest.mark.asyncio +async def test_verify_key_propagates_http_5xx(monkeypatch): + # A 5xx (e.g. DB outage) must NOT be masked as 401. + async def fake_user_api_key_auth(request, api_key): + raise HTTPException(status_code=503, detail="db unavailable") + + monkeypatch.setattr( + "litellm.proxy.logging_endpoints.callback_logs_endpoints.user_api_key_auth", + fake_user_api_key_auth, + ) + + body = VerifyKeyRequest(api_key="sk-key", route="/v1/realtime") + with pytest.raises(HTTPException) as exc_info: + await verify_key(body=body) + assert exc_info.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_verify_key_propagates_proxy_5xx(monkeypatch): + # A ProxyException carrying a 5xx code propagates too (not converted to 401). + async def fake_user_api_key_auth(request, api_key): + raise ProxyException( + message="internal", type="internal_error", param=None, code="500" + ) + + monkeypatch.setattr( + "litellm.proxy.logging_endpoints.callback_logs_endpoints.user_api_key_auth", + fake_user_api_key_auth, + ) + + body = VerifyKeyRequest(api_key="sk-key", route="/v1/realtime") + with pytest.raises(ProxyException): + await verify_key(body=body) + + def test_epoch_to_datetime_handles_float_and_fallback(): dt = CallbackLogsReplayer._epoch_to_datetime(1_700_000_000.5) assert dt.year == 2023 diff --git a/tests/test_litellm/proxy/rust_control_plane/__init__.py b/tests/test_litellm/proxy/rust_control_plane/__init__.py deleted file mode 100644 index c44841b4a73..00000000000 --- a/tests/test_litellm/proxy/rust_control_plane/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Tests for Rust control-plane endpoints.""" diff --git a/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py b/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py deleted file mode 100644 index 1e91e0332a1..00000000000 --- a/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py +++ /dev/null @@ -1,248 +0,0 @@ -"""Unit tests for the Rust control-plane auth endpoints.""" - -import pytest -from fastapi import HTTPException, Request - -from litellm.proxy._types import ProxyException, UserAPIKeyAuth -from litellm.proxy.rust_control_plane.auth_endpoints import ( - DATA_PLANE_KEY_ENV_VAR, - DATA_PLANE_KEY_HEADER, - VerifyKeyRequest, - _synthetic_request, - require_data_plane_key, - router, - verify_key, -) - - -def _make_request(headers: dict) -> Request: - """Build a minimal ASGI Request with the given headers.""" - raw_headers = [ - (k.lower().encode("latin-1"), v.encode("latin-1")) for k, v in headers.items() - ] - scope = { - "type": "http", - "method": "POST", - "path": "/v1/rust_control_plane/authentication", - "headers": raw_headers, - } - return Request(scope) - - -def test_require_data_plane_key_500_when_env_unset(monkeypatch): - monkeypatch.delenv(DATA_PLANE_KEY_ENV_VAR, raising=False) - request = _make_request({DATA_PLANE_KEY_HEADER: "anything"}) - with pytest.raises(HTTPException) as exc_info: - require_data_plane_key(request) - assert exc_info.value.status_code == 500 - assert exc_info.value.detail == "data-plane auth not configured" - - -def test_require_data_plane_key_500_when_env_empty(monkeypatch): - monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "") - request = _make_request({DATA_PLANE_KEY_HEADER: "anything"}) - with pytest.raises(HTTPException) as exc_info: - require_data_plane_key(request) - assert exc_info.value.status_code == 500 - - -def test_require_data_plane_key_401_when_header_missing(monkeypatch): - monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "secret-dp-key") - request = _make_request({}) - with pytest.raises(HTTPException) as exc_info: - require_data_plane_key(request) - assert exc_info.value.status_code == 401 - - -def test_require_data_plane_key_401_when_header_wrong(monkeypatch): - monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "secret-dp-key") - request = _make_request({DATA_PLANE_KEY_HEADER: "wrong-key"}) - with pytest.raises(HTTPException) as exc_info: - require_data_plane_key(request) - assert exc_info.value.status_code == 401 - - -def test_require_data_plane_key_does_not_accept_master_key(monkeypatch): - """The data-plane key must be a dedicated secret, not the master key.""" - monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "secret-dp-key") - monkeypatch.setenv("LITELLM_MASTER_KEY", "sk-master-1234") - request = _make_request({DATA_PLANE_KEY_HEADER: "sk-master-1234"}) - with pytest.raises(HTTPException) as exc_info: - require_data_plane_key(request) - assert exc_info.value.status_code == 401 - - -def test_require_data_plane_key_passes_when_correct(monkeypatch): - monkeypatch.setenv(DATA_PLANE_KEY_ENV_VAR, "secret-dp-key") - request = _make_request({DATA_PLANE_KEY_HEADER: "secret-dp-key"}) - # Should not raise. - assert require_data_plane_key(request) is None - - -def test_router_mounts_auth_verify_under_rust_control_plane(): - assert any( - getattr(route, "path", None) == "/v1/rust_control_plane/authentication" - for route in router.routes - ) - - -@pytest.mark.asyncio -async def test_synthetic_request_skips_budget_reservation(): - request = _synthetic_request( - route="/v1/realtime", - authorization_header="Bearer sk-test-key", - model="gpt-realtime", - ) - - assert request.url.path == "/v1/realtime" - assert request.state.skip_budget_reservation is True - assert (await request.json()) == {"model": "gpt-realtime"} - - -@pytest.mark.asyncio -async def test_verify_key_returns_model_dump(monkeypatch): - expected_auth = UserAPIKeyAuth( - api_key="hashed-key", user_id="user-123", max_budget=100.0 - ) - - captured = {} - - async def fake_user_api_key_auth(request, api_key): - captured["api_key"] = api_key - captured["request"] = request - return expected_auth - - monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", - fake_user_api_key_auth, - ) - - body = VerifyKeyRequest( - api_key="sk-test-key", route="/v1/realtime", model="gpt-realtime" - ) - result = await verify_key(body=body) - - # The key is forwarded WITH the Bearer prefix (user_api_key_auth strips it). - assert captured["api_key"] == "Bearer sk-test-key" - # Validation runs against a synthetic request carrying the gateway's route... - assert captured["request"].url.path == "/v1/realtime" - assert captured["request"].headers["authorization"] == "Bearer sk-test-key" - # ...and the requested model in the body, so model-access checks enforce it. - assert (await captured["request"].json())["model"] == "gpt-realtime" - assert result == expected_auth.model_dump(exclude_none=True, mode="json") - assert result["user_id"] == "user-123" - - -@pytest.mark.asyncio -async def test_verify_key_omits_model_when_absent(monkeypatch): - captured = {} - - async def fake_user_api_key_auth(request, api_key): - captured["request"] = request - return UserAPIKeyAuth(api_key="hashed-key") - - monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", - fake_user_api_key_auth, - ) - - body = VerifyKeyRequest(api_key="sk-test-key", route="/v1/realtime") - await verify_key(body=body) - # No model requested → empty body, not {"model": null}. - assert (await captured["request"].json()) == {} - - -@pytest.mark.asyncio -async def test_verify_key_does_not_double_prefix_existing_bearer(monkeypatch): - captured = {} - - async def fake_user_api_key_auth(request, api_key): - captured["api_key"] = api_key - captured["request"] = request - return UserAPIKeyAuth(api_key="hashed-key") - - monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", - fake_user_api_key_auth, - ) - - body = VerifyKeyRequest(api_key="Bearer sk-test-key", route="/v1/realtime") - await verify_key(body=body) - - assert captured["api_key"] == "Bearer sk-test-key" - assert captured["request"].headers["authorization"] == "Bearer sk-test-key" - - -@pytest.mark.asyncio -async def test_verify_key_401_on_proxy_exception(monkeypatch): - async def fake_user_api_key_auth(request, api_key): - raise ProxyException( - message="bad key", - type="auth_error", - param=None, - code="401", - ) - - monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", - fake_user_api_key_auth, - ) - - body = VerifyKeyRequest(api_key="sk-bad-key", route="/v1/realtime") - with pytest.raises(HTTPException) as exc_info: - await verify_key(body=body) - assert exc_info.value.status_code == 401 - assert exc_info.value.detail == "invalid api key" - - -@pytest.mark.asyncio -async def test_verify_key_401_on_http_exception(monkeypatch): - async def fake_user_api_key_auth(request, api_key): - raise HTTPException(status_code=403, detail="forbidden internals") - - monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", - fake_user_api_key_auth, - ) - - body = VerifyKeyRequest(api_key="sk-bad-key", route="/v1/realtime") - with pytest.raises(HTTPException) as exc_info: - await verify_key(body=body) - assert exc_info.value.status_code == 401 - # Internals must not leak. - assert exc_info.value.detail == "invalid api key" - - -@pytest.mark.asyncio -async def test_verify_key_propagates_http_5xx(monkeypatch): - # A 5xx (e.g. DB outage) must NOT be masked as 401 — operators need the real error. - async def fake_user_api_key_auth(request, api_key): - raise HTTPException(status_code=503, detail="db unavailable") - - monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", - fake_user_api_key_auth, - ) - - body = VerifyKeyRequest(api_key="sk-key", route="/v1/realtime") - with pytest.raises(HTTPException) as exc_info: - await verify_key(body=body) - assert exc_info.value.status_code == 503 - - -@pytest.mark.asyncio -async def test_verify_key_propagates_proxy_5xx(monkeypatch): - # A ProxyException carrying a 5xx code propagates too (not converted to 401). - async def fake_user_api_key_auth(request, api_key): - raise ProxyException( - message="internal", type="internal_error", param=None, code="500" - ) - - monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", - fake_user_api_key_auth, - ) - - body = VerifyKeyRequest(api_key="sk-key", route="/v1/realtime") - with pytest.raises(ProxyException): - await verify_key(body=body)