Refine seller store implementation

server/tests/test_storage_security.py

@@ -0,0 +1,105 @@
+"""远程图片下载的 SSRF 防护回归测试。"""
+from __future__ import annotations
+
+import sys
+import unittest
+from unittest.mock import patch
+from pathlib import Path
+from io import BytesIO
+
+import httpx
+from PIL import Image
+
+SERVER_DIR = Path(__file__).resolve().parents[1]
+if str(SERVER_DIR) not in sys.path:
+    sys.path.insert(0, str(SERVER_DIR))
+
+from services.storage import (
+    ResourceDownloadTimeoutError,
+    ImageValidationError,
+    UnsafeRemoteUrlError,
+    download_bytes,
+    normalize_image_bytes,
+    validate_remote_url,
+)
+
+
+class RemoteUrlValidationTests(unittest.IsolatedAsyncioTestCase):
+    async def test_allows_public_literal_ip(self) -> None:
+        await validate_remote_url("https://1.1.1.1/image.jpg")
+
+    async def test_rejects_local_and_private_targets(self) -> None:
+        for url in (
+            "http://localhost:8800/image.jpg",
+            "http://127.0.0.1/image.jpg",
+            "http://10.0.0.1/image.jpg",
+            "http://169.254.169.254/latest/meta-data/",
+            "http://[::1]/image.jpg",
+        ):
+            with self.subTest(url=url), self.assertRaises(UnsafeRemoteUrlError):
+                await validate_remote_url(url)
+
+    async def test_download_rejects_before_opening_a_local_connection(self) -> None:
+        with self.assertRaises(UnsafeRemoteUrlError):
+            await download_bytes("http://127.0.0.1:1/image.jpg")
+
+    async def test_redirect_target_is_validated_again(self) -> None:
+        async def handler(request: httpx.Request) -> httpx.Response:
+            self.assertEqual(request.url.host, "1.1.1.1")
+            return httpx.Response(
+                302,
+                headers={"location": "http://127.0.0.1:9000/private.jpg"},
+                request=request,
+            )
+
+        real_async_client = httpx.AsyncClient
+        with patch(
+            "services.storage.httpx.AsyncClient",
+            side_effect=lambda **kwargs: real_async_client(
+                transport=httpx.MockTransport(handler), **kwargs
+            ),
+        ), self.assertRaises(UnsafeRemoteUrlError):
+            await download_bytes("https://1.1.1.1/redirect.jpg")
+
+    async def test_http_client_timeout_is_normalized(self) -> None:
+        async def handler(request: httpx.Request) -> httpx.Response:
+            raise httpx.ReadTimeout("slow response", request=request)
+
+        real_async_client = httpx.AsyncClient
+        with patch(
+            "services.storage.httpx.AsyncClient",
+            side_effect=lambda **kwargs: real_async_client(
+                transport=httpx.MockTransport(handler), **kwargs
+            ),
+        ), self.assertRaises(ResourceDownloadTimeoutError):
+            await download_bytes("https://1.1.1.1/slow.jpg")
+
+
+class ImageNormalizationTests(unittest.TestCase):
+    def test_uses_real_format_and_corrects_exif_orientation(self) -> None:
+        source = Image.new("RGB", (30, 10), "red")
+        exif = Image.Exif()
+        exif[274] = 6  # 顺时针旋转 90°
+        raw = BytesIO()
+        source.save(raw, format="JPEG", exif=exif)
+
+        image = normalize_image_bytes(raw.getvalue(), max_pixels=1_000)
+
+        self.assertEqual(image.content_type, "image/jpeg")
+        self.assertEqual((image.width, image.height), (10, 30))
+        self.assertTrue(image.orientation_normalized)
+        with Image.open(BytesIO(image.data)) as stored:
+            self.assertNotEqual(stored.getexif().get(274), 6)
+
+    def test_rejects_invalid_content_and_pixel_bomb(self) -> None:
+        with self.assertRaises(ImageValidationError):
+            normalize_image_bytes(b"not an image")
+
+        raw = BytesIO()
+        Image.new("RGB", (11, 10), "blue").save(raw, format="PNG")
+        with self.assertRaises(ImageValidationError):
+            normalize_image_bytes(raw.getvalue(), max_pixels=100)
+
+
+if __name__ == "__main__":
+    unittest.main()