fix(serverless): block SSRF in job-input downloads (#533)

* fix(serverless): block SSRF in job-input downloads

Job input carries arbitrary URLs that the worker downloads. Without
restriction a job can point the worker at cloud instance-metadata
(169.254.169.254) or other non-public addresses (CWE-918).

Add runpod/serverless/utils/rp_ssrf.py: an http(s) scheme allowlist,
DNS resolution that fails closed on any non-global IP (link-local,
loopback, RFC1918, CGNAT, reserved, multicast, IPv6 ULA), a connection
adapter that pins the socket to the validated IP to defeat DNS
rebinding, per-hop redirect re-validation, and a streamed size cap.
Route both rp_download fetch sites through it, and stream file() to disk
instead of buffering the whole untrusted body in memory.

Blocks raise SSRFError (a ValueError, not a RequestException) so they
fail loudly rather than being retried or silently dropped. Configurable
via RUNPOD_ALLOW_PRIVATE_DOWNLOAD_URLS (escape hatch, off by default) and
RUNPOD_MAX_DOWNLOAD_BYTES (default 5 GiB).

* fix(serverless): close pinned SSRF sessions to prevent FD leaks

Address PR #533 review feedback.

- Tie each single-use SyncClientSession's lifecycle to its response via
  _bind_session_to_response, so closing the response (with-block or explicit)
  also closes the session; without it the session's pools/FDs leaked once per
  download in long-lived workers.
- Close the session on redirect hops and on request exceptions; make teardown
  best-effort so a cleanup failure cannot mask the caller's real error.
- Close response and session deterministically in the pinned-IP adapter test.
- Add tests for the request-exception path and the error-masking guarantee.

* fix(serverless): tolerate malformed Content-Length in downloads

Content-Length is remote-controlled, so int() on it can raise ValueError
and abort an otherwise valid download. In file() this was a new failure
path (main never parsed the header there); in download_files_from_urls
the ValueError also escaped both the backoff retry and the RequestException
handler, failing the whole batch.

Parse defensively via parse_content_length(), falling back to 0 so chunk
sizing still works. The size limit is unaffected: it is enforced while
streaming by iter_content_capped().

* fix(serverless): fall back across validated IPs when pinning

Pinning to ips[0] turned a multi-address host into a hard failure when the
first address is unroutable from the worker: a dual-stack host whose AAAA
sorts first is unreachable on a pod with no IPv6 route, where plain requests
would have tried the A record. Walk the validated addresses in resolution
order, advancing only on connection errors, so every candidate stays
pre-validated while the fallback behavior is preserved.

Also patch os.makedirs in the file() tests, which were creating a real
job_files/ directory in the working tree.

* fix(serverless): fail file() on HTTP error responses

file() saved whatever body came back, so a 4xx/5xx error page was written
as the downloaded file and, when the URL ended in .zip, handed to the
extractor. download_files_from_urls already called raise_for_status();
file() never did.

Behavior change: file() now raises requests.RequestException on an error
status instead of returning a dict pointing at the error page. It has no
in-tree callers and is not re-exported from serverless.utils, so this is
limited to direct importers.

* fix(serverless): bracket IPv6 literals in the pinned Host header

urlparse strips the brackets from an IPv6-literal URL host, so the adapter
built a malformed Host header ("2606::1:443" instead of "[2606::1]:443"),
which RFC 7230 requires and some servers reject. Extracted the authority
construction into _host_header_for() and covered every host form.

* fix(serverless): refuse proxied fetches that defeat IP pinning

Pinning only holds when this process opens the socket. Through a proxy the
request carries the hostname and the proxy resolves it, so the validated IP
is not what gets dialed and the DNS-rebind protection silently stops
applying while still appearing to be in force.

Refuse such a URL with SSRFError instead. NO_PROXY exclusions are honored,
since those hosts are fetched directly, and RUNPOD_ALLOW_PRIVATE_DOWNLOAD_URLS
allows the proxied fetch for deployments that accept the trade-off. Chosen
over trust_env=False, which would silently ignore a needed proxy and also
drop REQUESTS_CA_BUNDLE.

Author: deanq

runpod/serverless/utils/rp_download.py

@@ -16,7 +16,11 @@
 import backoff
 from requests import RequestException
 
-from runpod.http_client import SyncClientSession
+from runpod.serverless.utils.rp_ssrf import (
+    iter_content_capped,
+    max_download_bytes,
+    safe_get,
+)
 
 HEADERS = {"User-Agent": "runpod-python/0.0.0 (https://runpod.io; support@runpod.io)"}
 
@@ -33,6 +37,20 @@ def calculate_chunk_size(file_size: int) -> int:
     return 1024 * 1024 * 10  # 10 MB
 
 
+def parse_content_length(headers) -> int:
+    """
+    Reads Content-Length for chunk sizing, tolerating a missing or malformed value.
+
+    The header is remote-controlled, so a non-integer value must not abort an
+    otherwise valid download; 0 simply yields the smallest chunk size. The real
+    size limit is enforced while streaming by iter_content_capped().
+    """
+    try:
+        return int(headers.get("Content-Length", 0))
+    except (TypeError, ValueError):
+        return 0
+
+
 def extract_disposition_params(content_disposition: str) -> Dict[str, str]:
     parts = (p.strip() for p in content_disposition.split(";"))
 
@@ -57,7 +75,7 @@ def download_files_from_urls(job_id: str, urls: Union[str, List[str]]) -> List[s
 
     @backoff.on_exception(backoff.expo, RequestException, max_tries=3)
     def download_file(url: str, path_to_save: str) -> str:
-        with SyncClientSession().get(url, headers=HEADERS, stream=True, timeout=5) as response:
+        with safe_get(url, stream=True, timeout=5, headers=HEADERS) as response:
             response.raise_for_status()
             content_disposition = response.headers.get("Content-Disposition")
             file_extension = ""
@@ -69,14 +87,12 @@ def download_file(url: str, path_to_save: str) -> str:
             if not file_extension:
                 file_extension = os.path.splitext(urlparse(url).path)[1]
 
-            file_size = int(response.headers.get("Content-Length", 0))
-            chunk_size = calculate_chunk_size(file_size)
+            chunk_size = calculate_chunk_size(parse_content_length(response.headers))
 
-            # write the content in chunks to the file
+            # write the content in chunks to the file, aborting past the size cap
             with open(path_to_save + file_extension, "wb") as file_path:
-                for chunk in response.iter_content(chunk_size=chunk_size):
-                    if chunk:  # filter out keep-alive chunks
-                        file_path.write(chunk)
+                for chunk in iter_content_capped(response, chunk_size, max_download_bytes()):
+                    file_path.write(chunk)
 
             return file_extension
 
@@ -117,27 +133,35 @@ def file(file_url: str) -> dict:
     """
     os.makedirs("job_files", exist_ok=True)
 
-    download_response = SyncClientSession().get(file_url, headers=HEADERS, timeout=30)
+    with safe_get(file_url, stream=True, timeout=30, headers=HEADERS) as download_response:
+        # Fail on an error response rather than saving the error page as the
+        # file; for a .zip that body would go straight to the extractor.
+        download_response.raise_for_status()
+
+        content_disposition = download_response.headers.get("Content-Disposition")
 
-    content_disposition = download_response.headers.get("Content-Disposition")
+        original_file_name = ""
+        if content_disposition:
+            params = extract_disposition_params(content_disposition)
 
-    original_file_name = ""
-    if content_disposition:
-        params = extract_disposition_params(content_disposition)
+            original_file_name = params.get("filename", "")
 
-        original_file_name = params.get("filename", "")
+        if not original_file_name:
+            download_path = urlparse(file_url).path
+            original_file_name = os.path.basename(download_path)
 
-    if not original_file_name:
-        download_path = urlparse(file_url).path
-        original_file_name = os.path.basename(download_path)
+        file_type = os.path.splitext(original_file_name)[1].replace(".", "")
 
-    file_type = os.path.splitext(original_file_name)[1].replace(".", "")
+        file_name = f"{uuid.uuid4()}"
 
-    file_name = f"{uuid.uuid4()}"
+        output_file_path = os.path.join("job_files", f"{file_name}.{file_type}")
 
-    output_file_path = os.path.join("job_files", f"{file_name}.{file_type}")
-    with open(output_file_path, "wb") as output_file:
-        output_file.write(download_response.content)
+        # Stream to disk in chunks (aborting past the size cap) instead of
+        # buffering the entire untrusted body in memory.
+        chunk_size = calculate_chunk_size(parse_content_length(download_response.headers))
+        with open(output_file_path, "wb") as output_file:
+            for chunk in iter_content_capped(download_response, chunk_size, max_download_bytes()):
+                output_file.write(chunk)
 
     if file_type == "zip":
         unzipped_directory = os.path.join("job_files", file_name)