tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: simplify fetch_url allowlist + add update_connection tool

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/test_fetch_url.py
+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):
_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")