fix(oauth2): avoid redundant JWKS network fetches (#17891)

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

mypy.ini

@@ -57,6 +57,9 @@ ignore_missing_imports = True
 [mypy-grpc_status]
 ignore_missing_imports = True
 
+[mypy-jwt,jwt.*]
+ignore_missing_imports = True
+
 [mypy-ibis.*]
 ignore_missing_imports = True
 

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

@@ -124,25 +124,43 @@ 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.
 
     Returns:
         Mapping[str, Any]: The decoded token.
     """
+    if isinstance(id_token, bytes):
+        id_token = id_token.decode("utf-8")
+
     certs = _fetch_certs(request, certs_url)
 
     if "keys" in certs:
         try:
-            import jwt as jwt_lib  # type: ignore
+            import jwt as jwt_lib
+            from jwt.api_jwk import PyJWKSet
+            from jwt.exceptions import PyJWKClientError
         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 kid and key.key_id == kid and key.public_key_use in ("sig", None):
+                signing_key = key
+                break
+
+        if signing_key is None:
+            raise PyJWKClientError(
+                f'Unable to find a signing key that matches: "{kid}"'
+            )
+
         return jwt_lib.decode(
             id_token,
             signing_key.key,

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

@@ -86,28 +86,66 @@ 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.MagicMock()
+    mock_key.key_id = "mock-kid"
+    mock_key.public_key_use = "sig"
+    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,
     )
 
 
+@mock.patch("google.oauth2.id_token._fetch_certs", 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_missing_kid(
+    decode, get_unverified_header, py_jwk_set, _fetch_certs
+):
+    from jwt.exceptions import PyJWKClientError
+
+    certs_url = "abc123"
+    data = {"keys": [{"alg": "RS256"}]}
+    _fetch_certs.return_value = data
+    get_unverified_header.return_value = {"kid": "mock-kid"}
+
+    mock_key = mock.MagicMock()
+    mock_key.key_id = "different-kid"
+    mock_key.public_key_use = "sig"
+    mock_key.key = mock.sentinel.key
+    mock_key.algorithm_name = "mock-alg"
+    py_jwk_set.from_dict.return_value.keys = [mock_key]
+
+    with pytest.raises(
+        PyJWKClientError, match='Unable to find a signing key that matches: "mock-kid"'
+    ):
+        id_token.verify_token(
+            mock.sentinel.token, mock.sentinel.request, certs_url=certs_url
+        )
+
+
 @mock.patch("google.auth.jwt.decode", autospec=True)
 @mock.patch("google.oauth2.id_token._fetch_certs", autospec=True)
 def test_verify_token_args(_fetch_certs, decode):