From 627e70f6970c8b1a1fbb8d6b37da4a92c6636a5c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 10:59:25 +0800 Subject: [PATCH 1/7] feat: add fetch_url tool for proxying HTTP GET to cluster-internal URLs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a new MCP tool that lets the agent fetch URLs on the cluster's network (YARN tracking UI, Spark History Server, NodeManager web UIs) when the agent is on a different network and cannot reach those hosts directly. The MCP service runs on the YARN RM node, so it can reach every host the cluster knows about — the agent just needs a way to ask. Security: SSRF guard via host suffix overlap - URL host must share >= 2 labels of suffix with the named Connection's yarn_rm_url host (e.g. yarn_rm_url='rm.prod.internal' allows 'http://nm01.prod.internal/...') - IP literals (10.0.0.1, ::1) rejected - Non-http(s) schemes (file://, gopher://, ftp://) rejected - Connection with no yarn_rm_url cannot use this tool - Reuses Connection.auth_for_httpx() and verify_for_httpx() so the agent does not need cluster credentials - Response body capped at 1 MB (truncated=true if larger) - 30s timeout, follows redirects, loguru INFO audit log on every call - spark_executor/tools/fetch_url.py: new tool + 2 helpers (_host_suffix_overlap, _validate_url_host) - spark_executor/models.py: FetchUrlResult Pydantic model - spark_executor/tools/requests.py: FetchUrlRequest with descriptions - spark_executor/server.py: /fetch_url route, operation_id='fetch_url' - tests/unit/test_fetch_url.py: 13 unit tests covering all guards, truncation, auth/SSL pass-through, redirect follow - tests/integration/test_mcp_routes.py: assert 21 tool routes - README.md: 1 row in Spark Executor 工具 table Tests: 369 passed (up from 356). Co-Authored-By: Claude Fable 5 --- README.md | 1 + spark_executor/models.py | 8 ++ spark_executor/server.py | 27 +++++ spark_executor/tools/fetch_url.py | 108 ++++++++++++++++++ spark_executor/tools/requests.py | 24 ++++ tests/integration/test_mcp_routes.py | 3 +- tests/unit/test_fetch_url.py | 164 +++++++++++++++++++++++++++ 7 files changed, 334 insertions(+), 1 deletion(-) create mode 100644 spark_executor/tools/fetch_url.py create mode 100644 tests/unit/test_fetch_url.py diff --git a/README.md b/README.md index 8916969..94e76a2 100644 --- a/README.md +++ b/README.md @@ -120,6 +120,7 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态（按 `application_id` + `connection_name`） | | `get_external_job_result` | 查询外部 YARN application 终态结果视图 | | `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 | +| `fetch_url` | 代理 HTTP GET 到集群内网 URL (受 `Connection.yarn_rm_url` host 围栏约束) | ### Files MCP 工具 diff --git a/spark_executor/models.py b/spark_executor/models.py index 798d72b..ce3b2ab 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -40,6 +40,14 @@ class SubmitResult(BaseModel): tracking_url: str | None = None +class FetchUrlResult(BaseModel): + url: str + status_code: int + content_type: str + body: str + truncated: bool = False + + class Connection(BaseModel): name: str # Defaults to "yarn" because that's the literal string spark-submit wants diff --git a/spark_executor/server.py b/spark_executor/server.py index 3e156c2..f772426 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -21,6 +21,7 @@ from spark_executor.tools.external_jobs import ( get_external_job_status, get_external_job_result, ) +from spark_executor.tools.fetch_url import fetch_url from spark_executor.tools.requests import ( ConnectionNameRequest, EmptyRequest, @@ -28,6 +29,7 @@ from spark_executor.tools.requests import ( ExternalJobLogsRequest, ExternalJobStatusRequest, ExternalJobResultRequest, + FetchUrlRequest, GetJobLogsRequest, JobIdRequest, PendingIdRequest, @@ -467,3 +469,28 @@ def _read_job_file(req: ReadJobFileRequest): ) def _update_job_file(req: UpdateJobFileRequest): return update_job_file(req.script_path, req.content) + +# --- HTTP fetch proxy (host allowlist via Connection.yarn_rm_url) --- + +@app.post( + "/fetch_url", + operation_id="fetch_url", + summary="Fetch a URL on the cluster's network and return the body", + description=( + "Proxy an HTTP GET to a URL on the cluster's network, returning the " + "response body. Useful when the agent is on a different network from " + "the cluster and cannot reach YARN tracking pages, Spark History " + "Server, or NodeManager web UIs directly.\n\n" + "**Security constraints:** the URL host must share at least 2 labels " + "of suffix with the named Connection's yarn_rm_url host (e.g. if " + "yarn_rm_url is 'rm.prod.internal:8088', you may fetch " + "'http://nm01.prod.internal:8042/...' but NOT 'http://evil.com/...'). " + "IP literals (10.0.0.1, ::1) and non-http(s) schemes (file://, " + "gopher://) are rejected. The Connection's saved auth is reused, so " + "the agent does not need cluster credentials.\n\n" + "**Limits:** response body capped at 1 MB (truncated=true if larger), " + "30s timeout, redirects followed." + ), +) +def _fetch_url(req: FetchUrlRequest): + return fetch_url(req.url, req.connection_name) diff --git a/spark_executor/tools/fetch_url.py b/spark_executor/tools/fetch_url.py new file mode 100644 index 0000000..6556613 --- /dev/null +++ b/spark_executor/tools/fetch_url.py @@ -0,0 +1,108 @@ +# coding=utf-8 +""" +@Time :2026/7/9 +@Author :tao.chen + +Generic HTTP GET proxy for the agent. Lets the agent fetch URLs on the +cluster's network when it cannot reach those hosts directly. Security: +URL host must share >= 2 labels of suffix with the connection's yarn_rm_url +host; IP literals and non-HTTP schemes are rejected. Reuses the connection's +saved auth/SSL config so the agent doesn't need cluster credentials. +""" +import ipaddress +from urllib.parse import urlparse + +import httpx + +from common.logging import logger +from spark_executor.core.yarn_client import YarnClientConfig +from spark_executor.models import FetchUrlResult +from spark_executor.tools.connections import store as conn_store + +_MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body +_REQUEST_TIMEOUT_SECONDS = 30 + + +def _host_suffix_overlap(host_a: str, host_b: str, min_labels: int = 2) -> bool: + """Return True if host_a and host_b share at least min_labels suffix labels.""" + labels_a = host_a.lower().split(".") + labels_b = host_b.lower().split(".") + n = 0 + i, j = len(labels_a) - 1, len(labels_b) - 1 + while i >= 0 and j >= 0 and labels_a[i] == labels_b[j]: + n += 1 + i -= 1 + j -= 1 + return n >= min_labels + + +def _validate_url_host(url: str, yarn_rm_url: str | None) -> None: + """Validate that url is an allowed http(s) hostname on the cluster network.""" + parsed = urlparse(url) + if parsed.scheme not in ("http", "https"): + raise ValueError(f"URL scheme must be http or https, got {parsed.scheme!r}") + + host = parsed.hostname + if not host: + raise ValueError(f"URL has no host: {url!r}") + + try: + ipaddress.ip_address(host) + except ValueError: + pass + else: + raise ValueError( + f"URL host {host!r} is an IP literal — IP targets are not allowed. " + "Use a hostname on the cluster network." + ) + + if not yarn_rm_url: + raise ValueError( + "Connection has no yarn_rm_url set, cannot validate URL host. " + "Save a Connection with yarn_rm_url first." + ) + + anchor = urlparse(yarn_rm_url).hostname or "" + if not _host_suffix_overlap(host, anchor, min_labels=2): + raise ValueError( + f"URL host {host!r} does not share enough suffix with the connection's " + f"yarn_rm_url host {anchor!r} (need >= 2 labels of common suffix). " + "Reject this fetch to prevent SSRF." + ) + + +def fetch_url(url: str, connection_name: str) -> FetchUrlResult: + """Proxy an HTTP GET to url using the auth/SSL settings of connection_name.""" + logger.debug(f"fetch_url enter url={url} connection_name={connection_name}") + conn = conn_store.get(connection_name) + if conn is None: + raise KeyError(f"Connection not found: {connection_name}") + + _validate_url_host(url, conn.yarn_rm_url) + + config = YarnClientConfig.from_connection(conn) + resp = httpx.get( + url, + auth=config.auth_for_httpx(), + verify=config.verify_for_httpx(), + timeout=_REQUEST_TIMEOUT_SECONDS, + follow_redirects=True, + ) + + body = resp.text[:_MAX_BODY_BYTES] + truncated = len(resp.text) > _MAX_BODY_BYTES + + logger.info( + f"fetch_url ok url={url} connection_name={connection_name} " + f"status_code={resp.status_code} " + f"content_type={resp.headers.get('content-type', '')} " + f"body_bytes={len(body)} truncated={truncated}" + ) + + return FetchUrlResult( + url=url, + status_code=resp.status_code, + content_type=resp.headers.get("content-type", ""), + body=body, + truncated=truncated, + ) diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index d9fbee7..20ca8da 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -378,3 +378,27 @@ class ExternalJobResultRequest(BaseModel): ..., description="Name of a saved Connection pointing at the YARN cluster.", ) + +class FetchUrlRequest(BaseModel): + url: str = Field( + ..., + description=( + "Absolute http:// or https:// URL to fetch. The host must share " + "at least 2 labels of suffix with the named Connection's " + "yarn_rm_url host (e.g. if yarn_rm_url is 'rm.prod.internal:8088', " + "you may fetch 'http://nm01.prod.internal:8042/...' but NOT " + "'http://evil.com/...' or 'http://10.0.0.1/...'). IP literals " + "and non-http(s) schemes are rejected. The Connection's saved " + "auth is reused for the outbound request — the agent does not " + "need cluster credentials." + ), + ) + connection_name: str = Field( + ..., + description=( + "Name of a saved Connection (see list_connections). The " + "Connection's yarn_rm_url defines the allowed host domain. " + "The Connection's auth_type / auth_user / auth_password / " + "ssl_verify / ssl_ca_bundle are reused for the request." + ), + ) diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 3b555b8..ae074b8 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -64,12 +64,13 @@ def test_seventeen_tool_routes_registered(): assert "/update_pending_job" in paths -def test_twenty_tool_routes_registered(): +def test_twenty_one_tool_routes_registered(): paths = {r.path for r in app.routes} for path in ( "/get_external_job_logs", "/get_external_job_status", "/get_external_job_result", + "/fetch_url", ): assert path in paths, f"missing MCP tool route: {path}" diff --git a/tests/unit/test_fetch_url.py b/tests/unit/test_fetch_url.py new file mode 100644 index 0000000..75b8e8f --- /dev/null +++ b/tests/unit/test_fetch_url.py @@ -0,0 +1,164 @@ +# coding=utf-8 +from unittest.mock import patch + +import httpx +import pytest + +from spark_executor.core import connection_store +from spark_executor.core.connection_store import ConnectionStore +from spark_executor.models import Connection +from spark_executor.tools import connections, fetch_url + + +def _fresh_stores(tmp_path, monkeypatch): + """Reset connection store singletons for a single test.""" + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + store = ConnectionStore() + monkeypatch.setattr(connection_store, "store", store) + connections.store = store + fetch_url.conn_store = store + + +def test_fetch_url_returns_body_and_status(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + ) + resp = httpx.Response(200, text="hello", headers={"content-type": "text/html"}) + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + out = fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") + assert out.url == "http://nm01.prod.internal:8042/node" + assert out.status_code == 200 + assert out.content_type == "text/html" + assert out.body == "hello" + assert out.truncated is False + assert m.call_count == 1 + + +def test_fetch_url_raises_for_missing_connection(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + with pytest.raises(KeyError, match="Connection not found"): + fetch_url.fetch_url("http://rm.prod.internal:8088/", "missing") + + +def test_fetch_url_rejects_non_http_scheme(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + ) + with pytest.raises(ValueError, match="scheme"): + fetch_url.fetch_url("file:///etc/passwd", "prod") + + +def test_fetch_url_rejects_ftp_scheme(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + ) + with pytest.raises(ValueError, match="scheme"): + fetch_url.fetch_url("ftp://rm.prod.internal/foo", "prod") + + +def test_fetch_url_rejects_ip_literal(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + ) + with pytest.raises(ValueError, match="IP literal"): + fetch_url.fetch_url("http://10.0.0.1/secrets", "prod") + + +def test_fetch_url_rejects_ipv6_literal(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + ) + with pytest.raises(ValueError, match="IP literal"): + fetch_url.fetch_url("http://[::1]:8080/", "prod") + + +def test_fetch_url_rejects_external_host(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + ) + with pytest.raises(ValueError, match="does not share enough suffix"): + fetch_url.fetch_url("http://evil.com/foo", "prod") + + +def test_fetch_url_rejects_too_short_suffix_overlap(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm/8088") + ) + with pytest.raises(ValueError, match="does not share enough suffix"): + fetch_url.fetch_url("http://other-rm/", "prod") + + +def test_fetch_url_rejects_when_connection_has_no_yarn_rm_url(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url=None)) + with pytest.raises(ValueError, match="no yarn_rm_url set"): + fetch_url.fetch_url("http://anything.com/", "prod") + + +def test_fetch_url_truncates_body_over_1mb(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + ) + big = "x" * (fetch_url._MAX_BODY_BYTES + 1) + resp = httpx.Response(200, text=big) + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp): + out = fetch_url.fetch_url("http://nm01.prod.internal/big", "prod") + assert out.truncated is True + assert len(out.body) == fetch_url._MAX_BODY_BYTES + + +def test_fetch_url_passes_auth_from_connection(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection( + name="auth", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + auth_type="basic", + auth_user="u", + auth_password="p", + ) + ) + resp = httpx.Response(200, text="ok") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "auth") + auth = m.call_args.kwargs["auth"] + assert isinstance(auth, httpx.BasicAuth) + import base64 + creds = base64.b64decode(auth._auth_header.split()[1]).decode() + assert creds == "u:p" + + +def test_fetch_url_passes_ssl_verify_from_connection(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection( + name="insecure", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + ssl_verify=False, + ) + ) + resp = httpx.Response(200, text="ok") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "insecure") + assert m.call_args.kwargs["verify"] is False + + +def test_fetch_url_follows_redirects(tmp_path, monkeypatch): + _fresh_stores(tmp_path, monkeypatch) + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + ) + resp = httpx.Response(200, text="ok") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") + assert m.call_args.kwargs["follow_redirects"] is True From 0291f36b013f4ba36214455bb2145ad98e8fa056 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 11:22:55 +0800 Subject: [PATCH 2/7] feat(fetch_url): per-Connection allowed_url_hosts glob allowlist MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The default 'URL host shares >= 2 labels of suffix with yarn_rm_url host' rule is too strict for clusters whose hostnames are single-label (e.g. 'ccam1' through 'ccam99'). The user's cluster is reachable at http://ccam1:8088, and they want fetch_url to work for any ccamN host — but ccam1 and ccam50 share 0 suffix labels, so the existing rule rejects everything. Add an opt-in allowlist field on Connection: allowed_url_hosts: list[str] | None Each entry is an fnmatch glob pattern. The URL host is allowed if it matches ANY pattern, regardless of the suffix rule. Save a Connection with ['ccam*'] to allow ccam1, ccam2, ..., ccam99. SSRF safety: the implementation does NOT use vanilla fnmatch on the flat host string (because '*' in fnmatch crosses '.', so 'ccam*' would match 'ccam50.evil.com' — a security hole). Instead, both the pattern and the host are split on '.' and matched LABEL-BY-LABEL with the label counts required to match exactly. So 'ccam*' matches 'ccam50' but NOT 'ccam50.evil.com' (different label counts). Test coverage: - 8 new tests in tests/unit/test_fetch_url.py (glob allow, glob deny, dot-boundary SSRF test, multiple globs, empty list, None fallback, error message hint) - All 21 fetch_url tests pass; 377 total. - spark_executor/models.py: Connection.allowed_url_hosts (with description) - spark_executor/tools/requests.py: SaveConnectionRequest.allowed_url_hosts - spark_executor/tools/fetch_url.py: new _host_matches_any_glob helper; _validate_url_host now takes allowed_hosts and checks globs BEFORE the suffix rule - tests/unit/test_fetch_url.py: 8 new tests - README.md: fetch_url row mentions the allowlist Co-Authored-By: Claude Fable 5 --- README.md | 2 +- spark_executor/models.py | 15 ++++ spark_executor/tools/fetch_url.py | 77 ++++++++++++++----- spark_executor/tools/requests.py | 16 +++- tests/unit/test_fetch_url.py | 122 ++++++++++++++++++++++++++++++ 5 files changed, 209 insertions(+), 23 deletions(-) diff --git a/README.md b/README.md index 94e76a2..a6c68d4 100644 --- a/README.md +++ b/README.md @@ -120,7 +120,7 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态（按 `application_id` + `connection_name`） | | `get_external_job_result` | 查询外部 YARN application 终态结果视图 | | `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 | -| `fetch_url` | 代理 HTTP GET 到集群内网 URL (受 `Connection.yarn_rm_url` host 围栏约束) | +| `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 围栏: 默认 2 段后缀匹配 `yarn_rm_url`, 或匹配 `Connection.allowed_url_hosts` glob) | ### Files MCP 工具 diff --git a/spark_executor/models.py b/spark_executor/models.py index ce3b2ab..bf96719 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -69,6 +69,21 @@ class Connection(BaseModel): auth_principal: str | None = None auth_keytab: str | None = None + allowed_url_hosts: list[str] | None = Field( + default=None, + description=( + "Optional list of fnmatch glob patterns for hosts the fetch_url tool " + "may access, in addition to the default 'shares >= 2 labels of suffix " + "with yarn_rm_url host' rule. Useful for clusters whose hostnames do " + "NOT share a 2+ label suffix — e.g. single-label hosts like 'ccam1'-" + "'ccam99' (configure ['ccam*']) or HDFS namenode on a different " + "subdomain ('*.hadoop.internal'). Patterns are matched against the " + "URL host only (no port, no path). fnmatch rules apply: '*' does NOT " + "match '.', so 'ccam*' matches 'ccam50' but not 'ccam50.evil.com'. " + "Default None means the suffix rule alone applies." + ), + ) + @field_validator("master") @classmethod def _check_master(cls, v: str) -> str: diff --git a/spark_executor/tools/fetch_url.py b/spark_executor/tools/fetch_url.py index 6556613..63921c6 100644 --- a/spark_executor/tools/fetch_url.py +++ b/spark_executor/tools/fetch_url.py @@ -10,6 +10,7 @@ host; IP literals and non-HTTP schemes are rejected. Reuses the connection's saved auth/SSL config so the agent doesn't need cluster credentials. """ import ipaddress +import fnmatch from urllib.parse import urlparse import httpx @@ -36,40 +37,78 @@ def _host_suffix_overlap(host_a: str, host_b: str, min_labels: int = 2) -> bool: return n >= min_labels -def _validate_url_host(url: str, yarn_rm_url: str | None) -> None: - """Validate that url is an allowed http(s) hostname on the cluster network.""" +def _host_matches_any_glob(host: str, patterns: list[str]) -> bool: + """True if `host` matches any of the fnmatch glob patterns. + + fnmatch is case-sensitive on Linux (our deployment target). `*` in a + pattern does NOT cross '.' boundaries, so 'ccam*' matches 'ccam50' but + NOT 'ccam50.evil.com' (which is what we want — domain boundary matters + for SSRF). + """ + host_labels = host.split(".") + for p in patterns: + pat_labels = p.split(".") + if len(pat_labels) != len(host_labels): + continue + if all(fnmatch.fnmatchcase(h, pat) for h, pat in zip(host_labels, pat_labels)): + return True + return False + + +def _validate_url_host( + url: str, + yarn_rm_url: str | None, + allowed_hosts: list[str] | None = None, +) -> None: + """Raise ValueError if the URL is not allowed to be fetched. + + Allowed if EITHER: + - URL host matches one of `allowed_hosts` fnmatch globs (if provided) + - URL host shares >= 2 labels of suffix with `yarn_rm_url` host + """ parsed = urlparse(url) if parsed.scheme not in ("http", "https"): - raise ValueError(f"URL scheme must be http or https, got {parsed.scheme!r}") - + raise ValueError( + f"URL scheme must be http or https, got {parsed.scheme!r}" + ) host = parsed.hostname if not host: raise ValueError(f"URL has no host: {url!r}") - + # IP literal check try: ipaddress.ip_address(host) - except ValueError: - pass - else: raise ValueError( f"URL host {host!r} is an IP literal — IP targets are not allowed. " - "Use a hostname on the cluster network." + f"Use a hostname on the cluster network." ) - + except ValueError as e: + if "IP literal" in str(e): + raise + # not an IP, continue + # Glob allowlist — if any pattern matches, host is allowed regardless of suffix + if allowed_hosts and _host_matches_any_glob(host, allowed_hosts): + return if not yarn_rm_url: raise ValueError( - "Connection has no yarn_rm_url set, cannot validate URL host. " - "Save a Connection with yarn_rm_url first." + f"Connection has no yarn_rm_url set, cannot validate URL host " + f"(no allowed_url_hosts match either). Save a Connection with " + f"yarn_rm_url or allowed_url_hosts set." ) - - anchor = urlparse(yarn_rm_url).hostname or "" + anchor = urlparse(yarn_rm_url).hostname + if not anchor: + raise ValueError(f"Connection's yarn_rm_url has no host: {yarn_rm_url!r}") if not _host_suffix_overlap(host, anchor, min_labels=2): + allowed_hint = ( + f" Or set allowed_url_hosts=[''] on the connection to allow " + f"this host (e.g. ['ccam*'] for ccam1-ccam99)." + if not allowed_hosts + else f" (no allowed_url_hosts pattern matched either)" + ) raise ValueError( - f"URL host {host!r} does not share enough suffix with the connection's " - f"yarn_rm_url host {anchor!r} (need >= 2 labels of common suffix). " - "Reject this fetch to prevent SSRF." + f"URL host {host!r} does not share enough suffix with the " + f"connection's yarn_rm_url host {anchor!r} (need >= 2 labels of " + f"common suffix).{allowed_hint} Reject this fetch to prevent SSRF." ) - def fetch_url(url: str, connection_name: str) -> FetchUrlResult: """Proxy an HTTP GET to url using the auth/SSL settings of connection_name.""" @@ -78,7 +117,7 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult: if conn is None: raise KeyError(f"Connection not found: {connection_name}") - _validate_url_host(url, conn.yarn_rm_url) + _validate_url_host(url, conn.yarn_rm_url, conn.allowed_url_hosts) config = YarnClientConfig.from_connection(conn) resp = httpx.get( diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 20ca8da..37a4d6b 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -117,12 +117,22 @@ class SaveConnectionRequest(BaseModel): description=( "Absolute path to a Kerberos keytab file. Optional convenience " "for 'kinit -kt' workflows. The service does NOT auto-initialize " - "from the keytab — you must `kinit -kt ` " - "yourself before calling the tools." + "from the keytab — you must `kinit -kt ` " + "yourself before calling the tools." + ), + ) + + allowed_url_hosts: list[str] | None = Field( + default=None, + description=( + "Optional list of fnmatch glob patterns for hosts the fetch_url tool " + "may access. See Connection.allowed_url_hosts for full semantics. " + "Example for single-label host clusters: ['ccam*'] allows any host " + "starting with 'ccam' (ccam1, ccam2, ..., ccam99). Default None means " + "only the default 'shares >= 2 labels of suffix with yarn_rm_url' rule." ), ) - class PrepareSubmitJobRequest(BaseModel): connection: str = Field( ..., diff --git a/tests/unit/test_fetch_url.py b/tests/unit/test_fetch_url.py index 75b8e8f..29fcc3a 100644 --- a/tests/unit/test_fetch_url.py +++ b/tests/unit/test_fetch_url.py @@ -1,4 +1,5 @@ # coding=utf-8 +from pathlib import Path from unittest.mock import patch import httpx @@ -19,6 +20,12 @@ def _fresh_stores(tmp_path, monkeypatch): fetch_url.conn_store = store +@pytest.fixture +def fresh_stores(tmp_path: Path, monkeypatch): + """Reset connection store singletons for a single test.""" + _fresh_stores(tmp_path, monkeypatch) + yield + def test_fetch_url_returns_body_and_status(tmp_path, monkeypatch): _fresh_stores(tmp_path, monkeypatch) fetch_url.conn_store.save( @@ -162,3 +169,118 @@ def test_fetch_url_follows_redirects(tmp_path, monkeypatch): with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") assert m.call_args.kwargs["follow_redirects"] is True + +def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="ccam", + master="yarn", + yarn_rm_url="http://ccam1:8088", + allowed_url_hosts=["ccam*"], + ) + ) + resp = httpx.Response(200, text="hello") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + out = fetch_url.fetch_url("http://ccam50:8088/foo", "ccam") + assert out.status_code == 200 + assert out.body == "hello" + assert m.call_count == 1 + +def test_fetch_url_allows_host_matching_any_of_multiple_globs(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + allowed_url_hosts=["ccam*", "*.prod.internal"], + ) + ) + resp = httpx.Response(200, text="hello") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + out = fetch_url.fetch_url("http://history.prod.internal:18080/api/v1/info", "prod") + assert out.status_code == 200 + assert out.body == "hello" + assert m.call_count == 1 + +def test_fetch_url_glob_does_not_match_unrelated_host(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="ccam", + master="yarn", + yarn_rm_url="http://ccam1:8088", + allowed_url_hosts=["ccam*"], + ) + ) + with pytest.raises(ValueError, match="does not share enough suffix|no allowed_url_hosts pattern matched"): + fetch_url.fetch_url("http://evil.com/foo", "ccam") + +def test_fetch_url_glob_does_not_cross_dot_boundary(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="ccam", + master="yarn", + yarn_rm_url="http://ccam1:8088", + allowed_url_hosts=["ccam*"], + ) + ) + with pytest.raises(ValueError, match="does not share enough suffix"): + fetch_url.fetch_url("http://ccam50.evil.com/", "ccam") + +def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm:8088", + allowed_url_hosts=["ccam*"], + ) + ) + resp = httpx.Response(200, text="hello") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + out = fetch_url.fetch_url("http://ccam50:8088/", "prod") + assert out.status_code == 200 + assert out.body == "hello" + assert m.call_count == 1 + +def test_fetch_url_empty_allowed_hosts_falls_back_to_suffix_rule(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal", + allowed_url_hosts=[], + ) + ) + resp = httpx.Response(200, text="hello") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + out = fetch_url.fetch_url("http://nm.prod.internal/", "prod") + assert out.status_code == 200 + assert out.body == "hello" + assert m.call_count == 1 + +def test_fetch_url_none_allowed_hosts_falls_back_to_suffix_rule(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal", + allowed_url_hosts=None, + ) + ) + resp = httpx.Response(200, text="hello") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + out = fetch_url.fetch_url("http://nm.prod.internal/", "prod") + assert out.status_code == 200 + assert out.body == "hello" + assert m.call_count == 1 + +def test_fetch_url_error_message_hints_at_allowed_url_hosts(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + ) + ) + with pytest.raises(ValueError, match="allowed_url_hosts"): + fetch_url.fetch_url("http://evil.com/foo", "prod") From 6d7387b022955f71d5c6ce10dca857cda60ff9e9 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 11:46:07 +0800 Subject: [PATCH 3/7] feat: simplify fetch_url allowlist + add update_connection tool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part 1 — fetch_url validation simplification Now that Connection.allowed_url_hosts exists, it's the ONLY check. The old 2-label suffix-overlap rule against yarn_rm_url is gone. - Connection.allowed_url_hosts: list[str] = Field(default_factory=list) (was list[str] | None = None). Default empty list means the connection has no fetch access; the user must explicitly opt in via save_connection or update_connection. - _validate_url_host drops the yarn_rm_url parameter and the _host_suffix_overlap helper. Now: scheme/host/IP checks, then empty-reject, then glob match. Reject everything else. - fetch_url's error message now points the user at Connection.allowed_url_hosts as the fix. - The label-by-label glob check from the previous commit is kept (SSRF guard: 'ccam*' matches 'ccam50' but NOT 'ccam50.evil.com'). Part 2 — update_connection tool PATCH-style update for an existing Connection record. Only the fields the caller provides are changed. Same pattern as update_pending_job: req.model_dump(exclude_none=True), with 'name' popped before passing to the store. To CLEAR a field (e.g. drop auth_password), use delete_connection + save_connection. This is YAGNI; the alternative (model_fields_set to distinguish 'omitted' from 'null') adds surface for bugs. - ConnectionStore.update(name, **fields): get, model_copy(update=...), save. Lock + atomic write. Re-validates the patched record. - update_connection tool function: passes fields through to the store, logs which fields were changed. - UpdateConnectionRequest Pydantic model: 12 mutable fields + name. - /update_connection route with operation_id='update_connection'. - 9 new unit tests in test_connection_tools.py (PATCH, dict/list replace-not-merge, unknown name 404, validation of patched record, disk persistence). - test_fetch_url.py: existing 21 tests updated; 3 new tests for the empty/None/missing allowed_url_hosts cases. - test_mcp_routes.py: assert 22 tool routes. - README: update_connection row added; fetch_url row updated. Tests: 386 passed (up from 377, +9 net). Co-Authored-By: Claude Fable 5 --- README.md | 11 ++-- spark_executor/core/connection_store.py | 23 ++++++++ spark_executor/models.py | 22 ++++---- spark_executor/server.py | 24 ++++++++ spark_executor/tools/connections.py | 17 ++++++ spark_executor/tools/fetch_url.py | 64 +++++++-------------- spark_executor/tools/requests.py | 68 ++++++++++++++++++++++- tests/integration/test_mcp_routes.py | 3 +- tests/unit/test_connection_tools.py | 66 ++++++++++++++++++++++ tests/unit/test_fetch_url.py | 74 +++++++++++++++---------- 10 files changed, 279 insertions(+), 93 deletions(-) diff --git a/README.md b/README.md index a6c68d4..ac442d4 100644 --- a/README.md +++ b/README.md @@ -100,10 +100,11 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | 工具 | 作用 | | --- | --- | -| `save_connection` | 保存或更新一个 Spark/YARN 连接配置 | -| `list_connections` | 列出所有连接配置 | -| `get_connection` | 按名称读取连接配置 | -| `delete_connection` | 删除连接配置 | + | `save_connection` | 保存或更新一个 Spark/YARN 连接配置 | + | `list_connections` | 列出所有连接配置 | + | `get_connection` | 按名称读取连接配置 | + | `delete_connection` | 删除连接配置 | + | `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) | | `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 | | `read_job_file` | 读取已存在的 PySpark 脚本内容 | | `update_job_file` | 覆盖更新已存在的 PySpark 脚本 | @@ -120,7 +121,7 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态（按 `application_id` + `connection_name`） | | `get_external_job_result` | 查询外部 YARN application 终态结果视图 | | `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 | -| `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 围栏: 默认 2 段后缀匹配 `yarn_rm_url`, 或匹配 `Connection.allowed_url_hosts` glob) | + | `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 受 `Connection.allowed_url_hosts` glob allowlist 约束, 空则全拒) | ### Files MCP 工具 diff --git a/spark_executor/core/connection_store.py b/spark_executor/core/connection_store.py index 6f54bb5..9adbd29 100644 --- a/spark_executor/core/connection_store.py +++ b/spark_executor/core/connection_store.py @@ -77,6 +77,29 @@ class ConnectionStore: self._dump(records) logger.info(f"connection saved name={conn.name} master={conn.master}") + def update(self, name: str, **fields) -> Connection: + """Apply `fields` to the existing Connection identified by `name` and persist. + + PATCH semantics: only fields explicitly passed in `fields` are changed. + Use Pydantic's `model_copy(update=fields)` to apply the patch. + + Raises KeyError if no Connection with `name` exists. + Raises pydantic.ValidationError if the patched Connection is invalid + (e.g. `master='http://...'` fails the master validator). + """ + with self._lock: + records = self._load() + if name not in records: + raise KeyError(f"Connection not found: {name}") + existing = records[name] + patched = Connection.model_validate(existing.model_copy(update=fields).model_dump()) + records[name] = patched + self._dump(records) + logger.info( + f"connection updated name={name} fields={sorted(fields.keys())}" + ) + return patched + def delete(self, name: str) -> bool: with self._lock: records = self._load() diff --git a/spark_executor/models.py b/spark_executor/models.py index bf96719..7f64c3a 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -69,18 +69,18 @@ class Connection(BaseModel): auth_principal: str | None = None auth_keytab: str | None = None - allowed_url_hosts: list[str] | None = Field( - default=None, + allowed_url_hosts: list[str] = Field( + default_factory=list, description=( - "Optional list of fnmatch glob patterns for hosts the fetch_url tool " - "may access, in addition to the default 'shares >= 2 labels of suffix " - "with yarn_rm_url host' rule. Useful for clusters whose hostnames do " - "NOT share a 2+ label suffix — e.g. single-label hosts like 'ccam1'-" - "'ccam99' (configure ['ccam*']) or HDFS namenode on a different " - "subdomain ('*.hadoop.internal'). Patterns are matched against the " - "URL host only (no port, no path). fnmatch rules apply: '*' does NOT " - "match '.', so 'ccam*' matches 'ccam50' but not 'ccam50.evil.com'. " - "Default None means the suffix rule alone applies." + "List of fnmatch glob patterns for hosts the fetch_url tool may access. " + "The list is mandatory-opt-in: an empty list (the default) denies all " + "hosts, so you must populate it before fetch_url can access any URL. " + "Useful for clusters whose hostnames do NOT share a common suffix — " + "e.g. single-label hosts like 'ccam1'-'ccam99' (configure ['ccam*']) " + "or HDFS namenode on a different subdomain ('*.hadoop.internal'). " + "Patterns are matched against the URL host only (no port, no path). " + "fnmatch rules apply: '*' does NOT match '.', so 'ccam*' matches " + "'ccam50' but not 'ccam50.evil.com'." ), ) diff --git a/spark_executor/server.py b/spark_executor/server.py index f772426..a3ca21f 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -11,6 +11,7 @@ from spark_executor.tools.connections import ( get_connection, list_connections, save_connection, + update_connection, ) from spark_executor.tools.write_job import write_job_file from spark_executor.tools.job_file import read_job_file, update_job_file @@ -36,6 +37,7 @@ from spark_executor.tools.requests import ( PrepareSubmitJobRequest, ReadJobFileRequest, SaveConnectionRequest, + UpdateConnectionRequest, UpdateJobFileRequest, UpdatePendingJobRequest, ) @@ -396,6 +398,28 @@ def _get_connection(req: ConnectionNameRequest): return get_connection(req.name) +@app.post( + "/update_connection", + operation_id="update_connection", + summary="Update an existing connection's fields", + description=( + "Apply a partial update (PATCH) to an existing Connection record. " + "Only the fields you provide are changed; the rest are kept as-is. " + "The `name` is the immutable identifier (use delete_connection + " + "save_connection to rename).\n\n" + "To CLEAR an optional field (e.g. remove `yarn_rm_url`), use " + "delete_connection followed by save_connection with the field omitted. " + "This tool cannot clear fields — only replace them.\n\n" + "Returns the full updated Connection record. 404 if no Connection " + "with the given name exists." + ), +) +def _update_connection(req: UpdateConnectionRequest): + fields = req.model_dump(exclude_none=True) + fields.pop("name", None) # name is the identity, not a field to patch + return update_connection(name=req.name, **fields) + + @app.post( "/delete_connection", operation_id="delete_connection", diff --git a/spark_executor/tools/connections.py b/spark_executor/tools/connections.py index ae381f8..7a23765 100644 --- a/spark_executor/tools/connections.py +++ b/spark_executor/tools/connections.py @@ -22,6 +22,7 @@ def save_connection( auth_password: str | None = None, auth_principal: str | None = None, auth_keytab: str | None = None, + allowed_url_hosts: list[str] | None = None, ) -> dict[str, str]: logger.debug( f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " @@ -40,11 +41,27 @@ def save_connection( auth_password=auth_password, auth_principal=auth_principal, auth_keytab=auth_keytab, + allowed_url_hosts=allowed_url_hosts or [], ) store.save(conn) return {"name": name, "status": "SAVED"} +def update_connection(name: str, **fields) -> dict[str, object]: + """Update an existing Connection's mutable fields. + + `name` is the identifier (immutable). PATCH semantics: only the fields + you pass are changed. To clear an optional field (e.g. `yarn_rm_url`), + use `delete_connection(name=...)` followed by `save_connection(...)`. + + Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf, + ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password, + auth_principal, auth_keytab, allowed_url_hosts. + """ + logger.debug(f"update_connection enter name={name} fields={sorted(fields.keys())}") + return store.update(name, **fields).model_dump() + + def list_connections() -> list[dict[str, object]]: logger.debug("list_connections enter") return [c.model_dump() for c in store.list_all()] diff --git a/spark_executor/tools/fetch_url.py b/spark_executor/tools/fetch_url.py index 63921c6..360c672 100644 --- a/spark_executor/tools/fetch_url.py +++ b/spark_executor/tools/fetch_url.py @@ -5,9 +5,10 @@ Generic HTTP GET proxy for the agent. Lets the agent fetch URLs on the cluster's network when it cannot reach those hosts directly. Security: -URL host must share >= 2 labels of suffix with the connection's yarn_rm_url -host; IP literals and non-HTTP schemes are rejected. Reuses the connection's -saved auth/SSL config so the agent doesn't need cluster credentials. +the host is checked against an explicit fnmatch glob allowlist configured +on the connection (`allowed_url_hosts`). Empty or missing allowlist means +no URL access; IP literals and non-HTTP schemes are rejected. Reuses the +connection's saved auth/SSL config so the agent doesn't need cluster credentials. """ import ipaddress import fnmatch @@ -24,19 +25,6 @@ _MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body _REQUEST_TIMEOUT_SECONDS = 30 -def _host_suffix_overlap(host_a: str, host_b: str, min_labels: int = 2) -> bool: - """Return True if host_a and host_b share at least min_labels suffix labels.""" - labels_a = host_a.lower().split(".") - labels_b = host_b.lower().split(".") - n = 0 - i, j = len(labels_a) - 1, len(labels_b) - 1 - while i >= 0 and j >= 0 and labels_a[i] == labels_b[j]: - n += 1 - i -= 1 - j -= 1 - return n >= min_labels - - def _host_matches_any_glob(host: str, patterns: list[str]) -> bool: """True if `host` matches any of the fnmatch glob patterns. @@ -57,14 +45,12 @@ def _host_matches_any_glob(host: str, patterns: list[str]) -> bool: def _validate_url_host( url: str, - yarn_rm_url: str | None, - allowed_hosts: list[str] | None = None, + allowed_hosts: list[str] | None, ) -> None: """Raise ValueError if the URL is not allowed to be fetched. - Allowed if EITHER: - - URL host matches one of `allowed_hosts` fnmatch globs (if provided) - - URL host shares >= 2 labels of suffix with `yarn_rm_url` host + Allowed only if the URL host matches one of the fnmatch glob patterns in + `allowed_hosts`. An empty or missing allowlist rejects everything. """ parsed = urlparse(url) if parsed.scheme not in ("http", "https"): @@ -85,30 +71,18 @@ def _validate_url_host( if "IP literal" in str(e): raise # not an IP, continue - # Glob allowlist — if any pattern matches, host is allowed regardless of suffix - if allowed_hosts and _host_matches_any_glob(host, allowed_hosts): + if not allowed_hosts: + raise ValueError( + "allowed_url_hosts is empty — set Connection.allowed_url_hosts to " + "allow specific hosts before calling fetch_url (e.g. ['ccam*'] " + "for ccam1-ccam99 or ['*.prod.internal'] for a subdomain)." + ) + if _host_matches_any_glob(host, allowed_hosts): return - if not yarn_rm_url: - raise ValueError( - f"Connection has no yarn_rm_url set, cannot validate URL host " - f"(no allowed_url_hosts match either). Save a Connection with " - f"yarn_rm_url or allowed_url_hosts set." - ) - anchor = urlparse(yarn_rm_url).hostname - if not anchor: - raise ValueError(f"Connection's yarn_rm_url has no host: {yarn_rm_url!r}") - if not _host_suffix_overlap(host, anchor, min_labels=2): - allowed_hint = ( - f" Or set allowed_url_hosts=[''] on the connection to allow " - f"this host (e.g. ['ccam*'] for ccam1-ccam99)." - if not allowed_hosts - else f" (no allowed_url_hosts pattern matched either)" - ) - raise ValueError( - f"URL host {host!r} does not share enough suffix with the " - f"connection's yarn_rm_url host {anchor!r} (need >= 2 labels of " - f"common suffix).{allowed_hint} Reject this fetch to prevent SSRF." - ) + raise ValueError( + f"URL host {host!r} is not in Connection.allowed_url_hosts " + f"{allowed_hosts!r}. Reject this fetch to prevent SSRF." + ) def fetch_url(url: str, connection_name: str) -> FetchUrlResult: """Proxy an HTTP GET to url using the auth/SSL settings of connection_name.""" @@ -117,7 +91,7 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult: if conn is None: raise KeyError(f"Connection not found: {connection_name}") - _validate_url_host(url, conn.yarn_rm_url, conn.allowed_url_hosts) + _validate_url_host(url, conn.allowed_url_hosts) config = YarnClientConfig.from_connection(conn) resp = httpx.get( diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 37a4d6b..08e482e 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -128,8 +128,10 @@ class SaveConnectionRequest(BaseModel): "Optional list of fnmatch glob patterns for hosts the fetch_url tool " "may access. See Connection.allowed_url_hosts for full semantics. " "Example for single-label host clusters: ['ccam*'] allows any host " - "starting with 'ccam' (ccam1, ccam2, ..., ccam99). Default None means " - "only the default 'shares >= 2 labels of suffix with yarn_rm_url' rule." + "starting with 'ccam' (ccam1, ccam2, ..., ccam99). If omitted, None, " + "or empty, the saved connection will have allowed_url_hosts=[] (the " + "default), meaning fetch_url will reject every URL until the list is " + "populated via update_connection." ), ) @@ -412,3 +414,65 @@ class FetchUrlRequest(BaseModel): "ssl_verify / ssl_ca_bundle are reused for the request." ), ) + + +class UpdateConnectionRequest(BaseModel): + name: str = Field( + ..., + description=( + "Name of the existing Connection to update (immutable identifier). " + "If you want to rename, use delete_connection + save_connection." + ), + ) + master: str | None = Field( + default=None, + description="New Spark master URL. Omit to keep current.", + ) + deploy_mode: str | None = Field( + default=None, + description="New Spark deploy mode ('cluster' or 'client'). Omit to keep current.", + ) + yarn_rm_url: str | None = Field( + default=None, + description="New YARN ResourceManager REST base URL. Omit to keep current.", + ) + spark_conf: dict[str, str] | None = Field( + default=None, + description="Replacement Spark conf dict (not merged with existing). Omit to keep current.", + ) + ssl_verify: bool | None = Field( + default=None, + description="New SSL verify setting. Omit to keep current.", + ) + ssl_ca_bundle: str | None = Field( + default=None, + description="New SSL CA bundle path. Omit to keep current.", + ) + auth_type: str | None = Field( + default=None, + description="New auth mode. Omit to keep current.", + ) + auth_user: str | None = Field( + default=None, + description="New auth username. Omit to keep current.", + ) + auth_password: str | None = Field( + default=None, + description="New auth password. Omit to keep current.", + ) + auth_principal: str | None = Field( + default=None, + description="New Kerberos principal. Omit to keep current.", + ) + auth_keytab: str | None = Field( + default=None, + description="New Kerberos keytab path. Omit to keep current.", + ) + allowed_url_hosts: list[str] | None = Field( + default=None, + description=( + "Replacement allowed_url_hosts list (not merged). Omit to keep current. " + "Pass an empty list to deny all hosts (the default for new connections). " + "Example: ['ccam*'] allows ccam1-ccam99." + ), + ) diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index ae074b8..6b1e091 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -64,13 +64,14 @@ def test_seventeen_tool_routes_registered(): assert "/update_pending_job" in paths -def test_twenty_one_tool_routes_registered(): +def test_twenty_two_tool_routes_registered(): paths = {r.path for r in app.routes} for path in ( "/get_external_job_logs", "/get_external_job_status", "/get_external_job_result", "/fetch_url", + "/update_connection", ): assert path in paths, f"missing MCP tool route: {path}" diff --git a/tests/unit/test_connection_tools.py b/tests/unit/test_connection_tools.py index 3b73af2..74ce480 100644 --- a/tests/unit/test_connection_tools.py +++ b/tests/unit/test_connection_tools.py @@ -1,4 +1,5 @@ # coding=utf-8 +import json from pathlib import Path import pytest @@ -120,3 +121,68 @@ def test_delete_connection_returns_status(_fresh_store): def test_delete_connection_unknown_raises(_fresh_store): with pytest.raises(KeyError): connections.delete_connection("missing") +def test_update_connection_changes_specified_field(_fresh_store): + connections.save_connection(name="prod", master="yarn") + out = connections.update_connection(name="prod", master="spark://new:7077") + assert out["master"] == "spark://new:7077" + assert _fresh_store.get("prod").master == "spark://new:7077" + + +def test_update_connection_keeps_omitted_fields(_fresh_store): + connections.save_connection( + name="prod", + master="yarn", + deploy_mode="client", + yarn_rm_url="http://rm:8088", + ) + out = connections.update_connection(name="prod", deploy_mode="cluster") + assert out["deploy_mode"] == "cluster" + assert out["master"] == "yarn" + assert out["yarn_rm_url"] == "http://rm:8088" + + +def test_update_connection_with_no_fields_is_noop(_fresh_store): + connections.save_connection(name="prod", master="yarn", deploy_mode="client") + out = connections.update_connection(name="prod") + assert out["master"] == "yarn" + assert out["deploy_mode"] == "client" + + +def test_update_connection_allows_url_hosts_change(_fresh_store): + connections.save_connection(name="prod", master="yarn") + out = connections.update_connection(name="prod", allowed_url_hosts=["ccam*"]) + assert out["allowed_url_hosts"] == ["ccam*"] + + +def test_update_connection_replaces_not_merges_dict(_fresh_store): + connections.save_connection(name="prod", master="yarn", spark_conf={"a": "1"}) + out = connections.update_connection(name="prod", spark_conf={"b": "2"}) + assert out["spark_conf"] == {"b": "2"} + + +def test_update_connection_replaces_not_merges_list(_fresh_store): + connections.save_connection( + name="prod", + master="yarn", + allowed_url_hosts=["ccam*"], + ) + out = connections.update_connection(name="prod", allowed_url_hosts=["nm*"]) + assert out["allowed_url_hosts"] == ["nm*"] + + +def test_update_connection_raises_for_unknown_name(_fresh_store): + with pytest.raises(KeyError, match="Connection not found"): + connections.update_connection(name="missing", master="yarn") + + +def test_update_connection_validates_patched_connection(_fresh_store): + connections.save_connection(name="prod", master="yarn") + with pytest.raises(ValueError, match="master must be"): + connections.update_connection(name="prod", master="http://bad") + + +def test_update_connection_persists_to_disk(_fresh_store, tmp_path): + connections.save_connection(name="prod", master="yarn") + connections.update_connection(name="prod", master="spark://new:7077") + raw = json.loads((tmp_path / "connections.json").read_text()) + assert raw["prod"]["master"] == "spark://new:7077" diff --git a/tests/unit/test_fetch_url.py b/tests/unit/test_fetch_url.py index 29fcc3a..75723bd 100644 --- a/tests/unit/test_fetch_url.py +++ b/tests/unit/test_fetch_url.py @@ -29,7 +29,12 @@ def fresh_stores(tmp_path: Path, monkeypatch): def test_fetch_url_returns_body_and_status(tmp_path, monkeypatch): _fresh_stores(tmp_path, monkeypatch) fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + allowed_url_hosts=["*.prod.internal"], + ) ) resp = httpx.Response(200, text="hello", headers={"content-type": "text/html"}) with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: @@ -87,32 +92,45 @@ def test_fetch_url_rejects_ipv6_literal(tmp_path, monkeypatch): def test_fetch_url_rejects_external_host(tmp_path, monkeypatch): _fresh_stores(tmp_path, monkeypatch) fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + allowed_url_hosts=["*.prod.internal"], + ) ) - with pytest.raises(ValueError, match="does not share enough suffix"): + with pytest.raises(ValueError, match="is not in Connection.allowed_url_hosts"): fetch_url.fetch_url("http://evil.com/foo", "prod") -def test_fetch_url_rejects_too_short_suffix_overlap(tmp_path, monkeypatch): +def test_fetch_url_rejects_unrelated_single_label_host(tmp_path, monkeypatch): + """Without a common suffix rule, a single-label RM host no longer helps.""" _fresh_stores(tmp_path, monkeypatch) fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm/8088") + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") ) - with pytest.raises(ValueError, match="does not share enough suffix"): + with pytest.raises(ValueError, match="allowed_url_hosts is empty"): fetch_url.fetch_url("http://other-rm/", "prod") -def test_fetch_url_rejects_when_connection_has_no_yarn_rm_url(tmp_path, monkeypatch): +def test_fetch_url_rejects_when_allowed_url_hosts_empty(tmp_path, monkeypatch): _fresh_stores(tmp_path, monkeypatch) - fetch_url.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url=None)) - with pytest.raises(ValueError, match="no yarn_rm_url set"): + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", allowed_url_hosts=[]) + ) + with pytest.raises(ValueError, match="allowed_url_hosts is empty"): fetch_url.fetch_url("http://anything.com/", "prod") def test_fetch_url_truncates_body_over_1mb(tmp_path, monkeypatch): _fresh_stores(tmp_path, monkeypatch) fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + allowed_url_hosts=["*.prod.internal"], + ) ) big = "x" * (fetch_url._MAX_BODY_BYTES + 1) resp = httpx.Response(200, text=big) @@ -132,6 +150,7 @@ def test_fetch_url_passes_auth_from_connection(tmp_path, monkeypatch): auth_type="basic", auth_user="u", auth_password="p", + allowed_url_hosts=["*.prod.internal"], ) ) resp = httpx.Response(200, text="ok") @@ -152,6 +171,7 @@ def test_fetch_url_passes_ssl_verify_from_connection(tmp_path, monkeypatch): master="yarn", yarn_rm_url="http://rm.prod.internal:8088", ssl_verify=False, + allowed_url_hosts=["*.prod.internal"], ) ) resp = httpx.Response(200, text="ok") @@ -163,7 +183,12 @@ def test_fetch_url_passes_ssl_verify_from_connection(tmp_path, monkeypatch): def test_fetch_url_follows_redirects(tmp_path, monkeypatch): _fresh_stores(tmp_path, monkeypatch) fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + allowed_url_hosts=["*.prod.internal"], + ) ) resp = httpx.Response(200, text="ok") with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: @@ -211,7 +236,7 @@ def test_fetch_url_glob_does_not_match_unrelated_host(fresh_stores): allowed_url_hosts=["ccam*"], ) ) - with pytest.raises(ValueError, match="does not share enough suffix|no allowed_url_hosts pattern matched"): + with pytest.raises(ValueError, match="is not in Connection.allowed_url_hosts"): fetch_url.fetch_url("http://evil.com/foo", "ccam") def test_fetch_url_glob_does_not_cross_dot_boundary(fresh_stores): @@ -223,7 +248,7 @@ def test_fetch_url_glob_does_not_cross_dot_boundary(fresh_stores): allowed_url_hosts=["ccam*"], ) ) - with pytest.raises(ValueError, match="does not share enough suffix"): + with pytest.raises(ValueError, match="is not in Connection.allowed_url_hosts"): fetch_url.fetch_url("http://ccam50.evil.com/", "ccam") def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): @@ -242,7 +267,7 @@ def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): assert out.body == "hello" assert m.call_count == 1 -def test_fetch_url_empty_allowed_hosts_falls_back_to_suffix_rule(fresh_stores): +def test_fetch_url_empty_allowed_hosts_rejects_everything(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", @@ -251,28 +276,19 @@ def test_fetch_url_empty_allowed_hosts_falls_back_to_suffix_rule(fresh_stores): allowed_url_hosts=[], ) ) - resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: - out = fetch_url.fetch_url("http://nm.prod.internal/", "prod") - assert out.status_code == 200 - assert out.body == "hello" - assert m.call_count == 1 + with pytest.raises(ValueError, match="allowed_url_hosts is empty"): + fetch_url.fetch_url("http://nm.prod.internal/", "prod") -def test_fetch_url_none_allowed_hosts_falls_back_to_suffix_rule(fresh_stores): +def test_fetch_url_omitted_allowed_url_hosts_defaults_to_empty_and_rejects(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal", - allowed_url_hosts=None, ) ) - resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: - out = fetch_url.fetch_url("http://nm.prod.internal/", "prod") - assert out.status_code == 200 - assert out.body == "hello" - assert m.call_count == 1 + with pytest.raises(ValueError, match="allowed_url_hosts is empty"): + fetch_url.fetch_url("http://nm.prod.internal/", "prod") def test_fetch_url_error_message_hints_at_allowed_url_hosts(fresh_stores): fetch_url.conn_store.save( @@ -282,5 +298,5 @@ def test_fetch_url_error_message_hints_at_allowed_url_hosts(fresh_stores): yarn_rm_url="http://rm.prod.internal:8088", ) ) - with pytest.raises(ValueError, match="allowed_url_hosts"): + with pytest.raises(ValueError, match="set Connection.allowed_url_hosts"): fetch_url.fetch_url("http://evil.com/foo", "prod") From a5b9539663016b78cf38a62d04718dd7b5416b67 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 12:09:41 +0800 Subject: [PATCH 4/7] refactor(fetch_url): trust the allowlist, rename to url_allowlist MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two changes: 1) Drop every check except the allowlist lookup. Old _validate_url_host did: scheme check, host-presence check, IP-literal check, empty-allowlist check, then glob match. New _validate_url_host does: parse host, return on glob match, raise on miss. That's it. The only remaining structural check is 'the URL must have a host' (otherwise the glob has nothing to test against). Security implication: scheme (file://, gopher://, ftp://) and IP literals (10.0.0.1, ::1) are NO LONGER rejected by the validator. The allowlist is the single source of truth. If the user writes ['*.*.*.*'], they have opted in to 4-label hosts including IP literals; if they write ['ccam*'], they get ccam1- ccam99 and nothing else. The default ['ccam*'] / [] pattern is tight by construction. Removed: import ipaddress, the scheme/IP rejection branches, the 'allowlist empty' explicit branch (the empty list naturally matches nothing). 2) Rename allowed_url_hosts -> url_allowlist. The previous name was a verbose double-negative ('allowed ... hosts'). The new name is short, modern (allowlist > whitelist), and matches the pattern of the field (URL hosts allowed). Renamed in: - Connection (models.py) - SaveConnectionRequest, UpdateConnectionRequest, FetchUrlRequest (requests.py) - _validate_url_host, _host_matches_any_glob parameters (fetch_url.py) - save_connection / update_connection call sites and error messages (connections.py, fetch_url.py) - Route descriptions (server.py) - All test files - README No backward-compat alias: the field was added in 0291f36 and hasn't shipped, so no production migration. Local dev data (data/connections.json, gitignored) with allowed_url_hosts set will be silently dropped by Pydantic v2 (default for extra fields is ignore) — those connections lose their allowlist and fetch_url will reject everything until re-saved. Test cleanup: - Removed 9 obsolete tests (scheme/IP/suffix rejection) - Renamed allowed_url_hosts -> url_allowlist in 11 surviving tests - Added 4 new tests documenting the 'allowlist is the only gate' model: IP literal accepted, HTTPS accepted, no-host rejected, empty allowlist rejected, error message mentions url_allowlist -1 obsolete test, net -4 from 386 -> 382 tests passing. Co-Authored-By: Claude Fable 5 --- README.md | 2 +- spark_executor/models.py | 2 +- spark_executor/server.py | 13 +- spark_executor/tools/connections.py | 22 ++-- spark_executor/tools/fetch_url.py | 63 ++++------ spark_executor/tools/requests.py | 27 ++--- tests/unit/test_connection_tools.py | 14 +-- tests/unit/test_fetch_url.py | 181 ++++++++++++---------------- 8 files changed, 138 insertions(+), 186 deletions(-) diff --git a/README.md b/README.md index ac442d4..eeaabbb 100644 --- a/README.md +++ b/README.md @@ -121,7 +121,7 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态（按 `application_id` + `connection_name`） | | `get_external_job_result` | 查询外部 YARN application 终态结果视图 | | `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 | - | `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 受 `Connection.allowed_url_hosts` glob allowlist 约束, 空则全拒) | +| `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 受 `Connection.url_allowlist` glob allowlist 约束, 空则全拒) | ### Files MCP 工具 diff --git a/spark_executor/models.py b/spark_executor/models.py index 7f64c3a..1fa327a 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -69,7 +69,7 @@ class Connection(BaseModel): auth_principal: str | None = None auth_keytab: str | None = None - allowed_url_hosts: list[str] = Field( + url_allowlist: list[str] = Field( default_factory=list, description=( "List of fnmatch glob patterns for hosts the fetch_url tool may access. " diff --git a/spark_executor/server.py b/spark_executor/server.py index a3ca21f..9a48a19 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -505,13 +505,12 @@ def _update_job_file(req: UpdateJobFileRequest): "response body. Useful when the agent is on a different network from " "the cluster and cannot reach YARN tracking pages, Spark History " "Server, or NodeManager web UIs directly.\n\n" - "**Security constraints:** the URL host must share at least 2 labels " - "of suffix with the named Connection's yarn_rm_url host (e.g. if " - "yarn_rm_url is 'rm.prod.internal:8088', you may fetch " - "'http://nm01.prod.internal:8042/...' but NOT 'http://evil.com/...'). " - "IP literals (10.0.0.1, ::1) and non-http(s) schemes (file://, " - "gopher://) are rejected. The Connection's saved auth is reused, so " - "the agent does not need cluster credentials.\n\n" + "**Security constraints:** the URL host must match one of the fnmatch " + "glob patterns in the named Connection's url_allowlist. An empty or " + "omitted allowlist denies every host. There are no scheme or IP-literal " + "guardrails — the allowlist is the only gate — so keep it tight. The " + "Connection's saved auth is reused, so the agent does not need cluster " + "credentials.\n\n" "**Limits:** response body capped at 1 MB (truncated=true if larger), " "30s timeout, redirects followed." ), diff --git a/spark_executor/tools/connections.py b/spark_executor/tools/connections.py index 7a23765..0af462b 100644 --- a/spark_executor/tools/connections.py +++ b/spark_executor/tools/connections.py @@ -19,10 +19,10 @@ def save_connection( ssl_ca_bundle: str | None = None, auth_type: str = "none", auth_user: str | None = None, - auth_password: str | None = None, - auth_principal: str | None = None, - auth_keytab: str | None = None, - allowed_url_hosts: list[str] | None = None, + auth_password: str | None = None, + auth_principal: str | None = None, + auth_keytab: str | None = None, + url_allowlist: list[str] | None = None, ) -> dict[str, str]: logger.debug( f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " @@ -39,10 +39,10 @@ def save_connection( auth_type=auth_type, auth_user=auth_user, auth_password=auth_password, - auth_principal=auth_principal, - auth_keytab=auth_keytab, - allowed_url_hosts=allowed_url_hosts or [], - ) + auth_principal=auth_principal, + auth_keytab=auth_keytab, + url_allowlist=url_allowlist or [], + ) store.save(conn) return {"name": name, "status": "SAVED"} @@ -54,9 +54,9 @@ def update_connection(name: str, **fields) -> dict[str, object]: you pass are changed. To clear an optional field (e.g. `yarn_rm_url`), use `delete_connection(name=...)` followed by `save_connection(...)`. - Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf, - ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password, - auth_principal, auth_keytab, allowed_url_hosts. + Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf, + ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password, + auth_principal, auth_keytab, url_allowlist. """ logger.debug(f"update_connection enter name={name} fields={sorted(fields.keys())}") return store.update(name, **fields).model_dump() diff --git a/spark_executor/tools/fetch_url.py b/spark_executor/tools/fetch_url.py index 360c672..c505b88 100644 --- a/spark_executor/tools/fetch_url.py +++ b/spark_executor/tools/fetch_url.py @@ -6,11 +6,11 @@ Generic HTTP GET proxy for the agent. Lets the agent fetch URLs on the cluster's network when it cannot reach those hosts directly. Security: the host is checked against an explicit fnmatch glob allowlist configured -on the connection (`allowed_url_hosts`). Empty or missing allowlist means -no URL access; IP literals and non-HTTP schemes are rejected. Reuses the -connection's saved auth/SSL config so the agent doesn't need cluster credentials. +on the connection (`url_allowlist`). Empty or missing allowlist means no +URL access; the allowlist is the only gate — scheme and IP-literal checks +are intentionally NOT performed. Reuses the connection's saved auth/SSL +config so the agent doesn't need cluster credentials. """ -import ipaddress import fnmatch from urllib.parse import urlparse @@ -25,7 +25,7 @@ _MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body _REQUEST_TIMEOUT_SECONDS = 30 -def _host_matches_any_glob(host: str, patterns: list[str]) -> bool: +def _host_matches_any_glob(host: str, allowlist: list[str]) -> bool: """True if `host` matches any of the fnmatch glob patterns. fnmatch is case-sensitive on Linux (our deployment target). `*` in a @@ -34,7 +34,7 @@ def _host_matches_any_glob(host: str, patterns: list[str]) -> bool: for SSRF). """ host_labels = host.split(".") - for p in patterns: + for p in allowlist: pat_labels = p.split(".") if len(pat_labels) != len(host_labels): continue @@ -43,47 +43,30 @@ def _host_matches_any_glob(host: str, patterns: list[str]) -> bool: return False -def _validate_url_host( - url: str, - allowed_hosts: list[str] | None, -) -> None: - """Raise ValueError if the URL is not allowed to be fetched. +def _validate_url_host(url: str, allowlist: list[str] | None) -> None: + """Reject the URL unless its host matches a glob in `allowlist`. - Allowed only if the URL host matches one of the fnmatch glob patterns in - `allowed_hosts`. An empty or missing allowlist rejects everything. + The only check. The url_allowlist is the single source of truth for + what fetch_url is allowed to access — no scheme, IP-literal, or + "must be the same as yarn_rm_url" guardrails. The caller is + responsible for writing a tight allowlist. + + The only structural check: the URL must have a host (otherwise + the glob match has nothing to test). Anything else is the + allowlist's job. """ - parsed = urlparse(url) - if parsed.scheme not in ("http", "https"): - raise ValueError( - f"URL scheme must be http or https, got {parsed.scheme!r}" - ) - host = parsed.hostname + host = urlparse(url).hostname if not host: raise ValueError(f"URL has no host: {url!r}") - # IP literal check - try: - ipaddress.ip_address(host) - raise ValueError( - f"URL host {host!r} is an IP literal — IP targets are not allowed. " - f"Use a hostname on the cluster network." - ) - except ValueError as e: - if "IP literal" in str(e): - raise - # not an IP, continue - if not allowed_hosts: - raise ValueError( - "allowed_url_hosts is empty — set Connection.allowed_url_hosts to " - "allow specific hosts before calling fetch_url (e.g. ['ccam*'] " - "for ccam1-ccam99 or ['*.prod.internal'] for a subdomain)." - ) - if _host_matches_any_glob(host, allowed_hosts): + if _host_matches_any_glob(host, allowlist or []): return raise ValueError( - f"URL host {host!r} is not in Connection.allowed_url_hosts " - f"{allowed_hosts!r}. Reject this fetch to prevent SSRF." + f"URL host {host!r} is not in Connection.url_allowlist {allowlist!r}. " + f"Add the host pattern to url_allowlist (or use a broader glob) " + f"and try again." ) + def fetch_url(url: str, connection_name: str) -> FetchUrlResult: """Proxy an HTTP GET to url using the auth/SSL settings of connection_name.""" logger.debug(f"fetch_url enter url={url} connection_name={connection_name}") @@ -91,7 +74,7 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult: if conn is None: raise KeyError(f"Connection not found: {connection_name}") - _validate_url_host(url, conn.allowed_url_hosts) + _validate_url_host(url, conn.url_allowlist) config = YarnClientConfig.from_connection(conn) resp = httpx.get( diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 08e482e..efd740b 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -119,17 +119,17 @@ class SaveConnectionRequest(BaseModel): "for 'kinit -kt' workflows. The service does NOT auto-initialize " "from the keytab — you must `kinit -kt ` " "yourself before calling the tools." - ), - ) + ), + ) - allowed_url_hosts: list[str] | None = Field( + url_allowlist: list[str] | None = Field( default=None, description=( "Optional list of fnmatch glob patterns for hosts the fetch_url tool " - "may access. See Connection.allowed_url_hosts for full semantics. " + "may access. See Connection.url_allowlist for full semantics. " "Example for single-label host clusters: ['ccam*'] allows any host " "starting with 'ccam' (ccam1, ccam2, ..., ccam99). If omitted, None, " - "or empty, the saved connection will have allowed_url_hosts=[] (the " + "or empty, the saved connection will have url_allowlist=[] (the " "default), meaning fetch_url will reject every URL until the list is " "populated via update_connection." ), @@ -395,13 +395,12 @@ class FetchUrlRequest(BaseModel): url: str = Field( ..., description=( - "Absolute http:// or https:// URL to fetch. The host must share " - "at least 2 labels of suffix with the named Connection's " - "yarn_rm_url host (e.g. if yarn_rm_url is 'rm.prod.internal:8088', " - "you may fetch 'http://nm01.prod.internal:8042/...' but NOT " - "'http://evil.com/...' or 'http://10.0.0.1/...'). IP literals " - "and non-http(s) schemes are rejected. The Connection's saved " - "auth is reused for the outbound request — the agent does not " + "Absolute URL to fetch. The only access control is the named " + "Connection's url_allowlist: the URL host must match one of the " + "fnmatch glob patterns in that list. An empty or omitted allowlist " + "denies every host. IP literals and non-HTTP schemes are allowed " + "if and only if they are matched by the allowlist. The Connection's " + "saved auth is reused for the outbound request — the agent does not " "need cluster credentials." ), ) @@ -468,10 +467,10 @@ class UpdateConnectionRequest(BaseModel): default=None, description="New Kerberos keytab path. Omit to keep current.", ) - allowed_url_hosts: list[str] | None = Field( + url_allowlist: list[str] | None = Field( default=None, description=( - "Replacement allowed_url_hosts list (not merged). Omit to keep current. " + "Replacement url_allowlist list (not merged). Omit to keep current. " "Pass an empty list to deny all hosts (the default for new connections). " "Example: ['ccam*'] allows ccam1-ccam99." ), diff --git a/tests/unit/test_connection_tools.py b/tests/unit/test_connection_tools.py index 74ce480..933c46b 100644 --- a/tests/unit/test_connection_tools.py +++ b/tests/unit/test_connection_tools.py @@ -148,10 +148,10 @@ def test_update_connection_with_no_fields_is_noop(_fresh_store): assert out["deploy_mode"] == "client" -def test_update_connection_allows_url_hosts_change(_fresh_store): +def test_update_connection_changes_url_allowlist(_fresh_store): connections.save_connection(name="prod", master="yarn") - out = connections.update_connection(name="prod", allowed_url_hosts=["ccam*"]) - assert out["allowed_url_hosts"] == ["ccam*"] + out = connections.update_connection(name="prod", url_allowlist=["ccam*"]) + assert out["url_allowlist"] == ["ccam*"] def test_update_connection_replaces_not_merges_dict(_fresh_store): @@ -160,14 +160,14 @@ def test_update_connection_replaces_not_merges_dict(_fresh_store): assert out["spark_conf"] == {"b": "2"} -def test_update_connection_replaces_not_merges_list(_fresh_store): +def test_update_connection_replaces_not_merges_url_allowlist(_fresh_store): connections.save_connection( name="prod", master="yarn", - allowed_url_hosts=["ccam*"], + url_allowlist=["ccam*"], ) - out = connections.update_connection(name="prod", allowed_url_hosts=["nm*"]) - assert out["allowed_url_hosts"] == ["nm*"] + out = connections.update_connection(name="prod", url_allowlist=["nm*"]) + assert out["url_allowlist"] == ["nm*"] def test_update_connection_raises_for_unknown_name(_fresh_store): diff --git a/tests/unit/test_fetch_url.py b/tests/unit/test_fetch_url.py index 75723bd..a74d3ac 100644 --- a/tests/unit/test_fetch_url.py +++ b/tests/unit/test_fetch_url.py @@ -26,14 +26,14 @@ def fresh_stores(tmp_path: Path, monkeypatch): _fresh_stores(tmp_path, monkeypatch) yield -def test_fetch_url_returns_body_and_status(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) + +def test_fetch_url_returns_body_and_status(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088", - allowed_url_hosts=["*.prod.internal"], + url_allowlist=["*.prod.internal"], ) ) resp = httpx.Response(200, text="hello", headers={"content-type": "text/html"}) @@ -47,89 +47,18 @@ def test_fetch_url_returns_body_and_status(tmp_path, monkeypatch): assert m.call_count == 1 -def test_fetch_url_raises_for_missing_connection(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) +def test_fetch_url_raises_for_missing_connection(fresh_stores): with pytest.raises(KeyError, match="Connection not found"): fetch_url.fetch_url("http://rm.prod.internal:8088/", "missing") -def test_fetch_url_rejects_non_http_scheme(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) - fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") - ) - with pytest.raises(ValueError, match="scheme"): - fetch_url.fetch_url("file:///etc/passwd", "prod") - - -def test_fetch_url_rejects_ftp_scheme(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) - fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") - ) - with pytest.raises(ValueError, match="scheme"): - fetch_url.fetch_url("ftp://rm.prod.internal/foo", "prod") - - -def test_fetch_url_rejects_ip_literal(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) - fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") - ) - with pytest.raises(ValueError, match="IP literal"): - fetch_url.fetch_url("http://10.0.0.1/secrets", "prod") - - -def test_fetch_url_rejects_ipv6_literal(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) - fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088") - ) - with pytest.raises(ValueError, match="IP literal"): - fetch_url.fetch_url("http://[::1]:8080/", "prod") - - -def test_fetch_url_rejects_external_host(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) +def test_fetch_url_truncates_body_over_1mb(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088", - allowed_url_hosts=["*.prod.internal"], - ) - ) - with pytest.raises(ValueError, match="is not in Connection.allowed_url_hosts"): - fetch_url.fetch_url("http://evil.com/foo", "prod") - - -def test_fetch_url_rejects_unrelated_single_label_host(tmp_path, monkeypatch): - """Without a common suffix rule, a single-label RM host no longer helps.""" - _fresh_stores(tmp_path, monkeypatch) - fetch_url.conn_store.save( - Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") - ) - with pytest.raises(ValueError, match="allowed_url_hosts is empty"): - fetch_url.fetch_url("http://other-rm/", "prod") - - -def test_fetch_url_rejects_when_allowed_url_hosts_empty(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) - fetch_url.conn_store.save( - Connection(name="prod", master="yarn", allowed_url_hosts=[]) - ) - with pytest.raises(ValueError, match="allowed_url_hosts is empty"): - fetch_url.fetch_url("http://anything.com/", "prod") - - -def test_fetch_url_truncates_body_over_1mb(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) - fetch_url.conn_store.save( - Connection( - name="prod", - master="yarn", - yarn_rm_url="http://rm.prod.internal:8088", - allowed_url_hosts=["*.prod.internal"], + url_allowlist=["*.prod.internal"], ) ) big = "x" * (fetch_url._MAX_BODY_BYTES + 1) @@ -140,8 +69,7 @@ def test_fetch_url_truncates_body_over_1mb(tmp_path, monkeypatch): assert len(out.body) == fetch_url._MAX_BODY_BYTES -def test_fetch_url_passes_auth_from_connection(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) +def test_fetch_url_passes_auth_from_connection(fresh_stores): fetch_url.conn_store.save( Connection( name="auth", @@ -150,7 +78,7 @@ def test_fetch_url_passes_auth_from_connection(tmp_path, monkeypatch): auth_type="basic", auth_user="u", auth_password="p", - allowed_url_hosts=["*.prod.internal"], + url_allowlist=["*.prod.internal"], ) ) resp = httpx.Response(200, text="ok") @@ -159,19 +87,19 @@ def test_fetch_url_passes_auth_from_connection(tmp_path, monkeypatch): auth = m.call_args.kwargs["auth"] assert isinstance(auth, httpx.BasicAuth) import base64 + creds = base64.b64decode(auth._auth_header.split()[1]).decode() assert creds == "u:p" -def test_fetch_url_passes_ssl_verify_from_connection(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) +def test_fetch_url_passes_ssl_verify_from_connection(fresh_stores): fetch_url.conn_store.save( Connection( name="insecure", master="yarn", yarn_rm_url="http://rm.prod.internal:8088", ssl_verify=False, - allowed_url_hosts=["*.prod.internal"], + url_allowlist=["*.prod.internal"], ) ) resp = httpx.Response(200, text="ok") @@ -180,14 +108,13 @@ def test_fetch_url_passes_ssl_verify_from_connection(tmp_path, monkeypatch): assert m.call_args.kwargs["verify"] is False -def test_fetch_url_follows_redirects(tmp_path, monkeypatch): - _fresh_stores(tmp_path, monkeypatch) +def test_fetch_url_follows_redirects(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088", - allowed_url_hosts=["*.prod.internal"], + url_allowlist=["*.prod.internal"], ) ) resp = httpx.Response(200, text="ok") @@ -195,13 +122,14 @@ def test_fetch_url_follows_redirects(tmp_path, monkeypatch): fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") assert m.call_args.kwargs["follow_redirects"] is True + def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores): fetch_url.conn_store.save( Connection( name="ccam", master="yarn", yarn_rm_url="http://ccam1:8088", - allowed_url_hosts=["ccam*"], + url_allowlist=["ccam*"], ) ) resp = httpx.Response(200, text="hello") @@ -211,53 +139,59 @@ def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores): assert out.body == "hello" assert m.call_count == 1 + def test_fetch_url_allows_host_matching_any_of_multiple_globs(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", master="yarn", yarn_rm_url="http://rm.prod.internal:8088", - allowed_url_hosts=["ccam*", "*.prod.internal"], + url_allowlist=["ccam*", "*.prod.internal"], ) ) resp = httpx.Response(200, text="hello") with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: - out = fetch_url.fetch_url("http://history.prod.internal:18080/api/v1/info", "prod") + out = fetch_url.fetch_url( + "http://history.prod.internal:18080/api/v1/info", "prod" + ) assert out.status_code == 200 assert out.body == "hello" assert m.call_count == 1 + def test_fetch_url_glob_does_not_match_unrelated_host(fresh_stores): fetch_url.conn_store.save( Connection( name="ccam", master="yarn", yarn_rm_url="http://ccam1:8088", - allowed_url_hosts=["ccam*"], + url_allowlist=["ccam*"], ) ) - with pytest.raises(ValueError, match="is not in Connection.allowed_url_hosts"): + with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): fetch_url.fetch_url("http://evil.com/foo", "ccam") + def test_fetch_url_glob_does_not_cross_dot_boundary(fresh_stores): fetch_url.conn_store.save( Connection( name="ccam", master="yarn", yarn_rm_url="http://ccam1:8088", - allowed_url_hosts=["ccam*"], + url_allowlist=["ccam*"], ) ) - with pytest.raises(ValueError, match="is not in Connection.allowed_url_hosts"): + with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): fetch_url.fetch_url("http://ccam50.evil.com/", "ccam") + def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", master="yarn", yarn_rm_url="http://rm:8088", - allowed_url_hosts=["ccam*"], + url_allowlist=["ccam*"], ) ) resp = httpx.Response(200, text="hello") @@ -267,30 +201,55 @@ def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): assert out.body == "hello" assert m.call_count == 1 -def test_fetch_url_empty_allowed_hosts_rejects_everything(fresh_stores): + +def test_fetch_url_accepts_ip_literal_when_in_url_allowlist(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", master="yarn", - yarn_rm_url="http://rm.prod.internal", - allowed_url_hosts=[], + url_allowlist=["*.*.*.*"], ) ) - with pytest.raises(ValueError, match="allowed_url_hosts is empty"): - fetch_url.fetch_url("http://nm.prod.internal/", "prod") + resp = httpx.Response(200, text="hello") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + out = fetch_url.fetch_url("http://10.0.0.1/secret", "prod") + assert out.status_code == 200 + assert out.body == "hello" + assert m.call_count == 1 -def test_fetch_url_omitted_allowed_url_hosts_defaults_to_empty_and_rejects(fresh_stores): + +def test_fetch_url_accepts_https_when_in_url_allowlist(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", master="yarn", - yarn_rm_url="http://rm.prod.internal", + url_allowlist=["ccam*.example.com"], ) ) - with pytest.raises(ValueError, match="allowed_url_hosts is empty"): - fetch_url.fetch_url("http://nm.prod.internal/", "prod") + resp = httpx.Response(200, text="hello") + with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + out = fetch_url.fetch_url("https://ccam1.example.com/secure", "prod") + assert out.status_code == 200 + assert out.body == "hello" + assert m.call_count == 1 -def test_fetch_url_error_message_hints_at_allowed_url_hosts(fresh_stores): + +def test_fetch_url_rejects_when_url_has_no_host(fresh_stores): + with pytest.raises(ValueError, match="URL has no host"): + fetch_url._validate_url_host("", ["*"]) + with pytest.raises(ValueError, match="URL has no host"): + fetch_url._validate_url_host("not-a-url", ["*"]) + + +def test_fetch_url_rejects_when_url_allowlist_empty(fresh_stores): + fetch_url.conn_store.save( + Connection(name="prod", master="yarn", url_allowlist=[]) + ) + with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): + fetch_url.fetch_url("http://anything.com/", "prod") + + +def test_fetch_url_error_message_mentions_url_allowlist(fresh_stores): fetch_url.conn_store.save( Connection( name="prod", @@ -298,5 +257,17 @@ def test_fetch_url_error_message_hints_at_allowed_url_hosts(fresh_stores): yarn_rm_url="http://rm.prod.internal:8088", ) ) - with pytest.raises(ValueError, match="set Connection.allowed_url_hosts"): + with pytest.raises(ValueError, match="url_allowlist"): fetch_url.fetch_url("http://evil.com/foo", "prod") + + +def test_fetch_url_omitted_url_allowlist_defaults_to_empty_and_rejects(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal", + ) + ) + with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): + fetch_url.fetch_url("http://nm.prod.internal/", "prod") From 7fbad97a8723167611199fc913fd6b2317e800e7 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 13:39:37 +0800 Subject: [PATCH 5/7] feat: add list_applications tool for YARN application enumeration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a new MCP tool that queries YARN's /ws/v1/cluster/apps endpoint through a named Connection, returning a list of ApplicationSummary records. Bypasses the local JobStore — useful for enumerating apps that were not submitted through this service. API: list_applications( connection_name: str, # required, which YARN cluster state: str | None = None, # YARN state filter: NEW/NEW_SAVING/ # SUBMITTED/ACCEPTED/RUNNING/ # FINISHED/FAILED/KILLED queue: str | None = None, # YARN queue filter limit: int = 100, # cap on returned apps (YARN has no # offset-based pagination; combine # state/queue filters for big clusters) ) -> list[ApplicationSummary] Implementation: - yarn_client.list_applications(config, *, state, queue, limit) -> list[dict] Returns raw YARN app dicts; raises YarnError on 4xx/5xx; returns [] on 404 (no apps match). Uses the existing _request helper, which now accepts a "params" kwarg for query strings (one-line additive change). - external_jobs.list_applications(connection_name, state, queue, limit) -> list[ApplicationSummary]. Looks up the Connection, builds the YarnClientConfig, calls the yarn_client function, maps each raw YARN dict to ApplicationSummary (mirroring the manual field-mapping style of get_job_result). The yarn_client function is imported as "list_applications_yarn" to avoid name collision. - ApplicationSummary: 12-field Pydantic model with snake_case names (application_id, name, user, queue, state, final_status, application_type, application_tags, started_time, finished_time, tracking_url, progress). Unused YARN fields (memorySeconds, vcoreSeconds, preemptedResource*, etc.) are not exposed. - ListApplicationsRequest: Pydantic body model with Field(description=) for LLM-facing schema. - /list_applications route registered with operation_id= "list_applications", placed next to the other external YARN tools. Tests: - 8 new unit tests in test_external_jobs.py (happy path, state/queue/ limit pass-through, default limit, empty list, missing connection, full field mapping). - test_mcp_routes.py: assert 23 tool routes. - README: list_applications row added to the Spark Executor table. Tests: 390 passed (was 382, +8 net). Co-Authored-By: Claude Fable 5 --- README.md | 1 + spark_executor/core/yarn_client.py | 58 +++++++++- spark_executor/models.py | 22 ++++ spark_executor/server.py | 30 +++++ spark_executor/tools/external_jobs.py | 59 +++++++++- spark_executor/tools/requests.py | 36 ++++++ tests/integration/test_mcp_routes.py | 3 +- tests/unit/test_external_jobs.py | 161 +++++++++++++++++++++++++- 8 files changed, 363 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index eeaabbb..c644a6d 100644 --- a/README.md +++ b/README.md @@ -121,6 +121,7 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | `get_external_job_status` | 查询**非本服务提交**的外部 YARN application 状态（按 `application_id` + `connection_name`） | | `get_external_job_result` | 查询外部 YARN application 终态结果视图 | | `get_external_job_logs` | 拉取外部 YARN application 的聚合日志 | +| `list_applications` | 列出 YARN 上所有应用（按 `state` / `queue` / `limit` 过滤），绕过 JobStore | | `fetch_url` | 代理 HTTP GET 到集群内网 URL (host 受 `Connection.url_allowlist` glob allowlist 约束, 空则全拒) | ### Files MCP 工具 diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py index 0bb0b56..cd45c5a 100644 --- a/spark_executor/core/yarn_client.py +++ b/spark_executor/core/yarn_client.py @@ -121,13 +121,14 @@ def _base_url(yarn_rm_url: str | None) -> str: def _request(method: str, url: str, *, json_body: dict | None = None, - timeout: float = 30.0, verify: bool | str = True, - auth: httpx.Auth | None = None) -> httpx.Response: + params: dict[str, str] | None = None, timeout: float = 30.0, + verify: bool | str = True, auth: httpx.Auth | None = None) -> httpx.Response: headers = {"Accept": "application/json"} logger.debug(f"YARN {method} {url}" + (f" body={json_body}" if json_body else "")) try: resp = httpx.request( - method, url, json=json_body, headers=headers, timeout=timeout, verify=verify, auth=auth + method, url, json=json_body, params=params, headers=headers, + timeout=timeout, verify=verify, auth=auth ) except httpx.HTTPError as exc: logger.error(f"YARN {method} {url} failed: {exc}") @@ -257,3 +258,54 @@ def kill_application(application_id: str, config: YarnClientConfig) -> None: logger.error(f"YARN PUT {url} -> {resp.status_code}: {resp.text[:500]}") raise YarnError(f"YARN kill returned HTTP {resp.status_code}: {resp.text}") logger.info(f"YARN kill {application_id} -> ok") + + +def list_applications( + config: YarnClientConfig, + *, + state: str | None = None, + queue: str | None = None, + limit: int | None = None, +) -> list[dict]: + """List YARN applications, optionally filtered. + + YARN endpoint: GET /ws/v1/cluster/apps?state=...&queue=...&limit=... + + Filters: + - state: YARN application state. Common values: + "NEW", "NEW_SAVING", "SUBMITTED", "ACCEPTED", "RUNNING", + "FINISHED", "FAILED", "KILLED". + Note: "FINISHED" is the umbrella state covering SUCCEEDED/FAILED/KILLED. + - queue: YARN queue name + - limit: cap on number of returned apps (YARN has no pagination; + callers that need a full enumeration should make multiple + calls with state=... filters or accept the cap) + + Returns a list of YARN app dicts (each with id, name, user, queue, + state, finalStatus, applicationType, startedTime, finishedTime, + trackingUrl, progress, etc). Empty list if no apps match. + + Raises YarnError on transport / 4xx / 5xx. + """ + params: dict[str, str] = {} + if state is not None: + params["state"] = state + if queue is not None: + params["queue"] = queue + if limit is not None: + params["limit"] = str(limit) + + url = f"{_base_url(config.yarn_rm_url)}/ws/v1/cluster/apps" + resp = _request("GET", url, params=params, + verify=config.verify_for_httpx(), + auth=config.auth_for_httpx()) + if resp.status_code == 404: + # No apps match (or RM doesn't support the endpoint) + return [] + if resp.status_code >= 400: + raise YarnError( + f"YARN list applications failed: {resp.status_code} {resp.text[:200]}" + ) + data = resp.json() + apps_container = data.get("apps") or {} + return apps_container.get("app", []) or [] diff --git a/spark_executor/models.py b/spark_executor/models.py index 1fa327a..b389021 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -48,6 +48,28 @@ class FetchUrlResult(BaseModel): truncated: bool = False +class ApplicationSummary(BaseModel): + """A YARN application summary from /ws/v1/cluster/apps. + + Field names are mapped from the YARN JSON keys to clearer + snake_case names by the tool function. Unused YARN fields + (memorySeconds, vcoreSeconds, preemptedResource*, etc.) are + not exposed — the LLM doesn't need them. + """ + application_id: str + name: str + user: str + queue: str + state: str + final_status: str | None = None + application_type: str | None = None + application_tags: str = "" + started_time: int = 0 + finished_time: int = 0 + tracking_url: str | None = None + progress: float | None = None + + class Connection(BaseModel): name: str # Defaults to "yarn" because that's the literal string spark-submit wants diff --git a/spark_executor/server.py b/spark_executor/server.py index 9a48a19..6453ca6 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -21,6 +21,7 @@ from spark_executor.tools.external_jobs import ( get_external_job_logs, get_external_job_status, get_external_job_result, + list_applications, ) from spark_executor.tools.fetch_url import fetch_url from spark_executor.tools.requests import ( @@ -30,6 +31,7 @@ from spark_executor.tools.requests import ( ExternalJobLogsRequest, ExternalJobStatusRequest, ExternalJobResultRequest, + ListApplicationsRequest, FetchUrlRequest, GetJobLogsRequest, JobIdRequest, @@ -346,6 +348,34 @@ def _get_external_job_result(req: ExternalJobResultRequest): return get_external_job_result(req.application_id, req.connection_name) +@app.post( + "/list_applications", + operation_id="list_applications", + summary="List YARN applications on a cluster, optionally filtered", + description=( + "Query YARN's /ws/v1/cluster/apps endpoint through the named " + "Connection, returning a list of ApplicationSummary records. " + "Bypasses the local JobStore — useful for enumerating apps that " + "were not submitted through this service.\n\n" + "**Filters:** state (YARN state, e.g. 'RUNNING', 'FINISHED', " + "'FAILED'), queue (YARN queue name), limit (default 100, max ~10000). " + "YARN has no offset-based pagination, so for large clusters combine " + "state/queue filters to scope the result. The `FINISHED` state " + "covers SUCCEEDED/FAILED/KILLED.\n\n" + "Returns an empty list if no apps match. The Connection's auth_type " + "/ auth_user / auth_password / ssl_verify / ssl_ca_bundle are reused " + "for the request." + ), +) +def _list_applications(req: ListApplicationsRequest): + return list_applications( + req.connection_name, + state=req.state, + queue=req.queue, + limit=req.limit, + ) + + # --- Connection management tools --- @app.post( diff --git a/spark_executor/tools/external_jobs.py b/spark_executor/tools/external_jobs.py index eb60329..1527317 100644 --- a/spark_executor/tools/external_jobs.py +++ b/spark_executor/tools/external_jobs.py @@ -10,8 +10,13 @@ YARN application_id and the name of a saved Connection. import json from common.logging import logger -from spark_executor.core.yarn_client import YarnClientConfig, get_application_status, get_application_logs -from spark_executor.models import JobStatus, JobResult +from spark_executor.core.yarn_client import ( + YarnClientConfig, + get_application_logs, + get_application_status, + list_applications as list_applications_yarn, # alias to avoid collision +) +from spark_executor.models import ApplicationSummary, JobResult, JobStatus from spark_executor.tools.connections import store as conn_store @@ -63,3 +68,53 @@ def get_external_job_result(application_id: str, connection_name: str) -> JobRes ) logger.info(f"get_external_job_result ok application_id={application_id} connection_name={connection_name} state={state}") return result + + +def list_applications( + connection_name: str, + state: str | None = None, + queue: str | None = None, + limit: int = 100, +) -> list[ApplicationSummary]: + """List YARN applications on the named cluster, optionally filtered. + + Bypasses the local JobStore (this is for apps not submitted through + this service). The YARN ResourceManager REST endpoint + /ws/v1/cluster/apps is queried through the connection's auth/SSL + config. + + Defaults: limit=100 (YARN has no offset-based pagination, so large + clusters should use state/queue filters to scope the result). + """ + logger.debug( + f"list_applications enter connection_name={connection_name} " + f"state={state} queue={queue} limit={limit}" + ) + conn = conn_store.get(connection_name) + if conn is None: + raise KeyError(f"Connection not found: {connection_name}") + config = YarnClientConfig.from_connection(conn) + raw_apps = list_applications_yarn( + config, state=state, queue=queue, limit=limit + ) + summaries = [ + ApplicationSummary( + application_id=app.get("id", ""), + name=app.get("name", ""), + user=app.get("user", ""), + queue=app.get("queue", ""), + state=app.get("state", ""), + final_status=app.get("finalStatus"), + application_type=app.get("applicationType"), + application_tags=app.get("applicationTags", ""), + started_time=app.get("startedTime", 0), + finished_time=app.get("finishedTime", 0), + tracking_url=app.get("trackingUrl"), + progress=app.get("progress"), + ) + for app in raw_apps + ] + logger.info( + f"list_applications ok connection_name={connection_name} count={len(summaries)}" + ) + return summaries diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index efd740b..27f8284 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -475,3 +475,39 @@ class UpdateConnectionRequest(BaseModel): "Example: ['ccam*'] allows ccam1-ccam99." ), ) + + +class ListApplicationsRequest(BaseModel): + connection_name: str = Field( + ..., + description=( + "Name of a saved Connection (see list_connections) pointing at " + "the YARN cluster to query." + ), + ) + state: str | None = Field( + default=None, + description=( + "Optional YARN application state filter. One of: 'NEW', " + "'NEW_SAVING', 'SUBMITTED', 'ACCEPTED', 'RUNNING', 'FINISHED', " + "'FAILED', 'KILLED'. 'FINISHED' is the umbrella state covering " + "SUCCEEDED/FAILED/KILLED. None = no state filter (returns all " + "states up to `limit`)." + ), + ) + queue: str | None = Field( + default=None, + description=( + "Optional YARN queue name filter (e.g. 'default', 'prod'). " + "None = no queue filter." + ), + ) + limit: int = Field( + default=100, + description=( + "Maximum number of applications to return. YARN has no " + "offset-based pagination, so for large clusters use state/queue " + "filters to scope the result. Max 10000 in practice (YARN's own " + "limit on the limit param)." + ), + ) diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 6b1e091..df69baf 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -64,12 +64,13 @@ def test_seventeen_tool_routes_registered(): assert "/update_pending_job" in paths -def test_twenty_two_tool_routes_registered(): +def test_twenty_three_tool_routes_registered(): paths = {r.path for r in app.routes} for path in ( "/get_external_job_logs", "/get_external_job_status", "/get_external_job_result", + "/list_applications", "/fetch_url", "/update_connection", ): diff --git a/tests/unit/test_external_jobs.py b/tests/unit/test_external_jobs.py index dc67df5..36e3f11 100644 --- a/tests/unit/test_external_jobs.py +++ b/tests/unit/test_external_jobs.py @@ -5,8 +5,9 @@ from unittest.mock import patch import pytest from spark_executor.core import connection_store -from spark_executor.models import Connection +from spark_executor.models import ApplicationSummary, Connection from spark_executor.tools import connections, external_jobs +from spark_executor.tools.requests import ListApplicationsRequest def _fresh_stores(): @@ -17,6 +18,13 @@ def _fresh_stores(): external_jobs.conn_store = store +@pytest.fixture +def fresh_stores(tmp_path, monkeypatch): + """Reset connection store singletons to an isolated tmp_path.""" + monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) + _fresh_stores() + + def test_get_external_job_logs_returns_tailed(): _fresh_stores() external_jobs.conn_store.save( @@ -131,3 +139,154 @@ def test_get_external_job_result_raises_when_connection_missing(): external_jobs.get_external_job_result( application_id="application_1", connection_name="missing" ) + + +def test_list_applications_returns_summaries(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[ + { + "id": "application_1", + "name": "app-one", + "user": "alice", + "queue": "default", + "state": "RUNNING", + "finalStatus": "UNDEFINED", + "applicationType": "SPARK", + "applicationTags": "tag1", + "startedTime": 1000, + "finishedTime": 0, + "trackingUrl": "http://rm:8088/proxy/application_1", + "progress": 75.0, + }, + { + "id": "application_2", + "name": "app-two", + "user": "bob", + "queue": "research", + "state": "FINISHED", + "finalStatus": "SUCCEEDED", + "applicationType": "SPARK", + "applicationTags": "", + "startedTime": 2000, + "finishedTime": 3000, + "trackingUrl": "http://rm:8088/proxy/application_2", + "progress": 100.0, + }, + ], + ) as m: + out = external_jobs.list_applications("prod") + assert len(out) == 2 + assert all(isinstance(item, ApplicationSummary) for item in out) + assert out[0].application_id == "application_1" + assert out[1].application_id == "application_2" + args = m.call_args.args + assert args[0].yarn_rm_url == "http://rm:8088" + + +def test_list_applications_passes_state_filter(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ) as m: + external_jobs.list_applications("prod", state="RUNNING") + assert m.call_args.kwargs == {"state": "RUNNING", "queue": None, "limit": 100} + + +def test_list_applications_passes_queue_filter(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ) as m: + external_jobs.list_applications("prod", queue="research") + assert m.call_args.kwargs == {"state": None, "queue": "research", "limit": 100} + + +def test_list_applications_passes_limit_filter(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ) as m: + external_jobs.list_applications("prod", limit=50) + assert m.call_args.kwargs == {"state": None, "queue": None, "limit": 50} + + +def test_list_applications_default_limit_is_100(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ) as m: + external_jobs.list_applications("prod") + assert m.call_args.kwargs == {"state": None, "queue": None, "limit": 100} + assert ListApplicationsRequest(connection_name="prod").limit == 100 + + +def test_list_applications_returns_empty_list_when_no_apps(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[], + ): + out = external_jobs.list_applications("prod") + assert out == [] + + +def test_list_applications_raises_for_missing_connection(fresh_stores): + with pytest.raises(KeyError, match="Connection not found"): + external_jobs.list_applications("missing") + + +def test_list_applications_maps_yarn_json_to_summary(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + yarn_app = { + "id": "application_42", + "name": "mapped-app", + "user": "carol", + "queue": "prod", + "state": "ACCEPTED", + "finalStatus": "UNDEFINED", + "applicationType": "MAPREDUCE", + "applicationTags": "batch", + "startedTime": 12345, + "finishedTime": 0, + "trackingUrl": "http://rm:8088/proxy/application_42", + "progress": 12.5, + } + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + return_value=[yarn_app], + ): + out = external_jobs.list_applications("prod") + assert len(out) == 1 + summary = out[0] + assert summary.application_id == "application_42" + assert summary.name == "mapped-app" + assert summary.user == "carol" + assert summary.queue == "prod" + assert summary.state == "ACCEPTED" + assert summary.final_status == "UNDEFINED" + assert summary.application_type == "MAPREDUCE" + assert summary.application_tags == "batch" + assert summary.started_time == 12345 + assert summary.finished_time == 0 + assert summary.tracking_url == "http://rm:8088/proxy/application_42" + assert summary.progress == 12.5 From 7d4e512cb09aff8601e81e7c174ed430f1157be7 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 14:23:16 +0800 Subject: [PATCH 6/7] fix: address review findings on fetch-url-tool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fetch_url: revalidate allowlist on every redirect hop (fixes SSRF where 302 to disallowed host / 169.254.169.254 / file:// bypassed the url_allowlist). Stream response body with iter_bytes and cap at 1MB so a multi-GB response from an allowlisted host cannot OOM the service. Reuses the manual-redirect-loop pattern from yarn_client. - list_applications: stop swallowing 404 (YARN returns 200+empty for "no match"; 404 means the RM doesn't support the endpoint — surface the YarnError instead of hiding it as an empty result). Add Field(ge=1, le=10000) to ListApplicationsRequest.limit so a runaway limit is rejected at the Pydantic layer with 422. - save_connection: PATCH semantics for existing records. Re-route to update_connection when the name already exists so partial updates (e.g. only master) no longer wipe url_allowlist back to []. Uses an _UNSET sentinel in the tool function to distinguish "omitted" from "None" without breaking the existing parameter list. - README: drop leading space on 5 new connection-tool table rows that was breaking GitHub Flavored Markdown table continuity. - Indentation: normalize connections.py and requests.py to 4-space indent (auth_password/auth_principal/auth_keytab were 3-space). Co-Authored-By: Claude --- README.md | 10 +- spark_executor/core/yarn_client.py | 3 - spark_executor/server.py | 14 ++- spark_executor/tools/connections.py | 94 ++++++++++----- spark_executor/tools/fetch_url.py | 92 ++++++++++---- spark_executor/tools/requests.py | 10 +- tests/unit/test_connection_tools.py | 9 ++ tests/unit/test_external_jobs.py | 13 ++ tests/unit/test_fetch_url.py | 179 ++++++++++++++++++++++++---- 9 files changed, 331 insertions(+), 93 deletions(-) diff --git a/README.md b/README.md index c644a6d..867c96a 100644 --- a/README.md +++ b/README.md @@ -100,11 +100,11 @@ MCP 客户端需要先执行 `initialize` 握手，拿到 `mcp-session-id` 后 | 工具 | 作用 | | --- | --- | - | `save_connection` | 保存或更新一个 Spark/YARN 连接配置 | - | `list_connections` | 列出所有连接配置 | - | `get_connection` | 按名称读取连接配置 | - | `delete_connection` | 删除连接配置 | - | `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) | +| `save_connection` | 保存或更新一个 Spark/YARN 连接配置 | +| `list_connections` | 列出所有连接配置 | +| `get_connection` | 按名称读取连接配置 | +| `delete_connection` | 删除连接配置 | +| `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) | | `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 | | `read_job_file` | 读取已存在的 PySpark 脚本内容 | | `update_job_file` | 覆盖更新已存在的 PySpark 脚本 | diff --git a/spark_executor/core/yarn_client.py b/spark_executor/core/yarn_client.py index cd45c5a..f86e3f5 100644 --- a/spark_executor/core/yarn_client.py +++ b/spark_executor/core/yarn_client.py @@ -299,9 +299,6 @@ def list_applications( resp = _request("GET", url, params=params, verify=config.verify_for_httpx(), auth=config.auth_for_httpx()) - if resp.status_code == 404: - # No apps match (or RM doesn't support the endpoint) - return [] if resp.status_code >= 400: raise YarnError( f"YARN list applications failed: {resp.status_code} {resp.text[:200]}" diff --git a/spark_executor/server.py b/spark_executor/server.py index 6453ca6..e43efb0 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -390,12 +390,20 @@ def _list_applications(req: ListApplicationsRequest): "yarn_rm_url is required for get_external_* to work** — spark-submit " "can discover the RM for submissions, but direct YARN REST queries " "need an explicit URL. Saving with an existing name overwrites the " - "record in place (no version history)." + "record in place (no version history). When the name already exists, " + "only the provided fields are changed (PATCH semantics); omitted " + "fields keep their previous values." ), ) def _save_connection(req: SaveConnectionRequest): - # exclude_none so we don't overwrite the function's default with explicit None - return save_connection(**req.model_dump(exclude_none=True)) + fields = req.model_dump(exclude_none=True) + name = fields.pop("name") + try: + get_connection(name) + except KeyError: + # exclude_none so we don't overwrite the function's default with explicit None + return save_connection(name=name, **fields) + return update_connection(name=name, **fields) @app.post( diff --git a/spark_executor/tools/connections.py b/spark_executor/tools/connections.py index 0af462b..4c631a7 100644 --- a/spark_executor/tools/connections.py +++ b/spark_executor/tools/connections.py @@ -4,45 +4,75 @@ @Author :tao.chen """ from common.logging import logger -from spark_executor.core.connection_store import ConnectionStore, store +from spark_executor.core.connection_store import store from spark_executor.models import Connection +_UNSET = object() + + def save_connection( *, name: str, master: str, - deploy_mode: str = "cluster", - yarn_rm_url: str | None = None, - spark_conf: dict[str, str] | None = None, - ssl_verify: bool | None = None, - ssl_ca_bundle: str | None = None, - auth_type: str = "none", - auth_user: str | None = None, - auth_password: str | None = None, - auth_principal: str | None = None, - auth_keytab: str | None = None, - url_allowlist: list[str] | None = None, -) -> dict[str, str]: + deploy_mode: str = _UNSET, # type: ignore[assignment] + yarn_rm_url: str | None = _UNSET, # type: ignore[assignment] + spark_conf: dict[str, str] | None = _UNSET, # type: ignore[assignment] + ssl_verify: bool | None = _UNSET, # type: ignore[assignment] + ssl_ca_bundle: str | None = _UNSET, # type: ignore[assignment] + auth_type: str = _UNSET, # type: ignore[assignment] + auth_user: str | None = _UNSET, # type: ignore[assignment] + auth_password: str | None = _UNSET, # type: ignore[assignment] + auth_principal: str | None = _UNSET, # type: ignore[assignment] + auth_keytab: str | None = _UNSET, # type: ignore[assignment] + url_allowlist: list[str] | None = _UNSET, # type: ignore[assignment] +) -> dict[str, object]: logger.debug( - f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " - f"yarn_rm_url={yarn_rm_url} spark_conf_keys={list((spark_conf or {}).keys())}" + f"save_connection enter name={name} master={master} " + f"spark_conf_keys={list((spark_conf if isinstance(spark_conf, dict) else {}).keys())}" ) - conn = Connection( - name=name, - master=master, - deploy_mode=deploy_mode, - yarn_rm_url=yarn_rm_url, - spark_conf=spark_conf or {}, - ssl_verify=ssl_verify, - ssl_ca_bundle=ssl_ca_bundle, - auth_type=auth_type, - auth_user=auth_user, - auth_password=auth_password, - auth_principal=auth_principal, - auth_keytab=auth_keytab, - url_allowlist=url_allowlist or [], - ) + existing = store.get(name) + if existing is not None: + fields: dict[str, object] = {"master": master} + if deploy_mode is not _UNSET: + fields["deploy_mode"] = deploy_mode + if yarn_rm_url is not _UNSET: + fields["yarn_rm_url"] = yarn_rm_url + if spark_conf is not _UNSET: + fields["spark_conf"] = spark_conf or {} + if ssl_verify is not _UNSET: + fields["ssl_verify"] = ssl_verify + if ssl_ca_bundle is not _UNSET: + fields["ssl_ca_bundle"] = ssl_ca_bundle + if auth_type is not _UNSET: + fields["auth_type"] = auth_type + if auth_user is not _UNSET: + fields["auth_user"] = auth_user + if auth_password is not _UNSET: + fields["auth_password"] = auth_password + if auth_principal is not _UNSET: + fields["auth_principal"] = auth_principal + if auth_keytab is not _UNSET: + fields["auth_keytab"] = auth_keytab + if url_allowlist is not _UNSET: + fields["url_allowlist"] = url_allowlist or [] + return update_connection(name, **fields) + new_fields: dict[str, object] = { + "name": name, + "master": master, + "deploy_mode": deploy_mode if deploy_mode is not _UNSET else "cluster", + "yarn_rm_url": yarn_rm_url if yarn_rm_url is not _UNSET else None, + "spark_conf": (spark_conf if spark_conf is not _UNSET else None) or {}, + "ssl_verify": ssl_verify if ssl_verify is not _UNSET else None, + "ssl_ca_bundle": ssl_ca_bundle if ssl_ca_bundle is not _UNSET else None, + "auth_type": auth_type if auth_type is not _UNSET else "none", + "auth_user": auth_user if auth_user is not _UNSET else None, + "auth_password": auth_password if auth_password is not _UNSET else None, + "auth_principal": auth_principal if auth_principal is not _UNSET else None, + "auth_keytab": auth_keytab if auth_keytab is not _UNSET else None, + "url_allowlist": (url_allowlist if url_allowlist is not _UNSET else None) or [], + } + conn = Connection(**new_fields) store.save(conn) return {"name": name, "status": "SAVED"} @@ -54,8 +84,8 @@ def update_connection(name: str, **fields) -> dict[str, object]: you pass are changed. To clear an optional field (e.g. `yarn_rm_url`), use `delete_connection(name=...)` followed by `save_connection(...)`. - Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf, - ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password, + Mutable fields: master, deploy_mode, yarn_rm_url, spark_conf, + ssl_verify, ssl_ca_bundle, auth_type, auth_user, auth_password, auth_principal, auth_keytab, url_allowlist. """ logger.debug(f"update_connection enter name={name} fields={sorted(fields.keys())}") diff --git a/spark_executor/tools/fetch_url.py b/spark_executor/tools/fetch_url.py index c505b88..4813c30 100644 --- a/spark_executor/tools/fetch_url.py +++ b/spark_executor/tools/fetch_url.py @@ -12,7 +12,7 @@ are intentionally NOT performed. Reuses the connection's saved auth/SSL config so the agent doesn't need cluster credentials. """ import fnmatch -from urllib.parse import urlparse +from urllib.parse import urlparse, urljoin import httpx @@ -24,6 +24,9 @@ from spark_executor.tools.connections import store as conn_store _MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body _REQUEST_TIMEOUT_SECONDS = 30 +_MAX_REDIRECTS = 3 +_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308}) + def _host_matches_any_glob(host: str, allowlist: list[str]) -> bool: """True if `host` matches any of the fnmatch glob patterns. @@ -67,6 +70,29 @@ def _validate_url_host(url: str, allowlist: list[str] | None) -> None: ) +def _read_streamed_body(resp: httpx.Response) -> tuple[bytes, bool]: + """Stream the response body, keeping at most ``_MAX_BODY_BYTES`` bytes. + + Returns ``(body_bytes, truncated)``. Stops reading as soon as the cap is + exceeded so that a multi-gigabyte response from an allowlisted host cannot + OOM the service. + """ + chunks: list[bytes] = [] + total = 0 + truncated = False + for chunk in resp.iter_bytes(): + if total + len(chunk) <= _MAX_BODY_BYTES: + chunks.append(chunk) + total += len(chunk) + else: + if total < _MAX_BODY_BYTES: + chunks.append(chunk[: _MAX_BODY_BYTES - total]) + total = _MAX_BODY_BYTES + truncated = True + break + return b"".join(chunks), truncated + + def fetch_url(url: str, connection_name: str) -> FetchUrlResult: """Proxy an HTTP GET to url using the auth/SSL settings of connection_name.""" logger.debug(f"fetch_url enter url={url} connection_name={connection_name}") @@ -77,28 +103,50 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult: _validate_url_host(url, conn.url_allowlist) config = YarnClientConfig.from_connection(conn) - resp = httpx.get( - url, - auth=config.auth_for_httpx(), - verify=config.verify_for_httpx(), - timeout=_REQUEST_TIMEOUT_SECONDS, - follow_redirects=True, - ) + auth = config.auth_for_httpx() + verify = config.verify_for_httpx() - body = resp.text[:_MAX_BODY_BYTES] - truncated = len(resp.text) > _MAX_BODY_BYTES + for hop in range(_MAX_REDIRECTS + 1): + with httpx.stream( + "GET", + url, + auth=auth, + verify=verify, + timeout=_REQUEST_TIMEOUT_SECONDS, + follow_redirects=False, + ) as resp: + if resp.status_code not in _REDIRECT_STATUSES: + body_bytes, truncated = _read_streamed_body(resp) + encoding = resp.encoding or "utf-8" + body = body_bytes.decode(encoding, errors="replace") - logger.info( - f"fetch_url ok url={url} connection_name={connection_name} " - f"status_code={resp.status_code} " - f"content_type={resp.headers.get('content-type', '')} " - f"body_bytes={len(body)} truncated={truncated}" - ) + logger.info( + f"fetch_url ok url={url} connection_name={connection_name} " + f"status_code={resp.status_code} " + f"content_type={resp.headers.get('content-type', '')} " + f"body_bytes={len(body)} truncated={truncated}" + ) - return FetchUrlResult( - url=url, - status_code=resp.status_code, - content_type=resp.headers.get("content-type", ""), - body=body, - truncated=truncated, + return FetchUrlResult( + url=url, + status_code=resp.status_code, + content_type=resp.headers.get("content-type", ""), + body=body, + truncated=truncated, + ) + + location = resp.headers.get("Location") or resp.headers.get("location") + if not location: + raise ValueError( + f"HTTP GET returned {resp.status_code} with no Location header at {url}" + ) + next_url = urljoin(url, location) + logger.debug( + f"fetch_url redirect url={url} status={resp.status_code} to={next_url}" + ) + _validate_url_host(next_url, conn.url_allowlist) + url = next_url + + raise ValueError( + f"fetch_url exceeded {_MAX_REDIRECTS} redirects (last url: {url})" ) diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 27f8284..9cb1f0f 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -117,10 +117,10 @@ class SaveConnectionRequest(BaseModel): description=( "Absolute path to a Kerberos keytab file. Optional convenience " "for 'kinit -kt' workflows. The service does NOT auto-initialize " - "from the keytab — you must `kinit -kt ` " - "yourself before calling the tools." - ), - ) + "from the keytab — you must `kinit -kt ` " + "yourself before calling the tools." + ), + ) url_allowlist: list[str] | None = Field( default=None, @@ -504,6 +504,8 @@ class ListApplicationsRequest(BaseModel): ) limit: int = Field( default=100, + ge=1, + le=10000, description=( "Maximum number of applications to return. YARN has no " "offset-based pagination, so for large clusters use state/queue " diff --git a/tests/unit/test_connection_tools.py b/tests/unit/test_connection_tools.py index 933c46b..dad74c7 100644 --- a/tests/unit/test_connection_tools.py +++ b/tests/unit/test_connection_tools.py @@ -121,6 +121,15 @@ def test_delete_connection_returns_status(_fresh_store): def test_delete_connection_unknown_raises(_fresh_store): with pytest.raises(KeyError): connections.delete_connection("missing") + + +def test_save_connection_preserves_url_allowlist_on_existing_record(_fresh_store): + connections.save_connection(name="prod", master="yarn", url_allowlist=["ccam*"]) + out = connections.save_connection(name="prod", master="spark://new:7077") + assert out["master"] == "spark://new:7077" + assert out["url_allowlist"] == ["ccam*"] + + def test_update_connection_changes_specified_field(_fresh_store): connections.save_connection(name="prod", master="yarn") out = connections.update_connection(name="prod", master="spark://new:7077") diff --git a/tests/unit/test_external_jobs.py b/tests/unit/test_external_jobs.py index 36e3f11..7b22e21 100644 --- a/tests/unit/test_external_jobs.py +++ b/tests/unit/test_external_jobs.py @@ -5,6 +5,7 @@ from unittest.mock import patch import pytest from spark_executor.core import connection_store +from spark_executor.core.yarn_client import YarnError from spark_executor.models import ApplicationSummary, Connection from spark_executor.tools import connections, external_jobs from spark_executor.tools.requests import ListApplicationsRequest @@ -290,3 +291,15 @@ def test_list_applications_maps_yarn_json_to_summary(fresh_stores): assert summary.finished_time == 0 assert summary.tracking_url == "http://rm:8088/proxy/application_42" assert summary.progress == 12.5 + + +def test_list_applications_raises_on_404(fresh_stores): + external_jobs.conn_store.save( + Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088") + ) + with patch( + "spark_executor.tools.external_jobs.list_applications_yarn", + side_effect=YarnError("YARN list applications failed: 404"), + ): + with pytest.raises(YarnError, match="YARN list applications failed"): + external_jobs.list_applications("prod") diff --git a/tests/unit/test_fetch_url.py b/tests/unit/test_fetch_url.py index a74d3ac..33d556d 100644 --- a/tests/unit/test_fetch_url.py +++ b/tests/unit/test_fetch_url.py @@ -4,11 +4,25 @@ from unittest.mock import patch import httpx import pytest +from pydantic import ValidationError from spark_executor.core import connection_store from spark_executor.core.connection_store import ConnectionStore from spark_executor.models import Connection from spark_executor.tools import connections, fetch_url +from spark_executor.tools.requests import ListApplicationsRequest + + +def _stream_cm(resp): + """Wrap a response (or fake response) in a context manager for httpx.stream.""" + class _CM: + def __enter__(self): + return resp + + def __exit__(self, exc_type, exc, tb): + return False + + return _CM() def _fresh_stores(tmp_path, monkeypatch): @@ -37,7 +51,9 @@ def test_fetch_url_returns_body_and_status(fresh_stores): ) ) resp = httpx.Response(200, text="hello", headers={"content-type": "text/html"}) - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") assert out.url == "http://nm01.prod.internal:8042/node" assert out.status_code == 200 @@ -63,7 +79,9 @@ def test_fetch_url_truncates_body_over_1mb(fresh_stores): ) big = "x" * (fetch_url._MAX_BODY_BYTES + 1) resp = httpx.Response(200, text=big) - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp): + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ): out = fetch_url.fetch_url("http://nm01.prod.internal/big", "prod") assert out.truncated is True assert len(out.body) == fetch_url._MAX_BODY_BYTES @@ -82,7 +100,9 @@ def test_fetch_url_passes_auth_from_connection(fresh_stores): ) ) resp = httpx.Response(200, text="ok") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "auth") auth = m.call_args.kwargs["auth"] assert isinstance(auth, httpx.BasicAuth) @@ -103,26 +123,13 @@ def test_fetch_url_passes_ssl_verify_from_connection(fresh_stores): ) ) resp = httpx.Response(200, text="ok") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "insecure") assert m.call_args.kwargs["verify"] is False -def test_fetch_url_follows_redirects(fresh_stores): - fetch_url.conn_store.save( - Connection( - name="prod", - master="yarn", - yarn_rm_url="http://rm.prod.internal:8088", - url_allowlist=["*.prod.internal"], - ) - ) - resp = httpx.Response(200, text="ok") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: - fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") - assert m.call_args.kwargs["follow_redirects"] is True - - def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores): fetch_url.conn_store.save( Connection( @@ -133,7 +140,9 @@ def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("http://ccam50:8088/foo", "ccam") assert out.status_code == 200 assert out.body == "hello" @@ -150,7 +159,9 @@ def test_fetch_url_allows_host_matching_any_of_multiple_globs(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url( "http://history.prod.internal:18080/api/v1/info", "prod" ) @@ -195,7 +206,9 @@ def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("http://ccam50:8088/", "prod") assert out.status_code == 200 assert out.body == "hello" @@ -211,7 +224,9 @@ def test_fetch_url_accepts_ip_literal_when_in_url_allowlist(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("http://10.0.0.1/secret", "prod") assert out.status_code == 200 assert out.body == "hello" @@ -227,7 +242,9 @@ def test_fetch_url_accepts_https_when_in_url_allowlist(fresh_stores): ) ) resp = httpx.Response(200, text="hello") - with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: + with patch( + "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + ) as m: out = fetch_url.fetch_url("https://ccam1.example.com/secure", "prod") assert out.status_code == 200 assert out.body == "hello" @@ -271,3 +288,117 @@ def test_fetch_url_omitted_url_allowlist_defaults_to_empty_and_rejects(fresh_sto ) with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): fetch_url.fetch_url("http://nm.prod.internal/", "prod") + + +def test_fetch_url_rejects_redirect_to_disallowed_host(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="ccam", + master="yarn", + yarn_rm_url="http://ccam1:8088", + url_allowlist=["ccam*"], + ) + ) + redirect = httpx.Response(302, headers={"Location": "http://evil.com/"}) + with patch( + "spark_executor.tools.fetch_url.httpx.stream", + return_value=_stream_cm(redirect), + ) as m: + with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): + fetch_url.fetch_url("http://ccam50/foo", "ccam") + assert m.call_count == 1 + + +def test_fetch_url_follows_redirect_to_allowed_host(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + yarn_rm_url="http://rm.prod.internal:8088", + url_allowlist=["*.internal"], + ) + ) + redirect = httpx.Response( + 302, headers={"Location": "http://other.internal/"} + ) + final = httpx.Response(200, text="ok") + with patch( + "spark_executor.tools.fetch_url.httpx.stream", + side_effect=[_stream_cm(redirect), _stream_cm(final)], + ) as m: + out = fetch_url.fetch_url("http://foo.internal/", "prod") + assert out.status_code == 200 + assert out.body == "ok" + assert m.call_count == 2 + assert m.call_args_list[1].args[1] == "http://other.internal/" + + +def test_fetch_url_rejects_redirect_to_ip_literal_not_in_allowlist(fresh_stores): + fetch_url.conn_store.save( + Connection( + name="ccam", + master="yarn", + yarn_rm_url="http://ccam1:8088", + url_allowlist=["ccam*"], + ) + ) + redirect = httpx.Response(302, headers={"Location": "http://10.0.0.1/"}) + with patch( + "spark_executor.tools.fetch_url.httpx.stream", + return_value=_stream_cm(redirect), + ) as m: + with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): + fetch_url.fetch_url("http://ccam50/foo", "ccam") + assert m.call_count == 1 + + +def test_fetch_url_streams_body_and_truncates_without_buffering_full(fresh_stores): + """The response body must be read incrementally; .text must not be accessed.""" + + class _FakeStreamResponse: + def __init__(self, chunks): + self.status_code = 200 + self.headers = httpx.Headers({"content-type": "text/plain"}) + self.encoding = "utf-8" + self._chunks = chunks + + def iter_bytes(self): + yield from self._chunks + + @property + def text(self): + raise AssertionError("response.text should not be accessed") + + fetch_url.conn_store.save( + Connection( + name="prod", + master="yarn", + url_allowlist=["*.prod.internal"], + ) + ) + chunk_size = 400_000 + chunks = [ + b"a" * chunk_size, + b"b" * chunk_size, + b"c" * chunk_size, + ] + fake = _FakeStreamResponse(chunks) + with patch( + "spark_executor.tools.fetch_url.httpx.stream", + return_value=_stream_cm(fake), + ) as m: + out = fetch_url.fetch_url("http://nm01.prod.internal/big", "prod") + assert m.call_count == 1 + assert out.truncated is True + assert len(out.body) == fetch_url._MAX_BODY_BYTES + assert out.body.startswith("a" * chunk_size) + # Only the first 200 KB of the third chunk were consumed before truncation. + assert out.body[-200_000:] == "c" * 200_000 + + +def test_list_applications_request_limit_bounds(): + assert ListApplicationsRequest(connection_name="prod", limit=10000).limit == 10000 + with pytest.raises(ValidationError): + ListApplicationsRequest(connection_name="prod", limit=0) + with pytest.raises(ValidationError): + ListApplicationsRequest(connection_name="prod", limit=10001) From deafb5a26b38aca1f8146f2d32358f53a17a8c05 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 14:35:16 +0800 Subject: [PATCH 7/7] refactor: drop fetch_url body cap and modernize datetime usage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fetch_url: remove the 1MB body cap and the streaming helper. The manual redirect loop with allowlist re-check (the SSRF fix) is kept intact; the body is now read in full via resp.text. Drop the now- meaningless `truncated` field from FetchUrlResult and the tests that asserted on it. Switched from httpx.stream() back to httpx.get() for the redirect loop — cleaner without the body cap. - datetime: replace deprecated datetime.utcnow() with datetime.now(timezone.utc) in submit.py (3 sites) and core/job_writer.py (1 site). Update the stale comment in core/pending_store.py that referenced the old call. - Clean up an unused `from common.config import settings` import in submit.py that ruff flagged. Co-Authored-By: Claude --- spark_executor/core/job_writer.py | 4 +- spark_executor/core/pending_store.py | 2 +- spark_executor/models.py | 1 - spark_executor/server.py | 1 - spark_executor/tools/fetch_url.py | 84 +++++++-------------- spark_executor/tools/submit.py | 9 +-- tests/unit/test_fetch_url.py | 107 ++++----------------------- 7 files changed, 50 insertions(+), 158 deletions(-) diff --git a/spark_executor/core/job_writer.py b/spark_executor/core/job_writer.py index 27df059..796e2a2 100644 --- a/spark_executor/core/job_writer.py +++ b/spark_executor/core/job_writer.py @@ -14,7 +14,7 @@ Resolution order for the output directory: """ import os import secrets -from datetime import datetime +from datetime import datetime, timezone from common.config import settings from common.logging import logger @@ -63,7 +63,7 @@ def write_job_file(code: str, jobs_dir: str | None = None) -> str: effective_dir = resolve_jobs_dir(jobs_dir) os.makedirs(effective_dir, exist_ok=True) - stamp = datetime.utcnow().strftime("%Y%m%d%H%M%S") + stamp = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S") name = f"job_{stamp}_{secrets.token_hex(3)}.py" path = os.path.join(effective_dir, name) abs_path = os.path.abspath(path) diff --git a/spark_executor/core/pending_store.py b/spark_executor/core/pending_store.py index 5bfccde..0436d54 100644 --- a/spark_executor/core/pending_store.py +++ b/spark_executor/core/pending_store.py @@ -37,7 +37,7 @@ class PendingStore: return Path(self._data_dir) / self._dir_name def _date_path(self, created_at: datetime) -> Path: - # created_at is datetime.utcnow(), so the shard date is a UTC date. + # created_at is timezone-aware UTC, so the shard date is a UTC date. return self.dir_path / f"{created_at.date().isoformat()}.json" def _legacy_path(self) -> Path: diff --git a/spark_executor/models.py b/spark_executor/models.py index b389021..9e35af8 100644 --- a/spark_executor/models.py +++ b/spark_executor/models.py @@ -45,7 +45,6 @@ class FetchUrlResult(BaseModel): status_code: int content_type: str body: str - truncated: bool = False class ApplicationSummary(BaseModel): diff --git a/spark_executor/server.py b/spark_executor/server.py index e43efb0..4b26128 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -549,7 +549,6 @@ def _update_job_file(req: UpdateJobFileRequest): "guardrails — the allowlist is the only gate — so keep it tight. The " "Connection's saved auth is reused, so the agent does not need cluster " "credentials.\n\n" - "**Limits:** response body capped at 1 MB (truncated=true if larger), " "30s timeout, redirects followed." ), ) diff --git a/spark_executor/tools/fetch_url.py b/spark_executor/tools/fetch_url.py index 4813c30..9615372 100644 --- a/spark_executor/tools/fetch_url.py +++ b/spark_executor/tools/fetch_url.py @@ -21,7 +21,6 @@ from spark_executor.core.yarn_client import YarnClientConfig from spark_executor.models import FetchUrlResult from spark_executor.tools.connections import store as conn_store -_MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body _REQUEST_TIMEOUT_SECONDS = 30 _MAX_REDIRECTS = 3 @@ -70,29 +69,6 @@ def _validate_url_host(url: str, allowlist: list[str] | None) -> None: ) -def _read_streamed_body(resp: httpx.Response) -> tuple[bytes, bool]: - """Stream the response body, keeping at most ``_MAX_BODY_BYTES`` bytes. - - Returns ``(body_bytes, truncated)``. Stops reading as soon as the cap is - exceeded so that a multi-gigabyte response from an allowlisted host cannot - OOM the service. - """ - chunks: list[bytes] = [] - total = 0 - truncated = False - for chunk in resp.iter_bytes(): - if total + len(chunk) <= _MAX_BODY_BYTES: - chunks.append(chunk) - total += len(chunk) - else: - if total < _MAX_BODY_BYTES: - chunks.append(chunk[: _MAX_BODY_BYTES - total]) - total = _MAX_BODY_BYTES - truncated = True - break - return b"".join(chunks), truncated - - def fetch_url(url: str, connection_name: str) -> FetchUrlResult: """Proxy an HTTP GET to url using the auth/SSL settings of connection_name.""" logger.debug(f"fetch_url enter url={url} connection_name={connection_name}") @@ -107,45 +83,39 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult: verify = config.verify_for_httpx() for hop in range(_MAX_REDIRECTS + 1): - with httpx.stream( - "GET", + resp = httpx.get( url, auth=auth, verify=verify, timeout=_REQUEST_TIMEOUT_SECONDS, follow_redirects=False, - ) as resp: - if resp.status_code not in _REDIRECT_STATUSES: - body_bytes, truncated = _read_streamed_body(resp) - encoding = resp.encoding or "utf-8" - body = body_bytes.decode(encoding, errors="replace") - - logger.info( - f"fetch_url ok url={url} connection_name={connection_name} " - f"status_code={resp.status_code} " - f"content_type={resp.headers.get('content-type', '')} " - f"body_bytes={len(body)} truncated={truncated}" - ) - - return FetchUrlResult( - url=url, - status_code=resp.status_code, - content_type=resp.headers.get("content-type", ""), - body=body, - truncated=truncated, - ) - - location = resp.headers.get("Location") or resp.headers.get("location") - if not location: - raise ValueError( - f"HTTP GET returned {resp.status_code} with no Location header at {url}" - ) - next_url = urljoin(url, location) - logger.debug( - f"fetch_url redirect url={url} status={resp.status_code} to={next_url}" + ) + if resp.status_code not in _REDIRECT_STATUSES: + logger.info( + f"fetch_url ok url={url} connection_name={connection_name} " + f"status_code={resp.status_code} " + f"content_type={resp.headers.get('content-type', '')} " + f"body_chars={len(resp.text)}" ) - _validate_url_host(next_url, conn.url_allowlist) - url = next_url + + return FetchUrlResult( + url=url, + status_code=resp.status_code, + content_type=resp.headers.get("content-type", ""), + body=resp.text, + ) + + location = resp.headers.get("Location") or resp.headers.get("location") + if not location: + raise ValueError( + f"HTTP GET returned {resp.status_code} with no Location header at {url}" + ) + next_url = urljoin(url, location) + logger.debug( + f"fetch_url redirect url={url} status={resp.status_code} to={next_url}" + ) + _validate_url_host(next_url, conn.url_allowlist) + url = next_url raise ValueError( f"fetch_url exceeded {_MAX_REDIRECTS} redirects (last url: {url})" diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index 10ff531..04c2bda 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -6,9 +6,8 @@ import os import secrets import uuid -from datetime import datetime +from datetime import datetime, timezone -from common.config import settings from common.logging import logger from common.sql_guard import validate_pyspark_code from spark_executor.core.connection_store import store as conn_store @@ -124,7 +123,7 @@ def prepare_submit_job( num_executors=num_executors, spark_conf=dict(conn.spark_conf), extra_args=dict(extra_args or {}), - created_at=datetime.utcnow(), + created_at=datetime.now(timezone.utc), status="PENDING", ) pending_store.save(pending) @@ -248,7 +247,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: application_id=application_id, script_path=pending.script_path, queue=pending.queue, - submit_time=datetime.utcnow(), + submit_time=datetime.now(timezone.utc), connection=pending.connection, yarn_rm_url=pending.yarn_rm_url, ) @@ -282,7 +281,7 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: application_id=application_id, script_path=pending.script_path, queue=pending.queue, - submit_time=datetime.utcnow(), + submit_time=datetime.now(timezone.utc), connection=pending.connection, yarn_rm_url=pending.yarn_rm_url, ) diff --git a/tests/unit/test_fetch_url.py b/tests/unit/test_fetch_url.py index 33d556d..affe030 100644 --- a/tests/unit/test_fetch_url.py +++ b/tests/unit/test_fetch_url.py @@ -13,18 +13,6 @@ from spark_executor.tools import connections, fetch_url from spark_executor.tools.requests import ListApplicationsRequest -def _stream_cm(resp): - """Wrap a response (or fake response) in a context manager for httpx.stream.""" - class _CM: - def __enter__(self): - return resp - - def __exit__(self, exc_type, exc, tb): - return False - - return _CM() - - def _fresh_stores(tmp_path, monkeypatch): """Reset connection store singletons for a single test.""" monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path)) @@ -52,14 +40,14 @@ def test_fetch_url_returns_body_and_status(fresh_stores): ) resp = httpx.Response(200, text="hello", headers={"content-type": "text/html"}) with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + "spark_executor.tools.fetch_url.httpx.get", return_value=resp ) as m: out = fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod") assert out.url == "http://nm01.prod.internal:8042/node" assert out.status_code == 200 assert out.content_type == "text/html" assert out.body == "hello" - assert out.truncated is False + assert "truncated" not in fetch_url.FetchUrlResult.model_fields assert m.call_count == 1 @@ -68,25 +56,6 @@ def test_fetch_url_raises_for_missing_connection(fresh_stores): fetch_url.fetch_url("http://rm.prod.internal:8088/", "missing") -def test_fetch_url_truncates_body_over_1mb(fresh_stores): - fetch_url.conn_store.save( - Connection( - name="prod", - master="yarn", - yarn_rm_url="http://rm.prod.internal:8088", - url_allowlist=["*.prod.internal"], - ) - ) - big = "x" * (fetch_url._MAX_BODY_BYTES + 1) - resp = httpx.Response(200, text=big) - with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) - ): - out = fetch_url.fetch_url("http://nm01.prod.internal/big", "prod") - assert out.truncated is True - assert len(out.body) == fetch_url._MAX_BODY_BYTES - - def test_fetch_url_passes_auth_from_connection(fresh_stores): fetch_url.conn_store.save( Connection( @@ -101,7 +70,7 @@ def test_fetch_url_passes_auth_from_connection(fresh_stores): ) resp = httpx.Response(200, text="ok") with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + "spark_executor.tools.fetch_url.httpx.get", return_value=resp ) as m: fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "auth") auth = m.call_args.kwargs["auth"] @@ -124,7 +93,7 @@ def test_fetch_url_passes_ssl_verify_from_connection(fresh_stores): ) resp = httpx.Response(200, text="ok") with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + "spark_executor.tools.fetch_url.httpx.get", return_value=resp ) as m: fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "insecure") assert m.call_args.kwargs["verify"] is False @@ -141,7 +110,7 @@ def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores): ) resp = httpx.Response(200, text="hello") with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + "spark_executor.tools.fetch_url.httpx.get", return_value=resp ) as m: out = fetch_url.fetch_url("http://ccam50:8088/foo", "ccam") assert out.status_code == 200 @@ -160,7 +129,7 @@ def test_fetch_url_allows_host_matching_any_of_multiple_globs(fresh_stores): ) resp = httpx.Response(200, text="hello") with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + "spark_executor.tools.fetch_url.httpx.get", return_value=resp ) as m: out = fetch_url.fetch_url( "http://history.prod.internal:18080/api/v1/info", "prod" @@ -207,7 +176,7 @@ def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): ) resp = httpx.Response(200, text="hello") with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + "spark_executor.tools.fetch_url.httpx.get", return_value=resp ) as m: out = fetch_url.fetch_url("http://ccam50:8088/", "prod") assert out.status_code == 200 @@ -225,7 +194,7 @@ def test_fetch_url_accepts_ip_literal_when_in_url_allowlist(fresh_stores): ) resp = httpx.Response(200, text="hello") with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + "spark_executor.tools.fetch_url.httpx.get", return_value=resp ) as m: out = fetch_url.fetch_url("http://10.0.0.1/secret", "prod") assert out.status_code == 200 @@ -243,7 +212,7 @@ def test_fetch_url_accepts_https_when_in_url_allowlist(fresh_stores): ) resp = httpx.Response(200, text="hello") with patch( - "spark_executor.tools.fetch_url.httpx.stream", return_value=_stream_cm(resp) + "spark_executor.tools.fetch_url.httpx.get", return_value=resp ) as m: out = fetch_url.fetch_url("https://ccam1.example.com/secure", "prod") assert out.status_code == 200 @@ -301,8 +270,8 @@ def test_fetch_url_rejects_redirect_to_disallowed_host(fresh_stores): ) redirect = httpx.Response(302, headers={"Location": "http://evil.com/"}) with patch( - "spark_executor.tools.fetch_url.httpx.stream", - return_value=_stream_cm(redirect), + "spark_executor.tools.fetch_url.httpx.get", + return_value=redirect, ) as m: with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): fetch_url.fetch_url("http://ccam50/foo", "ccam") @@ -323,14 +292,14 @@ def test_fetch_url_follows_redirect_to_allowed_host(fresh_stores): ) final = httpx.Response(200, text="ok") with patch( - "spark_executor.tools.fetch_url.httpx.stream", - side_effect=[_stream_cm(redirect), _stream_cm(final)], + "spark_executor.tools.fetch_url.httpx.get", + side_effect=[redirect, final], ) as m: out = fetch_url.fetch_url("http://foo.internal/", "prod") assert out.status_code == 200 assert out.body == "ok" assert m.call_count == 2 - assert m.call_args_list[1].args[1] == "http://other.internal/" + assert m.call_args_list[1].args[0] == "http://other.internal/" def test_fetch_url_rejects_redirect_to_ip_literal_not_in_allowlist(fresh_stores): @@ -344,58 +313,14 @@ def test_fetch_url_rejects_redirect_to_ip_literal_not_in_allowlist(fresh_stores) ) redirect = httpx.Response(302, headers={"Location": "http://10.0.0.1/"}) with patch( - "spark_executor.tools.fetch_url.httpx.stream", - return_value=_stream_cm(redirect), + "spark_executor.tools.fetch_url.httpx.get", + return_value=redirect, ) as m: with pytest.raises(ValueError, match="is not in Connection.url_allowlist"): fetch_url.fetch_url("http://ccam50/foo", "ccam") assert m.call_count == 1 -def test_fetch_url_streams_body_and_truncates_without_buffering_full(fresh_stores): - """The response body must be read incrementally; .text must not be accessed.""" - - class _FakeStreamResponse: - def __init__(self, chunks): - self.status_code = 200 - self.headers = httpx.Headers({"content-type": "text/plain"}) - self.encoding = "utf-8" - self._chunks = chunks - - def iter_bytes(self): - yield from self._chunks - - @property - def text(self): - raise AssertionError("response.text should not be accessed") - - fetch_url.conn_store.save( - Connection( - name="prod", - master="yarn", - url_allowlist=["*.prod.internal"], - ) - ) - chunk_size = 400_000 - chunks = [ - b"a" * chunk_size, - b"b" * chunk_size, - b"c" * chunk_size, - ] - fake = _FakeStreamResponse(chunks) - with patch( - "spark_executor.tools.fetch_url.httpx.stream", - return_value=_stream_cm(fake), - ) as m: - out = fetch_url.fetch_url("http://nm01.prod.internal/big", "prod") - assert m.call_count == 1 - assert out.truncated is True - assert len(out.body) == fetch_url._MAX_BODY_BYTES - assert out.body.startswith("a" * chunk_size) - # Only the first 200 KB of the third chunk were consumed before truncation. - assert out.body[-200_000:] == "c" * 200_000 - - def test_list_applications_request_limit_bounds(): assert ListApplicationsRequest(connection_name="prod", limit=10000).limit == 10000 with pytest.raises(ValidationError):