set cookie in response directly

Author: IndominusByte

fastapi_jwt_auth/auth_jwt.py

@@ -315,12 +315,18 @@ def _get_csrf_token(self,encoded_token: str) -> str:
         """
         return self._verified_token(encoded_token)['csrf']
 
-    def set_access_cookies(self,encoded_access_token: str, max_age: Optional[int] = None) -> None:
+    def set_access_cookies(
+        self,
+        encoded_access_token: str,
+        response: Optional[Response] = None,
+        max_age: Optional[int] = None
+    ) -> None:
         """
         Configures the response to set access token in a cookie.
         this will also set the CSRF double submit values in a separate cookie
 
         :param encoded_access_token: The encoded access token to set in the cookies
+        :param response: The FastAPI response object to set the access cookies in
         :param max_age: The max age of the cookie value should be the number of seconds (integer)
         """
         if not self.jwt_in_cookies:
@@ -330,9 +336,13 @@ def set_access_cookies(self,encoded_access_token: str, max_age: Optional[int] =
 
         if max_age and not isinstance(max_age,int):
             raise TypeError("max_age must be a integer")
+        if response and not isinstance(response,Response):
+            raise TypeError("The response must be an object response FastAPI")
+
+        response = response or self._response
 
         # Set the access JWT in the cookie
-        self._response.set_cookie(
+        response.set_cookie(
             self._access_cookie_key,
             encoded_access_token,
             max_age=max_age or self._cookie_max_age,
@@ -345,7 +355,7 @@ def set_access_cookies(self,encoded_access_token: str, max_age: Optional[int] =
 
         # If enabled, set the csrf double submit access cookie
         if self._cookie_csrf_protect:
-            self._response.set_cookie(
+            response.set_cookie(
                 self._access_csrf_cookie_key,
                 self._get_csrf_token(encoded_access_token),
                 max_age=max_age or self._cookie_max_age,
@@ -356,12 +366,18 @@ def set_access_cookies(self,encoded_access_token: str, max_age: Optional[int] =
                 samesite=self._cookie_samesite
             )
 
-    def set_refresh_cookies(self, encoded_refresh_token: str, max_age: Optional[int] = None) -> None:
+    def set_refresh_cookies(
+        self,
+        encoded_refresh_token: str,
+        response: Optional[Response] = None,
+        max_age: Optional[int] = None
+    ) -> None:
         """
         Configures the response to set refresh token in a cookie.
         this will also set the CSRF double submit values in a separate cookie
 
         :param encoded_refresh_token: The encoded refresh token to set in the cookies
+        :param response: The FastAPI response object to set the refresh cookies in
         :param max_age: The max age of the cookie value should be the number of seconds (integer)
         """
         if not self.jwt_in_cookies:
@@ -371,9 +387,13 @@ def set_refresh_cookies(self, encoded_refresh_token: str, max_age: Optional[int]
 
         if max_age and not isinstance(max_age,int):
             raise TypeError("max_age must be a integer")
+        if response and not isinstance(response,Response):
+            raise TypeError("The response must be an object response FastAPI")
+
+        response = response or self._response
 
         # Set the refresh JWT in the cookie
-        self._response.set_cookie(
+        response.set_cookie(
             self._refresh_cookie_key,
             encoded_refresh_token,
             max_age=max_age or self._cookie_max_age,
@@ -386,7 +406,7 @@ def set_refresh_cookies(self, encoded_refresh_token: str, max_age: Optional[int]
 
         # If enabled, set the csrf double submit refresh cookie
         if self._cookie_csrf_protect:
-            self._response.set_cookie(
+            response.set_cookie(
                 self._refresh_csrf_cookie_key,
                 self._get_csrf_token(encoded_refresh_token),
                 max_age=max_age or self._cookie_max_age,
@@ -397,52 +417,68 @@ def set_refresh_cookies(self, encoded_refresh_token: str, max_age: Optional[int]
                 samesite=self._cookie_samesite
             )
 
-    def unset_jwt_cookies(self) -> None:
+    def unset_jwt_cookies(self,response: Optional[Response] = None) -> None:
         """
         Unset (delete) all jwt stored in a cookie
