fix(Screenshot): Impersonate current user

superset/charts/api.py

@@ -599,9 +599,10 @@ class ChartRestApi(BaseSupersetModelRestApi):
             cache_chart_thumbnail.delay(
                 current_user=get_current_user(),
                 chart_id=chart.id,
-                force=True,
                 window_size=window_size,
                 thumb_size=thumb_size,
+                force_current_user=True,
+                force=True,
             )
             return self.response(
                 202, cache_key=cache_key, chart_url=chart_url, image_url=image_url
@@ -718,6 +719,7 @@ class ChartRestApi(BaseSupersetModelRestApi):
             cache_chart_thumbnail.delay(
                 current_user=current_user,
                 chart_id=chart.id,
+                force_current_user=True,
                 force=True,
             )
             return self.response(202, message="OK Async")
@@ -734,6 +736,7 @@ class ChartRestApi(BaseSupersetModelRestApi):
             cache_chart_thumbnail.delay(
                 current_user=current_user,
                 chart_id=chart.id,
+                force_current_user=True,
                 force=True,
             )
             return self.response(202, message="OK Async")

superset/dashboards/api.py

@@ -932,6 +932,7 @@ class DashboardRestApi(BaseSupersetModelRestApi):
             cache_dashboard_thumbnail.delay(
                 current_user=current_user,
                 dashboard_id=dashboard.id,
+                force_current_user=True,
                 force=True,
             )
             return self.response(202, message="OK Async")
@@ -945,6 +946,7 @@ class DashboardRestApi(BaseSupersetModelRestApi):
             cache_dashboard_thumbnail.delay(
                 current_user=current_user,
                 dashboard_id=dashboard.id,
+                force_current_user=True,
                 force=True,
             )
             return self.response(202, message="OK Async")
@@ -1032,8 +1034,9 @@ class DashboardRestApi(BaseSupersetModelRestApi):
         ).run()
 
         dashboard_url = get_url_path("Superset.dashboard_permalink", key=permalink_key)
+        current_user = get_current_user()
         screenshot_obj = DashboardScreenshot(dashboard_url, dashboard.digest)
-        cache_key = screenshot_obj.cache_key(window_size, thumb_size)
+        cache_key = screenshot_obj.cache_key(window_size, thumb_size, current_user)
         image_url = get_url_path(
             "DashboardRestApi.screenshot", pk=dashboard.id, digest=cache_key
         )
@@ -1044,9 +1047,11 @@ class DashboardRestApi(BaseSupersetModelRestApi):
                 current_user=get_current_user(),
                 dashboard_id=dashboard.id,
                 dashboard_url=dashboard_url,
-                force=True,
+                cache_key=cache_key,
                 thumb_size=thumb_size,
                 window_size=window_size,
