diff --git a/nemo_deploy/multimodal/image_url_validator.py b/nemo_deploy/multimodal/image_url_validator.py new file mode 100644 index 000000000..ccfe17c29 --- /dev/null +++ b/nemo_deploy/multimodal/image_url_validator.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import base64 +import ipaddress +import socket +from io import BytesIO +from urllib.parse import urljoin, urlparse + +import urllib3 + +# Ranges that must never be reachable via a request-controlled image URL. +# 169.254.0.0/16 is the cloud IMDS range (AWS/GCP/Azure 169.254.169.254) — +# the primary SSRF target in cloud deployments. +_BLOCKED_NETWORKS = [ + ipaddress.ip_network("127.0.0.0/8"), # loopback — server's own local services + ipaddress.ip_network("10.0.0.0/8"), # RFC 1918 private + ipaddress.ip_network("172.16.0.0/12"), # RFC 1918 private + ipaddress.ip_network("192.168.0.0/16"), # RFC 1918 private + ipaddress.ip_network("169.254.0.0/16"), # link-local / cloud IMDS + ipaddress.ip_network("::1/128"), # IPv6 loopback + ipaddress.ip_network("fc00::/7"), # IPv6 unique-local + ipaddress.ip_network("fe80::/10"), # IPv6 link-local (IPv6 cloud IMDS equivalent) +] + + +def _is_blocked_ip(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: + """Return True if ip must not be reachable via an image URL fetch.""" + if any(ip in net for net in _BLOCKED_NETWORKS): + return True + return not ip.is_global or ip.is_multicast or ip.is_reserved or ip.is_unspecified + + +def _resolve_and_validate(hostname: str) -> list[tuple[int, str]]: + """Resolve hostname to every A/AAAA address and reject if any is blocked. + + Returns a list of (address_family, ip_string) tuples for all resolved + addresses. Validating every returned address (rather than just the + first) guards against DNS answers that mix a public address with a + private/link-local one (DNS rebinding / multi-answer SSRF). + """ + try: + infos = socket.getaddrinfo(hostname, None) + except socket.gaierror as exc: + raise ValueError(f"Cannot resolve image URL hostname '{hostname}': {exc}") from exc + resolved: list[tuple[int, str]] = [] + seen: set[str] = set() + for family, _, _, _, sockaddr in infos: + ip_str = sockaddr[0] + if ip_str in seen: + continue + seen.add(ip_str) + ip = ipaddress.ip_address(ip_str) + if _is_blocked_ip(ip): + raise ValueError( + f"Image URL hostname '{hostname}' resolves to a blocked address ({ip}). " + "Private, loopback, link-local, and other non-global addresses are not allowed." + ) + resolved.append((family, ip_str)) + if not resolved: + raise ValueError(f"Cannot resolve image URL hostname '{hostname}': no addresses returned.") + return resolved + + +def validate_image_url(url: str) -> None: + """Raise ValueError if url is not a safe http/https URL. + + Rejects file://, non-http(s) schemes, and URLs that resolve to + private/link-local/loopback ranges (SSRF guard). This is a point-in-time + check; use fetch_image_bytes_safely()/fetch_image_data_uri_safely() to + actually retrieve the image, since those pin validation to the + connection itself and guard redirects. + """ + parsed = urlparse(url) + if parsed.scheme not in ("http", "https"): + raise ValueError(f"Unsupported image URL scheme '{parsed.scheme}'. Only http and https are allowed.") + hostname = parsed.hostname + if not hostname: + raise ValueError("Image URL has no hostname.") + _resolve_and_validate(hostname) + + +def _pinned_request(url: str, hostname: str, ip: str, family: int, timeout: float) -> urllib3.HTTPResponse: + """Issue a GET request connected directly to ip, without following redirects. + + Connecting to the pre-validated IP directly (instead of letting the HTTP + client re-resolve hostname) closes the TOCTOU window between validation + and connection that DNS rebinding exploits. The Host header and TLS SNI + still use the original hostname so virtual hosting and certificate + validation behave normally. + """ + parsed = urlparse(url) + port = parsed.port or (443 if parsed.scheme == "https" else 80) + path = parsed.path or "/" + if parsed.query: + path = f"{path}?{parsed.query}" + pool_host = f"[{ip}]" if family == socket.AF_INET6 else ip + headers = {"Host": hostname} + if parsed.scheme == "https": + pool = urllib3.HTTPSConnectionPool( + pool_host, + port=port, + timeout=timeout, + assert_hostname=hostname, + server_hostname=hostname, + cert_reqs="CERT_REQUIRED", + ) + else: + pool = urllib3.HTTPConnectionPool(pool_host, port=port, timeout=timeout) + try: + return pool.request("GET", path, headers=headers, redirect=False, preload_content=True) + finally: + pool.close() + + +def fetch_image_bytes_safely(url: str, timeout: float = 5, max_redirects: int = 5) -> bytes: + """Safely fetch image bytes from url, guarding against SSRF. + + Unlike validate_image_url() followed by a separate request, this + resolves and validates every hostname (including redirect targets) + immediately before connecting, connects directly to the validated IP + address, and never auto-follows redirects — each hop is re-validated + from scratch. + """ + current_url = url + for _ in range(max_redirects + 1): + parsed = urlparse(current_url) + if parsed.scheme not in ("http", "https"): + raise ValueError(f"Unsupported image URL scheme '{parsed.scheme}'. Only http and https are allowed.") + hostname = parsed.hostname + if not hostname: + raise ValueError("Image URL has no hostname.") + resolved = _resolve_and_validate(hostname) + family, ip = resolved[0] + response = _pinned_request(current_url, hostname, ip, family, timeout) + if response.status in (301, 302, 303, 307, 308): + location = response.headers.get("Location") + if not location: + raise ValueError(f"Redirect response from '{current_url}' is missing a Location header.") + current_url = urljoin(current_url, location) + continue + if response.status >= 400: + raise ValueError(f"Failed to fetch image URL '{current_url}': HTTP {response.status}") + return response.data + raise ValueError(f"Exceeded max redirects ({max_redirects}) while fetching image URL '{url}'.") + + +def fetch_image_data_uri_safely(url: str, timeout: float = 5, max_redirects: int = 5) -> str: + """Safely fetch an image and return it as a base64 data URI. + + Useful for handing an image to code (e.g. third-party libraries) that + would otherwise perform its own, unguarded network fetch given a raw + URL — the network I/O happens here, under the SSRF guard, instead. + """ + from PIL import Image + + data = fetch_image_bytes_safely(url, timeout=timeout, max_redirects=max_redirects) + image_format = (Image.open(BytesIO(data)).format or "JPEG").lower() + encoded = base64.b64encode(data).decode("ascii") + return f"data:image/{image_format};base64,{encoded}" diff --git a/nemo_deploy/multimodal/megatron_multimodal_deployable.py b/nemo_deploy/multimodal/megatron_multimodal_deployable.py index 6f69a5177..21cab31e2 100644 --- a/nemo_deploy/multimodal/megatron_multimodal_deployable.py +++ b/nemo_deploy/multimodal/megatron_multimodal_deployable.py @@ -173,6 +173,15 @@ def process_image_input(self, image_source): if isinstance(self.inference_wrapped_model, QwenVLInferenceWrapper): from qwen_vl_utils import process_vision_info + from nemo_deploy.multimodal.image_url_validator import fetch_image_data_uri_safely + + # data: URIs are inline base64 and never trigger a network request. + # All other values are treated as URLs: fetch them ourselves under the SSRF + # guard and hand qwen_vl_utils a data URI, since it would otherwise perform + # its own unguarded network fetch for a raw URL. + if not image_source.startswith("data:"): + image_source = fetch_image_data_uri_safely(image_source) + messages = [ { "role": "user", diff --git a/nemo_deploy/multimodal/query_multimodal.py b/nemo_deploy/multimodal/query_multimodal.py index a729d0c27..64f91625f 100644 --- a/nemo_deploy/multimodal/query_multimodal.py +++ b/nemo_deploy/multimodal/query_multimodal.py @@ -17,7 +17,6 @@ from typing import List, Optional import numpy as np -import requests from nemo_deploy.utils import str_list2numpy from nemo_export_deploy_common.import_utils import ( @@ -100,8 +99,10 @@ def setup_media(self, input_media): raise UnavailableError(MISSING_PIL_MSG) if input_media.startswith("http") or input_media.startswith("https"): - response = requests.get(input_media, timeout=5) - media = Image.open(BytesIO(response.content)).convert("RGB") + from nemo_deploy.multimodal.image_url_validator import fetch_image_bytes_safely + + image_bytes = fetch_image_bytes_safely(input_media, timeout=5) + media = Image.open(BytesIO(image_bytes)).convert("RGB") else: media = Image.open(input_media).convert("RGB") return np.expand_dims(np.array(media), axis=0) diff --git a/tests/unit_tests/deploy/test_image_url_validator.py b/tests/unit_tests/deploy/test_image_url_validator.py new file mode 100644 index 000000000..fa9954e7e --- /dev/null +++ b/tests/unit_tests/deploy/test_image_url_validator.py @@ -0,0 +1,231 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import pathlib +import socket +from unittest.mock import MagicMock, patch + +import pytest + +# Load the validator directly by file path so we don't trigger nemo_deploy/__init__.py +# (which requires torch/triton). The module itself is pure stdlib + urllib3. +_validator_path = pathlib.Path(__file__).resolve().parents[3] / "nemo_deploy" / "multimodal" / "image_url_validator.py" +_spec = importlib.util.spec_from_file_location("image_url_validator", _validator_path) +_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_mod) +validate_image_url = _mod.validate_image_url +fetch_image_bytes_safely = _mod.fetch_image_bytes_safely + + +def _addrinfo(family, ip_str): + return (family, socket.SOCK_STREAM, 6, "", (ip_str, 0)) + + +def _mock_resolve(*ip_strs, family=socket.AF_INET): + """Return a patch for socket.getaddrinfo resolving to the given IP(s).""" + infos = [_addrinfo(family, ip) for ip in ip_strs] + return patch.object(_mod.socket, "getaddrinfo", return_value=infos) + + +class TestBlockedSchemes: + def test_file_scheme_rejected(self): + with pytest.raises(ValueError, match="scheme"): + validate_image_url("file:///etc/passwd") + + def test_ftp_scheme_rejected(self): + with pytest.raises(ValueError, match="scheme"): + validate_image_url("ftp://example.com/img.png") + + def test_no_scheme_rejected(self): + with pytest.raises(ValueError, match="scheme"): + validate_image_url("example.com/img.png") + + +class TestBlockedRanges: + def test_loopback_ipv4_rejected(self): + with _mock_resolve("127.0.0.1"): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://localhost/img.jpg") + + def test_loopback_other_subnet_rejected(self): + with _mock_resolve("127.1.2.3"): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://internal.local/img.jpg") + + def test_cloud_imds_rejected(self): + # 169.254.169.254 is the AWS/GCP/Azure metadata service + with _mock_resolve("169.254.169.254"): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://169.254.169.254/latest/meta-data/") + + def test_link_local_rejected(self): + with _mock_resolve("169.254.0.1"): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://169.254.0.1/img.jpg") + + def test_rfc1918_10_rejected(self): + with _mock_resolve("10.0.0.1"): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://internal.corp/img.jpg") + + def test_rfc1918_172_rejected(self): + with _mock_resolve("172.16.0.1"): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://internal.corp/img.jpg") + + def test_rfc1918_192_rejected(self): + with _mock_resolve("192.168.1.1"): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://192.168.1.1/img.jpg") + + +class TestIPv6: + def test_ipv6_loopback_rejected(self): + with _mock_resolve("::1", family=socket.AF_INET6): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://ipv6-loopback/img.jpg") + + def test_ipv6_link_local_rejected(self): + with _mock_resolve("fe80::1", family=socket.AF_INET6): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://ipv6-link-local/img.jpg") + + def test_ipv6_unique_local_rejected(self): + with _mock_resolve("fc00::1", family=socket.AF_INET6): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://ipv6-ula/img.jpg") + + def test_ipv6_global_allowed(self): + with _mock_resolve("2606:4700:4700::1111", family=socket.AF_INET6): + validate_image_url("http://ipv6-public/img.jpg") # must not raise + + +class TestMultipleDnsAnswers: + def test_mixed_public_and_blocked_answer_rejected(self): + """One public + one blocked answer for the same host must be rejected. + + Guards against DNS rebinding / multi-answer SSRF where only the + first-returned address would look safe. + """ + infos = [_addrinfo(socket.AF_INET, "93.184.216.34"), _addrinfo(socket.AF_INET, "169.254.169.254")] + with patch.object(_mod.socket, "getaddrinfo", return_value=infos): + with pytest.raises(ValueError, match="blocked address"): + validate_image_url("http://multi-answer.example.com/img.jpg") + + def test_all_public_answers_allowed(self): + infos = [_addrinfo(socket.AF_INET, "93.184.216.34"), _addrinfo(socket.AF_INET, "1.2.3.4")] + with patch.object(_mod.socket, "getaddrinfo", return_value=infos): + validate_image_url("http://multi-answer.example.com/img.jpg") # must not raise + + +class TestAllowedUrls: + def test_public_https_allowed(self): + with _mock_resolve("93.184.216.34"): # example.com + validate_image_url("https://example.com/image.jpg") # must not raise + + def test_public_http_allowed(self): + with _mock_resolve("1.2.3.4"): + validate_image_url("http://cdn.example.com/image.png") # must not raise + + +class TestNoHostname: + def test_url_without_hostname_rejected(self): + with pytest.raises(ValueError, match="hostname"): + validate_image_url("http:///image.jpg") + + def test_dns_failure_rejected(self): + with patch.object( + _mod.socket, + "getaddrinfo", + side_effect=socket.gaierror("Name or service not known"), + ): + with pytest.raises(ValueError, match="Cannot resolve"): + validate_image_url("http://nonexistent.invalid/img.jpg") + + +def _mock_pinned_response(status, data=b"", location=None): + response = MagicMock() + response.status = status + response.data = data + response.headers = {"Location": location} if location else {} + return response + + +class TestFetchImageBytesSafely: + def test_fetch_success(self): + with _mock_resolve("93.184.216.34"): + with patch.object(_mod, "_pinned_request", return_value=_mock_pinned_response(200, b"imgdata")): + result = fetch_image_bytes_safely("http://example.com/image.jpg") + assert result == b"imgdata" + + def test_fetch_rejects_redirect_to_blocked_address(self): + """A public URL redirecting to a blocked address must be rejected. + + Each redirect hop is independently resolved and validated, so the + redirect target's own DNS resolution (to a blocked IP) is what + trips the guard. + """ + + def fake_getaddrinfo(host, *_args, **_kwargs): + if host == "public.example.com": + return [_addrinfo(socket.AF_INET, "93.184.216.34")] + return [_addrinfo(socket.AF_INET, "169.254.169.254")] + + with patch.object(_mod.socket, "getaddrinfo", side_effect=fake_getaddrinfo): + with patch.object( + _mod, + "_pinned_request", + return_value=_mock_pinned_response(302, location="http://169.254.169.254/latest/meta-data/"), + ): + with pytest.raises(ValueError, match="blocked address"): + fetch_image_bytes_safely("http://public.example.com/image.jpg") + + def test_fetch_follows_redirect_chain_to_public_address(self): + def fake_getaddrinfo(host, *_args, **_kwargs): + return [_addrinfo(socket.AF_INET, "93.184.216.34")] + + responses = [ + _mock_pinned_response(302, location="http://public.example.com/final.jpg"), + _mock_pinned_response(200, data=b"final-bytes"), + ] + with patch.object(_mod.socket, "getaddrinfo", side_effect=fake_getaddrinfo): + with patch.object(_mod, "_pinned_request", side_effect=responses): + result = fetch_image_bytes_safely("http://public.example.com/image.jpg") + assert result == b"final-bytes" + + def test_fetch_raises_on_missing_location_header(self): + with _mock_resolve("93.184.216.34"): + with patch.object(_mod, "_pinned_request", return_value=_mock_pinned_response(302)): + with pytest.raises(ValueError, match="Location"): + fetch_image_bytes_safely("http://example.com/image.jpg") + + def test_fetch_raises_on_too_many_redirects(self): + def fake_getaddrinfo(host, *_args, **_kwargs): + return [_addrinfo(socket.AF_INET, "93.184.216.34")] + + with patch.object(_mod.socket, "getaddrinfo", side_effect=fake_getaddrinfo): + with patch.object( + _mod, + "_pinned_request", + return_value=_mock_pinned_response(302, location="http://public.example.com/next.jpg"), + ): + with pytest.raises(ValueError, match="redirects"): + fetch_image_bytes_safely("http://public.example.com/image.jpg", max_redirects=2) + + def test_fetch_raises_on_http_error_status(self): + with _mock_resolve("93.184.216.34"): + with patch.object(_mod, "_pinned_request", return_value=_mock_pinned_response(404)): + with pytest.raises(ValueError, match="404"): + fetch_image_bytes_safely("http://example.com/image.jpg") diff --git a/tests/unit_tests/deploy/test_megatron_multimodal_deployable.py b/tests/unit_tests/deploy/test_megatron_multimodal_deployable.py index 81a6c760e..2c72df9ee 100644 --- a/tests/unit_tests/deploy/test_megatron_multimodal_deployable.py +++ b/tests/unit_tests/deploy/test_megatron_multimodal_deployable.py @@ -587,6 +587,7 @@ def test_process_image_input_with_http_url(self, deployable): # HTTP URL as image source image_source = "https://example.com/image.jpg" + data_uri = "data:image/jpeg;base64,ZmFrZQ==" expected_image = Image.new("RGB", (100, 100)) with patch("nemo_deploy.multimodal.megatron_multimodal_deployable.QwenVLInferenceWrapper", mock_qwenvl_class): @@ -594,19 +595,27 @@ def test_process_image_input_with_http_url(self, deployable): with patch("nemo_deploy.multimodal.megatron_multimodal_deployable.isinstance") as mock_isinstance: mock_isinstance.return_value = True - with patch("qwen_vl_utils.process_vision_info") as mock_process: - mock_process.return_value = (expected_image, None) + with patch( + "nemo_deploy.multimodal.image_url_validator.fetch_image_data_uri_safely", + return_value=data_uri, + ) as mock_fetch: + with patch("qwen_vl_utils.process_vision_info") as mock_process: + mock_process.return_value = (expected_image, None) - result = deployable.process_image_input(image_source) + result = deployable.process_image_input(image_source) - # Verify process_vision_info was called with URL - call_args = mock_process.call_args[0][0] - assert len(call_args) == 1 - assert call_args[0]["role"] == "user" - assert call_args[0]["content"][0]["type"] == "image" - assert call_args[0]["content"][0]["image"] == image_source + # The raw URL must be fetched under the SSRF guard, never handed + # to process_vision_info directly. + mock_fetch.assert_called_once_with(image_source) - assert result == expected_image + # Verify process_vision_info was called with the resulting data URI + call_args = mock_process.call_args[0][0] + assert len(call_args) == 1 + assert call_args[0]["role"] == "user" + assert call_args[0]["content"][0]["type"] == "image" + assert call_args[0]["content"][0]["image"] == data_uri + + assert result == expected_image def test_process_image_input_with_unsupported_model(self, deployable): """Test process_image_input with unsupported model raises ValueError.""" diff --git a/tests/unit_tests/deploy/test_query_multimodal.py b/tests/unit_tests/deploy/test_query_multimodal.py index 4690d32bf..94934b036 100644 --- a/tests/unit_tests/deploy/test_query_multimodal.py +++ b/tests/unit_tests/deploy/test_query_multimodal.py @@ -63,12 +63,10 @@ def test_setup_media_image_local(self, query_multimodal, mock_image): assert result.shape[0] == 1 # Batch dimension os.unlink(mock_image) - @patch("requests.get") - def test_setup_media_image_url(self, mock_get, query_multimodal): - # Mock the response from requests.get - mock_response = MagicMock() - mock_response.content = b"fake_image_data" - mock_get.return_value = mock_response + @patch("nemo_deploy.multimodal.image_url_validator.fetch_image_bytes_safely") + def test_setup_media_image_url(self, mock_fetch, query_multimodal): + # Mock the safe fetch so no real network I/O happens + mock_fetch.return_value = b"fake_image_data" # Mock Image.open with patch("PIL.Image.open") as mock_image_open: