From 3920cf4dfe458b6d4982e2a64c162de21defa2a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 23:45:37 -0700 Subject: [PATCH] fix(image_handling): tell callers when the image host did not resolve instead of blaming the URL policy validate_url raises HostResolutionError, a SSRFError subclass, for the two DNS outcomes (lookup failed, no addresses). The image fetch helper maps that to a "host could not be resolved" message and keeps the user_url_allowed_hosts hint for the policy verdicts it can actually fix. --- .../prompt_templates/image_handling.py | 14 +++-- litellm/litellm_core_utils/url_utils.py | 8 ++- .../litellm_core_utils/test_image_handling.py | 56 +++++++++++-------- .../litellm_core_utils/test_url_utils.py | 19 ++++++- 4 files changed, 64 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index 4e7dfdb4017..99efbd13320 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -15,7 +15,7 @@ import litellm from litellm import verbose_logger from litellm.caching.caching import InMemoryCache from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB -from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError, async_safe_get, safe_get from litellm.types.llms.openai import AllMessageValues MAX_IMGS_IN_MEMORY: Final = 10 @@ -78,8 +78,12 @@ def _process_image_response(response: Response, url: str) -> str: return result -def _url_policy_rejection(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": - verbose_logger.warning("Image fetch of %s rejected by the URL policy: %s", url, verdict) +def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError": + verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict) + if isinstance(verdict, HostResolutionError): + return litellm.ImageFetchError( + f"Error: Unable to fetch image from URL. The image host could not be resolved. url={url}" + ) return litellm.ImageFetchError( "Error: Unable to fetch image from URL. The proxy's URL policy rejected this host; " f"an admin can allow it with `user_url_allowed_hosts` in general_settings. url={url}" @@ -108,7 +112,7 @@ async def async_convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise _url_policy_rejection(url, e) from e + raise _rejected_image_fetch(url, e) from e except Exception: pass raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}") @@ -136,7 +140,7 @@ def convert_url_to_base64(url: str) -> str: except litellm.ImageFetchError: raise except SSRFError as e: - raise _url_policy_rejection(url, e) from e + raise _rejected_image_fetch(url, e) from e except Exception as e: verbose_logger.exception(e) raise litellm.ImageFetchError( diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index fa070a648f5..f4388d5fd12 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -93,6 +93,10 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" +class HostResolutionError(SSRFError): + pass + + def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. @@ -324,10 +328,10 @@ def validate_url(url: str) -> tuple[str, str]: try: addrinfo: Final = socket.getaddrinfo(hostname, effective_port, proto=socket.IPPROTO_TCP) except socket.gaierror as e: - raise SSRFError(f"DNS resolution failed for '{hostname}': {e}") + raise HostResolutionError(f"DNS resolution failed for '{hostname}': {e}") if not addrinfo: - raise SSRFError(f"No addresses found for '{hostname}'") + raise HostResolutionError(f"No addresses found for '{hostname}'") if not is_allowlisted: for addrinfo_entry in addrinfo: diff --git a/tests/test_litellm/litellm_core_utils/test_image_handling.py b/tests/test_litellm/litellm_core_utils/test_image_handling.py index d8ccc9251dd..3a909298aba 100644 --- a/tests/test_litellm/litellm_core_utils/test_image_handling.py +++ b/tests/test_litellm/litellm_core_utils/test_image_handling.py @@ -17,7 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, convert_url_to_base64, ) -from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.litellm_core_utils.url_utils import HostResolutionError, SSRFError @pytest.fixture(autouse=True) @@ -115,9 +115,7 @@ class StreamingLargeImageClient: request=Request("GET", url), ) # Mock the iter_bytes method to return our generator - response.iter_bytes = lambda chunk_size=8192: generate_chunks( - size_bytes, chunk_size - ) + response.iter_bytes = lambda chunk_size=8192: generate_chunks(size_bytes, chunk_size) return response @@ -215,9 +213,7 @@ def test_streaming_download_handles_petabyte_file(monkeypatch): """ # Simulate a 1 petabyte file (1,000,000 GB) # Without streaming protection, this would cause OOM or hang indefinitely - client = StreamingLargeImageClient( - size_mb=1_000_000_000, include_content_length=False - ) + client = StreamingLargeImageClient(size_mb=1_000_000_000, include_content_length=False) monkeypatch.setattr(litellm, "module_level_client", client) with pytest.raises(litellm.ImageFetchError) as excinfo: @@ -429,20 +425,32 @@ async def test_async_inline_remote_media_cancels_the_other_fetches_when_one_fail _SSRF_VERDICTS = ( - "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " - "add the host to `user_url_allowed_hosts` in general_settings.", - "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known", - "No addresses found for 'internal.example'", + ( + SSRFError( + "URL targets a blocked address (10.0.0.8). If this is a legitimate internal service, " + "add the host to `user_url_allowed_hosts` in general_settings." + ), + "The proxy's URL policy rejected this host; an admin can allow it with `user_url_allowed_hosts`", + ), + ( + HostResolutionError( + "DNS resolution failed for 'internal.example': [Errno 8] nodename nor servname provided, or not known" + ), + "The image host could not be resolved", + ), + (HostResolutionError("No addresses found for 'internal.example'"), "The image host could not be resolved"), ) -def _assert_one_verdict_free_message(messages, url): - assert len(set(messages)) == 1 - assert "10.0.0.8" not in messages[0] - assert "DNS" not in messages[0] - assert "No addresses" not in messages[0] - assert "user_url_allowed_hosts" in messages[0] - assert url in messages[0] +def _assert_verdict_free_messages(messages, url): + for message, (verdict, expected_guidance) in zip(messages, _SSRF_VERDICTS, strict=True): + assert expected_guidance in message + assert url in message + assert "10.0.0.8" not in message + assert "DNS" not in message + assert "No addresses" not in message + if isinstance(verdict, HostResolutionError): + assert "user_url_allowed_hosts" not in message async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): @@ -450,11 +458,11 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict in _SSRF_VERDICTS: + for verdict, _ in _SSRF_VERDICTS: async def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) - raise SSRFError(verdict) + raise verdict monkeypatch.setattr(image_handling, "async_safe_get", block) with pytest.raises(litellm.ImageFetchError) as raised: @@ -462,7 +470,7 @@ async def test_async_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_r messages.append(raised.value.message) assert attempts == [url] * len(_SSRF_VERDICTS) - _assert_one_verdict_free_message(messages, url) + _assert_verdict_free_messages(messages, url) def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeypatch): @@ -470,11 +478,11 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp messages = [] url = f"http://internal.example/{uuid.uuid4()}.png" - for verdict in _SSRF_VERDICTS: + for verdict, _ in _SSRF_VERDICTS: def block(client, fetched_url, verdict=verdict, **kwargs): attempts.append(fetched_url) - raise SSRFError(verdict) + raise verdict monkeypatch.setattr(image_handling, "safe_get", block) with pytest.raises(litellm.ImageFetchError) as raised: @@ -482,7 +490,7 @@ def test_convert_url_to_base64_hides_the_ssrf_verdict_and_does_not_retry(monkeyp messages.append(raised.value.message) assert attempts == [url] * len(_SSRF_VERDICTS) - _assert_one_verdict_free_message(messages, url) + _assert_verdict_free_messages(messages, url) async def test_async_inline_remote_media_caps_in_flight_fetches_per_request(monkeypatch): diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index fccdc1a2a0a..151c7ceb5bd 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.url_utils import ( + HostResolutionError, SSRFError, _is_blocked_ip, assert_same_origin, @@ -161,10 +162,24 @@ class TestValidateUrl: assert "/path" in rewritten assert "key=value" in rewritten - def test_dns_failure_raises(self, mock_dns_failure): - with pytest.raises(SSRFError, match="DNS resolution failed"): + def test_dns_failure_raises_a_host_resolution_error(self, mock_dns_failure): + with pytest.raises(HostResolutionError, match="DNS resolution failed"): validate_url("http://this-domain-does-not-exist-xyz123.invalid/test") + def test_empty_resolution_raises_a_host_resolution_error(self, monkeypatch): + monkeypatch.setattr(url_utils.socket, "getaddrinfo", lambda *args, **kwargs: []) + with pytest.raises(HostResolutionError, match="No addresses found"): + validate_url("http://this-domain-resolves-to-nothing.invalid/test") + + def test_blocked_address_is_not_a_host_resolution_error(self, monkeypatch): + def fake(host, port, *a, **kw): + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("10.0.0.8", port or 80))] + + monkeypatch.setattr(url_utils.socket, "getaddrinfo", fake) + with pytest.raises(SSRFError, match="blocked address") as raised: + validate_url("http://internal.example/test") + assert not isinstance(raised.value, HostResolutionError) + def test_blocks_localhost_hostname(self, monkeypatch): def fake(host, port, *a, **kw): return [