fix(vision): reject non-image files and enforce website policy (salvage #1940) (#3845)

Three safety gaps in vision_analyze_tool:

1. Local files accepted without checking if they're actually images —
   a renamed text file would get base64-encoded and sent to the model.
   Now validates magic bytes (PNG, JPEG, GIF, BMP, WebP, SVG).

2. No website policy enforcement on image URLs — blocked domains could
   be fetched via the vision tool. Now checks before download.

3. No redirect check — if an allowed URL redirected to a blocked domain,
   the download would proceed. Now re-checks the final URL.

Fixed one test that needed _validate_image_url mocked to bypass DNS
resolution on the fake blocked.test domain (is_safe_url does DNS
checks that were added after the original PR).

Co-authored-by: GutSlabs <GutSlabs@users.noreply.github.com>
This commit is contained in:
Teknium
2026-03-29 20:55:04 -07:00
committed by GitHub
parent b60cfd6ce6
commit 5e67fc8c40
2 changed files with 113 additions and 1 deletions

View File

@@ -39,6 +39,7 @@ from urllib.parse import urlparse
import httpx
from agent.auxiliary_client import async_call_llm, extract_content_or_reasoning
from tools.debug_helpers import DebugSession
from tools.website_policy import check_website_access
logger = logging.getLogger(__name__)
@@ -76,6 +77,28 @@ def _validate_image_url(url: str) -> bool:
return True
def _detect_image_mime_type(image_path: Path) -> Optional[str]:
"""Return a MIME type when the file looks like a supported image."""
with image_path.open("rb") as f:
header = f.read(64)
if header.startswith(b"\x89PNG\r\n\x1a\n"):
return "image/png"
if header.startswith(b"\xff\xd8\xff"):
return "image/jpeg"
if header.startswith((b"GIF87a", b"GIF89a")):
return "image/gif"
if header.startswith(b"BM"):
return "image/bmp"
if len(header) >= 12 and header[:4] == b"RIFF" and header[8:12] == b"WEBP":
return "image/webp"
if image_path.suffix.lower() == ".svg":
head = image_path.read_text(encoding="utf-8", errors="ignore")[:4096].lower()
if "<svg" in head:
return "image/svg+xml"
return None
async def _download_image(image_url: str, destination: Path, max_retries: int = 3) -> Path:
"""
Download an image from a URL to a local destination (async) with retry logic.
@@ -115,6 +138,10 @@ async def _download_image(image_url: str, destination: Path, max_retries: int =
last_error = None
for attempt in range(max_retries):
try:
blocked = check_website_access(image_url)
if blocked:
raise PermissionError(blocked["message"])
# Download the image with appropriate headers using async httpx
# Enable follow_redirects to handle image CDNs that redirect (e.g., Imgur, Picsum)
# SSRF: event_hooks validates each redirect target against private IP ranges
@@ -131,6 +158,11 @@ async def _download_image(image_url: str, destination: Path, max_retries: int =
},
)
response.raise_for_status()
final_url = str(response.url)
blocked = check_website_access(final_url)
if blocked:
raise PermissionError(blocked["message"])
# Save the image content
destination.write_bytes(response.content)
@@ -257,6 +289,7 @@ async def vision_analyze_tool(
# Track whether we should clean up the file after processing.
# Local files (e.g. from the image cache) should NOT be deleted.
should_cleanup = True
detected_mime_type = None
try:
from tools.interrupt import is_interrupted
@@ -275,6 +308,9 @@ async def vision_analyze_tool(
should_cleanup = False # Don't delete cached/local files
elif _validate_image_url(image_url):
# Remote URL -- download to a temporary location
blocked = check_website_access(image_url)
if blocked:
raise PermissionError(blocked["message"])
logger.info("Downloading image from URL...")
temp_dir = Path("./temp_vision_images")
temp_image_path = temp_dir / f"temp_image_{uuid.uuid4()}.jpg"
@@ -289,10 +325,14 @@ async def vision_analyze_tool(
image_size_bytes = temp_image_path.stat().st_size
image_size_kb = image_size_bytes / 1024
logger.info("Image ready (%.1f KB)", image_size_kb)
detected_mime_type = _detect_image_mime_type(temp_image_path)
if not detected_mime_type:
raise ValueError("Only real image files are supported for vision analysis.")
# Convert image to base64 data URL
logger.info("Converting image to base64...")
image_data_url = _image_to_base64_data_url(temp_image_path)
image_data_url = _image_to_base64_data_url(temp_image_path, mime_type=detected_mime_type)
# Calculate size in KB for better readability
data_size_kb = len(image_data_url) / 1024
logger.info("Image converted to base64 (%.1f KB)", data_size_kb)