Merge pull request 'Only presign and accept images in the caller's own conversations' (#81) from fix/image-ownership into main

blueprints/conversation/__init__.py

@@ -138,9 +138,24 @@ conversation_blueprint = Blueprint(
 
 async def _get_owned_conversation(conversation_id, user_id) -> Conversation | None:
     """Fetch a conversation only if it belongs to the given user."""
+    try:
+        conversation_id = uuid.UUID(str(conversation_id))
+    except ValueError:
+        return None
     return await Conversation.get_or_none(id=conversation_id, user_id=user_id)
 
 
+async def _owns_storage_key(key: str, user_id) -> bool:
+    """True if `key` is stored under a conversation the user owns.
+
+    Every uploaded object lives under conversations/<conversation id>/.
+    """
+    parts = key.split("/")
+    if len(parts) < 3 or parts[0] != "conversations" or ".." in parts:
+        return False
+    return await _get_owned_conversation(parts[1], user_id) is not None
+
+
 def _build_user_message(
     query_text: str,
     image_block: dict | None = None,
@@ -221,6 +236,8 @@ async def upload_image():
 
     if not file or not conversation_id:
         return jsonify({"error": "file and conversation_id are required"}), 400
+    if await _get_owned_conversation(conversation_id, current_user_uuid) is None:
+        return jsonify({"error": "conversation not found"}), 404
 
     file_bytes = file.read()
     content_type = file.content_type or "image/jpeg"
@@ -326,6 +343,8 @@ async def upload_attachment():
 @conversation_blueprint.get("/image/<path:image_key>")
 @jwt_access_required
 async def serve_image(image_key: str):
+    if not await _owns_storage_key(image_key, get_jwt_identity()):
+        return jsonify({"error": "image not found"}), 404
     url = await s3_presigned_url(image_key)
     return jsonify({"url": url})
 

tests/unit/test_file_upload.py

@@ -301,3 +301,45 @@ class TestStreamQueryFileKey:
                 f"conversations/{CONV_ID}/files/u.txt"
             )
         assert name == "u.txt"
+
+
+class TestImageOwnership:
+    async def test_presigns_only_keys_in_owned_conversations(self, app, headers, s3):
+        presign = AsyncMock(return_value="https://s3.example/signed")
+        with patch.object(conv, "s3_presigned_url", presign):
+            client = app.test_client()
+            own = await client.get(
+                f"/api/conversation/image/conversations/{CONV_ID}/a.jpg",
+                headers=headers,
+            )
+            assert own.status_code == 200
+            for key in (
+                f"conversations/{OTHER_ID}/a.jpg",
+                f"conversations/{CONV_ID}/../{OTHER_ID}/a.jpg",
+                "conversations/not-a-uuid/a.jpg",
+                "beli/profile.jpg",
+            ):
+                resp = await client.get(
+                    f"/api/conversation/image/{key}", headers=headers
+                )
+                assert resp.status_code == 404, key
+        presign.assert_awaited_once_with(f"conversations/{CONV_ID}/a.jpg")
+
+    async def test_image_upload_requires_owned_conversation(self, app, headers, s3):
+        with patch.object(conv.blueprints.users.models.User, "get", AsyncMock()):
+            resp = await self._upload_image(app, headers)
+        assert resp.status_code == 404
+        s3.upload_image.assert_not_awaited()
+
+    @staticmethod
+    async def _upload_image(app, headers):
+        return await app.test_client().post(
+            "/api/conversation/upload-image",
+            headers=headers,
+            form={"conversation_id": OTHER_ID},
+            files={
+                "file": FileStorage(
+                    io.BytesIO(_png_bytes()), filename="a.png", content_type="image/png"
+                )
+            },
+        )