test: Add test for invalid kid that doesnt match on refreshing

Author: KShivendu

supertokens_python/recipe/session/jwks.py

@@ -136,6 +136,10 @@ def get_latest_keys(kid: Optional[str] = None) -> List[PyJWK]:
             if cached_jwks is not None:  # we found a valid JWKS
                 cached_keys = CachedKeys(cached_jwks)
                 log_debug_message("Returning JWKS from fetch")
-                return cached_keys.keys
+                matching_keys = find_matching_keys(get_cached_keys(), kid)
+                if matching_keys is not None:
+                    return matching_keys
+
+                raise Exception("No matching JWKS found")
 
     raise last_error

supertokens_python/utils.py

@@ -353,10 +353,10 @@ def __enter__(self):
             self.mutex.lock()
 
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
-        if exc_type is not None:
-            raise exc_type(exc_value).with_traceback(traceback)
-
         if self.read:
             self.mutex.r_unlock()
         else:
             self.mutex.unlock()
+
+        if exc_type is not None:
+            raise exc_type(exc_value).with_traceback(traceback)

tests/sessions/test_jwks.py

@@ -2,6 +2,7 @@
 import pytest
 import logging
 import threading
+import json
 import requests
 
 from typing import List, Any, Callable
@@ -29,6 +30,7 @@
     get_cached_keys,
     get_latest_keys,
 )
+from supertokens_python.utils import utf_base64encode
 
 from _pytest.logging import LogCaptureFixture
 
@@ -198,6 +200,26 @@ async def test_that_jwks_are_refresh_if_kid_is_unknown(caplog: LogCaptureFixture
 
     assert next(well_known_count) == 2  # no change
 
+    # use an access token with an invalid kid that doesn't match even on refreshing
+    _, payload, signature = tokens["accessToken"].split(".")
+    new_header = utf_base64encode(
+        json.dumps(
+            {"kid": "d-1234567890123", "typ": "JWT", "version": "3", "alg": "RS256"}
+        ),
+        urlsafe=False,
+    )
+    new_token = ".".join([new_header, payload, signature])
+
+    with pytest.raises(Exception) as e:
+        await get_session_without_request_response(
+            new_token, tokens.get("antiCsrfToken")
+        )
+
+    assert str(e.value) == "No matching JWKS found"
+    assert (
+        next(well_known_count) == 3
+    )  # kid not found in cache, so should have refreshed
+
 
 async def test_that_invalid_connection_uri_doesnot_throw_during_init_for_jwks():
     """This test makes sure that initialising SuperTokens and Session with an invalid connection uri does not