diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/PythonCodegenBase.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/PythonCodegenBase.scala index 7cd305bfb6e..bac14f0a8a0 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/PythonCodegenBase.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/huggingFace/codegen/PythonCodegenBase.scala @@ -128,6 +128,10 @@ object PythonCodegenBase { | # Hard cap on bytes pulled from an external (user/response-provided) URL. | MAX_REMOTE_FETCH_BYTES = 25 * 1024 * 1024 | + | # Redirect-chain bound for _fetch_remote_url. Redirects are followed + | # manually so every hop is re-validated (https-only, public address). + | MAX_REDIRECT_HOPS = 5 + | | def open(self): | # User-provided strings reach the operator via base64-encoded | # decode expressions so they cannot break Python syntax or @@ -704,13 +708,16 @@ object PythonCodegenBase { | # branches of _call_provider). | # ────────────────────────────────────────────────────────────────── | - | def _fetch_remote_url(self, url): - | '''Fetch an external URL with SSRF hardening. Returns (content_type, data). - | Enforces https-only, rejects private/loopback/link-local/reserved - | addresses (covers the 169.254.169.254 cloud-metadata endpoint), and - | caps the response at MAX_REMOTE_FETCH_BYTES. The address check runs - | before the request, so it mitigates but does not fully prevent DNS - | rebinding (requests re-resolves on connect). + | def _validate_remote_url(self, url): + | '''Validate one URL before it is fetched: https-only, has a host, + | and every address the host resolves to is a globally-routable + | public address (rejects private/loopback/link-local/reserved/ + | multicast/CGNAT addresses, covering the 169.254.169.254 + | cloud-metadata endpoint). Called on the original + | URL and again on every redirect hop, so a redirect can neither + | downgrade the scheme nor point at an internal address. The + | address check runs before the request, so it mitigates but does + | not fully prevent DNS rebinding (requests re-resolves on connect). | ''' | import ipaddress | import socket @@ -727,10 +734,43 @@ object PythonCodegenBase { | raise ValueError(f"Could not resolve host '{host}': {e}") | for info in addrinfos: | ip = ipaddress.ip_address(info[4][0]) - | if (ip.is_private or ip.is_loopback or ip.is_link_local - | or ip.is_reserved or ip.is_multicast or ip.is_unspecified): + | # Allowlist stance: require a globally-routable address. The + | # explicit predicates stay because is_global misses some ranges + | # (CPython reports multicast as global) — belt-and-suspenders, + | # and it also covers the CGNAT/shared range (100.64.0.0/10) + | # that the predicate list alone lets through. + | if (not ip.is_global or ip.is_private or ip.is_loopback + | or ip.is_link_local or ip.is_reserved or ip.is_multicast + | or ip.is_unspecified): | raise ValueError(f"Refusing to fetch from non-public address {ip}.") - | resp = requests.get(url, timeout=120, stream=True) + | + | def _fetch_remote_url(self, url): + | '''Fetch an external URL with SSRF hardening. Returns (content_type, data). + | Redirects are never followed automatically: each hop is re-validated + | by _validate_remote_url and the chain is bounded by + | MAX_REDIRECT_HOPS, so a redirect cannot escape the https-only / + | public-address checks. The body is capped at MAX_REMOTE_FETCH_BYTES. + | ''' + | from urllib.parse import urljoin as _urljoin + | current_url = url + | resp = None + | for _hop in range(self.MAX_REDIRECT_HOPS + 1): + | self._validate_remote_url(current_url) + | resp = requests.get(current_url, timeout=120, stream=True, allow_redirects=False) + | if resp.status_code in (301, 302, 303, 307, 308): + | location = resp.headers.get("Location", "") + | resp.close() + | if not location: + | raise ValueError("Redirect response has no Location header.") + | # Location may be relative; resolve it against the current + | # URL. The result is re-validated at the top of the loop. + | current_url = _urljoin(current_url, location) + | else: + | break + | else: + | raise ValueError( + | f"Too many redirects (more than {self.MAX_REDIRECT_HOPS}) while fetching remote URL." + | ) | resp.raise_for_status() | content_type = resp.headers.get("Content-Type", "") | total = 0 diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala index 83f6239903a..3db7576470b 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/huggingFace/HuggingFaceInferenceOpDescSpec.scala @@ -304,6 +304,41 @@ class HuggingFaceInferenceOpDescSpec extends AnyFlatSpec with Matchers { code should not include "open(audio_input" } + it should "re-validate every redirect hop in _fetch_remote_url instead of following blindly" in { + // requests follows redirects by default, which would skip the scheme/IP + // checks on the redirect target: a 302 to http://169.254.169.254/... or an + // internal host would be fetched. The helper must disable automatic + // redirects and re-run _validate_remote_url on each hop. + val code = makeDesc(task = "image-to-image", inputImageColumn = "img").generatePythonCode() + // Automatic redirect-following is off, and no redirect-following variant + // of the fetch remains anywhere in the helper. + code should include( + "resp = requests.get(current_url, timeout=120, stream=True, allow_redirects=False)" + ) + code should not include "resp = requests.get(url, timeout=120, stream=True)" + // The per-hop validator exists and runs BEFORE the request inside the loop. + code should include("def _validate_remote_url(self, url):") + val validateCall = code.indexOf("self._validate_remote_url(current_url)") + val fetchCall = code.indexOf("resp = requests.get(current_url") + validateCall should be > 0 + fetchCall should be > validateCall + // Every redirect status is intercepted; relative Location values are + // resolved against the current URL before re-validation. + code should include("if resp.status_code in (301, 302, 303, 307, 308):") + code should include("current_url = _urljoin(current_url, location)") + // Degenerate redirects fail closed: missing Location and unbounded chains. + code should include("Redirect response has no Location header.") + code should include("MAX_REDIRECT_HOPS = 5") + code should include("Too many redirects") + // The validator keeps the full pre-existing checks (https-only + public + // address) so each hop gets the same scrutiny as the original URL, and + // takes an allowlist stance (globally-routable only) that also blocks the + // CGNAT/shared range the predicate list alone misses. + code should include("""if parsed.scheme != "https":""") + code should include("not ip.is_global") + code should include("ip.is_multicast") + } + it should "treat pandas NA sentinels (NaN, pd.NA, NaT) as missing in _read_binary_value" in { // isinstance(value, float) only catches float('nan'); pd.NA / NaT are not // floats and previously fell through to be str()-ified into bytes. The