fix: add optional clock_skew parameter to token helpers (#333)

* fix: fix token clock skew issue

* lint

* fix unit tests

* change clock skew

* change clock skew to 10 instead of 60

* Made clock skew an optional arg

* added check for clock skew values

* make default clock skew = 0

* fix tests

* fix arg name

* rename arg

* add suggestion

Author: twishabansal

packages/toolbox-core/src/toolbox_core/auth_methods.py

@@ -41,6 +41,7 @@
 # --- Constants ---
 BEARER_TOKEN_PREFIX = "Bearer "
 CACHE_REFRESH_MARGIN = timedelta(seconds=60)
+DEFAULT_CLOCK_SKEW = 0
 
 _token_cache: Dict[str, Any] = {
     "token": None,
@@ -57,7 +58,7 @@ def _is_token_valid() -> bool:
     )
 
 
-def _update_cache(new_token: str) -> None:
+def _update_cache(new_token: str, clock_skew_in_seconds: int) -> None:
     """
     Validates a new token, extracts its expiry, and updates the cache.
 
@@ -71,7 +72,9 @@ def _update_cache(new_token: str) -> None:
         # verify_oauth2_token not only decodes but also validates the token's
         # signature and claims against Google's public keys.
         # It's a synchronous, CPU-bound operation, safe for async contexts.
-        claims = id_token.verify_oauth2_token(new_token, Request())
+        claims = id_token.verify_oauth2_token(
+            new_token, Request(), clock_skew_in_seconds=clock_skew_in_seconds
+        )
 
         expiry_timestamp = claims.get("exp")
         if not expiry_timestamp:
@@ -89,7 +92,15 @@ def _update_cache(new_token: str) -> None:
         raise ValueError(f"Failed to validate and cache the new token: {e}") from e
 
 
-def get_google_token_from_aud(audience: Optional[str] = None) -> str:
+def get_google_token_from_aud(
+    clock_skew_in_seconds: int = 0, audience: Optional[str] = None
+) -> str:
+    if clock_skew_in_seconds < 0 or clock_skew_in_seconds > 60:
+        raise ValueError(
+            f"Illegal clock_skew_in_seconds value: {clock_skew_in_seconds}. Must be between 0 and 60"
+            ", inclusive."
+        )
+
     if _is_token_valid():
         return BEARER_TOKEN_PREFIX + _token_cache["token"]
 
@@ -102,7 +113,7 @@ def get_google_token_from_aud(audience: Optional[str] = None) -> str:
     if hasattr(credentials, "id_token"):
         new_id_token = getattr(credentials, "id_token", None)
         if new_id_token:
-            _update_cache(new_id_token)
+            _update_cache(new_id_token, clock_skew_in_seconds)
             return BEARER_TOKEN_PREFIX + new_id_token
 
     if audience is None:
@@ -115,7 +126,7 @@ def get_google_token_from_aud(audience: Optional[str] = None) -> str:
     try:
         request = Request()
         new_token = id_token.fetch_id_token(request, audience)
-        _update_cache(new_token)
+        _update_cache(new_token, clock_skew_in_seconds)
         return BEARER_TOKEN_PREFIX + _token_cache["token"]
 
     except GoogleAuthError as e:
@@ -124,15 +135,20 @@ def get_google_token_from_aud(audience: Optional[str] = None) -> str:
         ) from e
 
 
-def get_google_id_token(audience: Optional[str] = None) -> Callable[[], str]:
+def get_google_id_token(
+    audience: Optional[str] = None, clock_skew_in_seconds: int = DEFAULT_CLOCK_SKEW
+) -> Callable[[], str]:
     """
     Returns a SYNC function that, when called, fetches a Google ID token.
     This function uses Application Default Credentials for local systems
     and standard google auth libraries for Google Cloud environments.
     It caches the token in memory.
 
     Args:
-        audience: The audience for the ID token (e.g., a service URL or client ID).
+        audience: The audience for the ID token (e.g., a service URL or client
+        ID).
+        clock_skew_in_seconds: The number of seconds to tolerate when checking the token.
+            Must be between 0-60. Defaults to 0.
 
     Returns:
         A function that when executed returns string in the format "Bearer <google_id_token>".
@@ -143,13 +159,13 @@ def get_google_id_token(audience: Optional[str] = None) -> Callable[[], str]:
     """
 
     def _token_getter() -> str:
-        return get_google_token_from_aud(audience)
+        return get_google_token_from_aud(clock_skew_in_seconds, audience)
 
     return _token_getter
 
 
 def aget_google_id_token(
-    audience: Optional[str] = None,
+    audience: Optional[str] = None, clock_skew_in_seconds: int = DEFAULT_CLOCK_SKEW
 ) -> Callable[[], Coroutine[Any, Any, str]]:
     """
     Returns an ASYNC function that, when called, fetches a Google ID token.
@@ -158,7 +174,10 @@ def aget_google_id_token(
     It caches the token in memory.
 
     Args:
-        audience: The audience for the ID token (e.g., a service URL or client ID).
+        audience: The audience for the ID token (e.g., a service URL or client
+        ID).
+        clock_skew_in_seconds: The number of seconds to tolerate when checking the token.
+            Must be between 0-60. Defaults to 0.
 
     Returns:
         An async function that when executed returns string in the format "Bearer <google_id_token>".
@@ -169,6 +188,8 @@ def aget_google_id_token(
     """
 
     async def _token_getter() -> str:
-        return await asyncio.to_thread(get_google_token_from_aud, audience)
+        return await asyncio.to_thread(
+            get_google_token_from_aud, clock_skew_in_seconds, audience
+        )
 
     return _token_getter

packages/toolbox-core/tests/test_auth_methods.py

@@ -154,7 +154,7 @@ def test_get_google_id_token_success_local_creds(
         mock_default.assert_called_once_with()
         mock_session.assert_called_once_with(mock_creds)
         mock_creds.refresh.assert_called_once_with(mock_request_instance)
-        mock_verify.assert_called_once_with(MOCK_ID_TOKEN, ANY)
+        mock_verify.assert_called_once_with(MOCK_ID_TOKEN, ANY, clock_skew_in_seconds=0)
         assert token == f"{auth_methods.BEARER_TOKEN_PREFIX}{MOCK_ID_TOKEN}"
         assert auth_methods._token_cache["token"] == MOCK_ID_TOKEN
         assert auth_methods._token_cache["expires_at"] == MOCK_EXPIRY_DATETIME