fix: Clear session changes for edge cases

Author: KShivendu

supertokens_python/framework/django/django_response.py

@@ -14,7 +14,7 @@
 import json
 from datetime import datetime
 from math import ceil
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional
 
 from supertokens_python.framework.response import BaseResponse
 
@@ -42,7 +42,7 @@ def set_cookie(
         value: str,
         expires: int,
         path: str = "/",
-        domain: Union[str, None] = None,
+        domain: Optional[str] = None,
         secure: bool = False,
         httponly: bool = False,
         samesite: str = "lax",
@@ -69,11 +69,14 @@ def set_status_code(self, status_code: int):
     def set_header(self, key: str, value: str):
         self.response[key] = value
 
-    def get_header(self, key: str):
+    def get_header(self, key: str) -> Optional[str]:
         if self.response.has_header(key):
             return self.response[key]
         return None
 
+    def remove_header(self, key: str):
+        del self.response[key]
+
     def set_json_content(self, content: Dict[str, Any]):
         if not self.response_sent:
             self.set_header("Content-Type", "application/json; charset=utf-8")

supertokens_python/framework/fastapi/fastapi_response.py

@@ -14,7 +14,7 @@
 import json
 from math import ceil
 from time import time
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional
 
 from supertokens_python.framework.response import BaseResponse
 
@@ -44,7 +44,7 @@ def set_cookie(
         value: str,
         expires: int,
         path: str = "/",
-        domain: Union[str, None] = None,
+        domain: Optional[str] = None,
         secure: bool = False,
         httponly: bool = False,
         samesite: str = "lax",
@@ -78,9 +78,12 @@ def set_cookie(
     def set_header(self, key: str, value: str):
         self.response.headers[key] = value
 
-    def get_header(self, key: str) -> Union[str, None]:
+    def get_header(self, key: str) -> Optional[str]:
         return self.response.headers.get(key, None)
 
+    def remove_header(self, key: str):
+        del self.response.headers[key]
+
     def set_status_code(self, status_code: int):
         if not self.status_set:
             self.response.status_code = status_code

supertokens_python/framework/flask/flask_response.py

@@ -12,7 +12,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 import json
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional
 
 from supertokens_python.framework.response import BaseResponse
 
@@ -40,7 +40,7 @@ def set_cookie(
         value: str,
         expires: int,
         path: str = "/",
-        domain: Union[str, None] = None,
+        domain: Optional[str] = None,
         secure: bool = False,
         httponly: bool = False,
         samesite: str = "lax",
@@ -59,9 +59,12 @@ def set_cookie(
     def set_header(self, key: str, value: str):
         self.response.headers.set(key, value)
 
-    def get_header(self, key: str) -> Union[None, str]:
+    def get_header(self, key: str) -> Optional[str]:
         return self.response.headers.get(key)
 
+    def remove_header(self, key: str):
+        del self.response.headers[key]
+
     def set_status_code(self, status_code: int):
         if not self.status_set:
             self.response.status_code = status_code

supertokens_python/framework/response.py

@@ -13,7 +13,7 @@
 # under the License.
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional
 
 
 class BaseResponse(ABC):
@@ -31,19 +31,23 @@ def set_cookie(
         #    max_age: Union[int, None] = None,
         expires: int,
         path: str = "/",
-        domain: Union[str, None] = None,
+        domain: Optional[str] = None,
         secure: bool = False,
         httponly: bool = False,
         samesite: str = "lax",
     ):
         pass
 
     @abstractmethod
-    def set_header(self, key: str, value: str):
+    def set_header(self, key: str, value: str) -> None:
         pass
 
     @abstractmethod
-    def get_header(self, key: str) -> Union[str, None]:
+    def get_header(self, key: str) -> Optional[str]:
+        pass
+
+    @abstractmethod
+    def remove_header(self, key: str) -> None:
         pass
 
     @abstractmethod

supertokens_python/recipe/session/cookie_and_header.py

@@ -97,6 +97,11 @@ def set_header(response: BaseResponse, key: str, value: str, allow_duplicate: bo
         response.set_header(key, value)
 
 
+def remove_header(response: BaseResponse, key: str):
+    if response.get_header(key) is not None:
+        response.remove_header(key)
+
+
 def get_cookie(request: BaseRequest, key: str):
     cookie_val = request.get_cookie(key)
     if cookie_val is None:
@@ -171,23 +176,31 @@ def get_rid_header(request: BaseRequest):
 
 
 def clear_session_from_all_token_transfer_methods(
-    response: BaseResponse, recipe: SessionRecipe, request: BaseRequest
+    response: BaseResponse, recipe: SessionRecipe
 ):
+    # We are clearing the session in all transfermethods to be sure to override cookies in case they have been already added to the response.
+    # This is done to handle the following use-case:
+    # If the app overrides signInPOST to check the ban status of the user after the original implementation and throwing an UNAUTHORISED error
+    # In this case: the SDK has attached cookies to the response, but none was sent with the request
+    # We can't know which to clear since we can't reliably query or remove the set-cookie header added to the response (causes issues in some frameworks, i.e.: hapi)
+    # The safe solution in this case is to overwrite all the response cookies/headers with an empty value, which is what we are doing here.
     for transfer_method in available_token_transfer_methods:
-        if get_token(request, "access", transfer_method) is not None:
-            _clear_session(response, recipe.config, transfer_method)
+        _clear_session(response, recipe.config, transfer_method)
 
 
 def _clear_session(
     response: BaseResponse,
     config: SessionConfig,
     transfer_method: TokenTransferMethod,
 ):
-    # If we can tell it's a cookie based session we are not clearing using headers
+    # If we can be specific about which transferMethod we want to clear, there is no reason to clear the other ones
     token_types: List[TokenType] = ["access", "refresh"]
     for token_type in token_types:
         _set_token(response, config, token_type, "", 0, transfer_method)
 
+    remove_header(
+        response, ANTI_CSRF_HEADER_KEY
+    )  # This can be added multiple times in some cases, but that should be OK
     set_header(response, FRONT_TOKEN_HEADER_SET_KEY, "remove", False)
     set_header(
         response, ACCESS_CONTROL_EXPOSE_HEADERS, FRONT_TOKEN_HEADER_SET_KEY, True

supertokens_python/recipe/session/interfaces.py

@@ -34,7 +34,9 @@
 from .utils import SessionConfig, TokenTransferMethod
 
 if TYPE_CHECKING:
-    from supertokens_python.framework import BaseRequest, BaseResponse
+    from supertokens_python.framework import BaseRequest
+
+from supertokens_python.framework import BaseResponse
 
 
 class SessionObj:

supertokens_python/recipe/session/utils.py

@@ -138,7 +138,7 @@ async def on_token_theft_detected(
         response: BaseResponse,
     ) -> BaseResponse:
         log_debug_message("Clearing tokens because of TOKEN_THEFT_DETECTED response")
-        clear_session_from_all_token_transfer_methods(response, recipe, request)
+        clear_session_from_all_token_transfer_methods(response, recipe)
         return await resolve(
             self.__on_token_theft_detected(request, session_handle, user_id, response)
         )
@@ -159,7 +159,7 @@ async def on_unauthorised(
     ):
         if do_clear_cookies:
             log_debug_message("Clearing tokens because of UNAUTHORISED response")
-            clear_session_from_all_token_transfer_methods(response, recipe, request)
+            clear_session_from_all_token_transfer_methods(response, recipe)
         return await resolve(self.__on_unauthorised(request, message, response))
 
     async def on_invalid_claim(

tests/Fastapi/test_fastapi.py

@@ -31,15 +31,20 @@
     get_session,
     refresh_session,
 )
+from supertokens_python.recipe.session.exceptions import UnauthorisedError
 from supertokens_python.recipe.session.framework.fastapi import verify_session
 from supertokens_python.recipe.session.interfaces import APIInterface
+from supertokens_python.recipe.session.interfaces import APIOptions as SessionAPIOptions
 from tests.utils import (
     TEST_DRIVER_CONFIG_ACCESS_TOKEN_PATH,
     TEST_DRIVER_CONFIG_COOKIE_DOMAIN,
     TEST_DRIVER_CONFIG_COOKIE_SAME_SITE,
     TEST_DRIVER_CONFIG_REFRESH_TOKEN_PATH,
+    assert_info_clears_tokens,
     clean_st,
     extract_all_cookies,
+    extract_info,
+    get_st_init_args,
     reset,
     setup_st,
     start_st,
@@ -107,6 +112,16 @@ async def custom_logout(request: Request):  # type: ignore
         await session.revoke_session()
         return {}  # type: ignore
 
+    @app.post("/create")
+    async def _create(request: Request):  # type: ignore
+        await create_new_session(request, "userId", {}, {})
+        return ""
+
+    @app.post("/create-throw")
+    async def _create_throw(request: Request):  # type: ignore
+        await create_new_session(request, "userId", {}, {})
+        raise UnauthorisedError("unauthorised")
+
     return TestClient(app)
 
 
@@ -524,3 +539,87 @@ def test_fastapi_root_path(fastapi_root_path: str):
     # The API should migrate (and return 404 here)
     response = test_client.get("/auth/signup/email/[email protected]")
     assert response.status_code == 404
+
+
+@mark.asyncio
+@mark.parametrize("token_transfer_method", ["cookie", "header"])
+async def test_should_clear_all_response_during_refresh_if_unauthorized(
+    driver_config_client: TestClient, token_transfer_method: str
+):
+    def override_session_apis(oi: APIInterface):
+        oi_refresh_post = oi.refresh_post
+
+        async def refresh_post(
+            api_options: SessionAPIOptions, user_context: Dict[str, Any]
+        ):
+            await oi_refresh_post(api_options, user_context)
+            raise UnauthorisedError("unauthorized", clear_tokens=True)
+
+        oi.refresh_post = refresh_post
+        return oi
+
+    init(
+        **get_st_init_args(
+            [
+                session.init(
+                    anti_csrf="VIA_TOKEN",
+                    override=session.InputOverrideConfig(apis=override_session_apis),
+                )
+            ]
+        )
+    )  # type: ignore
+    start_st()
+
+    res = driver_config_client.post(
+        "/create", headers={"st-auth-mode": token_transfer_method}
+    )
+    info = extract_info(res)
+
+    assert info["accessTokenFromAny"] is not None
+    assert info["refreshTokenFromAny"] is not None
+
+    headers: Dict[str, Any] = {}
+    cookies: Dict[str, Any] = {}
+
+    if token_transfer_method == "header":
+        headers.update({"authorization": f"Bearer {info['refreshTokenFromAny']}"})
+    else:
+        cookies.update(
+            {"sRefreshToken": info["refreshTokenFromAny"], "sIdRefreshToken": "asdf"}
+        )
+
+    if info["antiCsrf"] is not None:
+        headers.update({"anti-csrf": info["antiCsrf"]})
+
+    res = driver_config_client.post(
+        "/auth/session/refresh", headers=headers, cookies=cookies
+    )
+    info = extract_info(res)
+
+    assert res.status_code == 401
+    assert_info_clears_tokens(info, token_transfer_method)
+
+
+@mark.asyncio
+@mark.parametrize("token_transfer_method", ["cookie", "header"])
+async def test_revoking_session_after_create_new_session_with_throwing_unauthorized_error(
+    driver_config_client: TestClient, token_transfer_method: str
+):
+    init(
+        **get_st_init_args(
+            [
+                session.init(
+                    anti_csrf="VIA_TOKEN",
+                )
+            ]
+        )
+    )  # type: ignore
+    start_st()
+
+    res = driver_config_client.post(
+        "/create-throw", headers={"st-auth-mode": token_transfer_method}
+    )
+    info = extract_info(res)
+
+    assert res.status_code == 401
+    assert_info_clears_tokens(info, token_transfer_method)

tests/utils.py

@@ -239,6 +239,9 @@ def extract_info(response: Response) -> Dict[str, Any]:
     access_token = cookies.get("sAccessToken", {}).get("value")
     refresh_token = cookies.get("sRefreshToken", {}).get("value")
 
+    access_token_from_header = response.headers.get("st-access-token")
+    refresh_token_from_header = response.headers.get("st-refresh-token")
+
     return {
         **cookies,
         "accessToken": access_token,
@@ -247,11 +250,31 @@ def extract_info(response: Response) -> Dict[str, Any]:
         "status_code": response.status_code,
         "body": response.json(),
         "antiCsrf": response.headers.get("anti-csrf"),
-        "accessTokenFromHeader": response.headers.get("st-access-token"),
-        "refreshTokenFromHeader": response.headers.get("st-refresh-token"),
+        "accessTokenFromHeader": access_token_from_header,
+        "refreshTokenFromHeader": refresh_token_from_header,
+        "accessTokenFromAny": access_token_from_header or access_token,
+        "refreshTokenFromAny": refresh_token_from_header or refresh_token,
     }
 
 
+def assert_info_clears_tokens(info: Dict[str, Any], token_transfer_method: str):
+    if token_transfer_method == "cookie":
+        assert info["accessToken"] == ""
+        assert info["refreshToken"] == ""
+        assert info["sAccessToken"]["expires"] == "Thu, 01 Jan 1970 00:00:00 GMT"
+        assert info["sRefreshToken"]["expires"] == "Thu, 01 Jan 1970 00:00:00 GMT"
+        assert info["sAccessToken"]["domain"] == ""
+        assert info["sRefreshToken"]["domain"] == ""
+    elif token_transfer_method == "header":
+        assert info["accessTokenFromHeader"] == ""
+        assert info["refreshTokenFromHeader"] == ""
+    else:
+        raise Exception("unknown token transfer method: " + token_transfer_method)
+
+    assert info["frontToken"] == "remove"
+    assert info["antiCsrf"] is None
+
+
 def get_unix_timestamp(expiry: str):
     return int(
         datetime.strptime(expiry, "%a, %d %b %Y %H:%M:%S GMT")