tests: Add more tests and minor refactors

Author: KShivendu

CHANGELOG.md

@@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [unreleased]
 
+### Changes
+
+- Use `useStaticSigningKey` instead of `use_static_signing_key` in `create_jwt` function. This was a bug in the code.
+
 
 ## [0.14.3] - 2023-06-7
 

supertokens_python/recipe/session/access_token.py

@@ -52,22 +52,33 @@ def get_info_from_access_token(
 ):
     try:
         payload: Optional[Dict[str, Any]] = None
+        decode_algo = (
+            jwt_info.parsed_header["alg"]
+            if jwt_info.parsed_header is not None
+            else "RS256"
+        )
+
         if jwt_info.version >= 3:
             matching_key = jwk_client.get_matching_key_from_jwt(
                 jwt_info.raw_token_string
             )
             payload = jwt.decode(  # type: ignore
                 jwt_info.raw_token_string,
                 matching_key.key,  # type: ignore
-                algorithms=["RS256"],
+                algorithms=[decode_algo],
                 options={"verify_signature": True, "verify_exp": True},
             )
         else:
             # It won't have kid. So we'll have to try the token against all the keys from all the jwk_clients
             # If any of them work, we'll use that payload
             for k in jwk_client.get_latest_keys():
                 try:
-                    payload = jwt.decode(jwt_info.raw_token_string, k.key, algorithms=["RS256"])  # type: ignore
+                    payload = jwt.decode(  # type: ignore
+                        jwt_info.raw_token_string,
+                        k.key,  # type: ignore
+                        algorithms=[decode_algo],
+                        options={"verify_signature": True, "verify_exp": True},
+                    )
                     break
                 except DecodeError:
                     pass

supertokens_python/recipe/session/jwt.py

@@ -42,6 +42,7 @@ def __init__(
         payload: Dict[str, Any],
         signature: str,
         kid: Optional[str],
+        parsed_header: Optional[Dict[str, Any]] = None,
     ) -> None:
         self.version = version
         self.raw_token_string = raw_token_string
@@ -50,24 +51,26 @@ def __init__(
         self.payload = payload
         self.signature = signature
         self.kid = kid
+        self.parsed_header = parsed_header
 
 
 def parse_jwt_without_signature_verification(jwt: str) -> ParsedJWTInfo:
     splitted_input = jwt.split(".")
-    latest_access_token_version = 3
+    TOKEN_V3 = 3
     if len(splitted_input) != 3:
         raise Exception("invalid jwt")
 
     # V1 and V2 are functionally identical, plus all legacy tokens should be V2 now.
     # So we can assume these defaults:
     version = 2
     kid = None
+    parsed_header = None
     # V2 or older tokens didn't save the key id
     header, payload, signature = splitted_input
     # checking the header
     if header not in _allowed_headers:
         parsed_header = loads(utf_base64decode(header, True))
-        header_version = parsed_header.get("version", str(latest_access_token_version))
+        header_version = parsed_header.get("version", str(TOKEN_V3))
 
         try:
             version = int(header_version)
@@ -79,7 +82,7 @@ def parse_jwt_without_signature_verification(jwt: str) -> ParsedJWTInfo:
         if (
             parsed_header["typ"] != "JWT"
             or not isinstance(version, int)
-            or version < latest_access_token_version
+            or version < TOKEN_V3
             or kid is None
         ):
             raise Exception("JWT header mismatch")
@@ -94,4 +97,5 @@ def parse_jwt_without_signature_verification(jwt: str) -> ParsedJWTInfo:
         payload=loads(utf_base64decode(payload, True)),
         signature=signature,
         kid=kid,
+        parsed_header=parsed_header,
     )

supertokens_python/recipe/session/recipe_implementation.py

@@ -18,7 +18,7 @@
 
 from supertokens_python.logger import log_debug_message
 from supertokens_python.normalised_url_path import NormalisedURLPath
-from supertokens_python.utils import resolve
+from supertokens_python.utils import resolve, RWMutex, RWLockContext
 
 from ...types import MaybeAwaitable
 from . import session_functions
@@ -52,7 +52,6 @@
 
 
 from typing_extensions import TypedDict
-import threading
 from os import environ
 
 
@@ -66,61 +65,6 @@ class JWKSConfigType(TypedDict):
     "refresh_rate_limit": 500,
 }
 
-
-class RWMutex:
-    def __init__(self):
-        self._lock = threading.Lock()
-        self._readers = threading.Condition(self._lock)
-        self._writers = threading.Condition(self._lock)
-        self._reader_count = 0
-        self._writer_count = 0
-
-    def lock(self):
-        with self._lock:
-            while self._writer_count > 0 or self._reader_count > 0:
-                self._writers.wait()
-            self._writer_count += 1
-
-    def unlock(self):
-        with self._lock:
-            self._writer_count -= 1
-            self._readers.notify_all()
-            self._writers.notify_all()
-
-    def r_lock(self):
-        with self._lock:
-            while self._writer_count > 0:
-                self._readers.wait()
-            self._reader_count += 1
-
-    def r_unlock(self):
-        with self._lock:
-            self._reader_count -= 1
-            if self._reader_count == 0:
-                self._writers.notify_all()
-
-
-class RWLockContext:
-    def __init__(self, mutex: RWMutex, read: bool = True):
-        self.mutex = mutex
-        self.read = read
-
-    def __enter__(self):
-        if self.read:
-            self.mutex.r_lock()
-        else:
-            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()
-
-
 cached_jwk_client: Optional[JWKClient] = None
 mutex = RWMutex()
 

