Add safe check for Content-Type in photon (#2143)

api/api/utils/photon.py

@@ -37,6 +37,16 @@ def _get_file_extension_from_url(image_url: str) -> str:
     return ext[1:]  # remove the leading dot
 
 
+def _get_file_extension_from_content_type(content_type: str) -> str | None:
+    """
+    Return the image extension if present in the Response's content type
+    header.
+    """
+    if content_type and "/" in content_type:
+        return content_type.split("/")[1]
+    return None
+
+
 def check_image_type(image_url: str, media_obj) -> None:
     cache = django_redis.get_redis_connection("default")
     key = f"media:{media_obj.identifier}:thumb_type"
@@ -59,13 +69,14 @@ def check_image_type(image_url: str, media_obj) -> None:
                 f"type. {exc}"
             )
         else:
-            ext = response.headers["Content-Type"]
-            if ext and "/" in ext:
-                ext = ext.split("/")[1]
-                cache.set(key, ext)
+            if response.headers and "Content-Type" in response.headers:
+                content_type = response.headers["Content-Type"]
+                ext = _get_file_extension_from_content_type(content_type)
             else:
                 ext = None
 
+            cache.set(key, ext if ext else "unknown")
+
     if ext not in ALLOWED_TYPES:
         raise UnsupportedMediaType(ext)
 

api/test/unit/utils/test_photon.py

@@ -353,13 +353,20 @@ def test_check_image_type_raises_by_not_allowed_types(image_type):
 
 
 @pytest.mark.django_db
-def test_check_image_type_saves_image_type_to_cache(redis):
+@pytest.mark.parametrize(
+    "headers, expected_cache_val",
+    [
+        ({"Content-Type": "image/svg+xml"}, b"svg+xml"),
+        (None, b"unknown"),
+    ],
+)
+def test_check_image_type_saves_image_type_to_cache(redis, headers, expected_cache_val):
     image_url = TEST_IMAGE_URL.replace(".jpg", "")
     image = ImageFactory.create(url=image_url)
 
     with mock.patch("requests.head") as mock_head, pytest.raises(UnsupportedMediaType):
-        mock_head.return_value.headers = {"Content-Type": "image/svg+xml"}
+        mock_head.return_value.headers = headers
         check_image_type(image_url, image)
 
     key = f"media:{image.identifier}:thumb_type"
-    assert redis.get(key) == b"svg+xml"
+    assert redis.get(key) == expected_cache_val