From 8950a34f5e4ad7a7798453efa49856013c07e96e Mon Sep 17 00:00:00 2001 From: Pathways-on-Cloud Team Date: Wed, 5 Aug 2026 14:57:16 -0700 Subject: [PATCH] Fix unauthenticated profiling server vulnerability in pathwaysutils. PiperOrigin-RevId: 959886181 --- pathwaysutils/collect_profile.py | 5 +- pathwaysutils/profiling.py | 154 ++++++++++++--- pathwaysutils/test/profiling_test.py | 283 +++++++++++++++++++++++++-- 3 files changed, 393 insertions(+), 49 deletions(-) diff --git a/pathwaysutils/collect_profile.py b/pathwaysutils/collect_profile.py index 57c971e..1cf609b 100644 --- a/pathwaysutils/collect_profile.py +++ b/pathwaysutils/collect_profile.py @@ -67,7 +67,10 @@ def main() -> None: args = parser.parse_args() if profiling.collect_profile( - args.port, args.duration_ms, args.host, args.log_dir + args.port, + args.duration_ms, + args.host, + args.log_dir, ): _logger.info("Dumped profiling information in: %s", args.log_dir) else: diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 2a39b85..e31b1ca 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -20,6 +20,7 @@ import json import logging import os +import secrets import threading from typing import Any import urllib.parse @@ -32,7 +33,6 @@ import requests import uvicorn - _logger = logging.getLogger(__name__) @@ -45,6 +45,7 @@ class _ProfileState: profile_request: The mapping containing the profile request options. lock: A thread lock to protect access to the state. """ + executable: plugin_executable.PluginExecutable | None = None profile_request: Mapping[str, Any] | None = None lock: threading.Lock @@ -225,6 +226,40 @@ def _start_pathways_trace_from_profile_request( raise +def _validate_path(path: str) -> None: + """Validates that the path is a valid GCS path.""" + path_str = str(path) + if not path_str or not path_str.startswith("gs://"): + raise ValueError(f"Path must be a GCS path, got {path_str}") + + +def _validate_gcs_bucket(log_dir: str) -> None: + """Validates that log_dir is a valid GCS path and is allowed. + + It validates that the log_dir is in the list of allowed buckets specified + in the environment variable `PATHWAYS_PROFILING_ALLOWED_GCS_BUCKETS`. If the + environment variable is not set, it assumes all buckets are allowed. + + Args: + log_dir: The GCS path to validate. + + Raises: + ValueError: If the log_dir is not a valid GCS path or is not in allowed + buckets. + """ + _validate_path(log_dir) + + if env_buckets := os.environ.get("PATHWAYS_PROFILING_ALLOWED_GCS_BUCKETS"): + allowed_buckets = {b.strip() for b in env_buckets.split(",") if b.strip()} + if allowed_buckets and not any( + log_dir.startswith(bucket_name) for bucket_name in allowed_buckets + ): + raise ValueError( + f"GCS bucket '{log_dir}' is not in allowed buckets list: " + f"{allowed_buckets}" + ) + + def start_trace( log_dir: os.PathLike[str] | str, *, @@ -263,8 +298,7 @@ def start_trace( max_num_hosts: An optional integer to limit the number of hosts profiled (defaults to 1). """ - if not str(log_dir).startswith("gs://"): - raise ValueError(f"log_dir must be a GCS bucket path, got {log_dir}") + _validate_gcs_bucket(str(log_dir)) if create_perfetto_link or create_perfetto_trace: _logger.warning( @@ -275,8 +309,8 @@ def start_trace( if jax.version.__version_info__ < (0, 9, 2): if profiler_options is not None: _logger.warning( - "ProfileOptions are not supported until JAX 0.9.2 and will be omitted. " - "Some options can be specified via command line flags." + "ProfileOptions are not supported until JAX 0.9.2 and will be" + " omitted. Some options can be specified via command line flags." ) profiler_options = None else: @@ -297,19 +331,24 @@ def start_trace( _start_pathways_trace_from_profile_request(profile_request) - if jax.version.__version_info__ >= (0, 9, 2): - _original_start_trace( - log_dir=log_dir, - create_perfetto_link=create_perfetto_link, - create_perfetto_trace=create_perfetto_trace, - profiler_options=profiler_options, - ) - else: - _original_start_trace( - log_dir=log_dir, - create_perfetto_link=create_perfetto_link, - create_perfetto_trace=create_perfetto_trace, - ) + try: + if jax.version.__version_info__ >= (0, 9, 2): + _original_start_trace( + log_dir=log_dir, + create_perfetto_link=create_perfetto_link, + create_perfetto_trace=create_perfetto_trace, + profiler_options=profiler_options, + ) + else: + _original_start_trace( + log_dir=log_dir, + create_perfetto_link=create_perfetto_link, + create_perfetto_trace=create_perfetto_trace, + ) + except Exception: + with _profile_state.lock: + _profile_state.reset() + raise def stop_trace() -> None: @@ -336,13 +375,31 @@ def start_server(port: int, requires_backend: bool = True) -> None: to the server is returned because there is no `xla_client.profiler.ProfilerServer` to return. + The server will listen on the host specified by the environment variable + `PATHWAYS_PROFILING_SERVER_HOST`, or `0.0.0.0` if not set. + + The server will verify the token provided in the + `X-Auth-Token` or `Authorization` header against the value of the + environment variable `PATHWAYS_PROFILING_AUTH_TOKEN` (if set). If the token is + not provided or does not match the expected value (which should be equal to + `PATHWAYS_PROFILING_AUTH_TOKEN` used to start the server), the server will reject the + request with a 401 error. + Args: port: The port to start the server on. requires_backend: Unused in Pathways; accepted for parameter parity. """ del requires_backend - def server_loop(port: int): - _logger.debug("Starting JAX profiler server on port %s", port) + if allowed_buckets := os.environ.get("PATHWAYS_PROFILING_ALLOWED_GCS_BUCKETS"): + for bucket in allowed_buckets.split(","): + _validate_path(bucket.strip()) + host = os.environ.get("PATHWAYS_PROFILING_SERVER_HOST", "0.0.0.0") + token_to_verify = os.environ.get("PATHWAYS_PROFILING_AUTH_TOKEN") + + def server_loop(port: int, host: str, token_to_verify: str | None): + _logger.info( + "Starting Pathways profiler server on host %s port %s", host, port + ) app = fastapi.FastAPI() @dataclasses.dataclass @@ -350,22 +407,49 @@ class ProfilingConfig: duration_ms: int repository_path: str - @app.post("/profiling") + security = fastapi.security.HTTPBearer(auto_error=False) + + def verify_auth_token( + credentials: ( + fastapi.security.HTTPAuthorizationCredentials | None + ) = fastapi.Depends(security), + ) -> None: + if token_to_verify is None: + return + if not credentials or not secrets.compare_digest( + credentials.credentials, token_to_verify + ): + raise fastapi.HTTPException( + status_code=fastapi.status.HTTP_401_UNAUTHORIZED, + detail="Unauthorized: invalid or missing authentication token", + headers={"WWW-Authenticate": "Bearer"}, + ) + + @app.post("/profiling", dependencies=[fastapi.Depends(verify_auth_token)]) async def profiling(pc: ProfilingConfig) -> Mapping[str, str]: _logger.debug("Capturing profiling data for %s ms", pc.duration_ms) - _logger.debug("Writing profiling data to %s", pc.repository_path) - await asyncio.to_thread(jax.profiler.start_trace, pc.repository_path) - await asyncio.sleep(pc.duration_ms / 1e3) - await asyncio.to_thread(jax.profiler.stop_trace) + log_dir = pc.repository_path.strip() + _logger.debug("Writing profiling data to %s", log_dir) + _validate_gcs_bucket(log_dir) + + await asyncio.to_thread(start_trace, log_dir) + try: + await asyncio.sleep(pc.duration_ms / 1e3) + finally: + await asyncio.to_thread(stop_trace) + return {"response": "profiling completed"} - uvicorn.run(app, host="0.0.0.0", port=port, log_level="debug") + uvicorn.run(app, host=host, port=port, log_level="debug") global _profiler_thread if _profiler_thread is not None: raise RuntimeError("Only one profiler server can be active at a time.") - _profiler_thread = threading.Thread(target=server_loop, args=(port,)) + _profiler_thread = threading.Thread( + target=server_loop, + args=(port, host, token_to_verify), + ) _profiler_thread.start() @@ -386,6 +470,13 @@ def collect_profile( ) -> bool: """Collects a JAX profile and saves it to the specified directory. + This function sends a POST request to the Pathways profiler server running on + the specified host and port. The server will then collect the profile for the + specified duration and save it to the specified directory. + + Authentication is handled via the `Authorization` header with a Bearer token, which should contain + the value of the environment variable `PATHWAYS_PROFILING_AUTH_TOKEN`. + Args: port: The port on which the JAX profiler server is running. duration_ms: The duration in milliseconds for which to collect the profile. @@ -398,16 +489,19 @@ def collect_profile( Raises: ValueError: If the log_dir is not a GCS path. """ - if not str(log_dir).startswith("gs://"): - raise ValueError(f"log_dir must be a GCS bucket path, got {log_dir}") + _validate_path(str(log_dir)) request_json = { "duration_ms": duration_ms, "repository_path": log_dir, } + headers = {} + if effective_token := os.environ.get("PATHWAYS_PROFILING_AUTH_TOKEN"): + headers["Authorization"] = f"Bearer {effective_token}" + address = urllib.parse.urljoin(f"http://{host}:{port}", "profiling") try: - response = requests.post(address, json=request_json) + response = requests.post(address, json=request_json, headers=headers) response.raise_for_status() except requests.exceptions.RequestException: _logger.exception("Failed to collect profiling data") diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index 0909919..c146234 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -14,8 +14,10 @@ import json import logging -from unittest import mock from typing import Any +import os +from unittest import mock +import unittest from absl.testing import absltest from absl.testing import parameterized @@ -105,6 +107,7 @@ def test_collect_profile_port(self, port): "duration_ms": 1000, "repository_path": "gs://test_bucket/test_dir", }, + headers={}, ) @parameterized.parameters(1000, 1234) @@ -123,6 +126,7 @@ def test_collect_profile_duration_ms(self, duration_ms): "duration_ms": duration_ms, "repository_path": "gs://test_bucket/test_dir", }, + headers={}, ) @parameterized.parameters("127.0.0.1", "localhost", "192.168.1.1") @@ -141,6 +145,7 @@ def test_collect_profile_host(self, host): "duration_ms": 1000, "repository_path": "gs://test_bucket/test_dir", }, + headers={}, ) @parameterized.parameters( @@ -160,6 +165,7 @@ def test_collect_profile_log_dir(self, log_dir): "duration_ms": 1000, "repository_path": log_dir, }, + headers={}, ) @parameterized.parameters("/logs/test_log_dir", "relative_path/my_log_dir") @@ -216,9 +222,7 @@ def test_collect_profile_success(self): "not_a_gcs_path", ) def test_start_trace_log_dir_error(self, log_dir): - with self.assertRaisesRegex( - ValueError, "log_dir must be a GCS bucket path" - ): + with self.assertRaisesRegex(ValueError, "Path must be a GCS path"): profiling.start_trace(log_dir) def test_lock_released_on_success(self): @@ -238,9 +242,10 @@ def test_lock_released_on_start_failure(self): self.mock_plugin_executable_cls.return_value.call.return_value[1] ) mock_result.result.side_effect = RuntimeError("start failed") - with self.assertRaisesRegex( - RuntimeError, "start failed" - ), mock.patch.object(profiling._logger, "exception"): + with ( + self.assertRaisesRegex(RuntimeError, "start failed"), + mock.patch.object(profiling._logger, "exception"), + ): profiling.start_trace("gs://test_bucket/test_dir2") self.assertFalse(profiling._profile_state.lock.locked()) @@ -312,7 +317,9 @@ def test_start_trace_with_session_id_in_options(self): profiling.start_trace("gs://test_bucket/test_dir", profiler_options=options) expected_request = self._get_expected_profile_request( - "gs://test_bucket/test_dir", max_num_hosts=1, session_id="options_session" + "gs://test_bucket/test_dir", + max_num_hosts=1, + session_id="options_session", ) self.mock_plugin_executable_cls.assert_called_once_with( json.dumps(expected_request) @@ -323,7 +330,9 @@ def test_start_trace_with_session_id_in_options(self): self.assertEqual(call_args["log_dir"], "gs://test_bucket/test_dir") self.assertFalse(call_args["create_perfetto_link"]) self.assertFalse(call_args["create_perfetto_trace"]) - self.assertEqual(call_args["profiler_options"].session_id, "options_session") + self.assertEqual( + call_args["profiler_options"].session_id, "options_session" + ) def test_start_trace_no_toy_computation_second_time(self): profiling.start_trace("gs://test_bucket/test_dir") @@ -405,10 +414,56 @@ def test_start_server_starts_thread(self): mock.patch.object(profiling.threading, "Thread", autospec=True) ) profiling.start_server(9000) - mock_thread.assert_called_once_with(target=mock.ANY, args=(9000,)) + mock_thread.assert_called_once_with( + target=mock.ANY, args=(9000, "0.0.0.0", None) + ) mock_thread.return_value.start.assert_called_once() self.assertIsNotNone(profiling._profiler_thread) + @parameterized.named_parameters( + dict(testcase_name="unset", env_host=None, expected_host="0.0.0.0"), + dict(testcase_name="empty", env_host="", expected_host=""), + dict( + testcase_name="all_ipv4", + env_host="0.0.0.0", + expected_host="0.0.0.0", + ), + dict( + testcase_name="localhost", + env_host="127.0.0.1", + expected_host="127.0.0.1", + ), + dict( + testcase_name="public_ip", + env_host="192.15.2.4", + expected_host="192.15.2.4", + ), + dict( + testcase_name="private_ip", + env_host="10.0.0.3", + expected_host="10.0.0.3", + ), + dict(testcase_name="all_ipv6", env_host="[::]", expected_host="[::]"), + ) + def test_start_server_host_env_var( + self, env_host: str | None, expected_host: str + ): + mock_thread = self.enter_context( + mock.patch.object(profiling.threading, "Thread", autospec=True) + ) + env = dict(profiling.os.environ) + if env_host is not None: + env["PATHWAYS_PROFILING_SERVER_HOST"] = env_host + else: + env.pop("PATHWAYS_PROFILING_SERVER_HOST", None) + + with mock.patch.dict(profiling.os.environ, env, clear=True): + profiling.start_server(9000) + + mock_thread.assert_called_once_with( + target=mock.ANY, args=(9000, expected_host, None) + ) + def test_start_server_twice_raises_error(self): self.enter_context( mock.patch.object(profiling.threading, "Thread", autospec=True) @@ -526,7 +581,10 @@ def test_monkey_patched_start_server(self, profiler_module): profiler_module.start_server(1234, requires_backend=False) - mocks["start_server"].assert_called_once_with(1234, requires_backend=False) + mocks["start_server"].assert_called_once_with( + 1234, + requires_backend=False, + ) @parameterized.named_parameters( dict(testcase_name="jax_profiler", profiler_module=jax.profiler), @@ -598,12 +656,8 @@ def test_create_profile_request_with_options(self): "pwTraceOptions": { "enablePythonTracer": True, "advancedConfiguration": { - "tpu_num_chips_to_profile_per_task": { - "int64Value": 3 - }, - "tpu_num_sparse_core_tiles_to_trace": { - "int64Value": 5 - }, + "tpu_num_chips_to_profile_per_task": {"int64Value": 3}, + "tpu_num_sparse_core_tiles_to_trace": {"int64Value": 5}, "tpu_trace_mode": {"stringValue": "TRACE_COMPUTE"}, "tpu_num_sparse_cores_to_trace": {"int64Value": 1}, "tpu_enable_flag": {"boolValue": True}, @@ -653,7 +707,6 @@ def test_create_profile_request_with_options(self): }, ), ) - def test_start_pathways_trace_from_profile_request(self, profile_request): profiling._start_pathways_trace_from_profile_request(profile_request) @@ -715,6 +768,200 @@ def test_start_trace_compatibility_error(self): "gs://test_bucket/test_dir", profiler_options=options ) + @parameterized.named_parameters( + dict( + testcase_name="allowed_bucket", + log_dir="gs://bucket1/dir", + ), + ) + def test_validate_gcs_bucket_env_var_allowed(self, log_dir: str): + with mock.patch.dict( + profiling.os.environ, + {"PATHWAYS_PROFILING_ALLOWED_GCS_BUCKETS": "gs://bucket1,gs://bucket2"}, + ): + profiling._validate_gcs_bucket(log_dir) + + @parameterized.named_parameters( + dict( + testcase_name="disallowed_bucket", + log_dir="gs://bucket3/dir", + ), + ) + def test_validate_gcs_bucket_env_var_disallowed(self, log_dir: str): + with mock.patch.dict( + profiling.os.environ, + {"PATHWAYS_PROFILING_ALLOWED_GCS_BUCKETS": "gs://bucket1,gs://bucket2"}, + ): + with self.assertRaisesRegex(ValueError, "is not in allowed buckets list"): + profiling._validate_gcs_bucket(log_dir) + + def test_start_trace_rollback_on_original_failure(self): + self.mock_original_start_trace.side_effect = RuntimeError( + "original start trace error" + ) + with self.assertRaisesRegex(RuntimeError, "original start trace error"): + profiling.start_trace("gs://test_bucket/test_dir") + + self.assertIsNone(profiling._profile_state.executable) + self.assertFalse(profiling._profile_state.lock.locked()) + + def test_collect_profile_without_auth_token(self): + env = dict(profiling.os.environ) + env.pop("PATHWAYS_PROFILING_AUTH_TOKEN", None) + + with mock.patch.dict(profiling.os.environ, env, clear=True): + result = profiling.collect_profile( + port=8000, + duration_ms=1000, + host="127.0.0.1", + log_dir="gs://test_bucket/test_dir", + ) + + self.assertTrue(result) + self.mock_post.assert_called_once_with( + "http://127.0.0.1:8000/profiling", + json={ + "duration_ms": 1000, + "repository_path": "gs://test_bucket/test_dir", + }, + headers={}, + ) + + @parameterized.named_parameters( + dict( + testcase_name="valid_token", + env_token="secret_token", + expected_result=True, + http_error=None, + ), + dict( + testcase_name="http_error", + env_token="wrong_token", + expected_result=False, + http_error=requests.exceptions.HTTPError( + "401 Client Error: Unauthorized" + ), + ), + ) + def test_collect_profile_with_auth_token( + self, + env_token: str, + expected_result: bool, + http_error: Exception | None, + ): + env = dict(profiling.os.environ) + env["PATHWAYS_PROFILING_AUTH_TOKEN"] = env_token + + if http_error: + self.mock_post.return_value.raise_for_status.side_effect = http_error + + with mock.patch.dict(profiling.os.environ, env, clear=True): + result = profiling.collect_profile( + port=8000, + duration_ms=1000, + host="127.0.0.1", + log_dir="gs://test_bucket/test_dir", + ) + + self.assertEqual(result, expected_result) + self.mock_post.assert_called_once_with( + "http://127.0.0.1:8000/profiling", + json={ + "duration_ms": 1000, + "repository_path": "gs://test_bucket/test_dir", + }, + headers={"Authorization": f"Bearer {env_token}"}, + ) + + def _get_server_app(self) -> Any: + with ( + mock.patch.object(profiling.threading, "Thread") as mock_thread, + mock.patch.object(profiling.uvicorn, "run") as mock_uvicorn, + ): + profiling.start_server(9000) + server_loop_fn = mock_thread.call_args[1]["target"] + args = mock_thread.call_args[1]["args"] + server_loop_fn(*args) + return mock_uvicorn.call_args[0][0] + + @parameterized.named_parameters( + dict( + testcase_name="valid_token", + server_token="secret_token", + request_headers={"Authorization": "Bearer secret_token"}, + expected_status=200, + expected_detail_substring=None, + ), + dict( + testcase_name="missing_token", + server_token="secret_token", + request_headers=None, + expected_status=401, + expected_detail_substring=( + "Unauthorized: invalid or missing authentication token" + ), + ), + dict( + testcase_name="wrong_token", + server_token="secret_token", + request_headers={"Authorization": "Bearer invalid_token"}, + expected_status=401, + expected_detail_substring=( + "Unauthorized: invalid or missing authentication token" + ), + ), + dict( + testcase_name="token_when_not_needed", + server_token=None, + request_headers={"Authorization": "Bearer unneeded_token"}, + expected_status=200, + expected_detail_substring=None, + ), + dict( + testcase_name="no_token_when_not_needed", + server_token=None, + request_headers=None, + expected_status=200, + expected_detail_substring=None, + ), + ) + @unittest.skipIf( + os.environ.get("GITHUB_ACTIONS") == "true", + "Skipping FastAPI server test in GitHub CI", + ) + def test_server_auth( + self, + server_token: str | None, + request_headers: dict[str, str] | None, + expected_status: int, + expected_detail_substring: str | None, + ): + from fastapi import testclient + env = dict(profiling.os.environ) + if server_token is not None: + env["PATHWAYS_PROFILING_AUTH_TOKEN"] = server_token + else: + env.pop("PATHWAYS_PROFILING_AUTH_TOKEN", None) + + with ( + mock.patch.dict(profiling.os.environ, env, clear=True), + mock.patch.object(profiling, "start_trace"), + mock.patch.object(profiling, "stop_trace"), + mock.patch.object(profiling.asyncio, "sleep"), + ): + app = self._get_server_app() + client = testclient.TestClient(app) + response = client.post( + "/profiling", + json={"duration_ms": 100, "repository_path": "gs://test_bucket/dir"}, + headers=request_headers, + ) + self.assertEqual(response.status_code, expected_status) + if expected_status == 200: + self.assertEqual(response.json(), {"response": "profiling completed"}) + if expected_detail_substring: + self.assertIn(expected_detail_substring, response.json()["detail"]) + if __name__ == "__main__": absltest.main()