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.
@@ -141,8 +141,14 @@ def verify_token(
             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 = jwt_lib.PyJWKSet.from_dict(certs)
+        header = jwt_lib.get_unverified_header(id_token)
+        try:
+            signing_key = jwkset[header.get("kid")]
+        except KeyError as caught_exc:
+            raise ValueError(
+                "Token has an invalid kid"
+            ) from caught_exc
         return jwt_lib.decode(
             id_token,
             signing_key.key,

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

@@ -86,24 +86,28 @@ 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.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"}
     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)
+    
+    signing_key = py_jwk_set.from_dict.return_value.__getitem__.return_value
+    py_jwk_set.from_dict.return_value.__getitem__.assert_called_once_with("mock-kid")
     _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],
+        signing_key.key,
+        algorithms=[signing_key.algorithm_name],
         audience=None,
     )