feat: allow users to use jwk keys for verifying ID token (#1641)

* add support for jwk format

* update formatting

* formatting changes

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* remove pyjwt

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: harkamaljot

docs/requirements-docs.txt

@@ -2,4 +2,4 @@ cryptography
 sphinx-docstring-typing
 urllib3
 requests
-requests-oauthlib
+requests-oauthlib

google/oauth2/id_token.py

@@ -82,16 +82,17 @@ def _fetch_certs(request, certs_url):
     """Fetches certificates.
 
     Google-style cerificate endpoints return JSON in the format of
-    ``{'key id': 'x509 certificate'}``.
+    ``{'key id': 'x509 certificate'}`` or a certificate array according
+    to the JWK spec (see https://tools.ietf.org/html/rfc7517).
 
     Args:
         request (google.auth.transport.Request): The object used to make
             HTTP requests.
         certs_url (str): The certificate endpoint URL.
 
     Returns:
-        Mapping[str, str]: A mapping of public key ID to x.509 certificate
-            data.
+        Mapping[str, str] | Mapping[str, list]: A mapping of public keys
+        in x.509 or JWK spec.
     """
     response = request(certs_url, method="GET")
 
@@ -120,7 +121,8 @@ 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'}``.
+            ``{'key id': 'x509 certificate'}`` or a certificate array 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.
 
@@ -129,12 +131,28 @@ def verify_token(
     """
     certs = _fetch_certs(request, certs_url)
 
-    return jwt.decode(
-        id_token,
-        certs=certs,
-        audience=audience,
-        clock_skew_in_seconds=clock_skew_in_seconds,
-    )
+    if "keys" in certs:
+        try:
+            import jwt as jwt_lib  # type: ignore
+        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)
+        return jwt_lib.decode(
+            id_token,
+            signing_key.key,
+            algorithms=[signing_key.algorithm_name],
+            audience=audience,
+        )
+    else:
+        return jwt.decode(
+            id_token,
+            certs=certs,
+            audience=audience,
+            clock_skew_in_seconds=clock_skew_in_seconds,
+        )
 
 
 def verify_oauth2_token(id_token, request, audience=None, clock_skew_in_seconds=0):

setup.py

@@ -33,6 +33,7 @@
     "requests": "requests >= 2.20.0, < 3.0.0.dev0",
     "reauth": "pyu2f>=0.1.5",
     "enterprise_cert": ["cryptography", "pyopenssl"],
+    "pyjwt": ["pyjwt>=2.0", "cryptography>=38.0.3"],
 }
 
 with io.open("README.rst", "r") as fh:

system_tests/noxfile.py

@@ -162,7 +162,7 @@ def configure_cloud_sdk(session, application_default_credentials, project=False)
 # Test sesssions
 
 TEST_DEPENDENCIES_ASYNC = ["aiohttp", "pytest-asyncio", "nest-asyncio", "mock"]
-TEST_DEPENDENCIES_SYNC = ["pytest", "requests", "mock"]
+TEST_DEPENDENCIES_SYNC = ["pytest", "requests", "mock", "pyjwt"]
 PYTHON_VERSIONS_ASYNC = ["3.7"]
 PYTHON_VERSIONS_SYNC = ["3.7"]
 

testing/constraints-3.7.txt

@@ -11,3 +11,4 @@ setuptools==40.3.0
 rsa==3.1.4
 aiohttp==3.6.2
 requests==2.20.0
+pyjwt==2.0

testing/requirements.txt

@@ -8,6 +8,7 @@ pytest
 pytest-cov
 pytest-localserver
 pyu2f
+pyjwt
 requests
 urllib3
 cryptography < 39.0.0

tests/oauth2/test_id_token.py

@@ -78,6 +78,29 @@ 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.decode", autospec=True)
+def test_verify_token_jwk(decode, py_jwk, _fetch_certs):
+    certs_url = "abc123"
+    data = {"keys": [{"alg": "RS256"}]}
+    _fetch_certs.return_value = data
+    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
+    _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],
+        audience=None,
+    )
+
+
 @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):