+                force_current_user=True,
+                force=True,
             )
             return self.response(
                 202,
@@ -1102,10 +1107,16 @@ class DashboardRestApi(BaseSupersetModelRestApi):
         if not dashboard:
             return self.response_404()
 
+        decompressed_key = DashboardScreenshot.decompress_key(digest)
+        if (
+            not decompressed_key
+            or decompressed_key.get("current_user") != get_current_user()
+        ):
+            return self.response_403
+
         download_format = request.args.get("download_format", "png")
 
         # fetch the dashboard screenshot using the current user and cache if set
-
         if img := DashboardScreenshot.get_from_cache_key(thumbnail_cache, digest):
             if download_format == "pdf":
                 pdf_img = img.getvalue()

superset/tasks/thumbnails.py

@@ -24,6 +24,7 @@ from flask import current_app
 
 from superset import security_manager, thumbnail_cache
 from superset.extensions import celery_app
+from superset.tasks.types import ExecutorType
 from superset.tasks.utils import get_executor
 from superset.utils.core import override_user
 from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
@@ -40,6 +41,7 @@ def cache_chart_thumbnail(
     force: bool = False,
     window_size: Optional[WindowSize] = None,
     thumb_size: Optional[WindowSize] = None,
+    force_current_user: Optional[bool] = False,
 ) -> None:
     # pylint: disable=import-outside-toplevel
     from superset.models.slice import Slice
@@ -54,7 +56,9 @@ def cache_chart_thumbnail(
     url = get_url_path("Superset.slice", slice_id=chart.id)
     logger.info("Caching chart: %s", url)
     _, username = get_executor(
-        executor_types=current_app.config["THUMBNAIL_EXECUTE_AS"],
+        executor_types=[ExecutorType.CURRENT_USER]
+        if force_current_user
+        else current_app.config["THUMBNAIL_EXECUTE_AS"],
         model=chart,
         current_user=current_user,
     )
@@ -78,6 +82,7 @@ def cache_dashboard_thumbnail(
     force: bool = False,
     thumb_size: Optional[WindowSize] = None,
     window_size: Optional[WindowSize] = None,
+    force_current_user: Optional[bool] = False,
 ) -> None:
     # pylint: disable=import-outside-toplevel
     from superset.models.dashboard import Dashboard
@@ -90,7 +95,9 @@ def cache_dashboard_thumbnail(
 
     logger.info("Caching dashboard: %s", url)
     _, username = get_executor(
-        executor_types=current_app.config["THUMBNAIL_EXECUTE_AS"],
+        executor_types=[ExecutorType.CURRENT_USER]
+        if force_current_user
+        else current_app.config["THUMBNAIL_EXECUTE_AS"],
         model=dashboard,
         current_user=current_user,
     )
@@ -112,9 +119,11 @@ def cache_dashboard_screenshot(
     current_user: Optional[str],
     dashboard_id: int,
     dashboard_url: str,
-    force: bool = True,
+    cache_key: Optional[str] = None,
     thumb_size: Optional[WindowSize] = None,
     window_size: Optional[WindowSize] = None,
+    force_current_user: Optional[bool] = False,
+    force: bool = True,
 ) -> None:
     # pylint: disable=import-outside-toplevel
     from superset.models.dashboard import Dashboard
@@ -126,18 +135,23 @@ def cache_dashboard_screenshot(
     dashboard = Dashboard.get(dashboard_id)
 
     logger.info("Caching dashboard: %s", dashboard_url)
+
     _, username = get_executor(
-        executor_types=current_app.config["THUMBNAIL_EXECUTE_AS"],
+        executor_types=[ExecutorType.CURRENT_USER]
+        if force_current_user
+        else current_app.config["THUMBNAIL_EXECUTE_AS"],
         model=dashboard,
         current_user=current_user,
     )
+
     user = security_manager.find_user(username)
     with override_user(user):
         screenshot = DashboardScreenshot(dashboard_url, dashboard.digest)
         screenshot.compute_and_cache(
             user=user,
             cache=thumbnail_cache,
-            force=force,
+            cache_key=cache_key,
             window_size=window_size,
             thumb_size=thumb_size,
+            force=force,
         )

superset/utils/screenshots.py

@@ -16,14 +16,17 @@
 # under the License.
 from __future__ import annotations
 
+import base64
+import json
 import logging
+import zlib
 from io import BytesIO
 from typing import TYPE_CHECKING
 
 from flask import current_app
 
 from superset import feature_flag_manager
-from superset.utils.hashing import md5_sha_from_dict
+from superset.utils.hashing import md5_sha_from_dict, md5_sha_from_str
 from superset.utils.urls import modify_url_query
 from superset.utils.webdriver import (
     ChartStandaloneMode,
@@ -84,6 +87,7 @@ class BaseScreenshot:
             "window_size": window_size,
             "thumb_size": thumb_size,
         }
+
         return md5_sha_from_dict(args)
 
     def get_screenshot(
@@ -144,6 +148,7 @@ class BaseScreenshot:
         thumb_size: WindowSize | None = None,
         cache: Cache = None,
         force: bool = True,
+        cache_key: str | None = None,
     ) -> bytes | None:
         """
         Fetches the screenshot, computes the thumbnail and caches the result
@@ -155,7 +160,7 @@ class BaseScreenshot:
         :param force: Will force the computation even if it's already cached
         :return: Image payload
         """
-        cache_key = self.cache_key(window_size, thumb_size)
+        cache_key = cache_key or self.cache_key(window_size, thumb_size)
         window_size = window_size or self.window_size
         thumb_size = thumb_size or self.thumb_size
         if not force and cache and cache.get(cache_key):
@@ -251,3 +256,33 @@ class DashboardScreenshot(BaseScreenshot):
         super().__init__(url, digest)
         self.window_size = window_size or DEFAULT_DASHBOARD_WINDOW_SIZE
         self.thumb_size = thumb_size or DEFAULT_DASHBOARD_THUMBNAIL_SIZE
+
+    def cache_key(
+        self,
+        window_size: bool | WindowSize | None = None,
+        thumb_size: bool | WindowSize | None = None,
+        current_user=str,
+    ) -> str:
+        window_size = window_size or self.window_size
+        thumb_size = thumb_size or self.thumb_size
+        args = {
+            "digest": md5_sha_from_str(self.digest),
+            "thumbnail_type": "screenshot",
+            "target": self.thumbnail_type,
+            "window_size": window_size,
+            "thumb_size": thumb_size,
+            "current_user": current_user,
+        }
+        return self.compress_key(args)
+
+    def compress_key(self, obj: dict) -> str:
+        json_data = json.dumps(obj, separators=(",", ":"))
+        compressed_data = zlib.compress(json_data.encode("utf-8"), level=9)
+        compressed_key = base64.urlsafe_b64encode(compressed_data).decode("utf-8")
+        return compressed_key
+
+    @staticmethod
+    def decompress_key(encoded_data: str) -> dict:
+        compressed_data = base64.urlsafe_b64decode(encoded_data.encode("utf-8"))
+        json_data = zlib.decompress(compressed_data).decode("utf-8")
+        return json.loads(json_data)