Feat: override cookie domain from cookie utils (#447)

Author: bejito

flask_jwt_extended/utils.py

@@ -259,7 +259,7 @@ def get_csrf_token(encoded_token):
     return token["csrf"]
 
 
-def set_access_cookies(response, encoded_access_token, max_age=None):
+def set_access_cookies(response, encoded_access_token, max_age=None, domain=None):
     """
     Modifiy a Flask Response to set a cookie containing the access JWT.
     Also sets the corresponding CSRF cookies if ``JWT_CSRF_IN_COOKIES`` is ``True``
@@ -276,14 +276,20 @@ def set_access_cookies(response, encoded_access_token, max_age=None):
         ``JWT_SESSION_COOKIE`` option (see :ref:`Configuration Options`). Otherwise,
         it will use this as the cookies ``max-age`` and the JWT_SESSION_COOKIE option
         will be ignored. Values should be the number of seconds (as an integer).
+
+    :param domain:
+        The domain of the cookie. If this is None, it will use the
+        ``JWT_COOKIE_DOMAIN`` option (see :ref:`Configuration Options`). Otherwise,
+        it will use this as the cookies ``domain`` and the JWT_COOKIE_DOMAIN option
+        will be ignored.
     """
     response.set_cookie(
         config.access_cookie_name,
         value=encoded_access_token,
         max_age=max_age or config.cookie_max_age,
         secure=config.cookie_secure,
         httponly=True,
-        domain=config.cookie_domain,
+        domain=domain or config.cookie_domain,
         path=config.access_cookie_path,
         samesite=config.cookie_samesite,
     )
@@ -295,13 +301,13 @@ def set_access_cookies(response, encoded_access_token, max_age=None):
             max_age=max_age or config.cookie_max_age,
             secure=config.cookie_secure,
             httponly=False,
-            domain=config.cookie_domain,
+            domain=domain or config.cookie_domain,
             path=config.access_csrf_cookie_path,
             samesite=config.cookie_samesite,
         )
 
 
-def set_refresh_cookies(response, encoded_refresh_token, max_age=None):
+def set_refresh_cookies(response, encoded_refresh_token, max_age=None, domain=None):
     """
     Modifiy a Flask Response to set a cookie containing the refresh JWT.
     Also sets the corresponding CSRF cookies if ``JWT_CSRF_IN_COOKIES`` is ``True``
@@ -318,14 +324,20 @@ def set_refresh_cookies(response, encoded_refresh_token, max_age=None):
         ``JWT_SESSION_COOKIE`` option (see :ref:`Configuration Options`). Otherwise,
         it will use this as the cookies ``max-age`` and the JWT_SESSION_COOKIE option
         will be ignored. Values should be the number of seconds (as an integer).
+
+    :param domain:
+        The domain of the cookie. If this is None, it will use the
+        ``JWT_COOKIE_DOMAIN`` option (see :ref:`Configuration Options`). Otherwise,
+        it will use this as the cookies ``domain`` and the JWT_COOKIE_DOMAIN option
+        will be ignored.
     """
     response.set_cookie(
         config.refresh_cookie_name,
         value=encoded_refresh_token,
         max_age=max_age or config.cookie_max_age,
         secure=config.cookie_secure,
         httponly=True,
-        domain=config.cookie_domain,
+        domain=domain or config.cookie_domain,
         path=config.refresh_cookie_path,
         samesite=config.cookie_samesite,
     )
@@ -337,39 +349,45 @@ def set_refresh_cookies(response, encoded_refresh_token, max_age=None):
             max_age=max_age or config.cookie_max_age,
             secure=config.cookie_secure,
             httponly=False,
-            domain=config.cookie_domain,
+            domain=domain or config.cookie_domain,
             path=config.refresh_csrf_cookie_path,
             samesite=config.cookie_samesite,
         )
 
 
-def unset_jwt_cookies(response):
+def unset_jwt_cookies(response, domain=None):
     """
     Modifiy a Flask Response to delete the cookies containing access or refresh
     JWTs.  Also deletes the corresponding CSRF cookies if applicable.
 
     :param response:
         A Flask Response object
     """
-    unset_access_cookies(response)
-    unset_refresh_cookies(response)
+    unset_access_cookies(response, domain)
+    unset_refresh_cookies(response, domain)
 
 
-def unset_access_cookies(response):
+def unset_access_cookies(response, domain=None):
     """
     Modifiy a Flask Response to delete the cookie containing a refresh JWT.
     Also deletes the corresponding CSRF cookie if applicable.
 
     :param response:
         A Flask Response object
+
+    :param domain:
+        The domain of the cookie. If this is None, it will use the
+        ``JWT_COOKIE_DOMAIN`` option (see :ref:`Configuration Options`). Otherwise,
+        it will use this as the cookies ``domain`` and the JWT_COOKIE_DOMAIN option
+        will be ignored.
     """
     response.set_cookie(
         config.access_cookie_name,
         value="",
         expires=0,
         secure=config.cookie_secure,
         httponly=True,
-        domain=config.cookie_domain,
+        domain=domain or config.cookie_domain,
         path=config.access_cookie_path,
         samesite=config.cookie_samesite,
     )
@@ -381,27 +399,33 @@ def unset_access_cookies(response):
             expires=0,
             secure=config.cookie_secure,
             httponly=False,
