feat: image refactor

tests/api/user/test_base.py

@@ -4,6 +4,7 @@ from contextlib import contextmanager
 from unittest.mock import Mock, patch
 
 import pytest
+from PIL import Image
 from werkzeug.datastructures import FileStorage
 
 
@@ -112,7 +113,19 @@ class TestGetVectorStore:
 
 @pytest.mark.unit
 class TestHandleImageUpload:
-    pass
+    @staticmethod
+    def _image_bytes(
+        width: int = 1, height: int = 1, image_format: str = "PNG"
+    ) -> bytes:
+        buffer = io.BytesIO()
+        Image.new("RGB", (width, height), color="white").save(
+            buffer, format=image_format
+        )
+        return buffer.getvalue()
+
+    @classmethod
+    def _png_bytes(cls) -> bytes:
+        return cls._image_bytes()
 
     def test_returns_existing_url_when_no_file(self, flask_app):
         from application.api.user.base import handle_image_upload
@@ -135,7 +148,7 @@ class TestHandleImageUpload:
 
         with flask_app.test_request_context():
             mock_file = FileStorage(
-                stream=io.BytesIO(b"fake image data"), filename="test_image.png"
+                stream=io.BytesIO(self._png_bytes()), filename="test_image.png"
             )
             mock_request = Mock()
             mock_request.files = {"image": mock_file}
@@ -152,6 +165,23 @@ class TestHandleImageUpload:
             assert "user123" in url
             mock_storage.save_file.assert_called_once()
 
+    def test_uploads_image_with_non_ascii_basename(self, flask_app):
+        from application.api.user.base import handle_image_upload
+
+        with flask_app.test_request_context():
+            mock_file = FileStorage(
+                stream=io.BytesIO(self._png_bytes()), filename="上传.png"
+            )
+            mock_storage = Mock()
+
+            url, error = handle_image_upload(
+                Mock(files={"image": mock_file}), "", "user123", mock_storage
+            )
+
+            assert error is None
+            assert url.endswith("_avatar.png")
+            mock_storage.save_file.assert_called_once()
+
     def test_ignores_empty_filename(self, flask_app):
         from application.api.user.base import handle_image_upload
 
@@ -175,7 +205,9 @@ class TestHandleImageUpload:
         from application.api.user.base import handle_image_upload
 
         with flask_app.app_context():
-            mock_file = FileStorage(stream=io.BytesIO(b"data"), filename="test.png")
+            mock_file = FileStorage(
+                stream=io.BytesIO(self._png_bytes()), filename="test.png"
+            )
             mock_request = Mock()
             mock_request.files = {"image": mock_file}
             mock_storage = Mock()
@@ -189,6 +221,135 @@ class TestHandleImageUpload:
             assert error is not None
             assert error.status_code == 400
 
+    def test_rejects_non_image_content(self, flask_app):
+        from application.api.user.base import handle_image_upload
+
+        with flask_app.app_context():
+            mock_file = FileStorage(
+                stream=io.BytesIO(b"not an image"), filename="test.png"
+            )
+            mock_request = Mock(files={"image": mock_file})
+            mock_storage = Mock()
+
+            url, error = handle_image_upload(
+                mock_request, "old.jpg", "user123", mock_storage
+            )
+
+            assert url is None
+            assert error.status_code == 400
+            mock_storage.save_file.assert_not_called()
+
+    def test_accepts_image_at_encoded_byte_limit(self, flask_app):
+        from application.api.user.base import handle_image_upload
+
+        payload = self._png_bytes()
+        with patch(
+            "application.api.user.base.settings.AGENT_IMAGE_MAX_BYTES",
+            len(payload),
+        ), flask_app.app_context():
+            mock_file = FileStorage(
+                stream=io.BytesIO(payload), filename="at-limit.png"
+            )
+            mock_storage = Mock()
+
+            url, error = handle_image_upload(
+                Mock(files={"image": mock_file}), "", "user123", mock_storage
+            )
+
+            assert error is None
+            assert url.endswith("_at-limit.png")
+            mock_storage.save_file.assert_called_once()
+
+    def test_rejects_image_over_encoded_byte_limit(self, flask_app):
+        from application.api.user.base import handle_image_upload
+
+        payload = self._png_bytes()
+        with patch(
+            "application.api.user.base.settings.AGENT_IMAGE_MAX_BYTES",
+            len(payload) - 1,
+        ), flask_app.app_context():
+            mock_file = FileStorage(
+                stream=io.BytesIO(payload), filename="too-large.png"
+            )
+            mock_storage = Mock()
+
+            url, error = handle_image_upload(
+                Mock(files={"image": mock_file}), "old.png", "user123", mock_storage
+            )
+
+            assert url is None
+            assert error.status_code == 400
+            mock_storage.save_file.assert_not_called()
+
+    @pytest.mark.parametrize(
+        ("width", "height", "expected_status"),
+        [(4, 4, None), (5, 4, 400)],
+    )
+    def test_enforces_decoded_pixel_limit(
+        self, flask_app, width, height, expected_status
+    ):
+        from application.api.user.base import handle_image_upload
+
+        payload = self._image_bytes(width, height)
+        with patch(
+            "application.api.user.base.settings.AGENT_IMAGE_MAX_PIXELS", 16
+        ), flask_app.app_context():
+            mock_file = FileStorage(
+                stream=io.BytesIO(payload), filename="dimensions.png"
+            )
+            mock_storage = Mock()
+
+            url, error = handle_image_upload(
+                Mock(files={"image": mock_file}), "", "user123", mock_storage
+            )
+
+            if expected_status is None:
+                assert error is None
+                assert url
+                mock_storage.save_file.assert_called_once()
+            else:
+                assert url is None
+                assert error.status_code == expected_status
+                mock_storage.save_file.assert_not_called()
+
+    def test_rejects_image_whose_content_does_not_match_extension(
+        self, flask_app
+    ):
+        from application.api.user.base import handle_image_upload
+
+        with flask_app.app_context():
+            mock_file = FileStorage(
+                stream=io.BytesIO(self._png_bytes()), filename="disguised.jpg"
+            )
+            mock_storage = Mock()
+
+            url, error = handle_image_upload(
+                Mock(files={"image": mock_file}), "", "user123", mock_storage
+            )
+
+            assert url is None
+            assert error.status_code == 400
+            mock_storage.save_file.assert_not_called()
+
+    def test_sanitizes_user_directory_component(self, flask_app):
+        from application.api.user.base import handle_image_upload
+
+        with flask_app.app_context():
+            mock_file = FileStorage(
+                stream=io.BytesIO(self._png_bytes()), filename="test.png"
+            )
+            mock_request = Mock(files={"image": mock_file})
+            mock_storage = Mock()
+
+            url, error = handle_image_upload(
+                mock_request, "", "../../secrets", mock_storage
+            )
+
+            assert error is None
+            assert ".." not in url
+            assert url.startswith("inputs/secrets-")
+            assert "/attachments/" in url
+
 
 @pytest.mark.unit
 class TestRequireAgentDecorator:
@@ -321,4 +482,3 @@ class TestRequireAgentPgConn:
         with _patch_base_db(pg_conn), flask_app.app_context():
             result = fn(webhook_token="webhook-123")
         assert result["got"] == str(agent["id"])
-