feat: simplify fetch_url allowlist + add update_connection tool

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 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-07-09 11:46:07 +08:00
co-authored by Claude Fable 5
parent 0291f36b01
commit 6d7387b022
10 changed files with 279 additions and 93 deletions
+6 -5
View File
@@ -100,10 +100,11 @@ MCP 客户端需要先执行 `initialize` 握手,拿到 `mcp-session-id` 后
| 工具 | 作用 | | 工具 | 作用 |
| --- | --- | | --- | --- |
| `save_connection` | 保存或更新一个 Spark/YARN 连接配置 | | `save_connection` | 保存或更新一个 Spark/YARN 连接配置 |
| `list_connections` | 列出所有连接配置 | | `list_connections` | 列出所有连接配置 |
| `get_connection` | 按名称读取连接配置 | | `get_connection` | 按名称读取连接配置 |
| `delete_connection` | 删除连接配置 | | `delete_connection` | 删除连接配置 |
| `update_connection` | 部分更新一个已存在的连接 (PATCH 语义, 只改提供的字段) |
| `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 | | `write_job_file` | 将 LLM 已生成的 PySpark 代码写入服务端文件 |
| `read_job_file` | 读取已存在的 PySpark 脚本内容 | | `read_job_file` | 读取已存在的 PySpark 脚本内容 |
| `update_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_status` | 查询**非本服务提交**的外部 YARN application 状态(按 `application_id` + `connection_name` |
| `get_external_job_result` | 查询外部 YARN application 终态结果视图 | | `get_external_job_result` | 查询外部 YARN application 终态结果视图 |
| `get_external_job_logs` | 拉取外部 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 工具 ### Files MCP 工具
+23
View File
@@ -77,6 +77,29 @@ class ConnectionStore:
self._dump(records) self._dump(records)
logger.info(f"connection saved name={conn.name} master={conn.master}") 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: def delete(self, name: str) -> bool:
with self._lock: with self._lock:
records = self._load() records = self._load()
+11 -11
View File
@@ -69,18 +69,18 @@ class Connection(BaseModel):
auth_principal: str | None = None auth_principal: str | None = None
auth_keytab: str | None = None auth_keytab: str | None = None
allowed_url_hosts: list[str] | None = Field( allowed_url_hosts: list[str] = Field(
default=None, default_factory=list,
description=( description=(
"Optional list of fnmatch glob patterns for hosts the fetch_url tool " "List of fnmatch glob patterns for hosts the fetch_url tool may access. "
"may access, in addition to the default 'shares >= 2 labels of suffix " "The list is mandatory-opt-in: an empty list (the default) denies all "
"with yarn_rm_url host' rule. Useful for clusters whose hostnames do " "hosts, so you must populate it before fetch_url can access any URL. "
"NOT share a 2+ label suffix — e.g. single-label hosts like 'ccam1'-" "Useful for clusters whose hostnames do NOT share a common suffix — "
"'ccam99' (configure ['ccam*']) or HDFS namenode on a different " "e.g. single-label hosts like 'ccam1'-'ccam99' (configure ['ccam*']) "
"subdomain ('*.hadoop.internal'). Patterns are matched against the " "or HDFS namenode on a different subdomain ('*.hadoop.internal'). "
"URL host only (no port, no path). fnmatch rules apply: '*' does NOT " "Patterns are matched against the URL host only (no port, no path). "
"match '.', so 'ccam*' matches 'ccam50' but not 'ccam50.evil.com'. " "fnmatch rules apply: '*' does NOT match '.', so 'ccam*' matches "
"Default None means the suffix rule alone applies." "'ccam50' but not 'ccam50.evil.com'."
), ),
) )
+24
View File
@@ -11,6 +11,7 @@ from spark_executor.tools.connections import (
get_connection, get_connection,
list_connections, list_connections,
save_connection, save_connection,
update_connection,
) )
from spark_executor.tools.write_job import write_job_file from spark_executor.tools.write_job import write_job_file
from spark_executor.tools.job_file import read_job_file, update_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, PrepareSubmitJobRequest,
ReadJobFileRequest, ReadJobFileRequest,
SaveConnectionRequest, SaveConnectionRequest,
UpdateConnectionRequest,
UpdateJobFileRequest, UpdateJobFileRequest,
UpdatePendingJobRequest, UpdatePendingJobRequest,
) )
@@ -396,6 +398,28 @@ def _get_connection(req: ConnectionNameRequest):
return get_connection(req.name) 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( @app.post(
"/delete_connection", "/delete_connection",
operation_id="delete_connection", operation_id="delete_connection",
+17
View File
@@ -22,6 +22,7 @@ def save_connection(
auth_password: str | None = None, auth_password: str | None = None,
auth_principal: str | None = None, auth_principal: str | None = None,
auth_keytab: str | None = None, auth_keytab: str | None = None,
allowed_url_hosts: list[str] | None = None,
) -> dict[str, str]: ) -> dict[str, str]:
logger.debug( logger.debug(
f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} "
@@ -40,11 +41,27 @@ def save_connection(
auth_password=auth_password, auth_password=auth_password,
auth_principal=auth_principal, auth_principal=auth_principal,
auth_keytab=auth_keytab, auth_keytab=auth_keytab,
allowed_url_hosts=allowed_url_hosts or [],
) )
store.save(conn) store.save(conn)
return {"name": name, "status": "SAVED"} 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]]: def list_connections() -> list[dict[str, object]]:
logger.debug("list_connections enter") logger.debug("list_connections enter")
return [c.model_dump() for c in store.list_all()] return [c.model_dump() for c in store.list_all()]
+17 -43
View File
@@ -5,9 +5,10 @@
Generic HTTP GET proxy for the agent. Lets the agent fetch URLs on the 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: 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 the host is checked against an explicit fnmatch glob allowlist configured
host; IP literals and non-HTTP schemes are rejected. Reuses the connection's on the connection (`allowed_url_hosts`). Empty or missing allowlist means
saved auth/SSL config so the agent doesn't need cluster credentials. 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 ipaddress
import fnmatch import fnmatch
@@ -24,19 +25,6 @@ _MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body
_REQUEST_TIMEOUT_SECONDS = 30 _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: def _host_matches_any_glob(host: str, patterns: list[str]) -> bool:
"""True if `host` matches any of the fnmatch glob patterns. """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( def _validate_url_host(
url: str, url: str,
yarn_rm_url: str | None, allowed_hosts: list[str] | None,
allowed_hosts: list[str] | None = None,
) -> None: ) -> None:
"""Raise ValueError if the URL is not allowed to be fetched. """Raise ValueError if the URL is not allowed to be fetched.
Allowed if EITHER: Allowed only if the URL host matches one of the fnmatch glob patterns in
- URL host matches one of `allowed_hosts` fnmatch globs (if provided) `allowed_hosts`. An empty or missing allowlist rejects everything.
- URL host shares >= 2 labels of suffix with `yarn_rm_url` host
""" """
parsed = urlparse(url) parsed = urlparse(url)
if parsed.scheme not in ("http", "https"): if parsed.scheme not in ("http", "https"):
@@ -85,29 +71,17 @@ def _validate_url_host(
if "IP literal" in str(e): if "IP literal" in str(e):
raise raise
# not an IP, continue # not an IP, continue
# Glob allowlist — if any pattern matches, host is allowed regardless of suffix if not allowed_hosts:
if allowed_hosts and _host_matches_any_glob(host, 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 return
if not yarn_rm_url:
raise ValueError( raise ValueError(
f"Connection has no yarn_rm_url set, cannot validate URL host " f"URL host {host!r} is not in Connection.allowed_url_hosts "
f"(no allowed_url_hosts match either). Save a Connection with " f"{allowed_hosts!r}. Reject this fetch to prevent SSRF."
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=['<glob>'] 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."
) )
def fetch_url(url: str, connection_name: str) -> FetchUrlResult: def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
@@ -117,7 +91,7 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
if conn is None: if conn is None:
raise KeyError(f"Connection not found: {connection_name}") 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) config = YarnClientConfig.from_connection(conn)
resp = httpx.get( resp = httpx.get(
+66 -2
View File
@@ -128,8 +128,10 @@ class SaveConnectionRequest(BaseModel):
"Optional list of fnmatch glob patterns for hosts the fetch_url tool " "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.allowed_url_hosts for full semantics. "
"Example for single-label host clusters: ['ccam*'] allows any host " "Example for single-label host clusters: ['ccam*'] allows any host "
"starting with 'ccam' (ccam1, ccam2, ..., ccam99). Default None means " "starting with 'ccam' (ccam1, ccam2, ..., ccam99). If omitted, None, "
"only the default 'shares >= 2 labels of suffix with yarn_rm_url' rule." "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." "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."
),
)
+2 -1
View File
@@ -64,13 +64,14 @@ def test_seventeen_tool_routes_registered():
assert "/update_pending_job" in paths 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} paths = {r.path for r in app.routes}
for path in ( for path in (
"/get_external_job_logs", "/get_external_job_logs",
"/get_external_job_status", "/get_external_job_status",
"/get_external_job_result", "/get_external_job_result",
"/fetch_url", "/fetch_url",
"/update_connection",
): ):
assert path in paths, f"missing MCP tool route: {path}" assert path in paths, f"missing MCP tool route: {path}"
+66
View File
@@ -1,4 +1,5 @@
# coding=utf-8 # coding=utf-8
import json
from pathlib import Path from pathlib import Path
import pytest import pytest
@@ -120,3 +121,68 @@ def test_delete_connection_returns_status(_fresh_store):
def test_delete_connection_unknown_raises(_fresh_store): def test_delete_connection_unknown_raises(_fresh_store):
with pytest.raises(KeyError): with pytest.raises(KeyError):
connections.delete_connection("missing") 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"
+45 -29
View File
@@ -29,7 +29,12 @@ def fresh_stores(tmp_path: Path, monkeypatch):
def test_fetch_url_returns_body_and_status(tmp_path, monkeypatch): def test_fetch_url_returns_body_and_status(tmp_path, monkeypatch):
_fresh_stores(tmp_path, monkeypatch) _fresh_stores(tmp_path, monkeypatch)
fetch_url.conn_store.save( 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"}) 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.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): def test_fetch_url_rejects_external_host(tmp_path, monkeypatch):
_fresh_stores(tmp_path, monkeypatch) _fresh_stores(tmp_path, monkeypatch)
fetch_url.conn_store.save( 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") 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) _fresh_stores(tmp_path, monkeypatch)
fetch_url.conn_store.save( 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") 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) _fresh_stores(tmp_path, monkeypatch)
fetch_url.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url=None)) fetch_url.conn_store.save(
with pytest.raises(ValueError, match="no yarn_rm_url set"): 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") fetch_url.fetch_url("http://anything.com/", "prod")
def test_fetch_url_truncates_body_over_1mb(tmp_path, monkeypatch): def test_fetch_url_truncates_body_over_1mb(tmp_path, monkeypatch):
_fresh_stores(tmp_path, monkeypatch) _fresh_stores(tmp_path, monkeypatch)
fetch_url.conn_store.save( 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) big = "x" * (fetch_url._MAX_BODY_BYTES + 1)
resp = httpx.Response(200, text=big) 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_type="basic",
auth_user="u", auth_user="u",
auth_password="p", auth_password="p",
allowed_url_hosts=["*.prod.internal"],
) )
) )
resp = httpx.Response(200, text="ok") resp = httpx.Response(200, text="ok")
@@ -152,6 +171,7 @@ def test_fetch_url_passes_ssl_verify_from_connection(tmp_path, monkeypatch):
master="yarn", master="yarn",
yarn_rm_url="http://rm.prod.internal:8088", yarn_rm_url="http://rm.prod.internal:8088",
ssl_verify=False, ssl_verify=False,
allowed_url_hosts=["*.prod.internal"],
) )
) )
resp = httpx.Response(200, text="ok") 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): def test_fetch_url_follows_redirects(tmp_path, monkeypatch):
_fresh_stores(tmp_path, monkeypatch) _fresh_stores(tmp_path, monkeypatch)
fetch_url.conn_store.save( 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") 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.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*"], 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") fetch_url.fetch_url("http://evil.com/foo", "ccam")
def test_fetch_url_glob_does_not_cross_dot_boundary(fresh_stores): 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*"], 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") fetch_url.fetch_url("http://ccam50.evil.com/", "ccam")
def test_fetch_url_glob_match_does_not_require_suffix_overlap(fresh_stores): 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 out.body == "hello"
assert m.call_count == 1 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( fetch_url.conn_store.save(
Connection( Connection(
name="prod", name="prod",
@@ -251,28 +276,19 @@ def test_fetch_url_empty_allowed_hosts_falls_back_to_suffix_rule(fresh_stores):
allowed_url_hosts=[], allowed_url_hosts=[],
) )
) )
resp = httpx.Response(200, text="hello") with pytest.raises(ValueError, match="allowed_url_hosts is empty"):
with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: fetch_url.fetch_url("http://nm.prod.internal/", "prod")
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): def test_fetch_url_omitted_allowed_url_hosts_defaults_to_empty_and_rejects(fresh_stores):
fetch_url.conn_store.save( fetch_url.conn_store.save(
Connection( Connection(
name="prod", name="prod",
master="yarn", master="yarn",
yarn_rm_url="http://rm.prod.internal", yarn_rm_url="http://rm.prod.internal",
allowed_url_hosts=None,
) )
) )
resp = httpx.Response(200, text="hello") with pytest.raises(ValueError, match="allowed_url_hosts is empty"):
with patch("spark_executor.tools.fetch_url.httpx.get", return_value=resp) as m: fetch_url.fetch_url("http://nm.prod.internal/", "prod")
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): def test_fetch_url_error_message_hints_at_allowed_url_hosts(fresh_stores):
fetch_url.conn_store.save( 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", 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") fetch_url.fetch_url("http://evil.com/foo", "prod")