Merge pull request #9766 from BerriAI/litellm_add_auth_metrics_endpoint

[Security feature] Allow adding authentication on /metrics endpoints

litellm/__init__.py

@@ -123,6 +123,7 @@ callbacks: List[
 langfuse_default_tags: Optional[List[str]] = None
 langsmith_batch_size: Optional[int] = None
 prometheus_initialize_budget_metrics: Optional[bool] = False
+require_auth_for_metrics_endpoint: Optional[bool] = False
 argilla_batch_size: Optional[int] = None
 datadog_use_v1: Optional[bool] = False  # if you want to use v1 datadog logged payload
 gcs_pub_sub_use_v1: Optional[

litellm/integrations/prometheus.py

@@ -1721,6 +1721,36 @@ class PrometheusLogger(CustomLogger):
             return (end_time - start_time).total_seconds()
         return None
 
+    @staticmethod
+    def _mount_metrics_endpoint(premium_user: bool):
+        """
+        Mount the Prometheus metrics endpoint with optional authentication.
+
+        Args:
+            premium_user (bool): Whether the user is a premium user
+            require_auth (bool, optional): Whether to require authentication for the metrics endpoint.
+                                        Defaults to False.
+        """
+        from prometheus_client import make_asgi_app
+
+        from litellm._logging import verbose_proxy_logger
+        from litellm.proxy._types import CommonProxyErrors
+        from litellm.proxy.proxy_server import app
+
+        if premium_user is not True:
+            verbose_proxy_logger.warning(
+                f"Prometheus metrics are only available for premium users. {CommonProxyErrors.not_premium_user.value}"
+            )
+
+        # Create metrics ASGI app
+        metrics_app = make_asgi_app()
+
+        # Mount the metrics app to the app
+        app.mount("/metrics", metrics_app)
+        verbose_proxy_logger.debug(
+            "Starting Prometheus Metrics on /metrics (no authentication)"
+        )
+
 
 def prometheus_label_factory(
     supported_enum_labels: List[str],

litellm/proxy/common_utils/callback_utils.py

@@ -224,18 +224,9 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
             litellm.callbacks = imported_list  # type: ignore
 
         if "prometheus" in value:
-            if premium_user is not True:
-                verbose_proxy_logger.warning(
-                    f"Prometheus metrics are only available for premium users. {CommonProxyErrors.not_premium_user.value}"
-                )
-            from litellm.proxy.proxy_server import app
+            from litellm.integrations.prometheus import PrometheusLogger
 
-            verbose_proxy_logger.debug("Starting Prometheus Metrics on /metrics")
-            from prometheus_client import make_asgi_app
-
-            # Add prometheus asgi middleware to route /metrics requests
-            metrics_app = make_asgi_app()
-            app.mount("/metrics", metrics_app)
+            PrometheusLogger._mount_metrics_endpoint(premium_user)
     else:
         litellm.callbacks = [
             get_instance_fn(

litellm/proxy/middleware/prometheus_auth_middleware.py

@@ -0,0 +1,72 @@
+"""
+Prometheus Auth Middleware
+"""
+from fastapi import Request
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+
+import litellm
+from litellm.proxy._types import SpecialHeaders
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+
+class PrometheusAuthMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to authenticate requests to the metrics endpoint
+
+    By default, auth is not run on the metrics endpoint
+
+    Enabled by setting the following in proxy_config.yaml:
+
+    ```yaml
+    litellm_settings:
+        require_auth_for_metrics_endpoint: true
+    ```
+    """
+
+    async def dispatch(self, request: Request, call_next):
+        # Check if this is a request to the metrics endpoint
+
+        if self._is_prometheus_metrics_endpoint(request):
+            if self._should_run_auth_on_metrics_endpoint() is True:
+                try:
+                    await user_api_key_auth(
+                        request=request,
+                        api_key=request.headers.get(
+                            SpecialHeaders.openai_authorization.value
+                        )
+                        or "",
+                    )
+                except Exception as e:
+                    return JSONResponse(
+                        status_code=401,
+                        content=f"Unauthorized access to metrics endpoint: {getattr(e, 'message', str(e))}",
+                    )
+
+        # Process the request and get the response
+        response = await call_next(request)
+
+        return response
+
+    @staticmethod
+    def _is_prometheus_metrics_endpoint(request: Request):
+        try:
+            if "/metrics" in request.url.path:
+                return True
+            return False
+        except Exception:
+            return False
+
+    @staticmethod
+    def _should_run_auth_on_metrics_endpoint():
+        """
+        Returns True if auth should be run on the metrics endpoint
+
+        False by default, set to True in proxy_config.yaml to enable
+
+        ```yaml
+        litellm_settings:
+            require_auth_for_metrics_endpoint: true
+        ```
+        """
+        return litellm.require_auth_for_metrics_endpoint

litellm/proxy/proxy_config.yaml

@@ -4,7 +4,13 @@ model_list:
       model: openai/fake
       api_key: fake-key
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
+  - model_name: openai/gpt-4o
+    litellm_params:
+      model: openai/gpt-4o
+      api_key: fake-key
 
 litellm_settings:
+  require_auth_for_metrics_endpoint: true
+
   callbacks: ["prometheus"]
   service_callback: ["prometheus_system"]

litellm/proxy/proxy_server.py

@@ -249,6 +249,7 @@ from litellm.proxy.management_endpoints.ui_sso import (
 )
 from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router
 from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update
+from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware
 from litellm.proxy.openai_files_endpoints.files_endpoints import (
     router as openai_files_router,
 )
@@ -745,6 +746,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+app.add_middleware(PrometheusAuthMiddleware)
 
 from typing import Dict
 
@@ -1676,14 +1678,11 @@ class ProxyConfig:
                                 callback
                             )
                             if "prometheus" in callback:
-                                verbose_proxy_logger.debug(
-                                    "Starting Prometheus Metrics on /metrics"
+                                from litellm.integrations.prometheus import (
+                                    PrometheusLogger,
                                 )
-                                from prometheus_client import make_asgi_app
 
-                                # Add prometheus asgi middleware to route /metrics requests
-                                metrics_app = make_asgi_app()
-                                app.mount("/metrics", metrics_app)
+                                PrometheusLogger._mount_metrics_endpoint(premium_user)
                     print(  # noqa
                         f"{blue_color_code} Initialized Success Callbacks - {litellm.success_callback} {reset_color_code}"
                     )  # noqa

tests/litellm/proxy/middleware/test_prometheus_auth_middleware.py

@@ -0,0 +1,129 @@
+import json
+import os
+import sys
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+
+
+import pytest
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from fastapi.testclient import TestClient
+
+import litellm
+from litellm.proxy._types import SpecialHeaders
+from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware
+
+
+# Fake auth functions to simulate valid and invalid auth behavior.
+async def fake_valid_auth(request, api_key):
+    # Simulate valid authentication: do nothing (i.e. pass)
+    return
+
+
+async def fake_invalid_auth(request, api_key):
+    print("running fake invalid auth", request, api_key)
+    # Simulate invalid auth by raising an exception.
+    raise Exception("Invalid API key")
+
+
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+
+
+@pytest.fixture
+def app_with_middleware():
+    """Create a FastAPI app with the PrometheusAuthMiddleware and dummy endpoints."""
+    app = FastAPI()
+    # Add the PrometheusAuthMiddleware to the app.
+    app.add_middleware(PrometheusAuthMiddleware)
+
+    @app.get("/metrics")
+    async def metrics():
+        return {"msg": "metrics OK"}
+
+    # Also allow /metrics/ (trailing slash)
+    @app.get("/metrics/")
+    async def metrics_slash():
+        return {"msg": "metrics OK"}
+
+    @app.get("/chat/completions")
+    async def chat():
+        return {"msg": "chat completions OK"}
+
+    @app.get("/embeddings")
+    async def embeddings():
+        return {"msg": "embeddings OK"}
+
+    return app
+
+
+def test_valid_auth_metrics(app_with_middleware, monkeypatch):
+    """
+    Test that a request to /metrics (and /metrics/) with valid auth headers passes.
+    """
+    # Enable auth on metrics endpoints.
+    litellm.require_auth_for_metrics_endpoint = True
+    # Patch the auth function to simulate a valid authentication.
+    monkeypatch.setattr(
+        "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
+        fake_valid_auth,
+    )
+
+    client = TestClient(app_with_middleware)
+    headers = {SpecialHeaders.openai_authorization.value: "valid"}
+
+    # Test for /metrics (no trailing slash)
+    response = client.get("/metrics", headers=headers)
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "metrics OK"}
+
+    # Test for /metrics/ (with trailing slash)
+    response = client.get("/metrics/", headers=headers)
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "metrics OK"}
+
+
+def test_invalid_auth_metrics(app_with_middleware, monkeypatch):
+    """
+    Test that a request to /metrics with invalid auth headers fails with a 401.
+    """
+    litellm.require_auth_for_metrics_endpoint = True
+    # Patch the auth function to simulate a failed authentication.
+    monkeypatch.setattr(
+        "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
+        fake_invalid_auth,
+    )
+
+    client = TestClient(app_with_middleware)
+    headers = {SpecialHeaders.openai_authorization.value: "invalid"}
+
+    response = client.get("/metrics", headers=headers)
+    assert response.status_code == 401, response.text
+    assert "Unauthorized access to metrics endpoint" in response.text
+
+
+def test_no_auth_metrics_when_disabled(app_with_middleware, monkeypatch):
+    """
+    Test that when require_auth_for_metrics_endpoint is False, requests to /metrics
+    bypass the auth check.
+    """
+    litellm.require_auth_for_metrics_endpoint = False
+
+    # To ensure auth is not run, patch the auth function with one that will raise if called.
+    def should_not_be_called(*args, **kwargs):
+        raise Exception("Auth should not be called")
+
+    monkeypatch.setattr(
+        "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
+        should_not_be_called,
+    )
+
+    client = TestClient(app_with_middleware)
+    response = client.get("/metrics")
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "metrics OK"}