From 38467ea71a6b0f8757bd88e5065c3fceb0955fa2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Tue, 7 Jul 2026 22:08:12 -0700 Subject: [PATCH 1/2] 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). --- runpod/serverless/utils/rp_download.py | 48 +-- runpod/serverless/utils/rp_ssrf.py | 276 ++++++++++++++++++ .../test_utils/test_download.py | 138 ++++----- tests/test_serverless/test_utils/test_ssrf.py | 246 ++++++++++++++++ 4 files changed, 610 insertions(+), 98 deletions(-) create mode 100644 runpod/serverless/utils/rp_ssrf.py create mode 100644 tests/test_serverless/test_utils/test_ssrf.py diff --git a/runpod/serverless/utils/rp_download.py b/runpod/serverless/utils/rp_download.py index 5c889fd2..93a48e2b 100644 --- a/runpod/serverless/utils/rp_download.py +++ b/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)"} @@ -57,7 +61,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 = "" @@ -72,11 +76,10 @@ def download_file(url: str, path_to_save: str) -> str: file_size = int(response.headers.get("Content-Length", 0)) chunk_size = calculate_chunk_size(file_size) - # 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 +120,32 @@ 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: + 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. + file_size = int(download_response.headers.get("Content-Length", 0)) + chunk_size = calculate_chunk_size(file_size) + 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) diff --git a/runpod/serverless/utils/rp_ssrf.py b/runpod/serverless/utils/rp_ssrf.py new file mode 100644 index 00000000..37adf8dc --- /dev/null +++ b/runpod/serverless/utils/rp_ssrf.py @@ -0,0 +1,276 @@ +""" +PodWorker | serverless | utils | rp_ssrf.py + +SSRF-safe fetching of user-supplied (job-input) URLs. + +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). This module blocks any destination that is not a +globally-routable IP, pins connections to a pre-validated address to defeat DNS +rebinding, re-validates redirect hops, and caps download size. +""" + +import ipaddress +import os +import socket +from typing import Iterator, List, Optional +from urllib.parse import urljoin, urlparse + +from requests.adapters import HTTPAdapter + +from runpod.http_client import SyncClientSession + +_ALLOWED_SCHEMES = ("http", "https") +_REDIRECT_STATUS = frozenset((301, 302, 303, 307, 308)) +_DEFAULT_MAX_REDIRECTS = 5 + +# Size cap. Generous default that does not break legitimate image/zip/model +# downloads but bounds pathological abuse and the in-memory read in file(). +_MAX_BYTES_ENV = "RUNPOD_MAX_DOWNLOAD_BYTES" +_DEFAULT_MAX_BYTES = 5 * 1024 ** 3 # 5 GiB + +# CGNAT / shared address space (RFC 6598). Not reported by IPv4Address.is_private +# before CPython 3.12.4, so it is blocked explicitly rather than relied upon. +_EXTRA_BLOCKED_NETWORKS = [ipaddress.ip_network("100.64.0.0/10")] + +# Escape hatch for workers that legitimately fetch from a same-VPC / self-hosted +# endpoint. Off by default; scheme allowlist and size cap still apply when set. +_ALLOW_PRIVATE_ENV = "RUNPOD_ALLOW_PRIVATE_DOWNLOAD_URLS" + + +class SSRFError(ValueError): + """ + Raised when a URL is rejected as unsafe (non-public destination, blocked + scheme, or oversized body). + + Subclasses ValueError and deliberately NOT requests.RequestException so it + bypasses the download path's backoff retry and `except RequestException` + handler: a blocked internal URL must fail loudly, not be retried or silently + reported as an ordinary failed download. + """ + + +def _ssrf_protection_enabled() -> bool: + return os.environ.get(_ALLOW_PRIVATE_ENV, "").strip().lower() not in ( + "1", + "true", + "yes", + ) + + +def is_safe_address(ip: str) -> bool: + """ + Return True only for a globally-routable unicast IP address. + + Blocks loopback, link-local (incl. 169.254.169.254 metadata), RFC1918, + CGNAT, reserved, multicast, unspecified, IPv6 unique-local/link-local, and + IPv4-mapped IPv6 forms of any of the above. Unparseable input is unsafe. + """ + try: + address = ipaddress.ip_address(ip) + except ValueError: + return False + + # Unwrap IPv4-mapped IPv6 (e.g. ::ffff:10.0.0.1) and classify the real IPv4. + mapped = getattr(address, "ipv4_mapped", None) + if mapped is not None: + address = mapped + + if any(address in network for network in _EXTRA_BLOCKED_NETWORKS): + return False + + if ( + address.is_private + or address.is_loopback + or address.is_link_local + or address.is_reserved + or address.is_multicast + or address.is_unspecified + ): + return False + + return address.is_global + + +def resolve_and_validate(host: str, port: int) -> List[str]: + """ + Resolve `host` and return its IPs, failing closed on any unsafe address. + + Raises SSRFError if the host cannot be resolved, resolves to nothing, or + resolves to any non-global address (a host resolving to both a public and a + private IP is treated as hostile). Validation is skipped when the escape + hatch env var is set, but resolution still occurs so callers can pin. + """ + try: + infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM) + except socket.gaierror as err: + raise SSRFError(f"could not resolve host {host!r}: {err}") from err + + ips: List[str] = [] + for info in infos: + ip = info[4][0] + if ip not in ips: + ips.append(ip) + + if not ips: + raise SSRFError(f"no addresses resolved for host {host!r}") + + if _ssrf_protection_enabled(): + for ip in ips: + if not is_safe_address(ip): + raise SSRFError( + f"host {host!r} resolves to non-public address {ip}" + ) + + return ips + + +def max_download_bytes() -> int: + """Configured per-download byte cap; defaults to 5 GiB.""" + raw = os.environ.get(_MAX_BYTES_ENV) + if raw is None: + return _DEFAULT_MAX_BYTES + try: + value = int(raw) + except ValueError: + return _DEFAULT_MAX_BYTES + return value if value > 0 else _DEFAULT_MAX_BYTES + + +class PinnedIPAdapter(HTTPAdapter): + """ + Requests adapter that dials a pre-validated IP instead of re-resolving the + URL host, defeating DNS rebinding. + + The superclass builds the connection pool with its normal TLS/verification + configuration; this adapter only redirects the socket to the pinned IP and + keeps the TLS SNI, certificate hostname check, and Host header bound to the + original URL host. It never alters certificate verification. + """ + + def __init__(self, pinned_ip: str, **kwargs): + self._pinned_ip = pinned_ip + super().__init__(**kwargs) + + def send(self, request, **kwargs): + parsed = urlparse(request.url) + authority = parsed.hostname or "" + if parsed.port is not None: + authority = f"{authority}:{parsed.port}" + request.headers["Host"] = authority + return super().send(request, **kwargs) + + def _pin(self, pool, url: str): + parsed = urlparse(url) + if parsed.scheme == "https": + pool.assert_hostname = parsed.hostname + pool.conn_kw = dict(pool.conn_kw or {}) + pool.conn_kw["server_hostname"] = parsed.hostname + pool.host = self._pinned_ip + return pool + + def get_connection_with_tls_context(self, request, verify, proxies=None, cert=None): + pool = super().get_connection_with_tls_context( + request, verify, proxies=proxies, cert=cert + ) + return self._pin(pool, request.url) + + def get_connection(self, url, proxies=None): # requests < 2.32 fallback + pool = super().get_connection(url, proxies=proxies) + return self._pin(pool, url) + + +def _build_pinned_session(pinned_ip: str) -> SyncClientSession: + """A session whose http(s) connections dial only `pinned_ip`.""" + session = SyncClientSession() + adapter = PinnedIPAdapter(pinned_ip) + session.mount("http://", adapter) + session.mount("https://", adapter) + return session + + +def _enforce_content_length(response, max_bytes: Optional[int]) -> None: + if max_bytes is None: + return + declared = response.headers.get("Content-Length") + if declared is None: + return + try: + size = int(declared) + except ValueError: + return + if size > max_bytes: + response.close() + raise SSRFError(f"download exceeds size cap: {size} > {max_bytes} bytes") + + +def iter_content_capped(response, chunk_size: int, max_bytes: Optional[int]) -> Iterator[bytes]: + """ + Yield the response body in chunks, aborting with SSRFError if the running + total exceeds `max_bytes`. Guards against a lying/absent Content-Length. + """ + total = 0 + for chunk in response.iter_content(chunk_size=chunk_size): + if not chunk: + continue + total += len(chunk) + if max_bytes is not None and total > max_bytes: + response.close() + raise SSRFError(f"download exceeds size cap: >{max_bytes} bytes") + yield chunk + + +def safe_get( + url: str, + *, + stream: bool = True, + timeout: int = 30, + headers: Optional[dict] = None, + max_bytes: Optional[int] = None, + max_redirects: int = _DEFAULT_MAX_REDIRECTS, +): + """ + Fetch a user-supplied URL with SSRF protections. + + Enforces an http(s) scheme allowlist, resolves and validates the host to a + public IP, pins the connection to that IP, and re-validates every redirect + hop. Rejects a body whose declared size exceeds the cap. Returns the + (streamed) requests.Response for the caller to consume; use + iter_content_capped to enforce the cap while reading. Raises SSRFError on + any block. + """ + if max_bytes is None: + max_bytes = max_download_bytes() + + current = url + for _ in range(max_redirects + 1): + parsed = urlparse(current) + if parsed.scheme not in _ALLOWED_SCHEMES: + raise SSRFError(f"blocked URL scheme {parsed.scheme!r}: {current}") + if not parsed.hostname: + raise SSRFError(f"URL has no host: {current}") + + port = parsed.port or (443 if parsed.scheme == "https" else 80) + ips = resolve_and_validate(parsed.hostname, port) + + session = _build_pinned_session(ips[0]) + response = session.get( + current, + headers=headers, + allow_redirects=False, + stream=stream, + timeout=timeout, + ) + + if response.status_code in _REDIRECT_STATUS: + location = response.headers.get("Location") + response.close() + if not location: + raise SSRFError(f"redirect without Location header: {current}") + current = urljoin(current, location) + continue + + _enforce_content_length(response, max_bytes) + return response + + raise SSRFError(f"exceeded maximum redirects ({max_redirects}): {url}") diff --git a/tests/test_serverless/test_utils/test_download.py b/tests/test_serverless/test_utils/test_download.py index b7533f5c..c6bf6730 100644 --- a/tests/test_serverless/test_utils/test_download.py +++ b/tests/test_serverless/test_utils/test_download.py @@ -1,10 +1,10 @@ -"""Tests for runpod | serverless | modules | download.py""" +"""Tests for runpod | serverless | utils | rp_download.py""" # pylint: disable=R0903,W0613 import os import unittest -from unittest.mock import MagicMock, mock_open, patch +from unittest.mock import mock_open, patch from requests import RequestException @@ -23,51 +23,48 @@ JOB_ID = "job_123" -def mock_requests_get(*args, **kwargs): - """ - Mocks SyncClientSession.get - """ - headers = { - "Content-Disposition": 'attachment; filename="picture.jpg"', - "Content-Length": "1000", - } +class MockResponse: + """Stand-in for the streamed requests.Response returned by safe_get.""" - class MockResponse: - """Mocks SyncClientSession.get response""" + def __init__(self, content, status_code, headers=None): + self.content = content + self.status_code = status_code + self.headers = headers or {} - def __init__(self, content, status_code, headers=None): - """ - Mocks SyncClientSession.get response - """ - self.content = content - self.status_code = status_code - self.headers = headers or {} + def raise_for_status(self): + if 400 <= self.status_code < 600: + raise RequestException(f"Status code: {self.status_code}") - def raise_for_status(self): - """Mocks raise_for_status function""" - if 400 <= self.status_code < 600: - raise RequestException(f"Status code: {self.status_code}") + def iter_content(self, chunk_size=1024): + """A fresh generator each call (safe_get responses are streamed once).""" + content = self.content or b"" + for i in range(0, len(content), chunk_size): + yield content[i : min(i + chunk_size, len(content))] - def iter_content(self, chunk_size=1024): - """Mocks iter_content method""" - length = len(self.content) - for i in range(0, length, chunk_size): - yield self.content[i : min(i + chunk_size, length)] + def __enter__(self): + return self - def __enter__(self): - return self + def __exit__(self, *args): + return False - def __exit__(self, *args): - pass +def mock_safe_get(*args, **kwargs): + """Mocks rp_ssrf.safe_get for the download_files_from_urls tests.""" + headers = { + "Content-Disposition": 'attachment; filename="picture.jpg"', + "Content-Length": "1000", + } url = args[0] - # Check if the URL matches any of the URLs in URL_LIST if any(url.startswith(base_url) for base_url in URL_LIST): return MockResponse(b"nothing", 200, headers) - return MockResponse(None, 404) +def _called_urls(mock): + """URLs passed positionally to safe_get, order-independent.""" + return {call.args[0] for call in mock.call_args_list} + + class TestDownloadFilesFromUrls(unittest.TestCase): """Tests for download_files_from_urls""" @@ -81,7 +78,7 @@ def test_calculate_chunk_size(self): self.assertEqual(calculate_chunk_size(1024 * 1024 * 1024 * 10), 1024 * 1024 * 10) @patch("os.makedirs", return_value=None) - @patch("runpod.http_client.SyncClientSession.get", side_effect=mock_requests_get) + @patch("runpod.serverless.utils.rp_download.safe_get", side_effect=mock_safe_get) @patch("builtins.open", new_callable=mock_open) def test_download_files_from_urls(self, mock_open_file, mock_get, mock_makedirs): """ @@ -95,17 +92,12 @@ def test_download_files_from_urls(self, mock_open_file, mock_get, mock_makedirs) self.assertEqual(len(downloaded_files), len(urls)) - # Downloads run in parallel threads, so the order of get() calls is - # non-deterministic; assert the set of requested URLs instead of order. - requested_urls = {call.args[0] for call in mock_get.call_args_list} - self.assertEqual(requested_urls, set(urls)) + # download runs concurrently, so assert the set of fetched URLs (not order) + self.assertEqual(_called_urls(mock_get), set(urls)) - # executor.map preserves input order in results, so downloaded_files - # still aligns positionally with urls. for index in range(len(urls)): # Check that the file has the correct extension self.assertTrue(downloaded_files[index].endswith(".jpg")) - mock_open_file.assert_any_call(downloaded_files[index], "wb") mock_makedirs.assert_called_once_with(os.path.abspath(f"jobs/{JOB_ID}/downloaded_files"), exist_ok=True) @@ -124,7 +116,7 @@ def test_download_files_from_urls(self, mock_open_file, mock_get, mock_makedirs) ) @patch("os.makedirs", return_value=None) - @patch("runpod.http_client.SyncClientSession.get", side_effect=mock_requests_get) + @patch("runpod.serverless.utils.rp_download.safe_get", side_effect=mock_safe_get) @patch("builtins.open", new_callable=mock_open) def test_download_files_from_urls_signed(self, mock_open_file, mock_get, mock_makedirs): """ @@ -140,8 +132,8 @@ def test_download_files_from_urls_signed(self, mock_open_file, mock_get, mock_ma # Confirms that the same number of files were downloaded as urls provided self.assertEqual(len(downloaded_files), 1) - # Check that the url was called with SyncClientSession.get - self.assertIn(URL_LIST[1], mock_get.call_args_list[0][0]) + # Check that the signed url was fetched + self.assertEqual(_called_urls(mock_get), {URL_LIST[1]}) # Check that the file has the correct extension self.assertTrue(downloaded_files[0].endswith(".jpg")) @@ -151,82 +143,72 @@ def test_download_files_from_urls_signed(self, mock_open_file, mock_get, mock_ma class FileDownloaderTestCase(unittest.TestCase): - """Tests for file_downloader""" + """Tests for file()""" - @patch("runpod.serverless.utils.rp_download.SyncClientSession.get") + @patch("runpod.serverless.utils.rp_download.safe_get") @patch("builtins.open", new_callable=mock_open) def test_download_file(self, mock_file, mock_get): """ - Tests download_file + Tests file() """ - # Mock the response from SyncClientSession.get - mock_response = MagicMock() - mock_response.content = b"file content" - mock_response.headers = {"Content-Disposition": "filename=test_file.txt"} - mock_get.return_value = mock_response + mock_get.return_value = MockResponse( + b"file content", 200, {"Content-Disposition": "filename=test_file.txt"} + ) - # Call the function with a test URL result = file("http://test.com/test_file.txt") - # Check the result self.assertEqual(result["type"], "txt") self.assertEqual(result["original_name"], "test_file.txt") self.assertTrue(result["file_path"].endswith(".txt")) self.assertIsNone(result["extracted_path"]) - # Check that the file was written correctly + # Body is streamed to disk in a single chunk here mock_file().write.assert_called_once_with(b"file content") - @patch("runpod.serverless.utils.rp_download.SyncClientSession.get") + @patch("runpod.serverless.utils.rp_download.safe_get") @patch("builtins.open", new_callable=mock_open) def test_download_file_with_content_disposition(self, mock_file, mock_get): """ - Tests download_file using filename from Content-Disposition + Tests file() using filename from Content-Disposition """ - # Mock the response from SyncClientSession.get - mock_response = MagicMock() - mock_response.content = b"file content" - mock_response.headers = {"Content-Disposition": 'inline; filename="test_file.txt"'} - mock_get.return_value = mock_response + mock_get.return_value = MockResponse( + b"file content", 200, {"Content-Disposition": 'inline; filename="test_file.txt"'} + ) - # Call the function with a test URL result = file("http://test.com/file_without_extension") - # Check the result self.assertEqual(result["type"], "txt") self.assertEqual(result["original_name"], "test_file.txt") self.assertTrue(result["file_path"].endswith(".txt")) self.assertIsNone(result["extracted_path"]) - # Check that the file was written correctly mock_file().write.assert_called_once_with(b"file content") - @patch("runpod.serverless.utils.rp_download.SyncClientSession.get") + @patch("runpod.serverless.utils.rp_download.safe_get") @patch("builtins.open", new_callable=mock_open) @patch("runpod.serverless.utils.rp_download.zipfile.ZipFile") def test_download_zip_file(self, mock_zip, mock_file, mock_get): """ - Tests download_file with a zip file + Tests file() with a zip file """ - # Mock the response from SyncClientSession.get - mock_response = MagicMock() - mock_response.content = b"zip file content" - mock_response.headers = {"Content-Disposition": "filename=test_file.zip"} - mock_get.return_value = mock_response + mock_get.return_value = MockResponse( + b"zip file content", 200, {"Content-Disposition": "filename=test_file.zip"} + ) - # Call the function with a test URL result = file("http://test.com/test_file.zip") - # Check the result self.assertEqual(result["type"], "zip") self.assertEqual(result["original_name"], "test_file.zip") self.assertTrue(result["file_path"].endswith(".zip")) self.assertIsNotNone(result["extracted_path"]) - # Check that the file was written correctly mock_file().write.assert_called_once_with(b"zip file content") - # Check if no file name is provided - mock_response.headers = {"Content-Disposition": ""} + # Check if no file name is provided (falls back to URL basename) + mock_get.return_value = MockResponse(b"zip file content", 200, {"Content-Disposition": ""}) result = file("http://test.com/test_file.zip") self.assertEqual(result["original_name"], "test_file.zip") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_serverless/test_utils/test_ssrf.py b/tests/test_serverless/test_utils/test_ssrf.py new file mode 100644 index 00000000..f9092fd8 --- /dev/null +++ b/tests/test_serverless/test_utils/test_ssrf.py @@ -0,0 +1,246 @@ +"""Tests for runpod | serverless | utils | rp_ssrf.py""" + +import http.server +import os +import socket +import threading +import unittest +from unittest.mock import MagicMock, patch + +from requests import RequestException + +from runpod.serverless.utils.rp_ssrf import ( + SSRFError, + _build_pinned_session, + is_safe_address, + iter_content_capped, + max_download_bytes, + resolve_and_validate, + safe_get, +) + + +class _FakeResponse: + """Minimal stand-in for requests.Response used by safe_get.""" + + def __init__(self, status_code, headers=None): + self.status_code = status_code + self.headers = headers or {} + self.closed = False + + def close(self): + self.closed = True + + +def _addrinfo(*ips): + """Build a getaddrinfo-style return value for the given IPs.""" + infos = [] + for ip in ips: + family = socket.AF_INET6 if ":" in ip else socket.AF_INET + sockaddr = (ip, 443, 0, 0) if family == socket.AF_INET6 else (ip, 443) + infos.append((family, socket.SOCK_STREAM, 6, "", sockaddr)) + return infos + + +class TestIsSafeAddress(unittest.TestCase): + """Address classifier: only globally-routable IPs are safe.""" + + def test_blocks_non_global_and_allows_public(self): + blocked = [ + "169.254.169.254", # cloud metadata (link-local) + "169.254.0.1", # link-local + "127.0.0.1", # loopback + "0.0.0.0", # unspecified + "10.0.0.5", # RFC1918 + "172.16.0.1", # RFC1918 + "172.31.255.255", # RFC1918 + "192.168.1.1", # RFC1918 + "100.64.0.1", # CGNAT / shared address space + "224.0.0.1", # multicast + "240.0.0.1", # reserved + "::1", # IPv6 loopback + "fe80::1", # IPv6 link-local + "fc00::1", # IPv6 unique-local + "ff02::1", # IPv6 multicast + "::", # IPv6 unspecified + "::ffff:10.0.0.1", # IPv4-mapped IPv6 of a private addr + "::ffff:169.254.169.254", # IPv4-mapped metadata + "not-an-ip", # unparseable + ] + allowed = [ + "8.8.8.8", + "1.1.1.1", + "93.184.216.34", # example.com + "2606:2800:220:1:248:1893:25c8:1946", # public IPv6 + ] + + for ip in blocked: + self.assertFalse(is_safe_address(ip), f"{ip} should be blocked") + for ip in allowed: + self.assertTrue(is_safe_address(ip), f"{ip} should be allowed") + + +class TestResolveAndValidate(unittest.TestCase): + """DNS resolution fails closed on any non-global address.""" + + def test_ssrf_error_is_not_a_request_exception(self): + # Blocks must bypass the caller's `except RequestException` + backoff. + self.assertTrue(issubclass(SSRFError, ValueError)) + self.assertFalse(issubclass(SSRFError, RequestException)) + + @patch("runpod.serverless.utils.rp_ssrf.socket.getaddrinfo") + def test_returns_ips_when_all_public(self, mock_gai): + mock_gai.return_value = _addrinfo("8.8.8.8", "1.1.1.1") + self.assertEqual(resolve_and_validate("example.com", 443), ["8.8.8.8", "1.1.1.1"]) + + @patch("runpod.serverless.utils.rp_ssrf.socket.getaddrinfo") + def test_raises_on_private(self, mock_gai): + mock_gai.return_value = _addrinfo("169.254.169.254") + with self.assertRaises(SSRFError): + resolve_and_validate("metadata.attacker.test", 80) + + @patch("runpod.serverless.utils.rp_ssrf.socket.getaddrinfo") + def test_fails_closed_on_mixed_public_and_private(self, mock_gai): + # A host resolving to both a public and a private IP is hostile. + mock_gai.return_value = _addrinfo("8.8.8.8", "10.0.0.5") + with self.assertRaises(SSRFError): + resolve_and_validate("rebind.attacker.test", 443) + + @patch("runpod.serverless.utils.rp_ssrf.socket.getaddrinfo") + def test_raises_when_dns_fails(self, mock_gai): + mock_gai.side_effect = socket.gaierror("name resolution failed") + with self.assertRaises(SSRFError): + resolve_and_validate("nonexistent.attacker.test", 443) + + @patch.dict(os.environ, {"RUNPOD_ALLOW_PRIVATE_DOWNLOAD_URLS": "true"}) + @patch("runpod.serverless.utils.rp_ssrf.socket.getaddrinfo") + def test_escape_hatch_allows_private(self, mock_gai): + mock_gai.return_value = _addrinfo("10.0.0.5") + self.assertEqual(resolve_and_validate("internal.vpc.test", 80), ["10.0.0.5"]) + + +class TestSafeGet(unittest.TestCase): + """Scheme allowlist, redirect re-validation, and size-cap enforcement.""" + + def test_rejects_non_http_schemes(self): + for bad in ["file:///etc/passwd", "gopher://h/x", "ftp://h/f", "data:text/plain,x"]: + with self.assertRaises(SSRFError): + safe_get(bad) + + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") + @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate") + def test_redirect_to_metadata_is_blocked(self, mock_resolve, mock_session): + def resolve(host, _port): + if host == "example.com": + return ["93.184.216.34"] + raise SSRFError(f"blocked {host}") + + mock_resolve.side_effect = resolve + session = MagicMock() + session.get.return_value = _FakeResponse( + 302, {"Location": "http://169.254.169.254/latest/meta-data/"} + ) + mock_session.return_value = session + + with self.assertRaises(SSRFError): + safe_get("https://example.com/image.jpg") + + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") + @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"]) + def test_too_many_redirects(self, _mock_resolve, mock_session): + session = MagicMock() + session.get.return_value = _FakeResponse( + 302, {"Location": "https://example.com/loop"} + ) + mock_session.return_value = session + + with self.assertRaises(SSRFError): + safe_get("https://example.com/start", max_redirects=3) + + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") + @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"]) + def test_content_length_over_cap_rejected(self, _mock_resolve, mock_session): + session = MagicMock() + session.get.return_value = _FakeResponse(200, {"Content-Length": str(10 * 1024 * 1024)}) + mock_session.return_value = session + + with self.assertRaises(SSRFError): + safe_get("https://example.com/big.bin", max_bytes=1024) + + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") + @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"]) + def test_returns_response_on_success(self, _mock_resolve, mock_session): + response = _FakeResponse(200, {"Content-Length": "500"}) + session = MagicMock() + session.get.return_value = response + mock_session.return_value = session + + self.assertIs(safe_get("https://example.com/ok.jpg", max_bytes=4096), response) + + +class TestIterContentCapped(unittest.TestCase): + """Streaming byte counter aborts past the cap even when headers lie.""" + + def test_aborts_when_stream_exceeds_cap(self): + response = MagicMock() + response.iter_content.return_value = [b"x" * 100, b"y" * 100] + with self.assertRaises(SSRFError): + list(iter_content_capped(response, chunk_size=100, max_bytes=150)) + + def test_passes_through_under_cap(self): + response = MagicMock() + response.iter_content.return_value = [b"x" * 50, b"y" * 50] + self.assertEqual( + b"".join(iter_content_capped(response, chunk_size=100, max_bytes=1000)), + b"x" * 50 + b"y" * 50, + ) + + +class TestMaxDownloadBytes(unittest.TestCase): + def test_default_when_unset(self): + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("RUNPOD_MAX_DOWNLOAD_BYTES", None) + self.assertEqual(max_download_bytes(), 5 * 1024 ** 3) + + @patch.dict(os.environ, {"RUNPOD_MAX_DOWNLOAD_BYTES": "12345"}) + def test_env_override(self): + self.assertEqual(max_download_bytes(), 12345) + + +class _RecordingHandler(http.server.BaseHTTPRequestHandler): + received_host = None + + def do_GET(self): # noqa: N802 + type(self).received_host = self.headers.get("Host") + body = b"ok" + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): # silence test server logging + pass + + +class TestPinnedIPAdapter(unittest.TestCase): + """The socket must dial the pinned IP while the Host header keeps the URL host.""" + + def test_dials_pinned_ip_and_preserves_host_header(self): + server = http.server.HTTPServer(("127.0.0.1", 0), _RecordingHandler) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + # URL host is example.com (won't route to our server); pin to loopback. + session = _build_pinned_session("127.0.0.1") + response = session.get(f"http://example.com:{port}/path", timeout=5) + self.assertEqual(response.status_code, 200) + self.assertEqual(response.content, b"ok") + self.assertEqual(_RecordingHandler.received_host, f"example.com:{port}") + finally: + server.shutdown() + server.server_close() + + +if __name__ == "__main__": + unittest.main() From 6d4098b4531dbf2e6909dae38cedcc87eb8904bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Wed, 8 Jul 2026 11:32:30 -0700 Subject: [PATCH 2/2] 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. --- runpod/serverless/utils/rp_ssrf.py | 63 +++++++++++++++--- tests/test_serverless/test_utils/test_ssrf.py | 64 +++++++++++++++++-- 2 files changed, 114 insertions(+), 13 deletions(-) diff --git a/runpod/serverless/utils/rp_ssrf.py b/runpod/serverless/utils/rp_ssrf.py index 37adf8dc..357fe2c1 100644 --- a/runpod/serverless/utils/rp_ssrf.py +++ b/runpod/serverless/utils/rp_ssrf.py @@ -13,6 +13,7 @@ import ipaddress import os import socket +from contextlib import suppress from typing import Iterator, List, Optional from urllib.parse import urljoin, urlparse @@ -189,6 +190,42 @@ def _build_pinned_session(pinned_ip: str) -> SyncClientSession: return session +def _close_session_quietly(session: SyncClientSession) -> None: + """ + Close a pinned session best-effort, tolerating teardown errors. + + This runs only on cleanup paths (a failed request, a closing response) where + the caller's original exception is the meaningful one. A failure while + closing the session's connection pools must not mask that exception, so it is + suppressed rather than raised. Session.close() is idempotent and does no + actionable I/O, so nothing worth surfacing is lost. + """ + with suppress(Exception): + session.close() + + +def _bind_session_to_response(response, session: SyncClientSession) -> None: + """ + Tie `session` cleanup to `response.close()`. + + Each safe_get() call owns a single-use session; the caller only holds the + response. Closing a requests.Response releases its connection but not the + session (its adapters and connection pools), so without this the session + leaks one set of pooled sockets/FDs per download in a long-lived worker. + Wrapping close() makes `with safe_get(...) as r:` (and any direct + response.close()) tear down both. + """ + original_close = response.close + + def close_both(*args, **kwargs): + try: + original_close(*args, **kwargs) + finally: + _close_session_quietly(session) + + response.close = close_both + + def _enforce_content_length(response, max_bytes: Optional[int]) -> None: if max_bytes is None: return @@ -254,22 +291,32 @@ def safe_get( ips = resolve_and_validate(parsed.hostname, port) session = _build_pinned_session(ips[0]) - response = session.get( - current, - headers=headers, - allow_redirects=False, - stream=stream, - timeout=timeout, - ) + try: + response = session.get( + current, + headers=headers, + allow_redirects=False, + stream=stream, + timeout=timeout, + ) + except BaseException: + _close_session_quietly(session) + raise if response.status_code in _REDIRECT_STATUS: location = response.headers.get("Location") - response.close() + try: + response.close() + finally: + _close_session_quietly(session) if not location: raise SSRFError(f"redirect without Location header: {current}") current = urljoin(current, location) continue + # Bind before the cap check so a Content-Length breach (which closes + # the response) also tears down the session. + _bind_session_to_response(response, session) _enforce_content_length(response, max_bytes) return response diff --git a/tests/test_serverless/test_utils/test_ssrf.py b/tests/test_serverless/test_utils/test_ssrf.py index f9092fd8..370cb836 100644 --- a/tests/test_serverless/test_utils/test_ssrf.py +++ b/tests/test_serverless/test_utils/test_ssrf.py @@ -23,13 +23,16 @@ class _FakeResponse: """Minimal stand-in for requests.Response used by safe_get.""" - def __init__(self, status_code, headers=None): + def __init__(self, status_code, headers=None, close_exc=None): self.status_code = status_code self.headers = headers or {} self.closed = False + self._close_exc = close_exc def close(self): self.closed = True + if self._close_exc is not None: + raise self._close_exc def _addrinfo(*ips): @@ -145,6 +148,8 @@ def resolve(host, _port): with self.assertRaises(SSRFError): safe_get("https://example.com/image.jpg") + session.close.assert_called_once() # redirect hop must not leak its session + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"]) def test_too_many_redirects(self, _mock_resolve, mock_session): @@ -167,6 +172,8 @@ def test_content_length_over_cap_rejected(self, _mock_resolve, mock_session): with self.assertRaises(SSRFError): safe_get("https://example.com/big.bin", max_bytes=1024) + session.close.assert_called_once() # cap breach must close the pinned session + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"]) def test_returns_response_on_success(self, _mock_resolve, mock_session): @@ -177,6 +184,52 @@ def test_returns_response_on_success(self, _mock_resolve, mock_session): self.assertIs(safe_get("https://example.com/ok.jpg", max_bytes=4096), response) + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") + @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"]) + def test_closing_response_closes_session(self, _mock_resolve, mock_session): + """The session outlives safe_get() only until the caller closes the response.""" + response = _FakeResponse(200, {"Content-Length": "500"}) + session = MagicMock() + session.get.return_value = response + mock_session.return_value = session + + result = safe_get("https://example.com/ok.jpg", max_bytes=4096) + session.close.assert_not_called() # still open for the caller to stream + + result.close() + self.assertTrue(response.closed) # original close still runs + session.close.assert_called_once() # and the session is torn down with it + + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") + @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"]) + def test_request_error_closes_session(self, _mock_resolve, mock_session): + """A failed GET must tear down the session before the error propagates.""" + session = MagicMock() + session.get.side_effect = RequestException("boom") + mock_session.return_value = session + + with self.assertRaises(RequestException): + safe_get("https://example.com/x") + + session.close.assert_called_once() + + @patch("runpod.serverless.utils.rp_ssrf._build_pinned_session") + @patch("runpod.serverless.utils.rp_ssrf.resolve_and_validate", return_value=["93.184.216.34"]) + def test_session_close_error_does_not_mask_response_error(self, _mock_resolve, mock_session): + """Cleanup is best-effort: the caller's close error wins over a session-teardown error.""" + response = _FakeResponse(200, {"Content-Length": "500"}, close_exc=RuntimeError("primary")) + session = MagicMock() + session.get.return_value = response + session.close.side_effect = RuntimeError("secondary pool teardown") + mock_session.return_value = session + + result = safe_get("https://example.com/ok.jpg", max_bytes=4096) + with self.assertRaises(RuntimeError) as ctx: + result.close() + + self.assertEqual(str(ctx.exception), "primary") # response error surfaces + session.close.assert_called_once() # teardown was still attempted + class TestIterContentCapped(unittest.TestCase): """Streaming byte counter aborts past the cap even when headers lie.""" @@ -232,10 +285,11 @@ def test_dials_pinned_ip_and_preserves_host_header(self): thread.start() try: # URL host is example.com (won't route to our server); pin to loopback. - session = _build_pinned_session("127.0.0.1") - response = session.get(f"http://example.com:{port}/path", timeout=5) - self.assertEqual(response.status_code, 200) - self.assertEqual(response.content, b"ok") + with _build_pinned_session("127.0.0.1") as session, session.get( + f"http://example.com:{port}/path", timeout=5 + ) as response: + self.assertEqual(response.status_code, 200) + self.assertEqual(response.content, b"ok") self.assertEqual(_RecordingHandler.received_host, f"example.com:{port}") finally: server.shutdown()