diff --git a/superset/tasks/cache.py b/superset/tasks/cache.py index 1509b0d0624aa..94bc3cbb9c6b8 100644 --- a/superset/tasks/cache.py +++ b/superset/tasks/cache.py @@ -18,6 +18,7 @@ from typing import Any, Optional, Union from urllib import request from urllib.error import URLError +from urllib.parse import urlparse from celery.beat import SchedulingError from celery.utils.log import get_task_logger @@ -47,6 +48,23 @@ def get_payload(chart: Slice, dashboard: Optional[Dashboard] = None) -> dict[str return payload +def is_secure_url(url: str) -> bool: + """ + Validates if a URL is secure (uses HTTPS). + + :param url: The URL to validate. + :return: True if the URL uses HTTPS (secure), False if it uses HTTP (non-secure). + """ + try: + parsed_url = urlparse(url) + # Return True for HTTPS, False for HTTP + return parsed_url.scheme == "https" + except ValueError as exception: + # Log a warning with detailed context + logger.warning("Failed to parse URL '%s': %s", url, str(exception)) + return False + + class Strategy: # pylint: disable=too-few-public-methods """ A cache warm up strategy. @@ -220,10 +238,15 @@ def fetch_url(data: str, headers: dict[str, str]) -> dict[str, str]: """ result = {} try: + url = get_url_path("ChartRestApi.warm_up_cache") + + if is_secure_url(url): + logger.info("URL '%s' is secure. Adding Referer header.", url) + headers.update({"Referer": url}) + # Fetch CSRF token for API request headers.update(fetch_csrf_token(headers)) - url = get_url_path("ChartRestApi.warm_up_cache") logger.info("Fetching %s with payload %s", url, data) req = request.Request( url, data=bytes(data, "utf-8"), headers=headers, method="PUT" diff --git a/superset/tasks/utils.py b/superset/tasks/utils.py index 6fc799c4abc2c..5e3bc148082b9 100644 --- a/superset/tasks/utils.py +++ b/superset/tasks/utils.py @@ -133,7 +133,7 @@ def fetch_csrf_token( data = json.loads(body) res = {"X-CSRF-Token": data["result"]} if session_cookie is not None: - res["Cookie"] = session_cookie + res["Cookie"] = f"{session_cookie_name}={session_cookie}" return res logger.error("Error fetching CSRF token, status code: %s", response.status) diff --git a/tests/integration_tests/tasks/test_cache.py b/tests/integration_tests/tasks/test_cache.py index 6e8d3ffe03b4d..368cb1ebf0afb 100644 --- a/tests/integration_tests/tasks/test_cache.py +++ b/tests/integration_tests/tasks/test_cache.py @@ -22,17 +22,32 @@ @pytest.mark.parametrize( - "base_url", + "base_url, expected_referer", [ - "http://base-url", - "http://base-url/", + ("http://base-url", None), + ("http://base-url/", None), + ("https://base-url", "https://base-url/api/v1/chart/warm_up_cache"), + ("https://base-url/", "https://base-url/api/v1/chart/warm_up_cache"), + ], + ids=[ + "Without trailing slash (HTTP)", + "With trailing slash (HTTP)", + "Without trailing slash (HTTPS)", + "With trailing slash (HTTPS)", ], - ids=["Without trailing slash", "With trailing slash"], ) @mock.patch("superset.tasks.cache.fetch_csrf_token") @mock.patch("superset.tasks.cache.request.Request") @mock.patch("superset.tasks.cache.request.urlopen") -def test_fetch_url(mock_urlopen, mock_request_cls, mock_fetch_csrf_token, base_url): +@mock.patch("superset.tasks.cache.is_secure_url") +def test_fetch_url( + mock_is_secure_url, + mock_urlopen, + mock_request_cls, + mock_fetch_csrf_token, + base_url, + expected_referer, +): from superset.tasks.cache import fetch_url mock_request = mock.MagicMock() @@ -41,8 +56,17 @@ def test_fetch_url(mock_urlopen, mock_request_cls, mock_fetch_csrf_token, base_u mock_urlopen.return_value = mock.MagicMock() mock_urlopen.return_value.code = 200 + # Mock the URL validation to return True for HTTPS and False for HTTP + mock_is_secure_url.return_value = base_url.startswith("https") + initial_headers = {"Cookie": "cookie", "key": "value"} csrf_headers = initial_headers | {"X-CSRF-Token": "csrf_token"} + + # Conditionally add the Referer header and assert its presence + if expected_referer: + csrf_headers = csrf_headers | {"Referer": expected_referer} + assert csrf_headers["Referer"] == expected_referer + mock_fetch_csrf_token.return_value = csrf_headers app.config["WEBDRIVER_BASEURL"] = base_url @@ -51,13 +75,21 @@ def test_fetch_url(mock_urlopen, mock_request_cls, mock_fetch_csrf_token, base_u result = fetch_url(data, initial_headers) - assert data == result["success"] + expected_url = ( + f"{base_url}/api/v1/chart/warm_up_cache" + if not base_url.endswith("/") + else f"{base_url}api/v1/chart/warm_up_cache" + ) + mock_fetch_csrf_token.assert_called_once_with(initial_headers) + mock_request_cls.assert_called_once_with( - "http://base-url/api/v1/chart/warm_up_cache", + expected_url, # Use the dynamic URL based on base_url data=data_encoded, headers=csrf_headers, method="PUT", ) # assert the same Request object is used mock_urlopen.assert_called_once_with(mock_request, timeout=mock.ANY) + + assert data == result["success"] diff --git a/tests/integration_tests/tasks/test_utils.py b/tests/integration_tests/tasks/test_utils.py index b1213b78c85a0..29e5f38319cb9 100644 --- a/tests/integration_tests/tasks/test_utils.py +++ b/tests/integration_tests/tasks/test_utils.py @@ -26,8 +26,15 @@ [ "http://base-url", "http://base-url/", + "https://base-url", + "https://base-url/", + ], + ids=[ + "Without trailing slash (HTTP)", + "With trailing slash (HTTP)", + "Without trailing slash (HTTPS)", + "With trailing slash (HTTPS)", ], - ids=["Without trailing slash", "With trailing slash"], ) @mock.patch("superset.tasks.cache.request.Request") @mock.patch("superset.tasks.cache.request.urlopen") @@ -52,13 +59,19 @@ def test_fetch_csrf_token(mock_urlopen, mock_request_cls, base_url, app_context) result_headers = fetch_csrf_token(headers) + expected_url = ( + f"{base_url}/api/v1/security/csrf_token/" + if not base_url.endswith("/") + else f"{base_url}api/v1/security/csrf_token/" + ) + mock_request_cls.assert_called_with( - "http://base-url/api/v1/security/csrf_token/", + expected_url, headers=headers, method="GET", ) assert result_headers["X-CSRF-Token"] == "csrf_token" - assert result_headers["Cookie"] == "new_session_cookie" + assert result_headers["Cookie"] == "session=new_session_cookie" # Updated assertion # assert the same Request object is used mock_urlopen.assert_called_once_with(mock_request, timeout=mock.ANY)