supertokens_python/utils.py

@@ -16,6 +16,7 @@
 
 import asyncio
 import json
+import threading
 import warnings
 from base64 import urlsafe_b64decode, urlsafe_b64encode, b64encode, b64decode
 from math import floor
@@ -305,3 +306,57 @@ def get_top_level_domain_for_same_site_resolution(url: str) -> str:
         )
 
     return parsed_url.domain + "." + parsed_url.suffix  # type: ignore
+
+
+class RWMutex:
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._readers = threading.Condition(self._lock)
+        self._writers = threading.Condition(self._lock)
+        self._reader_count = 0
+        self._writer_count = 0
+
+    def lock(self):
+        with self._lock:
+            while self._writer_count > 0 or self._reader_count > 0:
+                self._writers.wait()
+            self._writer_count += 1
+
+    def unlock(self):
+        with self._lock:
+            self._writer_count -= 1
+            self._readers.notify_all()
+            self._writers.notify_all()
+
+    def r_lock(self):
+        with self._lock:
+            while self._writer_count > 0:
+                self._readers.wait()
+            self._reader_count += 1
+
+    def r_unlock(self):
+        with self._lock:
+            self._reader_count -= 1
+            if self._reader_count == 0:
+                self._writers.notify_all()
+
+
+class RWLockContext:
+    def __init__(self, mutex: RWMutex, read: bool = True):
+        self.mutex = mutex
+        self.read = read
+
+    def __enter__(self):
+        if self.read:
+            self.mutex.r_lock()
+        else:
+            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()

tests/sessions/test_jwks.py

@@ -462,6 +462,43 @@ async def test_session_verification_of_jwt_based_on_session_payload_with_check_d
     assert s_.get_user_id() == "userId"
 
 
+async def test_session_verification_of_jwt_with_dynamic_signing_key():
+    init(
+        **get_st_init_args(
+            recipe_list=[session.init(use_dynamic_access_token_signing_key=False)]
+        )
+    )
+    start_st()
+
+    s = await create_new_session_without_request_response("userId", {}, {})
+
+    payload = s.get_access_token_payload()
+    del payload["iat"]
+    del payload["exp"]
+    payload["tId"] = "public"  # tenant id
+
+    now = get_timestamp_ms()
+    jwt_expiry = now + 10 * 1000  # expiry jwt after 10sec
+
+    jwt_with_dynamic_key = await create_jwt(
+        payload, jwt_expiry, use_static_signing_key=False
+    )
+    assert isinstance(jwt_with_dynamic_key, CreateJwtOkResult)
+    try:
+        await get_session_without_request_response(jwt_with_dynamic_key.jwt)
+        assert False
+    except Exception:
+        pass
+
+    jwt_with_static_key = await create_jwt(
+        payload, jwt_expiry, use_static_signing_key=True
+    )
+    assert isinstance(jwt_with_static_key, CreateJwtOkResult)
+    s_ = await get_session_without_request_response(jwt_with_static_key.jwt)
+    assert s_ is not None
+    assert s_.get_user_id() == "userId"
+
+
 async def test_that_locking_for_jwks_cache_works(caplog: LogCaptureFixture):
     caplog.set_level(logging.DEBUG)
     not_returned_from_cache_count = get_log_occurence_count(

tests/test_utils.py

@@ -1,7 +1,11 @@
 from typing import Union, List, Any, Dict
 
 import pytest
+import threading
+
 from supertokens_python.utils import humanize_time, is_version_gte
+from supertokens_python.utils import RWMutex
+
 from tests.utils import is_subset
 
 
@@ -102,3 +106,62 @@ def test_is_subset(
         assert is_subset(d1, d2)
     else:
         assert not is_subset(d1, d2)
+
+
+class BankAccount:
+    def __init__(self):
+        self.balance = 0
+        self.mutex = RWMutex()
+
+    def deposit(self, amount: int):
+        self.mutex.lock()
+        self.balance += amount
+        self.mutex.unlock()
+
+    def withdraw(self, amount: int):
+        self.mutex.lock()
+        self.balance -= amount
+        self.mutex.unlock()
+
+    def get_balance(self):
+        self.mutex.r_lock()
+        balance = self.balance
+        self.mutex.r_unlock()
+        return balance
+
+
+def test_rw_mutex_writes():
+    account = BankAccount()
+    threads: List[threading.Thread] = []
+
+    # Create 10 deposit threads
+    for _ in range(10):
+        t = threading.Thread(target=account.deposit, args=(10,))
+        threads.append(t)
+
+    def balance_is_valid():
+        balance = account.get_balance()
+        assert balance % 5 == 0 and balance >= 0
+
+    # Create 15 balance checking threads
+    for _ in range(15):
+        t = threading.Thread(target=balance_is_valid)
+        threads.append(t)
+
+    # Create 10 withdraw threads
+    for _ in range(10):
+        t = threading.Thread(target=account.withdraw, args=(5,))
+        threads.append(t)
+
+    # Start all threads
+    for t in threads:
+        t.start()
+
+    # Wait for all threads to finish
+    for t in threads:
+        t.join()
+
+    # Check account balance
+    expected_balance = 10 * 10  # 10 threads depositing 10 each
+    expected_balance -= 10 * 5  # 10 threads withdrawing 5 each
+    assert account.get_balance() == expected_balance, "Incorrect account balance"