fix(oauth2): avoid redundant JWKS network fetches

Addresses a performance regression resulting in duplicate fetches of the JWKS endpoint by utilizing the already-fetched certs payload via PyJWKSet.from_dict instead of deferring to PyJWKClient.

Author: macastelaz

packages/google-auth/google/oauth2/id_token.py

@@ -124,7 +124,7 @@ def verify_token(
             intended for. If None then the audience is not verified.
         certs_url (str): The URL that specifies the certificates to use to
             verify the token. This URL should return JSON in the format of
-            ``{'key id': 'x509 certificate'}`` or a certificate array according to
+            ``{'key id': 'x509 certificate'}`` or a JWK Set according to
             the JWK spec (see https://tools.ietf.org/html/rfc7517).
         clock_skew_in_seconds (int): The clock skew used for `iat` and `exp`
             validation.
@@ -137,12 +137,24 @@ def verify_token(
     if "keys" in certs:
         try:
             import jwt as jwt_lib  # type: ignore
+            from jwt.api_jwk import PyJWKSet
         except ImportError as caught_exc:  # pragma: NO COVER
             raise ImportError(
                 "The pyjwt library is not installed, please install the pyjwt package to use the jwk certs format."
             ) from caught_exc
-        jwks_client = jwt_lib.PyJWKClient(certs_url)
-        signing_key = jwks_client.get_signing_key_from_jwt(id_token)
+        jwkset = PyJWKSet.from_dict(certs)
+        header = jwt_lib.get_unverified_header(id_token)
+        kid = header.get("kid")
+        
+        signing_key = None
+        for key in jwkset.keys:
+            if key.key_id == kid:
+                signing_key = key
+                break
+
+        if signing_key is None:
+            raise ValueError("Token has an invalid kid")
+        
         return jwt_lib.decode(
             id_token,
             signing_key.key,

packages/google-auth/tests/oauth2/test_id_token.py

@@ -86,24 +86,32 @@ def test_verify_token(_fetch_certs, decode):
 
 
 @mock.patch("google.oauth2.id_token._fetch_certs", autospec=True)
-@mock.patch("jwt.PyJWKClient", autospec=True)
+@mock.patch("jwt.api_jwk.PyJWKSet", autospec=True)
+@mock.patch("jwt.get_unverified_header", autospec=True)
 @mock.patch("jwt.decode", autospec=True)
-def test_verify_token_jwk(decode, py_jwk, _fetch_certs):
+def test_verify_token_jwk(decode, get_unverified_header, py_jwk_set, _fetch_certs):
     certs_url = "abc123"
     data = {"keys": [{"alg": "RS256"}]}
     _fetch_certs.return_value = data
+    get_unverified_header.return_value = {"kid": "mock-kid"}
+    
+    mock_key = mock.Mock()
+    mock_key.key_id = "mock-kid"
+    mock_key.key = mock.sentinel.key
+    mock_key.algorithm_name = "mock-alg"
+    py_jwk_set.from_dict.return_value.keys = [mock_key]
     result = id_token.verify_token(
         mock.sentinel.token, mock.sentinel.request, certs_url=certs_url
     )
     assert result == decode.return_value
-    py_jwk.assert_called_once_with(certs_url)
-    signing_key = py_jwk.return_value.get_signing_key_from_jwt
+    py_jwk_set.from_dict.assert_called_once_with(data)
+    get_unverified_header.assert_called_once_with(mock.sentinel.token)
+    
     _fetch_certs.assert_called_once_with(mock.sentinel.request, certs_url)
-    signing_key.assert_called_once_with(mock.sentinel.token)
     decode.assert_called_once_with(
         mock.sentinel.token,
-        signing_key.return_value.key,
-        algorithms=[signing_key.return_value.algorithm_name],
+        mock.sentinel.key,
+        algorithms=["mock-alg"],
         audience=None,
     )