fix(screenshots): validate cached screenshot image bytes on read and write

Reject stale/empty/corrupt cached screenshot payloads instead of serving
them: get_from_cache_key now treats a payload that claims a successful
screenshot but has empty or non-image bytes as a cache miss, and
compute_and_cache applies the same cheap header check before marking a
result as cached-success, falling back to ERROR status otherwise.

Co-Authored-By: Claude <noreply@anthropic.com>

superset/utils/screenshots.py

@@ -85,6 +85,26 @@ class ScreenshotCachePayloadType(TypedDict):
     status: str
 
 
+# Magic bytes for a cheap image sanity check. This is intentionally not a full
+# decode: it's meant to catch 0-byte/corrupt/blank payloads before they're
+# cached or served, not to validate the image is renderable.
+PNG_MAGIC_BYTES = b"\x89PNG\r\n\x1a\n"
+JPEG_MAGIC_BYTES = b"\xff\xd8\xff"
+
+
+def validate_screenshot_image(image: bytes | None) -> str | None:
+    """Cheaply validate screenshot bytes before they're cached or served.
+
+    :return: None if the bytes look like a usable image, otherwise a short
+        reason ("empty" or "undecodable") suitable for logging.
+    """
+    if not image:
+        return "empty"
+    if not image.startswith((PNG_MAGIC_BYTES, JPEG_MAGIC_BYTES)):
+        return "undecodable"
+    return None
+
+
 class ScreenshotCachePayload:
     def __init__(
         self,
@@ -147,6 +167,13 @@ class ScreenshotCachePayload:
     def get_status(self) -> str:
         return self.status.value
 
+    def get_invalid_image_reason(self) -> str | None:
+        """Reason this payload's image should not be served/cached, or None if
+        it passes validation (or it isn't claiming a successful screenshot)."""
+        if self.status != StatusValues.UPDATED:
+            return None
+        return validate_screenshot_image(self._image)
+
     def is_error_cache_ttl_expired(self) -> bool:
         error_cache_ttl = app.config["THUMBNAIL_ERROR_CACHE_TTL"]
         return (
@@ -258,6 +285,14 @@ class BaseScreenshot:
             elif isinstance(payload, dict):
                 payload = cast(ScreenshotCachePayloadType, payload)
                 payload = ScreenshotCachePayload.from_dict(payload)
+            if invalid_reason := payload.get_invalid_image_reason():
+                logger.warning(
+                    "Rejecting cached screenshot for %s: %s image payload; "
+                    "treating as a cache miss",
+                    cache_key,
+                    invalid_reason,
+                )
+                return None
             return payload
         logger.info("Failed at getting from cache: %s", cache_key)
         return None
@@ -326,15 +361,24 @@ class BaseScreenshot:
                         image = None
 
                 # Cache the result (success or error) to avoid immediate retries
-                if image:
+                invalid_reason = validate_screenshot_image(image)
+                if image and invalid_reason is None:
                     with event_logger.log_context(
                         f"screenshot.cache.{self.thumbnail_type}"
                     ):
                         cache_payload.update(image)
-                elif cache_payload.status != StatusValues.ERROR:
-                    # Only call error() if not already set — avoids overwriting
-                    # the timestamp recorded when the actual failure occurred above.
-                    cache_payload.error()
+                else:
+                    if invalid_reason:
+                        logger.warning(
+                            "Not caching screenshot result for %s: %s image payload",
+                            cache_key,
+                            invalid_reason,
+                        )
+                    if cache_payload.status != StatusValues.ERROR:
+                        # Only call error() if not already set — avoids overwriting
+                        # the timestamp recorded when the actual failure occurred
+                        # above.
+                        cache_payload.error()
 
                 logger.info("Caching thumbnail: %s", cache_key)
                 self.cache.set(cache_key, cache_payload.to_dict())

tests/unit_tests/utils/screenshot_test.py

@@ -33,6 +33,10 @@ from superset.utils.screenshots import (
 
 BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot"
 
+# A minimal valid PNG header, used wherever a test needs bytes that pass
+# ScreenshotCachePayload's image validation.
+FAKE_PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"fake-png-body"
+
 
 class MockCache:
     """A class to manage screenshot cache."""
@@ -92,7 +96,7 @@ def test_get_cache_key(app_context, screenshot_obj):
 def test_get_from_cache_key(mocker: MockerFixture, screenshot_obj):
     """get_from_cache_key should always return a ScreenshotCachePayload Object"""
     # backwards compatibility test for retrieving plain bytes
-    fake_bytes = b"fake_screenshot_data"
+    fake_bytes = FAKE_PNG_BYTES
     BaseScreenshot.cache = MockCache()
     BaseScreenshot.cache.set("key", fake_bytes)
     cache_payload = screenshot_obj.get_from_cache_key("key")
@@ -108,10 +112,10 @@ class TestComputeAndCache:
             BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None
         )
         get_screenshot = mocker.patch(
-            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"new_image_data"
+            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
         )
         resize_image = mocker.patch(
-            BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
+            BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES
         )
         BaseScreenshot.cache = MockCache()
         return {

tests/unit_tests/utils/test_screenshot_cache_fix.py

@@ -37,6 +37,10 @@ from superset.utils.screenshots import (
 BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot"
 DISTRIBUTED_LOCK_PATH = "superset.utils.screenshots.DistributedLock"
 
+# A minimal valid PNG header, used wherever a test needs bytes that pass
+# ScreenshotCachePayload's image validation.
+FAKE_PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"fake-png-body"
+
 
 class MockCache:
     """A class to manage screenshot cache for testing."""
@@ -83,11 +87,11 @@ class TestCacheOnlyOnSuccess:
         mocker.patch(DISTRIBUTED_LOCK_PATH)
         mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
         get_screenshot = mocker.patch(
-            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"image_data"
+            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
         )
         # Mock resize_image to avoid PIL errors with fake image data
         mocker.patch(
-            BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
+            BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES
         )
         BaseScreenshot.cache = MockCache()
         return get_screenshot
@@ -161,13 +165,15 @@ class TestCacheOnlyOnSuccess:
         screenshot_obj: BaseScreenshot,
         mock_user: MagicMock,
     ) -> None:
-        """Empty bytes from get_screenshot must set ERROR, not leave COMPUTING."""
+        """Empty bytes from get_screenshot must set ERROR, not leave COMPUTING,
+        and must log a WARNING that includes the cache key."""
         mocker.patch(DISTRIBUTED_LOCK_PATH)
         mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
         mocker.patch(
             BASE_SCREENSHOT_PATH + ".get_screenshot",
             return_value=b"",
         )
+        mock_logger = mocker.patch("superset.utils.screenshots.logger")
         BaseScreenshot.cache = MockCache()
 
         screenshot_obj.compute_and_cache(user=mock_user, force=True)
@@ -177,6 +183,43 @@ class TestCacheOnlyOnSuccess:
         assert cached_value is not None
         assert cached_value["status"] == "Error"
         assert cached_value.get("image") is None
+        assert any(
+            cache_key in call.args and "empty" in call.args
+            for call in mock_logger.warning.call_args_list
+        )
+
+    def test_cache_error_status_when_screenshot_returns_garbage_bytes(
+        self,
+        mocker: MockerFixture,
+        screenshot_obj: BaseScreenshot,
+        mock_user: MagicMock,
+    ) -> None:
+        """Non-empty bytes without a valid image header must set ERROR, not be
+        cached as a success, and must log a WARNING that includes the cache key."""
+        mocker.patch(DISTRIBUTED_LOCK_PATH)
+        mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
+        mocker.patch(
+            BASE_SCREENSHOT_PATH + ".get_screenshot",
+            return_value=b"this-is-not-a-real-image",
+        )
+        mocker.patch(
+            BASE_SCREENSHOT_PATH + ".resize_image",
+            return_value=b"this-is-not-a-real-image",
+        )
+        mock_logger = mocker.patch("superset.utils.screenshots.logger")
+        BaseScreenshot.cache = MockCache()
+
+        screenshot_obj.compute_and_cache(user=mock_user, force=True)
+
+        cache_key = screenshot_obj.get_cache_key()
+        cached_value = BaseScreenshot.cache.get(cache_key)
+        assert cached_value is not None
+        assert cached_value["status"] == "Error"
+        assert cached_value.get("image") is None
+        assert any(
+            cache_key in call.args and "undecodable" in call.args
+            for call in mock_logger.warning.call_args_list
+        )
 
     def test_computing_status_written_to_cache_early(
         self,
@@ -197,14 +240,14 @@ class TestCacheOnlyOnSuccess:
                 "Cache should be set to COMPUTING before screenshot starts"
             )
             assert cached_value["status"] == "Computing"
-            return b"image_data"
+            return FAKE_PNG_BYTES
 
         mocker.patch(
             BASE_SCREENSHOT_PATH + ".get_screenshot",
             side_effect=check_cache_during_screenshot,
         )
         mocker.patch(
-            BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
+            BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES
         )
 
         screenshot_obj.compute_and_cache(user=mock_user, force=True)
@@ -429,11 +472,11 @@ class TestIntegrationCacheBugFix:
         BaseScreenshot.cache.set(cache_key, stale_payload.to_dict())
 
         mocker.patch(
-            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"recovered_image"
+            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
         )
         # Mock resize to avoid PIL errors
         mocker.patch(
-            BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image"
+            BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES
         )
 
         # Should trigger task because COMPUTING is stale
@@ -482,3 +525,72 @@ class TestIntegrationCacheBugFix:
 
         assert payload._image == old_image
         assert payload.status == StatusValues.COMPUTING
+
+
+class TestReadSideImageValidation:
+    """A cached payload that claims a successful screenshot (status UPDATED)
+    but carries invalid image bytes must be served as a cache miss, not
+    returned to the caller — this is what the dashboard/chart screenshot
+    endpoints call to fetch bytes to serve."""
+
+    def test_zero_byte_image_is_treated_as_cache_miss(
+        self, mocker: MockerFixture, screenshot_obj: BaseScreenshot
+    ) -> None:
+        mock_logger = mocker.patch("superset.utils.screenshots.logger")
+        BaseScreenshot.cache = MockCache()
+        cache_key = screenshot_obj.get_cache_key()
+        stale_payload = ScreenshotCachePayload(image=b"", status=StatusValues.UPDATED)
+        BaseScreenshot.cache.set(cache_key, stale_payload.to_dict())
+
+        result = screenshot_obj.get_from_cache_key(cache_key)
+
+        assert result is None
+        assert any(
+            cache_key in call.args and "empty" in call.args
+            for call in mock_logger.warning.call_args_list
+        )
+
+    def test_garbage_bytes_image_is_treated_as_cache_miss(
+        self, mocker: MockerFixture, screenshot_obj: BaseScreenshot
+    ) -> None:
+        mock_logger = mocker.patch("superset.utils.screenshots.logger")
+        BaseScreenshot.cache = MockCache()
+        cache_key = screenshot_obj.get_cache_key()
+        garbage_payload = ScreenshotCachePayload(image=b"not-an-image-at-all")
+        BaseScreenshot.cache.set(cache_key, garbage_payload.to_dict())
+
+        result = screenshot_obj.get_from_cache_key(cache_key)
+
+        assert result is None
+        assert any(
+            cache_key in call.args and "undecodable" in call.args
+            for call in mock_logger.warning.call_args_list
+        )
+
+    def test_valid_image_is_served_normally(
+        self, screenshot_obj: BaseScreenshot
+    ) -> None:
+        BaseScreenshot.cache = MockCache()
+        cache_key = screenshot_obj.get_cache_key()
+        valid_payload = ScreenshotCachePayload(image=FAKE_PNG_BYTES)
+        BaseScreenshot.cache.set(cache_key, valid_payload.to_dict())
+
+        result = screenshot_obj.get_from_cache_key(cache_key)
+
+        assert result is not None
+        assert result.get_image().read() == FAKE_PNG_BYTES
+
+    def test_pending_status_with_no_image_is_not_rejected(
+        self, screenshot_obj: BaseScreenshot
+    ) -> None:
+        """Non-UPDATED statuses (e.g. PENDING/COMPUTING) aren't claiming a
+        successful screenshot, so they should be returned as-is."""
+        BaseScreenshot.cache = MockCache()
+        cache_key = screenshot_obj.get_cache_key()
+        pending_payload = ScreenshotCachePayload(status=StatusValues.PENDING)
+        BaseScreenshot.cache.set(cache_key, pending_payload.to_dict())
+
+        result = screenshot_obj.get_from_cache_key(cache_key)
+
+        assert result is not None
+        assert result.status == StatusValues.PENDING