From 440a6cb4c4d16db1e748140b06d3fd4493afa2ad Mon Sep 17 00:00:00 2001 From: truongsontung Date: Sun, 23 Aug 2026 12:46:05 +0700 Subject: [PATCH 1/4] fix: prevent SSRF in image decoder URL fetch Add _is_safe_url() to validate URLs against private/loopback/metadata IP ranges before fetching, preventing SSRF attacks via dataset image URLs. Fixes #4293 --- .../hf_datasets/multimodal/utils/image.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/torchtitan/hf_datasets/multimodal/utils/image.py b/torchtitan/hf_datasets/multimodal/utils/image.py index b6325257f7..20b8e5843f 100644 --- a/torchtitan/hf_datasets/multimodal/utils/image.py +++ b/torchtitan/hf_datasets/multimodal/utils/image.py @@ -10,8 +10,11 @@ vision encoder. """ +import ipaddress import math +import socket from collections.abc import Callable +from urllib.parse import urlparse import einops as E import requests @@ -26,6 +29,38 @@ from torchtitan.tools.logging import logger +_PRIVATE_IP_RANGES = [ + ipaddress.ip_network("10.0.0.0/8"), + ipaddress.ip_network("172.16.0.0/12"), + ipaddress.ip_network("192.168.0.0/16"), + ipaddress.ip_network("127.0.0.0/8"), + ipaddress.ip_network("169.254.0.0/16"), + ipaddress.ip_network("0.0.0.0/8"), + ipaddress.ip_network("224.0.0.0/4"), +] + + +def _is_safe_url(url: str) -> bool: + """Check if URL is safe from SSRF by blocking private/loopback/metadata IPs.""" + try: + parsed = urlparse(url) + hostname = parsed.hostname + if not hostname: + return False + if parsed.scheme not in ("http", "https"): + return False + addrinfo = socket.getaddrinfo(hostname, None) + for family, _, _, _, sockaddr in addrinfo: + ip_str = sockaddr[0] + ip = ipaddress.ip_address(ip_str) + for network in _PRIVATE_IP_RANGES: + if ip in network: + return False + return True + except (socket.gaierror, socket.herror, OSError, ValueError): + return False + + def _decode_image(image: str | bytes | Image.Image) -> torch.Tensor: """Decode an image to a (C, H, W) uint8 RGB tensor. @@ -33,6 +68,8 @@ def _decode_image(image: str | bytes | Image.Image) -> torch.Tensor: falls back to TVF.pil_to_tensor for PIL Image inputs. """ if isinstance(image, str) and image.startswith("http"): + if not _is_safe_url(image): + raise ValueError(f"URL not allowed (SSRF protection): {image}") response = requests.get(image, timeout=10) image = response.content if isinstance(image, bytes): From 23f6df36d71757ea4d1a9fb7471588bdce67db6c Mon Sep 17 00:00:00 2001 From: truongsontung Date: Sun, 23 Aug 2026 13:01:38 +0700 Subject: [PATCH 2/4] fix: improved SSRF protection with redirect blocking and IPv6 support Address review feedback: - Use ipaddress.is_private/is_loopback/is_link_local/is_multicast/is_unspecified (covers IPv6) - Add _SSRFProtectedAdapter to block redirects to unsafe IPs - Resolve all IP addresses and check each Fixes #4293 --- .../hf_datasets/multimodal/utils/image.py | 51 +++++++++++++------ 1 file changed, 35 insertions(+), 16 deletions(-) diff --git a/torchtitan/hf_datasets/multimodal/utils/image.py b/torchtitan/hf_datasets/multimodal/utils/image.py index 20b8e5843f..8c24f1931e 100644 --- a/torchtitan/hf_datasets/multimodal/utils/image.py +++ b/torchtitan/hf_datasets/multimodal/utils/image.py @@ -26,22 +26,24 @@ from PIL import Image +from requests.adapters import HTTPAdapter + from torchtitan.tools.logging import logger -_PRIVATE_IP_RANGES = [ - ipaddress.ip_network("10.0.0.0/8"), - ipaddress.ip_network("172.16.0.0/12"), - ipaddress.ip_network("192.168.0.0/16"), - ipaddress.ip_network("127.0.0.0/8"), - ipaddress.ip_network("169.254.0.0/16"), - ipaddress.ip_network("0.0.0.0/8"), - ipaddress.ip_network("224.0.0.0/4"), -] +def _is_blocked_ip(ip: ipaddress.ip_address) -> bool: + """Return True if the IP is private, loopback, link-local, multicast, or unspecified.""" + return ( + ip.is_private + or ip.is_loopback + or ip.is_link_local + or ip.is_multicast + or ip.is_unspecified + ) def _is_safe_url(url: str) -> bool: - """Check if URL is safe from SSRF by blocking private/loopback/metadata IPs.""" + """Check if URL is safe from SSRF by resolving and checking all IP addresses.""" try: parsed = urlparse(url) hostname = parsed.hostname @@ -49,18 +51,32 @@ def _is_safe_url(url: str) -> bool: return False if parsed.scheme not in ("http", "https"): return False + # Resolve hostname and check ALL resolved IPs addrinfo = socket.getaddrinfo(hostname, None) + resolved_ips: list[ipaddress.ip_address] = [] for family, _, _, _, sockaddr in addrinfo: - ip_str = sockaddr[0] - ip = ipaddress.ip_address(ip_str) - for network in _PRIVATE_IP_RANGES: - if ip in network: - return False + ip = ipaddress.ip_address(socket.inet_ntop(family, sockaddr[4])) + resolved_ips.append(ip) + if not resolved_ips: + return False + if any(_is_blocked_ip(ip) for ip in resolved_ips): + return False return True except (socket.gaierror, socket.herror, OSError, ValueError): return False +class _SSRFProtectedAdapter(HTTPAdapter): + """HTTPAdapter that blocks redirects to private/loopback/metadata IPs.""" + + def resolve_redirects(self, resp, req, stream=False, timeout=None, **kwargs): + location = resp.headers.get("Location") + if location: + if not _is_safe_url(location): + raise requests.exceptions.InvalidURL(f"Blocked redirect to unsafe URL: {location}") + return super().resolve_redirects(resp, req, stream=stream, timeout=timeout, **kwargs) + + def _decode_image(image: str | bytes | Image.Image) -> torch.Tensor: """Decode an image to a (C, H, W) uint8 RGB tensor. @@ -70,7 +86,10 @@ def _decode_image(image: str | bytes | Image.Image) -> torch.Tensor: if isinstance(image, str) and image.startswith("http"): if not _is_safe_url(image): raise ValueError(f"URL not allowed (SSRF protection): {image}") - response = requests.get(image, timeout=10) + session = requests.Session() + session.mount("http://", _SSRFProtectedAdapter()) + session.mount("https://", _SSRFProtectedAdapter()) + response = session.get(image, timeout=10, allow_redirects=True) image = response.content if isinstance(image, bytes): raw = torch.frombuffer(bytearray(image), dtype=torch.uint8) From 7522c6054d22d46cf482b6112f060fcdf2eab040 Mon Sep 17 00:00:00 2001 From: truongsontung Date: Sun, 23 Aug 2026 13:09:29 +0700 Subject: [PATCH 3/4] fix: use Session subclass for redirect protection (resolve_redirects is on Session, not HTTPAdapter) Address Claude re-review feedback: _SSRFProtectedAdapter was dead code because resolve_redirects is a Session method, not HTTPAdapter. Switched to _SSRFProtectedSession subclass that intercepts redirects. --- .../hf_datasets/multimodal/utils/image.py | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/torchtitan/hf_datasets/multimodal/utils/image.py b/torchtitan/hf_datasets/multimodal/utils/image.py index 8c24f1931e..51bfb148f0 100644 --- a/torchtitan/hf_datasets/multimodal/utils/image.py +++ b/torchtitan/hf_datasets/multimodal/utils/image.py @@ -14,7 +14,7 @@ import math import socket from collections.abc import Callable -from urllib.parse import urlparse +from urllib.parse import urlparse, urljoin import einops as E import requests @@ -26,8 +26,6 @@ from PIL import Image -from requests.adapters import HTTPAdapter - from torchtitan.tools.logging import logger @@ -51,7 +49,6 @@ def _is_safe_url(url: str) -> bool: return False if parsed.scheme not in ("http", "https"): return False - # Resolve hostname and check ALL resolved IPs addrinfo = socket.getaddrinfo(hostname, None) resolved_ips: list[ipaddress.ip_address] = [] for family, _, _, _, sockaddr in addrinfo: @@ -66,15 +63,16 @@ def _is_safe_url(url: str) -> bool: return False -class _SSRFProtectedAdapter(HTTPAdapter): - """HTTPAdapter that blocks redirects to private/loopback/metadata IPs.""" +class _SSRFProtectedSession(requests.Session): + """Session that blocks redirects to private/loopback/metadata IPs.""" - def resolve_redirects(self, resp, req, stream=False, timeout=None, **kwargs): + def resolve_redirects(self, resp, req, **kwargs): location = resp.headers.get("Location") if location: - if not _is_safe_url(location): - raise requests.exceptions.InvalidURL(f"Blocked redirect to unsafe URL: {location}") - return super().resolve_redirects(resp, req, stream=stream, timeout=timeout, **kwargs) + full_url = urljoin(resp.url, location) + if not _is_safe_url(full_url): + raise requests.exceptions.InvalidURL(f"Blocked redirect to unsafe URL: {full_url}") + return super().resolve_redirects(resp, req, **kwargs) def _decode_image(image: str | bytes | Image.Image) -> torch.Tensor: @@ -86,9 +84,7 @@ def _decode_image(image: str | bytes | Image.Image) -> torch.Tensor: if isinstance(image, str) and image.startswith("http"): if not _is_safe_url(image): raise ValueError(f"URL not allowed (SSRF protection): {image}") - session = requests.Session() - session.mount("http://", _SSRFProtectedAdapter()) - session.mount("https://", _SSRFProtectedAdapter()) + session = _SSRFProtectedSession() response = session.get(image, timeout=10, allow_redirects=True) image = response.content if isinstance(image, bytes): From e9b6c894cd41a308da0114158205fcde267a15b0 Mon Sep 17 00:00:00 2001 From: truongsontung Date: Mon, 24 Aug 2026 13:30:06 +0700 Subject: [PATCH 4/4] fix: validate every redirect hop in SSRF-protected session Claude review feedback: previous override only checked first redirect hop. Rewrote resolve_redirects to manually loop each hop with _is_safe_url() validation, preventing multi-hop redirect bypass. --- .../hf_datasets/multimodal/utils/image.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/torchtitan/hf_datasets/multimodal/utils/image.py b/torchtitan/hf_datasets/multimodal/utils/image.py index 51bfb148f0..a1c5f207eb 100644 --- a/torchtitan/hf_datasets/multimodal/utils/image.py +++ b/torchtitan/hf_datasets/multimodal/utils/image.py @@ -67,12 +67,24 @@ class _SSRFProtectedSession(requests.Session): """Session that blocks redirects to private/loopback/metadata IPs.""" def resolve_redirects(self, resp, req, **kwargs): - location = resp.headers.get("Location") - if location: + # Override to validate EVERY redirect hop, not just the first. + # requests.Session.resolve_redirects is a generator that internally + # loops the entire redirect chain via self.send + get_redirect_target. + # We bypass that loop and validate each hop ourselves. + while True: + location = resp.headers.get("Location") + if not location: + return + # Build fully-qualified redirect URL full_url = urljoin(resp.url, location) if not _is_safe_url(full_url): - raise requests.exceptions.InvalidURL(f"Blocked redirect to unsafe URL: {full_url}") - return super().resolve_redirects(resp, req, **kwargs) + raise requests.exceptions.InvalidURL( + f"Blocked redirect to unsafe URL: {full_url}" + ) + prepared_request = self.prepare_request( + requests.Request("GET", full_url).prepare() + ) + resp = self.send(prepared_request, allow_redirects=False, timeout=10) def _decode_image(image: str | bytes | Image.Image) -> torch.Tensor: