Added mimetype to validate extension against content-type #81

ral-facilities · Jan 9, 2025 · 9e3d72c · 9e3d72c
1 parent af5e70d
commit 9e3d72c
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/object_storage_api/services/image.py b/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(),

diff --git a/test/unit/services/test_image.py b/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,23 @@ 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 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'")