Add optional bearer auth to metrics endpoint (#2460)

Author: un-def

src/dstack/_internal/server/routers/prometheus.py

@@ -1,3 +1,4 @@
+import os
 from typing import Annotated
 
 from fastapi import APIRouter, Depends
@@ -6,12 +7,16 @@
 
 from dstack._internal.server import settings
 from dstack._internal.server.db import get_session
+from dstack._internal.server.security.permissions import OptionalServiceAccount
 from dstack._internal.server.services import prometheus
 from dstack._internal.server.utils.routers import error_not_found
 
+_auth = OptionalServiceAccount(os.getenv("DSTACK_PROMETHEUS_AUTH_TOKEN"))
+
 router = APIRouter(
     tags=["prometheus"],
     default_response_class=PlainTextResponse,
+    dependencies=[Depends(_auth)],
 )
 
 

src/dstack/_internal/server/security/permissions.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import Annotated, Optional, Tuple
 
 from fastapi import Depends, HTTPException, Security
 from fastapi.security import HTTPBearer
@@ -99,6 +99,24 @@ async def __call__(
         return await get_project_member(session, project_name, token.credentials)
 
 
+class OptionalServiceAccount:
+    def __init__(self, token: Optional[str]) -> None:
+        self._token = token
+
+    async def __call__(
+        self,
+        token: Annotated[
+            Optional[HTTPAuthorizationCredentials], Security(HTTPBearer(auto_error=False))
+        ],
+    ) -> None:
+        if self._token is None:
+            return
+        if token is None:
+            raise error_forbidden()
+        if token.credentials != self._token:
+            raise error_invalid_token()
+
+
 async def get_project_member(
     session: AsyncSession, project_name: str, token: str
 ) -> Tuple[UserModel, ProjectModel]:

src/tests/_internal/server/routers/test_prometheus.py

@@ -28,6 +28,7 @@
     create_repo,
     create_run,
     create_user,
+    get_auth_headers,
     get_instance_offer_with_availability,
     get_job_provisioning_data,
     get_job_runtime_data,
@@ -38,6 +39,7 @@
 @pytest.fixture
 def enable_metrics(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr("dstack._internal.server.settings.ENABLE_PROMETHEUS_METRICS", True)
+    monkeypatch.setattr("dstack._internal.server.routers.prometheus._auth._token", None)
 
 
 FAKE_NOW = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)
@@ -289,6 +291,25 @@ async def test_returns_404_if_not_enabled(
         response = await client.get("/metrics")
         assert response.status_code == 404
 
+    @pytest.mark.parametrize("token", [None, "foo"])
+    async def test_returns_403_if_not_authenticated(
+        self, monkeypatch: pytest.MonkeyPatch, client: AsyncClient, token: Optional[str]
+    ):
+        monkeypatch.setattr("dstack._internal.server.routers.prometheus._auth._token", "secret")
+        if token is not None:
+            headers = get_auth_headers(token)
+        else:
+            headers = None
+        response = await client.get("/metrics", headers=headers)
+        assert response.status_code == 403
+
+    async def test_returns_200_if_token_is_valid(
+        self, monkeypatch: pytest.MonkeyPatch, client: AsyncClient
+    ):
+        monkeypatch.setattr("dstack._internal.server.routers.prometheus._auth._token", "secret")
+        response = await client.get("/metrics", headers=get_auth_headers("secret"))
+        assert response.status_code == 200
+
 
 async def _create_project(session: AsyncSession, name: str, user: UserModel) -> ProjectModel:
     project = await create_project(session=session, owner=user, name=name)