From f6c574edd85014bd4ebf77ec6c68fe75b745d14f Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Fri, 31 Jul 2026 15:13:59 -0700 Subject: [PATCH] fix(screenshots): validate cached screenshot image bytes on read and write (#42120) Co-authored-by: Claude --- superset/utils/screenshots.py | 58 +++++++- tests/unit_tests/utils/screenshot_test.py | 10 +- .../utils/test_screenshot_cache_fix.py | 126 +++++++++++++++++- 3 files changed, 179 insertions(+), 15 deletions(-) diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py index 750ec54b597..32b7f2bda23 100644 --- a/superset/utils/screenshots.py +++ b/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 ( @@ -263,6 +290,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 @@ -331,15 +366,28 @@ class BaseScreenshot: image = None # Cache the result (success or error) to avoid immediate retries - if image: + invalid_reason = validate_screenshot_image(image) + # `image and` is redundant at runtime (validate_screenshot_image + # only returns None for truthy, well-formed bytes) but mypy can't + # infer that image is non-None from invalid_reason being None + # across the function-call boundary, so it's kept for narrowing. + 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()) diff --git a/tests/unit_tests/utils/screenshot_test.py b/tests/unit_tests/utils/screenshot_test.py index b7f7de6032e..c57df50484c 100644 --- a/tests/unit_tests/utils/screenshot_test.py +++ b/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 { diff --git a/tests/unit_tests/utils/test_screenshot_cache_fix.py b/tests/unit_tests/utils/test_screenshot_cache_fix.py index 4ffe07711c8..133c8c7e4d8 100644 --- a/tests/unit_tests/utils/test_screenshot_cache_fix.py +++ b/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