fix: validate PNGs at shared image resolver

This commit is contained in:
CannibalKush
2026-08-13 08:32:11 +02:00
committed by Teknium
parent ad0e83062f
commit 765142d94f
3 changed files with 45 additions and 1 deletions

View File

@@ -15,8 +15,14 @@ from unittest.mock import patch
import pytest
PNG = b"\x89PNG\r\n\x1a\n" + b"\x00" * 64
# Minimal valid 1x1 PNG bytes. Resolver validation requires a decodable fixture.
PNG = base64.b64decode(
b"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
)
JPEG = b"\xff\xd8\xff" + b"\x00" * 64
CORRUPT_PNG = base64.b64decode(
b"iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAIAAAACUFjqAAAAFElEQVR4nGP8z8Dwn4EIwESJ5gAAVQ4CH1evYJQAAAAASUVORK5CYII="
)
def _reload(monkeypatch, hermes_home: Path):
@@ -57,6 +63,16 @@ class TestDataUrl:
await isrc.resolve_image_source(
f"data:text/plain;base64,{b64}", isrc.ResolveContext())
@pytest.mark.asyncio
async def test_corrupt_png_rejected_at_resolver_boundary(self, tmp_path, monkeypatch):
"""A PNG-shaped but undecodable payload never becomes a resolved image."""
isrc = _reload(monkeypatch, tmp_path / "hermes")
monkeypatch.setenv("TERMINAL_ENV", "local")
img = tmp_path / "corrupt.png"
img.write_bytes(CORRUPT_PNG)
with pytest.raises(isrc.NotAnImage):
await isrc.resolve_image_source(str(img), isrc.ResolveContext())
class TestLocalBackend:
@pytest.mark.asyncio

View File

@@ -28,6 +28,11 @@ _TINY_PNG = base64.b64decode(
b"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
)
# PNG-shaped but undecodable; resolver/native fast path must reject it.
_CORRUPT_PNG = base64.b64decode(
b"iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAIAAAACUFjqAAAAFElEQVR4nGP8z8Dwn4EIwESJ5gAAVQ4CH1evYJQAAAAASUVORK5CYII="
)
def _animated_gif_bytes(colors, *, size=(4, 4)):
from PIL import Image
@@ -261,6 +266,19 @@ class TestVisionAnalyzeNative:
assert isinstance(result, dict)
assert result.get("_multimodal") is True
def test_corrupt_png_rejected_before_native_embed(self, tmp_path):
"""Header-only PNG bytes must not enter conversation history."""
img = tmp_path / "bad.png"
img.write_bytes(_CORRUPT_PNG)
result = asyncio.get_event_loop().run_until_complete(
_vision_analyze_native(str(img), "what is this?")
)
assert isinstance(result, str)
parsed = json.loads(result)
assert parsed.get("success") is False
assert "multimodal" not in parsed
assert "recognized image" in parsed.get("error", "")
def test_oversized_image_resized_under_embed_cap(self, tmp_path):
"""Regression for the wedged-session incident (May 2026).

View File

@@ -33,6 +33,7 @@ import contextlib
import asyncio
import json
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO
import logging
import os
import uuid
@@ -251,6 +252,15 @@ def _detect_image_mime_type_from_bytes(data: bytes) -> Optional[str]:
"""
header = data[:64]
if header.startswith(b"\x89PNG\r\n\x1a\n"):
# Magic bytes alone are insufficient: native vision history is
# immutable, so reject corrupt PNGs before they can be embedded.
try:
from PIL import Image
with Image.open(BytesIO(data)) as image:
image.verify()
except Exception:
return None
return "image/png"
if header.startswith(b"\xff\xd8\xff"):
return "image/jpeg"