fix(Embedded): Skip CSRF validation for dashboard download endpoints (#31798)

superset/dashboards/api.py

@@ -19,20 +19,19 @@ import functools
 import logging
 from datetime import datetime
 from io import BytesIO
-from typing import Any, Callable, cast, Optional
+from typing import Any, Callable, cast
 from zipfile import is_zipfile, ZipFile
 
 from flask import g, redirect, request, Response, send_file, url_for
 from flask_appbuilder import permission_name
 from flask_appbuilder.api import expose, protect, rison, safe
-from flask_appbuilder.hooks import before_request
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from flask_babel import gettext, ngettext
 from marshmallow import ValidationError
 from werkzeug.wrappers import Response as WerkzeugResponse
 from werkzeug.wsgi import FileWrapper
 
-from superset import db, is_feature_enabled, thumbnail_cache
+from superset import db, thumbnail_cache
 from superset.charts.schemas import ChartEntityResponseSchema
 from superset.commands.dashboard.copy import CopyDashboardCommand
 from superset.commands.dashboard.create import CreateDashboardCommand
@@ -124,6 +123,7 @@ from superset.views.base_api import (
     requires_form_data,
     requires_json,
     statsd_metrics,
+    validate_feature_flags,
 )
 from superset.views.error_handling import handle_api_exception
 from superset.views.filters import (
@@ -160,18 +160,6 @@ def with_dashboard(
 class DashboardRestApi(BaseSupersetModelRestApi):
     datamodel = SQLAInterface(Dashboard)
 
-    @before_request(only=["thumbnail", "cache_dashboard_screenshot", "screenshot"])
-    def ensure_thumbnails_enabled(self) -> Optional[Response]:
-        if not is_feature_enabled("THUMBNAILS"):
-            return self.response_404()
-        return None
-
-    @before_request(only=["cache_dashboard_screenshot", "screenshot"])
-    def ensure_screenshots_enabled(self) -> Optional[Response]:
-        if not is_feature_enabled("ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS"):
-            return self.response_404()
-        return None
-
     include_route_methods = RouteMethod.REST_MODEL_VIEW_CRUD_SET | {
         RouteMethod.EXPORT,
         RouteMethod.IMPORT,
@@ -1038,6 +1026,7 @@ class DashboardRestApi(BaseSupersetModelRestApi):
         return response
 
     @expose("/<pk>/thumbnail/<digest>/", methods=("GET",))
+    @validate_feature_flags(["THUMBNAILS"])
     @protect()
     @safe
     @rison(thumbnail_query_schema)
@@ -1141,6 +1130,7 @@ class DashboardRestApi(BaseSupersetModelRestApi):
         )
 
     @expose("/<pk>/cache_dashboard_screenshot/", methods=("POST",))
+    @validate_feature_flags(["THUMBNAILS", "ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS"])
     @protect()
     @rison(screenshot_query_schema)
     @safe
@@ -1241,6 +1231,7 @@ class DashboardRestApi(BaseSupersetModelRestApi):
         return trigger_celery()
 
     @expose("/<pk>/screenshot/<digest>/", methods=("GET",))
+    @validate_feature_flags(["THUMBNAILS", "ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS"])
     @protect()
     @safe
     @statsd_metrics

superset/views/base_api.py

@@ -31,6 +31,7 @@ from marshmallow import fields, Schema
 from sqlalchemy import and_, distinct, func
 from sqlalchemy.orm.query import Query
 
+from superset import is_feature_enabled
 from superset.exceptions import InvalidPayloadFormatError
 from superset.extensions import db, event_logger, security_manager, stats_logger_manager
 from superset.models.core import FavStar
@@ -130,6 +131,29 @@ def statsd_metrics(f: Callable[..., Any]) -> Callable[..., Any]:
     return functools.update_wrapper(wraps, f)
 
 
+def validate_feature_flags(
+    feature_flags: list[str],
+) -> Callable[[Callable[..., Response]], Callable[..., Response]]:
+    """
+    A decorator to check if all given feature flags are enabled.
+
+    :param feature_flags: List of feature flag names to be checked.
+    """
+
+    def decorate(f: Callable[..., Response]) -> Callable[..., Response]:
+        @functools.wraps(f)
+        def wrapper(
+            self: BaseSupersetModelRestApi, *args: Any, **kwargs: Any
+        ) -> Response:
+            if not all(is_feature_enabled(flag) for flag in feature_flags):
+                return self.response_404()
+            return f(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorate
+
+
 class RelatedFieldFilter:
     # data class to specify what filter to use on a /related endpoint
     # pylint: disable=too-few-public-methods

tests/integration_tests/dashboards/api_tests.py

@@ -41,6 +41,7 @@ from superset.utils import json
 
 from tests.integration_tests.base_api_tests import ApiOwnersTestCaseMixin
 from tests.integration_tests.base_tests import SupersetTestCase
+from tests.integration_tests.conftest import with_feature_flags
 from tests.integration_tests.constants import (
     ADMIN_USERNAME,
     ALPHA_USERNAME,
@@ -3027,10 +3028,9 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
         uri = f"/api/v1/dashboard/{dashboard_id}/screenshot/{cache_key}/?download_format={download_format}"  # noqa: E501
         return self.client.get(uri)
 
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
     @pytest.mark.usefixtures("create_dashboard_with_tag")
-    @patch("superset.dashboards.api.is_feature_enabled")
+    def test_cache_dashboard_screenshot_success(self):
-    def test_cache_dashboard_screenshot_success(self, is_feature_enabled):
-        is_feature_enabled.return_value = True
         self.login(ADMIN_USERNAME)
         dashboard = (
             db.session.query(Dashboard)
@@ -3040,10 +3040,9 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
         response = self._cache_screenshot(dashboard.id)
         assert response.status_code == 202
 
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
     @pytest.mark.usefixtures("create_dashboard_with_tag")
-    @patch("superset.dashboards.api.is_feature_enabled")
+    def test_cache_dashboard_screenshot_dashboard_validation(self):
-    def test_cache_dashboard_screenshot_dashboard_validation(self, is_feature_enabled):
-        is_feature_enabled.return_value = True
         self.login(ADMIN_USERNAME)
         dashboard = (
             db.session.query(Dashboard)
@@ -3059,25 +3058,21 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
         response = self._cache_screenshot(dashboard.id, invalid_payload)
         assert response.status_code == 400
 
-    @patch("superset.dashboards.api.is_feature_enabled")
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
-    def test_cache_dashboard_screenshot_dashboard_not_found(self, is_feature_enabled):
+    def test_cache_dashboard_screenshot_dashboard_not_found(self):
-        is_feature_enabled.return_value = True
         self.login(ADMIN_USERNAME)
         non_existent_id = 999
         response = self._cache_screenshot(non_existent_id)
         assert response.status_code == 404
 
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
     @pytest.mark.usefixtures("create_dashboard_with_tag")
     @patch("superset.dashboards.api.cache_dashboard_screenshot")
     @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
-    @patch("superset.dashboards.api.is_feature_enabled")
+    def test_screenshot_success_png(self, mock_get_cache, mock_cache_task):
-    def test_screenshot_success_png(
-        self, is_feature_enabled, mock_get_cache, mock_cache_task
-    ):
         """
         Validate screenshot returns png
         """
-        is_feature_enabled.return_value = True
         self.login(ADMIN_USERNAME)
         mock_cache_task.return_value = None
         mock_get_cache.return_value = BytesIO(b"fake image data")
@@ -3096,18 +3091,17 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
         assert response.mimetype == "image/png"
         assert response.data == b"fake image data"
 
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
     @pytest.mark.usefixtures("create_dashboard_with_tag")
     @patch("superset.dashboards.api.cache_dashboard_screenshot")
     @patch("superset.dashboards.api.build_pdf_from_screenshots")
     @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
-    @patch("superset.dashboards.api.is_feature_enabled")
     def test_screenshot_success_pdf(
-        self, is_feature_enabled, mock_get_from_cache, mock_build_pdf, mock_cache_task
+        self, mock_get_from_cache, mock_build_pdf, mock_cache_task
     ):
         """
         Validate screenshot can return pdf.
         """
-        is_feature_enabled.return_value = True
         self.login(ADMIN_USERNAME)
         mock_cache_task.return_value = None
         mock_get_from_cache.return_value = BytesIO(b"fake image data")
@@ -3127,14 +3121,11 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
         assert response.mimetype == "application/pdf"
         assert response.data == b"fake pdf data"
 
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
     @pytest.mark.usefixtures("create_dashboard_with_tag")
     @patch("superset.dashboards.api.cache_dashboard_screenshot")
     @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
-    @patch("superset.dashboards.api.is_feature_enabled")
+    def test_screenshot_not_in_cache(self, mock_get_cache, mock_cache_task):
-    def test_screenshot_not_in_cache(
-        self, is_feature_enabled, mock_get_cache, mock_cache_task
-    ):
-        is_feature_enabled.return_value = True
         self.login(ADMIN_USERNAME)
         mock_cache_task.return_value = None
         mock_get_cache.return_value = None
@@ -3151,22 +3142,18 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
         response = self._get_screenshot(dashboard.id, cache_key, "pdf")
         assert response.status_code == 404
 
-    @patch("superset.dashboards.api.is_feature_enabled")
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
-    def test_screenshot_dashboard_not_found(self, is_feature_enabled):
+    def test_screenshot_dashboard_not_found(self):
-        is_feature_enabled.return_value = True
         self.login(ADMIN_USERNAME)
         non_existent_id = 999
         response = self._get_screenshot(non_existent_id, "some_cache_key", "png")
         assert response.status_code == 404
 
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
     @pytest.mark.usefixtures("create_dashboard_with_tag")
     @patch("superset.dashboards.api.cache_dashboard_screenshot")
     @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
-    @patch("superset.dashboards.api.is_feature_enabled")
+    def test_screenshot_invalid_download_format(self, mock_get_cache, mock_cache_task):
-    def test_screenshot_invalid_download_format(
-        self, is_feature_enabled, mock_get_cache, mock_cache_task
-    ):
-        is_feature_enabled.return_value = True
         self.login(ADMIN_USERNAME)
         mock_cache_task.return_value = None
         mock_get_cache.return_value = BytesIO(b"fake png data")
@@ -3184,10 +3171,41 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
         response = self._get_screenshot(dashboard.id, cache_key, "invalid")
         assert response.status_code == 404
 
+    @with_feature_flags(THUMBNAILS=False, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
     @pytest.mark.usefixtures("create_dashboard_with_tag")
-    @patch("superset.dashboards.api.is_feature_enabled")
+    def test_cache_dashboard_screenshot_feature_thumbnails_ff_disabled(self):
-    def test_cache_dashboard_screenshot_feature_disabled(self, is_feature_enabled):
+        self.login(ADMIN_USERNAME)
-        is_feature_enabled.return_value = False
+
+        dashboard = (
+            db.session.query(Dashboard)
+            .filter(Dashboard.dashboard_title == "dash with tag")
+            .first()
+        )
+
+        assert dashboard is not None
+
+        response = self._cache_screenshot(dashboard.id)
+        assert response.status_code == 404
+
+    @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=False)
+    @pytest.mark.usefixtures("create_dashboard_with_tag")
+    def test_cache_dashboard_screenshot_feature_screenshot_ff_disabled(self):
+        self.login(ADMIN_USERNAME)
+
+        dashboard = (
+            db.session.query(Dashboard)
+            .filter(Dashboard.dashboard_title == "dash with tag")
+            .first()
+        )
+
+        assert dashboard is not None
+
+        response = self._cache_screenshot(dashboard.id)
+        assert response.status_code == 404
+
+    @with_feature_flags(THUMBNAILS=False, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=False)
+    @pytest.mark.usefixtures("create_dashboard_with_tag")
+    def test_cache_dashboard_screenshot_feature_both_ff_disabled(self):
         self.login(ADMIN_USERNAME)
 
         dashboard = (