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:
Claude
2026-07-09 14:23:16 +08:00
co-authored by Claude
parent 7fbad97a87
commit 7d4e512cb0
9 changed files with 331 additions and 93 deletions
-3
View File
@@ -299,9 +299,6 @@ def list_applications(
resp = _request("GET", url, params=params, resp = _request("GET", url, params=params,
verify=config.verify_for_httpx(), verify=config.verify_for_httpx(),
auth=config.auth_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: if resp.status_code >= 400:
raise YarnError( raise YarnError(
f"YARN list applications failed: {resp.status_code} {resp.text[:200]}" f"YARN list applications failed: {resp.status_code} {resp.text[:200]}"
+10 -2
View File
@@ -390,12 +390,20 @@ def _list_applications(req: ListApplicationsRequest):
"yarn_rm_url is required for get_external_* to work** — spark-submit " "yarn_rm_url is required for get_external_* to work** — spark-submit "
"can discover the RM for submissions, but direct YARN REST queries " "can discover the RM for submissions, but direct YARN REST queries "
"need an explicit URL. Saving with an existing name overwrites the " "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): 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 # 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( @app.post(
+60 -30
View File
@@ -4,45 +4,75 @@
@Author :tao.chen @Author :tao.chen
""" """
from common.logging import logger 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 from spark_executor.models import Connection
_UNSET = object()
def save_connection( def save_connection(
*, *,
name: str, name: str,
master: str, master: str,
deploy_mode: str = "cluster", deploy_mode: str = _UNSET, # type: ignore[assignment]
yarn_rm_url: str | None = None, yarn_rm_url: str | None = _UNSET, # type: ignore[assignment]
spark_conf: dict[str, str] | None = None, spark_conf: dict[str, str] | None = _UNSET, # type: ignore[assignment]
ssl_verify: bool | None = None, ssl_verify: bool | None = _UNSET, # type: ignore[assignment]
ssl_ca_bundle: str | None = None, ssl_ca_bundle: str | None = _UNSET, # type: ignore[assignment]
auth_type: str = "none", auth_type: str = _UNSET, # type: ignore[assignment]
auth_user: str | None = None, auth_user: str | None = _UNSET, # type: ignore[assignment]
auth_password: str | None = None, auth_password: str | None = _UNSET, # type: ignore[assignment]
auth_principal: str | None = None, auth_principal: str | None = _UNSET, # type: ignore[assignment]
auth_keytab: str | None = None, auth_keytab: str | None = _UNSET, # type: ignore[assignment]
url_allowlist: list[str] | None = None, url_allowlist: list[str] | None = _UNSET, # type: ignore[assignment]
) -> dict[str, str]: ) -> dict[str, object]:
logger.debug( logger.debug(
f"save_connection enter name={name} master={master} deploy_mode={deploy_mode} " f"save_connection enter name={name} master={master} "
f"yarn_rm_url={yarn_rm_url} spark_conf_keys={list((spark_conf or {}).keys())}" f"spark_conf_keys={list((spark_conf if isinstance(spark_conf, dict) else {}).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 [],
) )
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) store.save(conn)
return {"name": name, "status": "SAVED"} return {"name": name, "status": "SAVED"}
+58 -10
View File
@@ -12,7 +12,7 @@ are intentionally NOT performed. Reuses the connection's saved auth/SSL
config so the agent doesn't need cluster credentials. config so the agent doesn't need cluster credentials.
""" """
import fnmatch import fnmatch
from urllib.parse import urlparse from urllib.parse import urlparse, urljoin
import httpx 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 _MAX_BODY_BYTES = 1_000_000 # 1 MB cap on response body
_REQUEST_TIMEOUT_SECONDS = 30 _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: def _host_matches_any_glob(host: str, allowlist: list[str]) -> bool:
"""True if `host` matches any of the fnmatch glob patterns. """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: def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
"""Proxy an HTTP GET to url using the auth/SSL settings of connection_name.""" """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}") 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) _validate_url_host(url, conn.url_allowlist)
config = YarnClientConfig.from_connection(conn) config = YarnClientConfig.from_connection(conn)
resp = httpx.get( auth = config.auth_for_httpx()
url, verify = config.verify_for_httpx()
auth=config.auth_for_httpx(),
verify=config.verify_for_httpx(),
timeout=_REQUEST_TIMEOUT_SECONDS,
follow_redirects=True,
)
body = resp.text[:_MAX_BODY_BYTES] for hop in range(_MAX_REDIRECTS + 1):
truncated = len(resp.text) > _MAX_BODY_BYTES 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( logger.info(
f"fetch_url ok url={url} connection_name={connection_name} " 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, body=body,
truncated=truncated, 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})"
)
+2
View File
@@ -504,6 +504,8 @@ class ListApplicationsRequest(BaseModel):
) )
limit: int = Field( limit: int = Field(
default=100, default=100,
ge=1,
le=10000,
description=( description=(
"Maximum number of applications to return. YARN has no " "Maximum number of applications to return. YARN has no "
"offset-based pagination, so for large clusters use state/queue " "offset-based pagination, so for large clusters use state/queue "
+9
View File
@@ -121,6 +121,15 @@ 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_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): def test_update_connection_changes_specified_field(_fresh_store):
connections.save_connection(name="prod", master="yarn") connections.save_connection(name="prod", master="yarn")
out = connections.update_connection(name="prod", master="spark://new:7077") out = connections.update_connection(name="prod", master="spark://new:7077")
+13
View File
@@ -5,6 +5,7 @@ from unittest.mock import patch
import pytest import pytest
from spark_executor.core import connection_store 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.models import ApplicationSummary, Connection
from spark_executor.tools import connections, external_jobs from spark_executor.tools import connections, external_jobs
from spark_executor.tools.requests import ListApplicationsRequest 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.finished_time == 0
assert summary.tracking_url == "http://rm:8088/proxy/application_42" assert summary.tracking_url == "http://rm:8088/proxy/application_42"
assert summary.progress == 12.5 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
View File
@@ -4,11 +4,25 @@ from unittest.mock import patch
import httpx import httpx
import pytest import pytest
from pydantic import ValidationError
from spark_executor.core import connection_store from spark_executor.core import connection_store
from spark_executor.core.connection_store import ConnectionStore from spark_executor.core.connection_store import ConnectionStore
from spark_executor.models import Connection from spark_executor.models import Connection
from spark_executor.tools import connections, fetch_url 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): 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"}) 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") out = fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "prod")
assert out.url == "http://nm01.prod.internal:8042/node" assert out.url == "http://nm01.prod.internal:8042/node"
assert out.status_code == 200 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) big = "x" * (fetch_url._MAX_BODY_BYTES + 1)
resp = httpx.Response(200, text=big) 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") out = fetch_url.fetch_url("http://nm01.prod.internal/big", "prod")
assert out.truncated is True assert out.truncated is True
assert len(out.body) == fetch_url._MAX_BODY_BYTES 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") 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") fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "auth")
auth = m.call_args.kwargs["auth"] auth = m.call_args.kwargs["auth"]
assert isinstance(auth, httpx.BasicAuth) 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") 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") fetch_url.fetch_url("http://nm01.prod.internal:8042/node", "insecure")
assert m.call_args.kwargs["verify"] is False 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): def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores):
fetch_url.conn_store.save( fetch_url.conn_store.save(
Connection( Connection(
@@ -133,7 +140,9 @@ def test_fetch_url_allows_host_matching_glob_pattern(fresh_stores):
) )
) )
resp = httpx.Response(200, text="hello") 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") out = fetch_url.fetch_url("http://ccam50:8088/foo", "ccam")
assert out.status_code == 200 assert out.status_code == 200
assert out.body == "hello" 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") 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( out = fetch_url.fetch_url(
"http://history.prod.internal:18080/api/v1/info", "prod" "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") 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") out = fetch_url.fetch_url("http://ccam50:8088/", "prod")
assert out.status_code == 200 assert out.status_code == 200
assert out.body == "hello" 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") 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") out = fetch_url.fetch_url("http://10.0.0.1/secret", "prod")
assert out.status_code == 200 assert out.status_code == 200
assert out.body == "hello" 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") 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") out = fetch_url.fetch_url("https://ccam1.example.com/secure", "prod")
assert out.status_code == 200 assert out.status_code == 200
assert out.body == "hello" 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"): with pytest.raises(ValueError, match="is not in Connection.url_allowlist"):
fetch_url.fetch_url("http://nm.prod.internal/", "prod") 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)