Fail closed for server and local endpoint authentication

cluster/dev-values.yaml

@@ -45,6 +45,7 @@ pctasks:
 
     dev:
       local_dev_endpoints_url: "http://local-dev-endpoints:8512"
+      local_dev_endpoints_token: "local-dev-endpoints-token"
 
       azurite:
         enabled: true

deployment/helm/published/pctasks-server/templates/deployment.yaml

@@ -85,6 +85,8 @@ spec:
               value: "{{ .Values.pctasks.run.dev.enabled }}"
             - name: PCTASKS_RUN__LOCAL_DEV_ENDPOINTS_URL
               value: "{{ .Values.pctasks.run.dev.local_dev_endpoints_url }}"
+            - name: PCTASKS_RUN__LOCAL_DEV_ENDPOINTS_TOKEN
+              value: "{{ .Values.pctasks.run.dev.local_dev_endpoints_token }}"
 
             # Argo
             - name: PCTASKS_RUN__ARGO_HOST

deployment/helm/published/pctasks-server/values.yaml

@@ -113,6 +113,7 @@ pctasks:
 
     dev:
       local_dev_endpoints_url: ""
+      local_dev_endpoints_token: ""
 
       azurite:
         enabled: false

docker-compose.console.yml

@@ -19,6 +19,7 @@ services:
       # Dev executor settings
       - PCTASKS_RUN__TASK_RUNNER_TYPE=local
       - PCTASKS_RUN__LOCAL_DEV_ENDPOINTS_URL=http://local-dev-endpoints:8512
+      - PCTASKS_RUN__LOCAL_DEV_ENDPOINTS_TOKEN=local-dev-endpoints-token
       - PCTASKS_RUN__LOCAL_SECRETS=true
       - PCTASKS_RUN__WORKFLOW_RUNNER_TYPE=local
 

docker-compose.yml

@@ -72,6 +72,7 @@ services:
       # Dev executor settings
       - PCTASKS_RUN__TASK_RUNNER_TYPE=${PCTASKS_RUN__TASK_RUNNER_TYPE:-local}
       - PCTASKS_RUN__LOCAL_DEV_ENDPOINTS_URL=http://local-dev-endpoints:8512
+      - PCTASKS_RUN__LOCAL_DEV_ENDPOINTS_TOKEN=local-dev-endpoints-token
       - PCTASKS_RUN__LOCAL_SECRETS=${PCTASKS_RUN__LOCAL_SECRETS:-true}
       - PCTASKS_RUN__WORKFLOW_RUNNER_TYPE=local
 
@@ -142,6 +143,7 @@ services:
       - AZURITE_STORAGE_ACCOUNT=devstoreaccount1
       # Local dev secrets file
       - DEV_SECRETS_FILE=/opt/src/dev-secrets.yaml
+      - PCTASKS_RUN__LOCAL_DEV_ENDPOINTS_TOKEN=local-dev-endpoints-token
       # Cosmos DB
       - PCTASKS_COSMOSDB__TEST_CONTAINER_SUFFIX
 

pctasks/dev/pctasks/dev/local_dev_endpoints.py

@@ -7,25 +7,31 @@
 
 import logging
 import os
+import secrets
 import time
 from pathlib import Path
 from threading import Lock
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Optional
 from uuid import uuid1
 
 import yaml
 from cachetools import LRUCache
-from fastapi import BackgroundTasks, FastAPI, Request, Response
+from fastapi import (
+    BackgroundTasks,
+    Depends,
+    FastAPI,
+    Header,
+    HTTPException,
+    Request,
+    Response,
+)
 from fastapi.responses import JSONResponse, PlainTextResponse
 from yaml import Loader
 
 from pctasks.cli.cli import pctasks_cmd
 from pctasks.core.models.run import TaskRunStatus
 from pctasks.run.models import TaskPollResult
 
-app = FastAPI()
-
-
 logger = logging.getLogger(__name__)
 
 _task_cache_lock = Lock()