+
+        :param response: The FastAPI response object to delete the JWT cookies in.
         """
-        self.unset_access_cookies()
-        self.unset_refresh_cookies()
+        self.unset_access_cookies(response)
+        self.unset_refresh_cookies(response)
 
-    def unset_access_cookies(self) -> None:
+    def unset_access_cookies(self,response: Optional[Response] = None) -> None:
         """
         Remove access token and access CSRF double submit from the response cookies
+
+        :param response: The FastAPI response object to delete the access cookies in.
         """
         if not self.jwt_in_cookies:
             raise RuntimeWarning(
                 "unset_access_cookies() called without 'authjwt_token_location' configured to use cookies"
             )
 
-        self._response.delete_cookie(
+        if response and not isinstance(response,Response):
+            raise TypeError("The response must be an object response FastAPI")
+
+        response = response or self._response
+
+        response.delete_cookie(
             self._access_cookie_key,
             path=self._access_cookie_path,
             domain=self._cookie_domain
         )
 
         if self._cookie_csrf_protect:
-            self._response.delete_cookie(
+            response.delete_cookie(
                 self._access_csrf_cookie_key,
                 path=self._access_csrf_cookie_path,
                 domain=self._cookie_domain
             )
 
-    def unset_refresh_cookies(self) -> None:
+    def unset_refresh_cookies(self,response: Optional[Response] = None) -> None:
         """
         Remove refresh token and refresh CSRF double submit from the response cookies
