tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
0291f36b013f4ba36214455bb2145ad98e8fa056
mcp-server/tests/unit/test_fetch_url.py
T
ClaudeandClaude Fable 5 0291f36b01 feat(fetch_url): per-Connection allowed_url_hosts glob allowlist 
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 <noreply@anthropic.com>
2026-07-09 11:22:55 +08:00

287 lines
11 KiB
Python
Raw Blame History

# coding=utf-8
from pathlib import Path
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
@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(
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
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")
Reference in New Issue View Git Blame Copy Permalink