Add validation of content_type to image entity (#95248)

homeassistant/components/image/__init__.py

@@ -58,12 +58,23 @@ class Image:
     content: bytes
 
 
+class ImageContentTypeError(HomeAssistantError):
+    """Error with the content type while loading an image."""
+
+
+def valid_image_content_type(content_type: str) -> str:
+    """Validate the assigned content type is one of an image."""
+    if content_type.split("/", 1)[0] != "image":
+        raise ImageContentTypeError
+    return content_type
+
+
 async def _async_get_image(image_entity: ImageEntity, timeout: int) -> Image:
     """Fetch image from an image entity."""
-    with suppress(asyncio.CancelledError, asyncio.TimeoutError):
+    with suppress(asyncio.CancelledError, asyncio.TimeoutError, ImageContentTypeError):
         async with async_timeout.timeout(timeout):
             if image_bytes := await image_entity.async_image():
-                content_type = image_entity.content_type
+                content_type = valid_image_content_type(image_entity.content_type)
                 image = Image(content_type, image_bytes)
                 return image
 

tests/components/image/conftest.py

@@ -36,6 +36,21 @@ class MockImageEntity(image.ImageEntity):
         return b"Test"
 
 
+class MockImageEntityInvalidContentType(image.ImageEntity):
+    """Mock image entity."""
+
+    _attr_name = "Test"
+
+    async def async_added_to_hass(self):
+        """Set the update time and assign and incorrect content type."""
+        self._attr_content_type = "text/json"
+        self._attr_image_last_updated = dt_util.utcnow()
+
+    async def async_image(self) -> bytes | None:
+        """Return bytes of image."""
+        return b"Test"
+
+
 class MockURLImageEntity(image.ImageEntity):
     """Mock image entity."""
 
@@ -125,7 +140,9 @@ def config_flow_fixture(hass: HomeAssistant) -> Generator[None, None, None]:
 
 
 @pytest.fixture(name="mock_image_config_entry")
-async def mock_image_config_entry_fixture(hass: HomeAssistant, config_flow: None):
+async def mock_image_config_entry_fixture(
+    hass: HomeAssistant, config_flow: None
+) -> ConfigEntry:
     """Initialize a mock image config_entry."""
 
     async def async_setup_entry_init(
@@ -166,7 +183,7 @@ async def mock_image_config_entry_fixture(hass: HomeAssistant, config_flow: None
 
 
 @pytest.fixture(name="mock_image_platform")
-async def mock_image_platform_fixture(hass: HomeAssistant):
+async def mock_image_platform_fixture(hass: HomeAssistant) -> None:
     """Initialize a mock image platform."""
     mock_integration(hass, MockModule(domain="test"))
     mock_platform(hass, "test.image", MockImagePlatform([MockImageEntity(hass)]))

tests/components/image/test_init.py

@@ -9,11 +9,13 @@ import pytest
 import respx
 
 from homeassistant.components import image
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
 from .conftest import (
     MockImageEntity,
+    MockImageEntityInvalidContentType,
     MockImageNoStateEntity,
     MockImagePlatform,
     MockImageSyncEntity,
@@ -26,7 +28,7 @@ from tests.typing import ClientSessionGenerator
 
 @pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
 async def test_state(
-    hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform
+    hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform: None
 ) -> None:
     """Test image state."""
     state = hass.states.get("image.test")
@@ -41,7 +43,9 @@ async def test_state(
 
 @pytest.mark.freeze_time("2023-04-01 00:00:00+00:00")
 async def test_config_entry(
-    hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_config_entry
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    mock_image_config_entry: ConfigEntry,
 ) -> None:
     """Test setting up an image platform from a config entry."""
     state = hass.states.get("image.test")
@@ -99,8 +103,35 @@ async def test_no_state(
     }
 
 
+async def test_no_valid_content_type(
+    hass: HomeAssistant, hass_client: ClientSessionGenerator
+) -> None:
+    """Test invalid content type."""
+    mock_integration(hass, MockModule(domain="test"))
+    mock_platform(
+        hass, "test.image", MockImagePlatform([MockImageEntityInvalidContentType(hass)])
+    )
+    assert await async_setup_component(
+        hass, image.DOMAIN, {"image": {"platform": "test"}}
+    )
+    await hass.async_block_till_done()
+
+    client = await hass_client()
+
+    state = hass.states.get("image.test")
+    # assert state.state == "unknown"
+    access_token = state.attributes["access_token"]
+    assert state.attributes == {
+        "access_token": access_token,
+        "entity_picture": f"/api/image_proxy/image.test?token={access_token}",
+        "friendly_name": "Test",
+    }
+    resp = await client.get(f"/api/image_proxy/image.test?token={access_token}")
+    assert resp.status == HTTPStatus.INTERNAL_SERVER_ERROR
+
+
 async def test_fetch_image_authenticated(
-    hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform
+    hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform: None
 ) -> None:
     """Test fetching an image with an authenticated client."""
     client = await hass_client()
@@ -115,7 +146,7 @@ async def test_fetch_image_authenticated(
 
 
 async def test_fetch_image_fail(
-    hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform
+    hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_image_platform: None
 ) -> None:
     """Test fetching an image with an authenticated client."""
     client = await hass_client()
@@ -147,7 +178,7 @@ async def test_fetch_image_sync(
 async def test_fetch_image_unauthenticated(
     hass: HomeAssistant,
     hass_client_no_auth: ClientSessionGenerator,
-    mock_image_platform,
+    mock_image_platform: None,
 ) -> None:
     """Test fetching an image with an unauthenticated client."""
     client = await hass_client_no_auth()