-            domain=config.cookie_domain,
+            domain=domain or config.cookie_domain,
             path=config.access_csrf_cookie_path,
             samesite=config.cookie_samesite,
         )
 
 
-def unset_refresh_cookies(response):
+def unset_refresh_cookies(response, domain=None):
     """
     Modifiy a Flask Response to delete the cookie containing an access JWT.
     Also deletes the corresponding CSRF cookie if applicable.
 
     :param response:
         A Flask Response object
+
+    :param domain:
+        The domain of the cookie. If this is None, it will use the
+        ``JWT_COOKIE_DOMAIN`` option (see :ref:`Configuration Options`). Otherwise,
+        it will use this as the cookies ``domain`` and the JWT_COOKIE_DOMAIN option
+        will be ignored.
     """
     response.set_cookie(
         config.refresh_cookie_name,
         value="",
         expires=0,
         secure=config.cookie_secure,
         httponly=True,
-        domain=config.cookie_domain,
+        domain=domain or config.cookie_domain,
         path=config.refresh_cookie_path,
         samesite=config.cookie_samesite,
     )
@@ -413,7 +437,7 @@ def unset_refresh_cookies(response):
             expires=0,
             secure=config.cookie_secure,
             httponly=False,
-            domain=config.cookie_domain,
+            domain=domain or config.cookie_domain,
             path=config.refresh_csrf_cookie_path,
             samesite=config.cookie_samesite,
         )

tests/test_cookies.py

@@ -1,6 +1,7 @@
 import pytest
 from flask import Flask
 from flask import jsonify
+from flask import request
 
 from flask_jwt_extended import create_access_token
 from flask_jwt_extended import create_refresh_token
@@ -35,34 +36,39 @@ def app():
 
     @app.route("/access_token", methods=["GET"])
     def access_token():
+        domain = request.args.get("domain")
         resp = jsonify(login=True)
         access_token = create_access_token("username")
-        set_access_cookies(resp, access_token)
+        set_access_cookies(resp, access_token, domain=domain)
         return resp
 
     @app.route("/refresh_token", methods=["GET"])
     def refresh_token():
+        domain = request.args.get("domain")
         resp = jsonify(login=True)
         refresh_token = create_refresh_token("username")
-        set_refresh_cookies(resp, refresh_token)
+        set_refresh_cookies(resp, refresh_token, domain=domain)
         return resp
 
     @app.route("/delete_tokens", methods=["GET"])
     def delete_tokens():
+        domain = request.args.get("domain")
         resp = jsonify(logout=True)
-        unset_jwt_cookies(resp)
+        unset_jwt_cookies(resp, domain=domain)
         return resp
 
     @app.route("/delete_access_tokens", methods=["GET"])
     def delete_access_tokens():
+        domain = request.args.get("domain")
         resp = jsonify(access_revoked=True)
-        unset_access_cookies(resp)
+        unset_access_cookies(resp, domain=domain)
         return resp
 
     @app.route("/delete_refresh_tokens", methods=["GET"])
     def delete_refresh_tokens():
+        domain = request.args.get("domain")
         resp = jsonify(refresh_revoked=True)
-        unset_refresh_cookies(resp)
+        unset_refresh_cookies(resp, domain=domain)
         return resp
 
     @app.route("/protected", methods=["GET"])
@@ -494,3 +500,54 @@ def test_jwt_optional_with_csrf_enabled(app):
     response = test_client.post("/optional_post_protected")
     assert response.status_code == 401
     assert response.get_json() == {"msg": "Missing CSRF token"}
+
+
+@pytest.mark.parametrize(
+    "options",
+    [
+        (
+            "/access_token",
+            "/delete_access_tokens",
+            "access_token_cookie",
+            "csrf_access_token",
+        ),
+        (
+            "/refresh_token",
+            "/delete_refresh_tokens",
+            "refresh_token_cookie",
+            "csrf_refresh_token",
+        ),
+    ],
+)
+def test_override_domain_option(app, options):
+    auth_url, delete_url, auth_cookie_name, csrf_cookie_name = options
+    domain = "yolo.com"
+
+    test_client = app.test_client()
+    app.config["JWT_COOKIE_DOMAIN"] = "test.com"
+
+    # Test set access cookies with custom domain
+    response = test_client.get(f"{auth_url}?domain={domain}")
+    cookies = response.headers.getlist("Set-Cookie")
+    assert len(cookies) == 2  # JWT and CSRF value
+
+    access_cookie = _get_cookie_from_response(response, auth_cookie_name)
+    assert access_cookie is not None
+    assert access_cookie["domain"] == domain
+
+    access_csrf_cookie = _get_cookie_from_response(response, csrf_cookie_name)
+    assert access_csrf_cookie is not None
+    assert access_csrf_cookie["domain"] == domain
+
+    # Test unset access cookies with custom domain
+    response = test_client.get(f"{delete_url}?domain={domain}")
+    cookies = response.headers.getlist("Set-Cookie")
+    assert len(cookies) == 2  # JWT and CSRF value
+
+    access_cookie = _get_cookie_from_response(response, auth_cookie_name)
+    assert access_cookie is not None
+    assert access_cookie["domain"] == domain
+
+    access_csrf_cookie = _get_cookie_from_response(response, csrf_cookie_name)
+    assert access_csrf_cookie is not None
+    assert access_csrf_cookie["domain"] == domain