@@ -35,6 +41,25 @@
 WAIT_AND_FAIL_TAG = "wait_and_fail"
 
 DEV_SECRETS_FILE_ENV_VAR = "DEV_SECRETS_FILE"
+LOCAL_DEV_ENDPOINTS_TOKEN_ENV_VAR = "PCTASKS_RUN__LOCAL_DEV_ENDPOINTS_TOKEN"
+
+
+def authenticate_local_request(
+    authorization: Optional[str] = Header(default=None),
+) -> None:
+    token = os.getenv(LOCAL_DEV_ENDPOINTS_TOKEN_ENV_VAR)
+    if not token:
+        logger.error("Local dev endpoints token is not configured")
+        raise HTTPException(status_code=503, detail="Authentication is not configured")
+
+    expected = f"Bearer {token}"
+    if not authorization or not secrets.compare_digest(authorization, expected):
+        raise HTTPException(
+            status_code=401, detail="Invalid authentication credentials"
+        )
+
+
+app = FastAPI(dependencies=[Depends(authenticate_local_request)])
 
 # Execute tasks
 

pctasks/dev/tests/test_local_dev_endpoints.py

@@ -0,0 +1,72 @@
+from pathlib import Path
+
+import pytest
+from fastapi.testclient import TestClient
+
+from pctasks.dev.local_dev_endpoints import (
+    FAIL_SUBMIT_TAG,
+    LOCAL_DEV_ENDPOINTS_TOKEN_ENV_VAR,
+    app,
+)
+
+TOKEN = "test-local-token"
+AUTH_HEADERS = {"Authorization": f"Bearer {TOKEN}"}
+
+
+def test_endpoints_fail_closed_without_configured_token(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv(LOCAL_DEV_ENDPOINTS_TOKEN_ENV_VAR, raising=False)
+
+    with TestClient(app) as client:
+        assert client.post("/execute", json={"args": []}).status_code == 503
+        assert client.get("/poll/missing").status_code == 503
+        assert client.get("/secrets/test").status_code == 503
+
+
+def test_endpoints_reject_missing_or_invalid_token(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv(LOCAL_DEV_ENDPOINTS_TOKEN_ENV_VAR, TOKEN)
+
+    with TestClient(app) as client:
+        assert client.post("/execute", json={"args": []}).status_code == 401
+        assert client.get("/poll/missing").status_code == 401
+        assert client.get("/secrets/test").status_code == 401
+        assert (
+            client.get(
+                "/secrets/test", headers={"Authorization": "Bearer invalid"}
+            ).status_code
+            == 401
+        )
+
+
+def test_authenticated_execute_and_poll(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv(LOCAL_DEV_ENDPOINTS_TOKEN_ENV_VAR, TOKEN)
+
+    with TestClient(app) as client:
+        response = client.post(
+            "/execute",
+            headers=AUTH_HEADERS,
+            json={"args": [], "tags": {FAIL_SUBMIT_TAG: "true"}},
+        )
+        assert response.status_code == 200
+        task_id = response.json()["id"]
+
+        poll_response = client.get(f"/poll/{task_id}", headers=AUTH_HEADERS)
+        assert poll_response.status_code == 200
+
+
+def test_authenticated_secret_access(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    monkeypatch.setenv(LOCAL_DEV_ENDPOINTS_TOKEN_ENV_VAR, TOKEN)
+    secrets_file = tmp_path / "secrets.yaml"
+    secrets_file.write_text("test: secret-value\n")
+    monkeypatch.setenv("DEV_SECRETS_FILE", str(secrets_file))
+
+    with TestClient(app) as client:
+        response = client.get("/secrets/test", headers=AUTH_HEADERS)
+
+    assert response.status_code == 200
+    assert response.text == "secret-value"

pctasks/run/pctasks/run/secrets/local.py

@@ -18,7 +18,12 @@ def get_secret(self, name: str) -> str:
             local_dev_endpoints_url = self.settings.local_dev_endpoints_url
             if local_dev_endpoints_url:
                 resp = requests.get(
-                    os.path.join(local_dev_endpoints_url, f"secrets/{name}")
+                    os.path.join(local_dev_endpoints_url, f"secrets/{name}"),
+                    headers={
+                        "Authorization": (
+                            f"Bearer {self.settings.local_dev_endpoints_token}"
+                        )
+                    },
                 )
                 if resp.status_code == 200:
                     result = resp.text

pctasks/run/pctasks/run/settings.py

@@ -63,6 +63,7 @@ def section_name(cls) -> str:
 
     # Dev
     local_dev_endpoints_url: Optional[str] = None
+    local_dev_endpoints_token: Optional[str] = None
     local_secrets: bool = False
 
     notification_queue: NotificationQueueConnStrConfig = Field()
@@ -145,6 +146,12 @@ def keyvault_url_validator(self) -> Self:
 
     @model_validator(mode="after")
     def _task_runner_type_validator(self) -> Self:
+        if self.local_dev_endpoints_url and not self.local_dev_endpoints_token:
+            raise ValueError(
+                "Must specify local_dev_endpoints_token when "
+                "local_dev_endpoints_url is configured."
+            )
+
         if self.task_runner_type == TaskRunnerType.LOCAL:
             if self.local_dev_endpoints_url is None:
                 raise ValueError(

pctasks/run/pctasks/run/task/__init__.py

@@ -12,7 +12,10 @@ def get_task_runner(settings: Optional[RunSettings] = None) -> TaskRunner:
 
     if settings.task_runner_type == TaskRunnerType.LOCAL:
         assert settings.local_dev_endpoints_url  # Checked during settings validation
-        return LocalTaskRunner(settings.local_dev_endpoints_url)
+        assert settings.local_dev_endpoints_token  # Checked during settings validation
+        return LocalTaskRunner(
+            settings.local_dev_endpoints_url, settings.local_dev_endpoints_token
+        )
     elif settings.task_runner_type == TaskRunnerType.BATCH:
         return BatchTaskRunner(settings)
     elif settings.task_runner_type == TaskRunnerType.ARGO:

pctasks/run/pctasks/run/task/local.py

@@ -26,8 +26,9 @@ class LocalTaskRunner(TaskRunner):
     See the local-dev-endpoints service in the development environment.
     """
 
-    def __init__(self, local_dev_endpoints_url: str) -> None:
+    def __init__(self, local_dev_endpoints_url: str, token: str) -> None:
         self.local_dev_endpoints_url = local_dev_endpoints_url
+        self.headers = {"Authorization": f"Bearer {token}"}
 
     def __enter__(self) -> "LocalTaskRunner":
         return self
@@ -65,7 +66,11 @@ def submit_tasks(
                 args.extend(["--account-url", task_input_blob_config.account_url])
 
             data = json.dumps({"args": args, "tags": task_tags or {}}).encode("utf-8")
-            resp = requests.post(self.local_dev_endpoints_url + "/execute", data=data)
+            resp = requests.post(
+                self.local_dev_endpoints_url + "/execute",
+                data=data,
+                headers=self.headers,
+            )
             if resp.status_code == 200:
                 results.append(SuccessfulTaskSubmitResult(task_runner_id=resp.json()))
             else:
@@ -82,7 +87,8 @@ def poll_task(
     ) -> TaskPollResult:
         try:
             resp = requests.get(
-                self.local_dev_endpoints_url + f"/poll/{runner_id['id']}"
+                self.local_dev_endpoints_url + f"/poll/{runner_id['id']}",
+                headers=self.headers,
             )
             if resp.status_code == 200:
                 return TaskPollResult.model_validate(resp.json())

pctasks/run/tests/secrets/test_local.py

@@ -0,0 +1,37 @@
+from types import SimpleNamespace
+from typing import Dict, cast
+
+import pytest
+
+from pctasks.run.secrets.local import LocalSecretsProvider
+from pctasks.run.settings import RunSettings
+
+TOKEN = "test-local-token"
+
+
+class FakeResponse:
+    status_code = 200
+    text = "secret-value"
+
+
+def test_secret_request_sends_authentication_header(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    captured_headers: Dict[str, str] = {}
+
+    def get(url: str, headers: Dict[str, str]) -> FakeResponse:
+        captured_headers.update(headers)
+        return FakeResponse()
+
+    monkeypatch.setattr("pctasks.run.secrets.local.requests.get", get)
+    settings = cast(
+        RunSettings,
+        SimpleNamespace(
+            local_dev_endpoints_url="http://local-dev-endpoints:8512",
+            local_dev_endpoints_token=TOKEN,
+        ),
+    )
+    provider = LocalSecretsProvider(settings)
+
+    assert provider.get_secret("test") == "secret-value"
+    assert captured_headers == {"Authorization": f"Bearer {TOKEN}"}

pctasks/run/tests/task/test_local.py

@@ -0,0 +1,62 @@
+from types import SimpleNamespace
+from typing import Dict
+
+import pytest
+
+from pctasks.run.task.local import LocalTaskRunner
+
+TOKEN = "test-local-token"
+EXPECTED_HEADERS = {"Authorization": f"Bearer {TOKEN}"}
+
+
+class FakeResponse:
+    status_code = 200
+    text = ""
+
+    def __init__(self, payload: Dict[str, str]) -> None:
+        self.payload = payload
+
+    def json(self) -> Dict[str, str]:
+        return self.payload
+
+
+def test_submit_sends_authentication_header(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    captured_headers: Dict[str, str] = {}
+
+    def post(url: str, data: bytes, headers: Dict[str, str]) -> FakeResponse:
+        captured_headers.update(headers)
+        return FakeResponse({"id": "task-id"})
+
+    monkeypatch.setattr("pctasks.run.task.local.requests.post", post)
+    prepared_task = SimpleNamespace(
+        task_input_blob_config=SimpleNamespace(
+            uri="blob://account/container/input.json",
+            sas_token="token",
+            account_url=None,
+        ),
+        task_data=SimpleNamespace(tags={}),
+    )
+
+    runner = LocalTaskRunner("http://local-dev-endpoints:8512", TOKEN)
+    results = runner.submit_tasks([prepared_task])  # type: ignore[list-item]
+
+    assert results[0].success
+    assert captured_headers == EXPECTED_HEADERS
+
+
+def test_poll_sends_authentication_header(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured_headers: Dict[str, str] = {}
+
+    def get(url: str, headers: Dict[str, str]) -> FakeResponse:
+        captured_headers.update(headers)
+        return FakeResponse({"task_status": "completed"})
+
+    monkeypatch.setattr("pctasks.run.task.local.requests.get", get)
+
+    runner = LocalTaskRunner("http://local-dev-endpoints:8512", TOKEN)
+    result = runner.poll_task({"id": "task-id"}, previous_poll_count=0)
+
+    assert result.task_status.value == "completed"
+    assert captured_headers == EXPECTED_HEADERS

pctasks/server/pctasks/server/request.py

@@ -1,4 +1,5 @@
 import logging
+import secrets
 from typing import Dict, Optional, Union
 from urllib.parse import urlparse
 
@@ -62,12 +63,16 @@ def is_authenticated(self) -> bool:
         if self.dev:
             return self.has_subscription or self.has_authorization
 
-        if self.access_key:
-            if not self.request_access_key == self.access_key:
-                logger.warning("Request made with mismatched access key")
-                return False
-        else:
+        if not self.access_key:
             logger.warning("Access key is unset in non-dev environment!")
+            return False
+
+        request_access_key = self.request_access_key
+        if not request_access_key or not secrets.compare_digest(
+            request_access_key, self.access_key
+        ):
+            logger.warning("Request made with mismatched access key")
+            return False
 
         has_subscription_key = (
             self.has_subscription