From 1a3c1b8b391784ea767a55e8198a6b99834a4c17 Mon Sep 17 00:00:00 2001 From: InfinityPacer <160988576+InfinityPacer@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:56:46 +0800 Subject: [PATCH] fix(system): derive image proxy MIME from content (#6202) --- app/api/endpoints/system.py | 10 +- app/helper/image.py | 91 ++++++++++--- tests/test_system_image_proxy.py | 220 +++++++++++++++++++++++++++++++ tests/test_system_nettest.py | 12 +- 4 files changed, 310 insertions(+), 23 deletions(-) create mode 100644 tests/test_system_image_proxy.py diff --git a/app/api/endpoints/system.py b/app/api/endpoints/system.py index 1728df66..9536306d 100644 --- a/app/api/endpoints/system.py +++ b/app/api/endpoints/system.py @@ -582,23 +582,27 @@ async def fetch_image( ): return None - content = await ImageHelper().async_fetch_image( + image_result = await ImageHelper().async_fetch_image_with_mime_type( url=fetch_url, proxy=proxy, use_cache=use_cache, cookies=cookies, ) - if content: + if image_result: + content, media_type = image_result + # 检查 If-None-Match etag = HashUtils.md5(content) headers = RequestUtils.generate_cache_headers(etag, max_age=86400 * 7) + headers["Content-Type"] = media_type + headers["X-Content-Type-Options"] = "nosniff" if if_none_match == etag: return Response(status_code=304, headers=headers) # 返回缓存图片 return Response( content=content, - media_type=UrlUtils.get_mime_type(fetch_url, "image/jpeg"), + media_type=media_type, headers=headers, ) return None diff --git a/app/helper/image.py b/app/helper/image.py index f76b3088..52d09758 100644 --- a/app/helper/image.py +++ b/app/helper/image.py @@ -188,16 +188,31 @@ class ImageHelper(metaclass=Singleton): return cache_path.as_posix() @staticmethod - def _validate_image(content: bytes) -> bool: - """验证图片""" + def get_image_mime_type(content: bytes, verify: bool = True) -> Optional[str]: + """ + 根据图片内容返回 Pillow 识别的图片 MIME 类型。 + + 外部响应在写入缓存前需要完整校验；已校验的缓存只需读取格式头。 + 非图片或可脚本化的 MIME 类型不作为图片代理响应。 + """ if not content: - return False + return None try: - Image.open(io.BytesIO(content)).verify() - return True - except Exception as e: - logger.warn(f"Invalid image format: {e}") - return False + with Image.open(io.BytesIO(content)) as image: + image_format = (image.format or "").upper() + if verify: + image.verify() + mime_type = Image.MIME.get(image_format) + if ( + not mime_type + or not mime_type.startswith("image/") + or mime_type == "image/svg+xml" + ): + return None + return mime_type + except Exception as err: + logger.warning(f"Invalid image format: {err}") + return None @staticmethod def _get_request_params(url: str, proxy: Optional[bool], cookies: Optional[str | dict]) -> dict: @@ -224,6 +239,26 @@ class ImageHelper(metaclass=Singleton): """ 获取图片（同步版本） """ + result = self.fetch_image_with_mime_type( + url=url, + proxy=proxy, + use_cache=use_cache, + cookies=cookies, + ) + return result[0] if result else None + + def fetch_image_with_mime_type( + self, + url: str, + proxy: Optional[bool] = None, + use_cache: bool = True, + cookies: Optional[str | dict] = None, + ) -> Optional[tuple[bytes, str]]: + """ + 同步获取图片及其内容识别 MIME 类型。 + + 网络响应在写入缓存前完整验证一次；缓存命中仅重新识别格式头。 + """ if not url: return None @@ -233,7 +268,9 @@ class ImageHelper(metaclass=Singleton): if use_cache: content = self.file_cache.get(cache_path, region="images") if content: - return content + mime_type = self.get_image_mime_type(content, verify=False) + if mime_type: + return content, mime_type # 请求远程图片 params = self._get_request_params(url, proxy, cookies) @@ -243,13 +280,13 @@ class ImageHelper(metaclass=Singleton): return None content = response.content - # 验证图片 - if not self._validate_image(content): + mime_type = self.get_image_mime_type(content) + if not mime_type: return None # 保存缓存 self.file_cache.set(cache_path, content, region="images") - return content + return content, mime_type async def async_fetch_image( self, @@ -260,6 +297,26 @@ class ImageHelper(metaclass=Singleton): """ 获取图片（异步版本） """ + result = await self.async_fetch_image_with_mime_type( + url=url, + proxy=proxy, + use_cache=use_cache, + cookies=cookies, + ) + return result[0] if result else None + + async def async_fetch_image_with_mime_type( + self, + url: str, + proxy: Optional[bool] = None, + use_cache: bool = True, + cookies: Optional[str | dict] = None, + ) -> Optional[tuple[bytes, str]]: + """ + 异步获取图片及其内容识别 MIME 类型。 + + 网络响应在写入缓存前完整验证一次；缓存命中仅重新识别格式头。 + """ if not url: return None @@ -269,7 +326,9 @@ class ImageHelper(metaclass=Singleton): if use_cache: content = await self.async_file_cache.get(cache_path, region="images") if content: - return content + mime_type = self.get_image_mime_type(content, verify=False) + if mime_type: + return content, mime_type # 请求远程图片 params = self._get_request_params(url, proxy, cookies) @@ -279,10 +338,10 @@ class ImageHelper(metaclass=Singleton): return None content = response.content - # 验证图片 - if not self._validate_image(content): + mime_type = self.get_image_mime_type(content) + if not mime_type: return None # 保存缓存 await self.async_file_cache.set(cache_path, content, region="images") - return content + return content, mime_type diff --git a/tests/test_system_image_proxy.py b/tests/test_system_image_proxy.py new file mode 100644 index 00000000..0eb80320 --- /dev/null +++ b/tests/test_system_image_proxy.py @@ -0,0 +1,220 @@ +import asyncio +import io +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from PIL import Image + +from app.api.endpoints import system as system_endpoint +from app.helper.image import ImageHelper + + +def _image_bytes(image_format: str, trailing: bytes = b"") -> bytes: + buffer = io.BytesIO() + Image.new("RGB", (2, 2), color=(32, 96, 160)).save(buffer, format=image_format) + return buffer.getvalue() + trailing + + +@pytest.mark.parametrize( + ("image_format", "expected_mime"), + [ + ("PNG", "image/png"), + ("JPEG", "image/jpeg"), + ("GIF", "image/gif"), + ("WEBP", "image/webp"), + ("PCX", "image/x-pcx"), + ("PPM", "image/x-portable-anymap"), + ], +) +def test_get_image_mime_type_uses_pillow_detected_format( + image_format: str, + expected_mime: str, +): + assert ImageHelper.get_image_mime_type(_image_bytes(image_format)) == expected_mime + + +def test_get_image_mime_type_rejects_non_image_pillow_mime(): + assert ImageHelper.get_image_mime_type(_image_bytes("EPS")) is None + + +def test_get_image_mime_type_rejects_scriptable_svg_mime(): + with patch.dict(Image.MIME, {"PNG": "image/svg+xml"}): + assert ImageHelper.get_image_mime_type(_image_bytes("PNG")) is None + + +def test_fetch_image_with_mime_type_only_reads_cached_format_header(): + content = _image_bytes("PNG") + image_helper = ImageHelper() + + with patch.object( + image_helper.file_cache, + "get", + return_value=content, + ), patch.object( + image_helper, + "get_image_mime_type", + return_value="image/png", + ) as get_mime_type: + result = image_helper.fetch_image_with_mime_type( + "https://images.example/wallpaper.png" + ) + + assert result == (content, "image/png") + get_mime_type.assert_called_once_with(content, verify=False) + + +def test_fetch_image_with_mime_type_validates_network_content_once(): + content = _image_bytes("PNG") + image_helper = ImageHelper() + response = Mock(status_code=200, content=content) + request = Mock() + request.get_res.return_value = response + + with patch.object( + image_helper.file_cache, + "get", + return_value=None, + ), patch.object( + image_helper.file_cache, + "set", + ), patch( + "app.helper.image.RequestUtils", + return_value=request, + ), patch.object( + image_helper, + "get_image_mime_type", + return_value="image/png", + ) as get_mime_type: + result = image_helper.fetch_image_with_mime_type( + "https://images.example/wallpaper.png" + ) + + assert result == (content, "image/png") + get_mime_type.assert_called_once_with(content) + + +def test_async_fetch_image_with_mime_type_only_reads_cached_format_header(): + content = _image_bytes("PNG") + image_helper = ImageHelper() + + with patch.object( + image_helper.async_file_cache, + "get", + new=AsyncMock(return_value=content), + ), patch.object( + image_helper, + "get_image_mime_type", + return_value="image/png", + ) as get_mime_type: + result = asyncio.run( + image_helper.async_fetch_image_with_mime_type( + "https://images.example/wallpaper.png" + ) + ) + + assert result == (content, "image/png") + get_mime_type.assert_called_once_with(content, verify=False) + + +def test_async_fetch_image_with_mime_type_validates_network_content_once(): + content = _image_bytes("PNG") + image_helper = ImageHelper() + response = Mock(status_code=200, content=content) + request = Mock() + request.get_res = AsyncMock(return_value=response) + + with patch.object( + image_helper.async_file_cache, + "get", + new=AsyncMock(return_value=None), + ), patch.object( + image_helper.async_file_cache, + "set", + new=AsyncMock(), + ), patch( + "app.helper.image.AsyncRequestUtils", + return_value=request, + ), patch.object( + image_helper, + "get_image_mime_type", + return_value="image/png", + ) as get_mime_type: + result = asyncio.run( + image_helper.async_fetch_image_with_mime_type( + "https://images.example/wallpaper.png" + ) + ) + + assert result == (content, "image/png") + get_mime_type.assert_called_once_with(content) + + +def test_fetch_image_does_not_trust_active_url_suffix(): + content = _image_bytes("PNG", b"") + image_helper = Mock() + image_helper.async_fetch_image_with_mime_type = AsyncMock( + return_value=(content, "image/png") + ) + + with patch.object( + system_endpoint.SecurityUtils, + "is_safe_image_url_async", + new=AsyncMock(return_value=True), + ), patch.object(system_endpoint, "ImageHelper", return_value=image_helper): + response = asyncio.run( + system_endpoint.fetch_image( + url="https://images.example/wallpaper.html", + allowed_domains={"images.example"}, + ) + ) + + assert response is not None + assert response.headers["content-type"] == "image/png" + assert response.headers["x-content-type-options"] == "nosniff" + assert response.body == content + + +def test_fetch_image_rejects_unverified_content(): + image_helper = Mock() + image_helper.async_fetch_image_with_mime_type = AsyncMock(return_value=None) + + with patch.object( + system_endpoint.SecurityUtils, + "is_safe_image_url_async", + new=AsyncMock(return_value=True), + ), patch.object(system_endpoint, "ImageHelper", return_value=image_helper): + response = asyncio.run( + system_endpoint.fetch_image( + url="https://images.example/wallpaper.png", + allowed_domains={"images.example"}, + ) + ) + + assert response is None + + +def test_fetch_image_adds_nosniff_to_not_modified_response(): + content = _image_bytes("JPEG") + image_helper = Mock() + image_helper.async_fetch_image_with_mime_type = AsyncMock( + return_value=(content, "image/jpeg") + ) + etag = system_endpoint.HashUtils.md5(content) + + with patch.object( + system_endpoint.SecurityUtils, + "is_safe_image_url_async", + new=AsyncMock(return_value=True), + ), patch.object(system_endpoint, "ImageHelper", return_value=image_helper): + response = asyncio.run( + system_endpoint.fetch_image( + url="https://images.example/wallpaper.jpg", + if_none_match=etag, + allowed_domains={"images.example"}, + ) + ) + + assert response is not None + assert response.status_code == 304 + assert response.headers["content-type"] == "image/jpeg" + assert response.headers["x-content-type-options"] == "nosniff" diff --git a/tests/test_system_nettest.py b/tests/test_system_nettest.py index b9fbd890..4244a33a 100644 --- a/tests/test_system_nettest.py +++ b/tests/test_system_nettest.py @@ -88,7 +88,9 @@ class NettestSecurityTest(unittest.TestCase): image_url = "http://192.168.1.50:8096/System/Info/Public" signed_url = system_endpoint.SecurityUtils.sign_url(image_url) image_helper = Mock() - image_helper.async_fetch_image = AsyncMock(return_value=b"image-bytes") + image_helper.async_fetch_image_with_mime_type = AsyncMock( + return_value=(b"image-bytes", "image/jpeg") + ) with patch.object(system_endpoint, "ImageHelper", return_value=image_helper), patch.object( system_endpoint.HashUtils, "md5", return_value="etag", create=True @@ -103,7 +105,7 @@ class NettestSecurityTest(unittest.TestCase): ) self.assertEqual(resp.status_code, 200) - image_helper.async_fetch_image.assert_awaited_once_with( + image_helper.async_fetch_image_with_mime_type.assert_awaited_once_with( url=image_url, proxy=None, use_cache=False, @@ -133,7 +135,9 @@ class NettestSecurityTest(unittest.TestCase): 图片代理在域名白名单命中后，可按配置放行指定非公网解析网段。 """ image_helper = Mock() - image_helper.async_fetch_image = AsyncMock(return_value=b"image-bytes") + image_helper.async_fetch_image_with_mime_type = AsyncMock( + return_value=(b"image-bytes", "image/jpeg") + ) with patch.object(system_endpoint, "ImageHelper", return_value=image_helper), patch.object( system_endpoint.HashUtils, "md5", return_value="etag", create=True @@ -161,7 +165,7 @@ class NettestSecurityTest(unittest.TestCase): ) self.assertEqual(resp.status_code, 200) - image_helper.async_fetch_image.assert_awaited_once_with( + image_helper.async_fetch_image_with_mime_type.assert_awaited_once_with( url="https://img1.doubanio.com/poster.webp", proxy=None, use_cache=False,