Prevent file extension and content type mismatch

object_storage_api/services/image.py

@@ -4,6 +4,7 @@
 """
 
 import logging
+import mimetypes
 from typing import Annotated, Optional
 
 from bson import ObjectId
@@ -58,6 +59,12 @@ def create(self, image_metadata: ImagePostMetadataSchema, upload_file: UploadFil
         # Upload the full size image to object storage
         object_key = self._image_store.upload(image_id, image_metadata, upload_file)
 
+        expected_file_type = mimetypes.guess_type(upload_file.filename)[0]
+        if expected_file_type != upload_file.content_type:
+            raise InvalidObjectIdError(
+                f"File extension `{upload_file.filename}` does not match content type `{upload_file.content_type}`"
+            )
+
         try:
             image_in = ImageIn(
                 **image_metadata.model_dump(),

object_storage_api/stores/image.py

@@ -18,7 +18,12 @@ class ImageStore:
     Store for managing images in an S3 object store.
     """
 
-    def upload(self, image_id: str, image_metadata: ImagePostMetadataSchema, upload_file: UploadFile) -> str:
+    def upload(
+        self,
+        image_id: str,
+        image_metadata: ImagePostMetadataSchema,
+        upload_file: UploadFile,
+    ) -> str:
         """
         Uploads a given image to object storage.
 
@@ -46,13 +51,17 @@ def create_presigned_get(self, image: ImageOut) -> str:
         :param image: `ImageOut` model of the image.
         :return: Presigned url to get the image.
         """
-        logger.info("Generating presigned url to get image with object key: %s from the object store", image.object_key)
+        logger.info(
+            "Generating presigned url to get image with object key: %s from the object store",
+            image.object_key,
+        )
         response = s3_client.generate_presigned_url(
             "get_object",
             Params={
                 "Bucket": object_storage_config.bucket_name.get_secret_value(),
                 "Key": image.object_key,
                 "ResponseContentDisposition": f'inline; filename="{image.file_name}"',
+                "ResponseContentType": "application/octet-stream",
             },
             ExpiresIn=object_storage_config.presigned_url_expiry_seconds,
         )

test/unit/services/test_image.py

@@ -62,16 +62,18 @@ class CreateDSL(ImageServiceDSL):
     _created_image: ImageMetadataSchema
     _create_exception: pytest.ExceptionInfo
 
-    def mock_create(self, image_post_metadata_data: dict) -> None:
+    def mock_create(self, image_post_metadata_data: dict, filename: str) -> None:
         """
         Mocks repo & store methods appropriately to test the `create` service method.
 
         :param image_post_metadata_data: Dictionary containing the image metadata data as would be required for an
                                          `ImagePostMetadataSchema`.
+        :filename: Filename of the image.
         """
 
         self._image_post_metadata = ImagePostMetadataSchema(**image_post_metadata_data)
-        self._upload_file = UploadFile(MagicMock(), size=100, filename="test.png", headers=MagicMock())
+        header = {"content-type": "image/png"}
+        self._upload_file = UploadFile(MagicMock(), size=100, filename=filename, headers=header)
 
         self._expected_image_id = ObjectId()
         self.mock_object_id.return_value = self._expected_image_id
@@ -151,14 +153,24 @@ class TestCreate(CreateDSL):
     def test_create(self):
         """Test creating an image."""
 
-        self.mock_create(IMAGE_POST_METADATA_DATA_ALL_VALUES)
+        self.mock_create(IMAGE_POST_METADATA_DATA_ALL_VALUES, "test.png")
         self.call_create()
         self.check_create_success()
 
+    def test_create_with_invalid_file(self):
+        """Test creating an image with an inconsistent file extension and content type."""
+
+        self.mock_create(IMAGE_POST_METADATA_DATA_ALL_VALUES, "test.jpeg")
+        self.call_create_expecting_error(InvalidObjectIdError)
+        self.check_create_failed_with_exception(
+            f"File extension `{self._upload_file.filename}` does not match "
+            f"content type `{self._upload_file.content_type}`"
+        )
+
     def test_create_with_invalid_entity_id(self):
         """Test creating an image with an invalid `entity_id`."""
 
-        self.mock_create({**IMAGE_POST_METADATA_DATA_ALL_VALUES, "entity_id": "invalid-id"})
+        self.mock_create({**IMAGE_POST_METADATA_DATA_ALL_VALUES, "entity_id": "invalid-id"}, "test.png")
         self.call_create_expecting_error(InvalidObjectIdError)
         self.check_create_failed_with_exception("Invalid ObjectId value 'invalid-id'")
 

test/unit/stores/test_image.py

@@ -151,6 +151,7 @@ def check_create_presigned_get_success(self) -> None:
                 "Bucket": object_storage_config.bucket_name.get_secret_value(),
                 "Key": self._image_out.object_key,
                 "ResponseContentDisposition": f'inline; filename="{self._image_out.file_name}"',
+                "ResponseContentType": "application/octet-stream",
             },
             ExpiresIn=object_storage_config.presigned_url_expiry_seconds,
         )