fix: prevent SSRF in image decoder URL fetch

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, urljoin
 
 import einops as E
 import requests
@@ -26,14 +29,75 @@
 from torchtitan.tools.logging import logger
 
 
+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 resolving and checking all IP addresses."""
+    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)
+        resolved_ips: list[ipaddress.ip_address] = []
+        for family, _, _, _, sockaddr in addrinfo:
+            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 _SSRFProtectedSession(requests.Session):
+    """Session that blocks redirects to private/loopback/metadata IPs."""
+
+    def resolve_redirects(self, resp, req, **kwargs):
+        # 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}"
+                )
+            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:
     """Decode an image to a (C, H, W) uint8 RGB tensor.
 
     Uses torchvision.io.decode_image for bytes/paths (faster SIMD decode),
     falls back to TVF.pil_to_tensor for PIL Image inputs.
     """
     if isinstance(image, str) and image.startswith("http"):
-        response = requests.get(image, timeout=10)
+        if not _is_safe_url(image):
+            raise ValueError(f"URL not allowed (SSRF protection): {image}")
+        session = _SSRFProtectedSession()
+        response = session.get(image, timeout=10, allow_redirects=True)
         image = response.content
     if isinstance(image, bytes):
         raw = torch.frombuffer(bytearray(image), dtype=torch.uint8)