diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py index d251734..356f884 100644 --- a/blueprints/conversation/__init__.py +++ b/blueprints/conversation/__init__.py @@ -149,9 +149,24 @@ def _persist_in_background(coro) -> asyncio.Task: 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//. + """ + 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, @@ -253,6 +268,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" @@ -358,6 +375,8 @@ async def upload_attachment(): @conversation_blueprint.get("/image/") @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}) diff --git a/tests/unit/test_file_upload.py b/tests/unit/test_file_upload.py index c85613a..4677dc6 100644 --- a/tests/unit/test_file_upload.py +++ b/tests/unit/test_file_upload.py @@ -296,3 +296,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" + ) + }, + )