From c4e2a8b0bd4a7d4c8b56b8b8aa666a8121c13a0f Mon Sep 17 00:00:00 2001 From: MB Date: Wed, 25 Aug 2021 17:58:05 +0200 Subject: [PATCH] Feat: override cookie domain from cookie utils --- flask_jwt_extended/utils.py | 54 +++++++++++++++++++++--------- tests/test_cookies.py | 67 ++++++++++++++++++++++++++++++++++--- 2 files changed, 101 insertions(+), 20 deletions(-) diff --git a/flask_jwt_extended/utils.py b/flask_jwt_extended/utils.py index 48ae99ba..a670263d 100644 --- a/flask_jwt_extended/utils.py +++ b/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,6 +276,12 @@ 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, @@ -283,7 +289,7 @@ 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=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,6 +324,12 @@ 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, @@ -325,7 +337,7 @@ 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=True, - domain=config.cookie_domain, + domain=domain or config.cookie_domain, path=config.refresh_cookie_path, samesite=config.cookie_samesite, ) @@ -337,13 +349,13 @@ 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. @@ -351,17 +363,23 @@ def unset_jwt_cookies(response): :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, @@ -369,7 +387,7 @@ def unset_access_cookies(response): 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,19 +399,25 @@ 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, @@ -401,7 +425,7 @@ def unset_refresh_cookies(response): 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, ) diff --git a/tests/test_cookies.py b/tests/test_cookies.py index 9a8434ee..9cc33247 100644 --- a/tests/test_cookies.py +++ b/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