fix: address review findings on fetch-url-tool
- 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 <noreply@anthropic.com>
This commit is contained in:
@@ -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 脚本 |
|
||||
|
||||
@@ -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]}"
|
||||
|
||||
@@ -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):
|
||||
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(**req.model_dump(exclude_none=True))
|
||||
return save_connection(name=name, **fields)
|
||||
return update_connection(name=name, **fields)
|
||||
|
||||
|
||||
@app.post(
|
||||
|
||||
@@ -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())}"
|
||||
)
|
||||
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 [],
|
||||
f"save_connection enter name={name} master={master} "
|
||||
f"spark_conf_keys={list((spark_conf if isinstance(spark_conf, dict) else {}).keys())}"
|
||||
)
|
||||
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"}
|
||||
|
||||
|
||||
@@ -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,16 +103,22 @@ 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} "
|
||||
@@ -102,3 +134,19 @@ def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
|
||||
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})"
|
||||
)
|
||||
|
||||
@@ -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 "
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
|
||||
+155
-24
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user