From d15196b5197d89270469faabaece0dd26cf1f4b6 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 16 Apr 2026 04:30:14 +0000 Subject: [PATCH] fix(proxy): add safe_get/async_safe_get with redirect validation Add safe_get() and async_safe_get() helpers that validate each redirect hop before following. For HTTPS, rely on TLS certificate binding instead of URL rewriting. Simplify call sites to use the new helpers. --- .../prompt_templates/image_handling.py | 20 +----- litellm/litellm_core_utils/token_counter.py | 11 +-- .../mcp_server/openapi_to_mcp_generator.py | 9 +-- litellm/proxy/common_utils/url_utils.py | 70 +++++++++++++++++-- litellm/rag/ingestion/base_ingestion.py | 9 +-- 5 files changed, 73 insertions(+), 46 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index c0727699167..c036ec23ddb 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -10,7 +10,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.proxy.common_utils.url_utils import SSRFError, validate_url +from litellm.proxy.common_utils.url_utils import async_safe_get, safe_get MAX_IMGS_IN_MEMORY = 10 @@ -82,17 +82,10 @@ async def async_convert_url_to_base64(url: str) -> str: if cached_result: return cached_result - # Resolve DNS once, validate IPs, rewrite URL to validated IP - validated_url, original_host = validate_url(url) - client = litellm.module_level_aclient for _ in range(3): try: - response = await client.get( - validated_url, - headers={"Host": original_host}, - follow_redirects=False, - ) + response = await async_safe_get(client, url) return _process_image_response(response, url) except litellm.ImageFetchError: raise @@ -114,17 +107,10 @@ def convert_url_to_base64(url: str) -> str: if cached_result: return cached_result - # Resolve DNS once, validate IPs, rewrite URL to validated IP - validated_url, original_host = validate_url(url) - client = litellm.module_level_client for _ in range(3): try: - response = client.get( - validated_url, - headers={"Host": original_host}, - follow_redirects=False, - ) + response = safe_get(client, url) return _process_image_response(response, url) except litellm.ImageFetchError: raise diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e2d2a56c698..ad2691156a5 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -30,7 +30,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.default_encoding import encoding as default_encoding from litellm.llms.custom_httpx.http_handler import _get_httpx_client -from litellm.proxy.common_utils.url_utils import validate_url +from litellm.proxy.common_utils.url_utils import safe_get from litellm.types.llms.anthropic import ( AnthropicMessagesToolResultParam, AnthropicMessagesToolUseParam, @@ -212,14 +212,9 @@ def get_image_dimensions( """ img_data = None try: - # Try to open as URL — validate and pin to resolved IP - validated_url, original_host = validate_url(data) + # Try to open as URL with SSRF protection client = _get_httpx_client() - response = client.get( - validated_url, - headers={"Host": original_host}, - follow_redirects=False, - ) + response = safe_get(client, data) img_data = response.read() except Exception: # If not URL, assume it's base64 diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 68f52f34395..d6b2cb86b26 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -15,7 +15,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.proxy.common_utils.url_utils import validate_url +from litellm.proxy.common_utils.url_utils import async_safe_get from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -75,13 +75,8 @@ def load_openapi_spec(filepath: str) -> Dict[str, Any]: async def load_openapi_spec_async(filepath: str) -> Dict[str, Any]: if filepath.startswith("http://") or filepath.startswith("https://"): - validated_url, original_host = validate_url(filepath) client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - r = await client.get( - validated_url, - headers={"Host": original_host}, - follow_redirects=False, - ) + r = await async_safe_get(client, filepath) r.raise_for_status() return r.json() diff --git a/litellm/proxy/common_utils/url_utils.py b/litellm/proxy/common_utils/url_utils.py index 97fee0c0965..0583cc3a42b 100644 --- a/litellm/proxy/common_utils/url_utils.py +++ b/litellm/proxy/common_utils/url_utils.py @@ -4,16 +4,15 @@ URL validation for user-controlled URLs. Use validate_url() before fetching any URL that originates from user input (image_url, file_url, spec_path, etc.) to prevent SSRF attacks. -The function resolves DNS once, validates all IPs, and rewrites the URL -to connect to the validated IP directly — no TOCTOU gap, no DNS rebinding. -Callers should also set follow_redirects=False to prevent redirect-based -SSRF bypasses. +validate_url() resolves DNS once, validates all IPs, and rewrites the +URL to connect to the validated IP directly — no TOCTOU gap, no DNS +rebinding. Redirects are followed manually with validation at each hop. """ import ipaddress import socket from ipaddress import ip_address, ip_network -from typing import Optional, Tuple +from typing import Any, Optional, Tuple, Union from urllib.parse import urlparse, urlunparse _BLOCKED_NETWORKS = [ @@ -105,12 +104,18 @@ def validate_url(url: str) -> Tuple[str, str]: "provider configuration instead of a user-supplied URL." ) - # Rewrite URL to connect to the first validated IP + # For HTTPS, TLS certificate validation binds the connection to the + # hostname — DNS rebinding can't redirect to a different server because + # the cert wouldn't match. Return the original URL. + if parsed.scheme == "https": + return url, hostname + + # For HTTP, rewrite URL to connect to the validated IP directly + # to prevent DNS rebinding (no TLS to bind the connection). validated_ip = addrinfo[0][4][0] is_ipv6 = addrinfo[0][0] == socket.AF_INET6 ip_host = f"[{validated_ip}]" if is_ipv6 else validated_ip - # Reconstruct netloc with IP instead of hostname if port: new_netloc = f"{ip_host}:{port}" else: @@ -121,3 +126,54 @@ def validate_url(url: str) -> Tuple[str, str]: ) return rewritten, hostname + + +_MAX_REDIRECTS = 10 + + +def safe_get(client: Any, url: str, **kwargs: Any) -> Any: + """ + Fetch a user-supplied URL with SSRF protection on every redirect hop. + + Validates the initial URL and each redirect target before making the + request. No DNS rebinding (resolve-and-rewrite). No redirect bypass + (each hop validated). No breaking change for legitimate CDN redirects. + + Args: + client: An httpx.Client or httpx.AsyncClient (sync version). + url: The user-supplied URL. + **kwargs: Additional kwargs passed to client.get(). + + Returns: + The final httpx.Response. + """ + kwargs.pop("follow_redirects", None) + for _ in range(_MAX_REDIRECTS): + validated_url, original_host = validate_url(url) + response = client.get( + validated_url, + headers={**kwargs.pop("headers", {}), "Host": original_host}, + follow_redirects=False, + **kwargs, + ) + if not response.is_redirect or response.next_request is None: + return response + url = str(response.next_request.url) + raise SSRFError("Too many redirects") + + +async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: + """Async version of safe_get.""" + kwargs.pop("follow_redirects", None) + for _ in range(_MAX_REDIRECTS): + validated_url, original_host = validate_url(url) + response = await client.get( + validated_url, + headers={**kwargs.pop("headers", {}), "Host": original_host}, + follow_redirects=False, + **kwargs, + ) + if not response.is_redirect or response.next_request is None: + return response + url = str(response.next_request.url) + raise SSRFError("Too many redirects") diff --git a/litellm/rag/ingestion/base_ingestion.py b/litellm/rag/ingestion/base_ingestion.py index 1d868d35c7a..6f139764a85 100644 --- a/litellm/rag/ingestion/base_ingestion.py +++ b/litellm/rag/ingestion/base_ingestion.py @@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.proxy.common_utils.url_utils import validate_url +from litellm.proxy.common_utils.url_utils import async_safe_get from litellm.rag.ingestion.file_parsers import extract_text_from_pdf from litellm.rag.text_splitters import RecursiveCharacterTextSplitter from litellm.types.rag import RAGIngestOptions, RAGIngestResponse @@ -112,13 +112,8 @@ class BaseRAGIngestion(ABC): return filename, file_content, content_type, None if file_url: - validated_url, original_host = validate_url(file_url) http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.RAG) - response = await http_client.get( - validated_url, - headers={"Host": original_host}, - follow_redirects=False, - ) + response = await async_safe_get(http_client, file_url) response.raise_for_status() file_content = response.content filename = file_url.split("/")[-1] or "document"