From 974ce3981bdc2109309becdedeb52535e5953ed3 Mon Sep 17 00:00:00 2001 From: Derek Tan Date: Wed, 2 Sep 2026 15:01:42 +0000 Subject: [PATCH] Fail closed for server and local endpoint authentication --- cluster/dev-values.yaml | 1 + .../pctasks-server/templates/deployment.yaml | 2 + .../helm/published/pctasks-server/values.yaml | 1 + docker-compose.console.yml | 1 + docker-compose.yml | 2 + .../dev/pctasks/dev/local_dev_endpoints.py | 35 +++++++-- pctasks/dev/tests/test_local_dev_endpoints.py | 72 +++++++++++++++++++ pctasks/run/pctasks/run/secrets/local.py | 7 +- pctasks/run/pctasks/run/settings.py | 7 ++ pctasks/run/pctasks/run/task/__init__.py | 5 +- pctasks/run/pctasks/run/task/local.py | 12 +++- pctasks/run/tests/secrets/test_local.py | 37 ++++++++++ pctasks/run/tests/task/test_local.py | 62 ++++++++++++++++ pctasks/server/pctasks/server/request.py | 15 ++-- pctasks/server/tests/test_request.py | 19 +++++ 15 files changed, 263 insertions(+), 15 deletions(-) create mode 100644 pctasks/dev/tests/test_local_dev_endpoints.py create mode 100644 pctasks/run/tests/secrets/test_local.py create mode 100644 pctasks/run/tests/task/test_local.py diff --git a/cluster/dev-values.yaml b/cluster/dev-values.yaml index 870cee4d8..2bcea1b23 100644 --- a/cluster/dev-values.yaml +++ b/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 diff --git a/deployment/helm/published/pctasks-server/templates/deployment.yaml b/deployment/helm/published/pctasks-server/templates/deployment.yaml index 520c0e6bd..faa0156d8 100644 --- a/deployment/helm/published/pctasks-server/templates/deployment.yaml +++ b/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 diff --git a/deployment/helm/published/pctasks-server/values.yaml b/deployment/helm/published/pctasks-server/values.yaml index 279ef5e11..fab20d15d 100644 --- a/deployment/helm/published/pctasks-server/values.yaml +++ b/deployment/helm/published/pctasks-server/values.yaml @@ -113,6 +113,7 @@ pctasks: dev: local_dev_endpoints_url: "" + local_dev_endpoints_token: "" azurite: enabled: false diff --git a/docker-compose.console.yml b/docker-compose.console.yml index 542cb61e4..196404e82 100644 --- a/docker-compose.console.yml +++ b/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 diff --git a/docker-compose.yml b/docker-compose.yml index 7957126d7..da13390e4 100644 --- a/docker-compose.yml +++ b/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 diff --git a/pctasks/dev/pctasks/dev/local_dev_endpoints.py b/pctasks/dev/pctasks/dev/local_dev_endpoints.py index a659bc09a..485aacaf9 100644 --- a/pctasks/dev/pctasks/dev/local_dev_endpoints.py +++ b/pctasks/dev/pctasks/dev/local_dev_endpoints.py @@ -7,15 +7,24 @@ 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 @@ -23,9 +32,6 @@ 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 diff --git a/pctasks/dev/tests/test_local_dev_endpoints.py b/pctasks/dev/tests/test_local_dev_endpoints.py new file mode 100644 index 000000000..db2c30a16 --- /dev/null +++ b/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" diff --git a/pctasks/run/pctasks/run/secrets/local.py b/pctasks/run/pctasks/run/secrets/local.py index 6aee14bd2..07ffc3e4e 100644 --- a/pctasks/run/pctasks/run/secrets/local.py +++ b/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 diff --git a/pctasks/run/pctasks/run/settings.py b/pctasks/run/pctasks/run/settings.py index be8c55f26..b0181a1d4 100644 --- a/pctasks/run/pctasks/run/settings.py +++ b/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( diff --git a/pctasks/run/pctasks/run/task/__init__.py b/pctasks/run/pctasks/run/task/__init__.py index 5bf8ec839..2d4555db8 100644 --- a/pctasks/run/pctasks/run/task/__init__.py +++ b/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: diff --git a/pctasks/run/pctasks/run/task/local.py b/pctasks/run/pctasks/run/task/local.py index c474f8d08..a7729aabf 100644 --- a/pctasks/run/pctasks/run/task/local.py +++ b/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()) diff --git a/pctasks/run/tests/secrets/test_local.py b/pctasks/run/tests/secrets/test_local.py new file mode 100644 index 000000000..7dc97fd8d --- /dev/null +++ b/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}"} diff --git a/pctasks/run/tests/task/test_local.py b/pctasks/run/tests/task/test_local.py new file mode 100644 index 000000000..c795d6670 --- /dev/null +++ b/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 diff --git a/pctasks/server/pctasks/server/request.py b/pctasks/server/pctasks/server/request.py index 9a35aa3fc..5f6d70db2 100644 --- a/pctasks/server/pctasks/server/request.py +++ b/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 diff --git a/pctasks/server/tests/test_request.py b/pctasks/server/tests/test_request.py index 81cfa7b57..a46893f7b 100644 --- a/pctasks/server/tests/test_request.py +++ b/pctasks/server/tests/test_request.py @@ -1,4 +1,5 @@ import pytest +from fastapi import Request from fastapi.testclient import TestClient from pctasks.core.utils import ignore_ssl_warnings @@ -9,6 +10,7 @@ HAS_SUBSCRIPTION_HEADER, SUBSCRIPTION_KEY_HEADER, USER_EMAIL_HEADER, + ParsedRequest, ) from pctasks.server.settings import ServerSettings @@ -87,3 +89,20 @@ def test_unauthorized_requests(): with ignore_ssl_warnings(): response = client.get("/workflows", headers=headers) assert response.status_code == 401 + + +def test_unset_access_key_fails_closed() -> None: + request = Request( + { + "type": "http", + "headers": [ + (HAS_AUTHORIZATION_HEADER.lower().encode(), b"true"), + (USER_EMAIL_HEADER.lower().encode(), b"attacker@example.com"), + ], + } + ) + parsed_request = ParsedRequest(request) + parsed_request.dev = False + parsed_request.access_key = None + + assert not parsed_request.is_authenticated