fix: cache-warmup fails (#31173)

Co-authored-by: Sivarajan Narayanan <narayanan_sivarajan@apple.com>
Co-authored-by: Pat Heard <patrick.heard@cds-snc.ca>

UPDATING.md

@@ -28,6 +28,7 @@ assists people when migrating to a new version.
 - [29798](https://github.com/apache/superset/pull/29798) Since 3.1.0, the intial schedule for an alert or report was mistakenly offset by the specified timezone's relation to UTC. The initial schedule should now begin at the correct time.
 - [30021](https://github.com/apache/superset/pull/30021) The `dev` layer in our Dockerfile no long includes firefox binaries, only Chromium to reduce bloat/docker-build-time.
 - [30099](https://github.com/apache/superset/pull/30099) Translations are no longer included in the default docker image builds. If your environment requires translations, you'll want to set the docker build arg `BUILD_TRANSACTION=true`.
+- [31173](https://github.com/apache/superset/pull/31173) Modified `fetch_csrf_token` to align with HTTP standards, particularly regarding how cookies are handled. If you encounter any issues related to CSRF functionality, please report them as a new issue and reference this PR for context.
 
 ### Potential Downtime
 

superset/tasks/cache.py

@@ -33,7 +33,7 @@ from superset.tasks.utils import fetch_csrf_token
 from superset.utils import json
 from superset.utils.date_parser import parse_human_datetime
 from superset.utils.machine_auth import MachineAuthProvider
-from superset.utils.urls import get_url_path
+from superset.utils.urls import get_url_path, is_secure_url
 
 logger = get_task_logger(__name__)
 logger.setLevel(logging.INFO)
@@ -220,10 +220,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"

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)

superset/utils/urls.py

@@ -16,6 +16,7 @@
 # under the License.
 import urllib
 from typing import Any
+from urllib.parse import urlparse
 
 from flask import current_app, url_for
 
@@ -50,3 +51,14 @@ def modify_url_query(url: str, **kwargs: Any) -> str:
         f"{k}={urllib.parse.quote(str(v[0]))}" for k, v in params.items()
     )
     return urllib.parse.urlunsplit(parts)
+
+
+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).
+    """
+    parsed_url = urlparse(url)
+    return parsed_url.scheme == "https"

tests/integration_tests/tasks/test_cache.py

@@ -22,17 +22,32 @@ from tests.integration_tests.test_app import app
 
 
 @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"]

tests/integration_tests/tasks/test_utils.py

@@ -26,8 +26,15 @@ from tests.integration_tests.test_app import app
     [
         "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)