review-fix(vision): restore _validate_image_url/_image_url_shape_ok (public on main) + SSRF/localhost-block tests

ethernet8023: TestValidateImageUrl (incl. localhost block) was deleted with the sync
validator. Both names restored with BASE semantics; the async validator (the live
download path) now shares _image_url_shape_ok and gets its own localhost/malformed test.
This commit is contained in:
Teknium
2026-09-03 09:37:18 -07:00
parent cb9b7c36f3
commit 9c7238572e
2 changed files with 69 additions and 2 deletions

View File

@@ -1,5 +1,6 @@
"""Tests for tools/vision_tools.py — URL validation, type hints, error logging."""
import asyncio
import base64
import json
import logging
@@ -12,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from tools.vision_tools import (
_validate_image_url,
_handle_vision_analyze,
_determine_mime_type,
_image_to_base64_data_url,
@@ -40,6 +42,57 @@ VALID_PNG = (
_RESOLVES = [(2, 1, 6, "", ("93.184.216.34", 0))]
# ---------------------------------------------------------------------------
# _validate_image_url — urlparse-based validation
# ---------------------------------------------------------------------------
class TestValidateImageUrl:
"""Tests for URL validation, including urlparse-based netloc check."""
def test_accepts_valid_http_and_https_urls(self):
with patch("tools.url_safety.socket.getaddrinfo", return_value=_RESOLVES):
assert _validate_image_url("https://example.com/image.jpg") is True
assert _validate_image_url("http://cdn.example.org/photo.png") is True
# CDN endpoints that redirect to images should still pass.
assert _validate_image_url("https://cdn.example.com/abcdef123") is True
assert _validate_image_url("https://img.example.com/pic?w=200&h=200") is True
assert _validate_image_url("http://example.com:8080/image.png") is True
assert _validate_image_url("https://example.com/") is True
def test_localhost_url_blocked_by_ssrf(self):
"""localhost URLs are blocked by SSRF protection."""
assert _validate_image_url("http://localhost:8080/image.png") is False
def test_rejects_malformed_and_non_string_inputs(self):
# http:// alone has no network location — urlparse catches this.
assert _validate_image_url("http://") is False
assert _validate_image_url("https://") is False
assert _validate_image_url("http:") is False
assert _validate_image_url("") is False
assert _validate_image_url(" ") is False
assert _validate_image_url(None) is False
assert _validate_image_url(12345) is False
assert _validate_image_url(True) is False
assert _validate_image_url(["https://example.com"]) is False
def test_async_validator_shares_shape_and_ssrf_gate(self):
"""The live download path uses the async validator; it must block
localhost / malformed input exactly like the sync one."""
from tools.vision_tools import _validate_image_url_async
async def _run():
assert await _validate_image_url_async("http://localhost:8080/image.png") is False
assert await _validate_image_url_async("http://127.0.0.1/image.png") is False
assert await _validate_image_url_async("http://") is False
assert await _validate_image_url_async(None) is False
with patch("tools.url_safety.socket.getaddrinfo", return_value=_RESOLVES):
assert await _validate_image_url_async("https://cdn.example.com/abcdef123") is True
asyncio.run(_run())
# ---------------------------------------------------------------------------
# _determine_mime_type
# ---------------------------------------------------------------------------

View File

@@ -119,9 +119,23 @@ async def _run_encode_on_cpu_executor(fn, *args, **kwargs):
return await loop.run_in_executor(_vision_cpu_executor, functools.partial(fn, *args, **kwargs))
def _image_url_shape_ok(url: str) -> bool:
"""HTTP(S) shape check only (scheme, netloc). No DNS. Extension-less CDN URLs pass."""
return bool(isinstance(url, str) and url.startswith(("http://", "https://")) and urlparse(url).netloc)
def _validate_image_url(url: str) -> bool:
"""Validate image URL for sync callers and tests (SSRF via sync DNS check)."""
if not _image_url_shape_ok(url):
return False
# Block private/internal addresses to prevent SSRF
from tools.url_safety import is_safe_url
return is_safe_url(url)
async def _validate_image_url_async(url: str) -> bool:
"""HTTP(S) shape check (extension-less CDN URLs pass) + SSRF guard, DNS off the event loop."""
if not (isinstance(url, str) and url.startswith(("http://", "https://")) and urlparse(url).netloc):
"""Shape check + SSRF guard with DNS off the event loop."""
if not _image_url_shape_ok(url):
return False
from tools.url_safety import async_is_safe_url
return await async_is_safe_url(url)