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:
@@ -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 工具
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
+11
-11
@@ -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'."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()]
|
||||
|
||||
@@ -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=['<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."
|
||||
)
|
||||
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(
|
||||
|
||||
@@ -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."
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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}"
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user