+
+        :param response: The FastAPI response object to delete the refresh cookies in.
         """
         if not self.jwt_in_cookies:
             raise RuntimeWarning(
                 "unset_refresh_cookies() called without 'authjwt_token_location' configured to use cookies"
             )
 
-        self._response.delete_cookie(
+        if response and not isinstance(response,Response):
+            raise TypeError("The response must be an object response FastAPI")
+
+        response = response or self._response
+
+        response.delete_cookie(
             self._refresh_cookie_key,
             path=self._refresh_cookie_path,
             domain=self._cookie_domain
         )
 
         if self._cookie_csrf_protect:
-            self._response.delete_cookie(
+            response.delete_cookie(
                 self._refresh_csrf_cookie_key,
                 path=self._refresh_csrf_cookie_path,
                 domain=self._cookie_domain

tests/test_cookies.py

@@ -24,23 +24,52 @@ def all_token(Authorize: AuthJWT = Depends()):
         Authorize.set_refresh_cookies(refresh_token)
         return {"msg":"all token"}
 
+    @app.get('/all-token-response')
+    def all_token_response(Authorize: AuthJWT = Depends()):
+        access_token = Authorize.create_access_token(subject=1,fresh=True)
+        refresh_token = Authorize.create_refresh_token(subject=1)
+        response = JSONResponse(content={"msg":"all token"})
+        Authorize.set_access_cookies(access_token,response)
+        Authorize.set_refresh_cookies(refresh_token,response)
+        return response
+
     @app.get('/access-token')
     def access_token(Authorize: AuthJWT = Depends()):
         access_token = Authorize.create_access_token(subject=1)
         Authorize.set_access_cookies(access_token)
         return {"msg":"access token"}
 
+    @app.get('/access-token-response')
+    def access_token_response(Authorize: AuthJWT = Depends()):
+        access_token = Authorize.create_access_token(subject=1)
+        response = JSONResponse(content={"msg":"access token"})
+        Authorize.set_access_cookies(access_token,response)
+        return response
+
     @app.get('/refresh-token')
     def refresh_token(Authorize: AuthJWT = Depends()):
         refresh_token = Authorize.create_refresh_token(subject=1)
         Authorize.set_refresh_cookies(refresh_token)
         return {"msg":"refresh token"}
 
+    @app.get('/refresh-token-response')
+    def refresh_token_response(Authorize: AuthJWT = Depends()):
+        refresh_token = Authorize.create_refresh_token(subject=1)
+        response = JSONResponse(content={"msg":"refresh token"})
+        Authorize.set_refresh_cookies(refresh_token,response)
+        return response
+
     @app.get('/unset-all-token')
     def unset_all_token(Authorize: AuthJWT = Depends()):
         Authorize.unset_jwt_cookies()
         return {"msg":"unset all token"}
 
+    @app.get('/unset-all-token-response')
+    def unset_all_token_response(Authorize: AuthJWT = Depends()):
+        response = JSONResponse(content={"msg":"unset all token"})
+        Authorize.unset_jwt_cookies(response)
+        return response
+
     @app.get('/unset-access-token')
     def unset_access_token(Authorize: AuthJWT = Depends()):
         Authorize.unset_access_cookies()
@@ -96,7 +125,29 @@ def get_cookie_location():
     with pytest.raises(TypeError,match=r"max_age"):
         Authorize.set_refresh_cookies(token,max_age="string")
 
-@pytest.mark.parametrize("url",["/access-token","/refresh-token"])
+def test_set_unset_cookies_not_valid_type_response(Authorize):
+    @AuthJWT.load_config
+    def get_cookie_location():
+        return [("authjwt_token_location",{'cookies'}),("authjwt_secret_key","secret")]
+
+    token = Authorize.create_access_token(subject=1)
+
+    with pytest.raises(TypeError,match=r"response"):
+        Authorize.set_access_cookies(token,response={"msg":"hello"})
+
+    with pytest.raises(TypeError,match=r"response"):
+        Authorize.set_refresh_cookies(token,response={"msg":"hello"})
+
+    with pytest.raises(TypeError,match=r"response"):
+        Authorize.unset_jwt_cookies({"msg":"hello"})
+
+    with pytest.raises(TypeError,match=r"response"):
+        Authorize.unset_access_cookies({"msg":"hello"})
+
+    with pytest.raises(TypeError,match=r"response"):
+        Authorize.unset_refresh_cookies({"msg":"hello"})
+
+@pytest.mark.parametrize("url",["/access-token","/refresh-token","/access-token-response","/refresh-token-response"])
 def test_set_cookie_csrf_protect_false(url,client):
     @AuthJWT.load_config
     def get_cookie_location():
@@ -110,7 +161,7 @@ def get_cookie_location():
     response = client.get(url)
     assert response.cookies.get("csrf_{}_token".format(cookie_key)) is None
 
-@pytest.mark.parametrize("url",["/access-token","/refresh-token"])
+@pytest.mark.parametrize("url",["/access-token","/refresh-token","/access-token-response","/refresh-token-response"])
 def test_set_cookie_csrf_protect_true(url,client):
     @AuthJWT.load_config
     def get_cookie_location():
@@ -140,6 +191,26 @@ def get_cookie_location():
     assert response.cookies.get("refresh_token_cookie") is None
     assert response.cookies.get("csrf_refresh_token") is None
 
+def test_unset_all_cookie_response(client):
+    @AuthJWT.load_config
+    def get_cookie_location():
+        return [("authjwt_token_location",{'cookies'}),("authjwt_secret_key","secret")]
+
+    response = client.get('/all-token-response')
+    assert response.cookies.get("access_token_cookie") is not None
+    assert response.cookies.get("csrf_access_token") is not None
+
+    assert response.cookies.get("refresh_token_cookie") is not None
+    assert response.cookies.get("csrf_refresh_token") is not None
+
+    response = client.get('/unset-all-token-response')
+
+    assert response.cookies.get("access_token_cookie") is None
+    assert response.cookies.get("csrf_access_token") is None
+
+    assert response.cookies.get("refresh_token_cookie") is None
+    assert response.cookies.get("csrf_refresh_token") is None
+
 def test_custom_cookie_key(client):
     @AuthJWT.load_config
     def get_cookie_location():