adding jwt_leeway param to allow configuring the leeway param on PyJWT jwt.decode method

Author: Th3R3p0

workos/_base_client.py

@@ -26,6 +26,7 @@ class BaseClient(ClientConfiguration):
     _base_url: str
     _client_id: str
     _request_timeout: int
+    _jwt_leeway: float
 
     def __init__(
         self,
@@ -34,6 +35,7 @@ def __init__(
         client_id: Optional[str],
         base_url: Optional[str] = None,
         request_timeout: Optional[int] = None,
+        jwt_leeway: float = 0,
     ) -> None:
         api_key = api_key or os.getenv("WORKOS_API_KEY")
         if api_key is None:
@@ -65,6 +67,8 @@ def __init__(
             else int(os.getenv("WORKOS_REQUEST_TIMEOUT", DEFAULT_REQUEST_TIMEOUT))
         )
 
+        self._jwt_leeway = jwt_leeway
+
     @property
     @abstractmethod
     def audit_logs(self) -> AuditLogsModule: ...
@@ -127,3 +131,7 @@ def client_id(self) -> str:
     @property
     def request_timeout(self) -> int:
         return self._request_timeout
+
+    @property
+    def jwt_leeway(self) -> float:
+        return self._jwt_leeway

workos/async_client.py

@@ -30,12 +30,14 @@ def __init__(
         client_id: Optional[str] = None,
         base_url: Optional[str] = None,
         request_timeout: Optional[int] = None,
+        jwt_leeway: float = 0,
     ):
         super().__init__(
             api_key=api_key,
             client_id=client_id,
             base_url=base_url,
             request_timeout=request_timeout,
+            jwt_leeway=jwt_leeway,
         )
         self._http_client = AsyncHTTPClient(
             api_key=self._api_key,

workos/client.py

@@ -30,12 +30,14 @@ def __init__(
         client_id: Optional[str] = None,
         base_url: Optional[str] = None,
         request_timeout: Optional[int] = None,
+        jwt_leeway: float = 0,
     ):
         super().__init__(
             api_key=api_key,
             client_id=client_id,
             base_url=base_url,
             request_timeout=request_timeout,
+            jwt_leeway=jwt_leeway,
         )
         self._http_client = SyncHTTPClient(
             api_key=self._api_key,

workos/session.py

@@ -35,6 +35,7 @@ class SessionModule(Protocol):
     cookie_password: str
     jwks: PyJWKClient
     jwk_algorithms: List[str]
+    jwt_leeway: float
 
     def __init__(
         self,
@@ -43,6 +44,7 @@ def __init__(
         client_id: str,
         session_data: str,
         cookie_password: str,
+        jwt_leeway: float = 0,
     ) -> None:
         # If the cookie password is not provided, throw an error
         if cookie_password is None or cookie_password == "":
@@ -52,6 +54,7 @@ def __init__(
         self.client_id = client_id
         self.session_data = session_data
         self.cookie_password = cookie_password
+        self.jwt_leeway = jwt_leeway
 
         self.jwks = _get_jwks_client(self.user_management.get_jwks_url())
 
@@ -91,13 +94,13 @@ def authenticate(
                 signing_key.key,
                 algorithms=self.jwk_algorithms,
                 options={"verify_aud": False},
+                leeway=self.jwt_leeway,
             )
         except jwt.exceptions.InvalidTokenError:
             return AuthenticateWithSessionCookieErrorResponse(
                 authenticated=False,
                 reason=AuthenticateWithSessionCookieFailureReason.INVALID_JWT,
             )
-
         return AuthenticateWithSessionCookieSuccessResponse(
             authenticated=True,
             session_id=decoded["sid"],
@@ -137,6 +140,20 @@ def get_logout_url(self, return_to: Optional[str] = None) -> str:
         )
         return str(result)
 
+    def _is_valid_jwt(self, token: str) -> bool:
+        try:
+            signing_key = self.jwks.get_signing_key_from_jwt(token)
+            jwt.decode(
+                token,
+                signing_key.key,
+                algorithms=self.jwk_algorithms,
+                options={"verify_aud": False},
+                leeway=self.jwt_leeway,
+            )
+            return True
+        except jwt.exceptions.InvalidTokenError:
+            return False
+
     @staticmethod
     def seal_data(data: Dict[str, Any], key: str) -> str:
         fernet = Fernet(key)
@@ -163,6 +180,7 @@ def __init__(
         client_id: str,
         session_data: str,
         cookie_password: str,
+        jwt_leeway: float = 0,
     ) -> None:
         # If the cookie password is not provided, throw an error
         if cookie_password is None or cookie_password == "":
@@ -172,6 +190,7 @@ def __init__(
         self.client_id = client_id
         self.session_data = session_data
         self.cookie_password = cookie_password
+        self.jwt_leeway = jwt_leeway
 
         self.jwks = _get_jwks_client(self.user_management.get_jwks_url())
 
@@ -224,6 +243,7 @@ def refresh(
                 signing_key.key,
                 algorithms=self.jwk_algorithms,
                 options={"verify_aud": False},
+                leeway=self.jwt_leeway,
             )
 
             return RefreshWithSessionCookieSuccessResponse(
@@ -255,6 +275,7 @@ def __init__(
         client_id: str,
         session_data: str,
         cookie_password: str,
+        jwt_leeway: float = 0,
     ) -> None:
         # If the cookie password is not provided, throw an error
         if cookie_password is None or cookie_password == "":
@@ -264,6 +285,7 @@ def __init__(
         self.client_id = client_id
         self.session_data = session_data
         self.cookie_password = cookie_password
+        self.jwt_leeway = jwt_leeway
 
         self.jwks = _get_jwks_client(self.user_management.get_jwks_url())
 
@@ -316,6 +338,7 @@ async def refresh(
                 signing_key.key,
                 algorithms=self.jwk_algorithms,
                 options={"verify_aud": False},
+                leeway=self.jwt_leeway,
             )
 
             return RefreshWithSessionCookieSuccessResponse(

workos/user_management.py

@@ -891,6 +891,7 @@ def load_sealed_session(
             client_id=self._http_client.client_id,
             session_data=sealed_session,
             cookie_password=cookie_password,
+            jwt_leeway=self._client_configuration.jwt_leeway,
         )
 
     def get_user(self, user_id: str) -> User:
@@ -1531,6 +1532,7 @@ async def load_sealed_session(
             client_id=self._http_client.client_id,
             session_data=sealed_session,
             cookie_password=cookie_password,
+            jwt_leeway=self._client_configuration.jwt_leeway,
         )
 
     async def get_user(self, user_id: str) -> User: