fix: skip DNS resolution for base64 data in token counter, add unit tests

Check URL scheme before calling safe_get in token counter to avoid
unnecessary DNS resolution on base64-encoded image data.

Add 14 unit tests for validate_url covering blocked networks, scheme
validation, URL rewriting, and DNS failure handling.

litellm/litellm_core_utils/token_counter.py

@@ -211,13 +211,15 @@ def get_image_dimensions(
         Tuple[int, int]: The width and height of the image.
     """
     img_data = None
-    try:
-        # Try to open as URL with SSRF protection
-        client = _get_httpx_client()
-        response = safe_get(client, data)
-        img_data = response.read()
-    except Exception:
-        # If not URL, assume it's base64
+    if data.startswith(("http://", "https://")):
+        try:
+            client = _get_httpx_client()
+            response = safe_get(client, data)
+            img_data = response.read()
+        except Exception:
+            pass
+    if img_data is None:
+        # Not a URL or fetch failed — assume base64
         _header, encoded = data.split(",", 1)
         img_data = base64.b64decode(encoded)
 

tests/test_litellm/proxy/common_utils/test_url_utils.py

@@ -0,0 +1,64 @@
+import pytest
+
+from litellm.proxy.common_utils.url_utils import SSRFError, validate_url
+
+
+class TestValidateUrl:
+    def test_blocks_loopback(self):
+        with pytest.raises(SSRFError):
+            validate_url("http://127.0.0.1/test")
+
+    def test_blocks_imds(self):
+        with pytest.raises(SSRFError):
+            validate_url("http://169.254.169.254/latest/meta-data/")
+
+    def test_blocks_rfc1918_class_a(self):
+        with pytest.raises(SSRFError):
+            validate_url("http://10.0.1.5:8080/v1/completions")
+
+    def test_blocks_rfc1918_class_b(self):
+        with pytest.raises(SSRFError):
+            validate_url("http://172.16.0.1/")
+
+    def test_blocks_rfc1918_class_c(self):
+        with pytest.raises(SSRFError):
+            validate_url("http://192.168.1.1/")
+
+    def test_blocks_file_scheme(self):
+        with pytest.raises(SSRFError):
+            validate_url("file:///etc/passwd")
+
+    def test_blocks_ftp_scheme(self):
+        with pytest.raises(SSRFError):
+            validate_url("ftp://internal.host/data")
+
+    def test_blocks_no_hostname(self):
+        with pytest.raises(SSRFError):
+            validate_url("http:///path")
+
+    def test_allows_public_https(self):
+        rewritten, host = validate_url("https://example.com/image.png")
+        assert host == "example.com"
+        assert rewritten == "https://example.com/image.png"
+
+    def test_rewrites_public_http_to_ip(self):
+        rewritten, host = validate_url("http://example.com/image.png")
+        assert host == "example.com"
+        assert "example.com" not in rewritten
+
+    def test_preserves_path_and_query(self):
+        rewritten, host = validate_url("http://example.com/path?key=value")
+        assert "/path" in rewritten
+        assert "key=value" in rewritten
+
+    def test_dns_failure_raises(self):
+        with pytest.raises(SSRFError, match="DNS resolution failed"):
+            validate_url("http://this-domain-does-not-exist-xyz123.invalid/test")
+
+    def test_blocks_localhost_hostname(self):
+        with pytest.raises(SSRFError):
+            validate_url("http://localhost/")
+
+    def test_blocks_ipv6_loopback(self):
+        with pytest.raises(SSRFError):
+            validate_url("http://[::1]/")