Updated JWT validation to new library (GoogleCloudPlatform#4393)

engelke · web-flow · commit 41c173ead304 · 2020-07-31T10:29:37.000-07:00
* Updated JWT validation to new library

* Fix changed variable name

* Debug audience test value

* Debug audience in test environment

* More verbose assertion

* Clearer assertion

* Clearer assertion

* Specified certs for validation
diff --git a/iap/README.md b/iap/README.md
@@ -70,12 +70,6 @@ service account private key can impersonate that account!
 
 ## Using validate_jwt
 
-`validate_jwt` is not compatible with App Engine standard environment;
-use App Engine's Users API instead. (See `app_engine_app` for an example
-of how to do this.)
-
-For all other environments:
-
 1. Install the libraries listed in `requirements.txt`, e.g. by running:
    ```
    virtualenv/bin/pip install -r requirements.txt
diff --git a/iap/example_gce_backend.py b/iap/example_gce_backend.py
@@ -30,7 +30,7 @@ def root():
     if jwt is None:
         return 'Unauthorized request.'
     user_id, user_email, error_str = (
-        validate_jwt.validate_iap_jwt_from_compute_engine(
+        validate_jwt.validate_iap_jwt(
             jwt, CLOUD_PROJECT_ID, BACKEND_SERVICE_ID))
     if error_str:
         return 'Error: {}'.format(error_str)
diff --git a/iap/iap_test.py b/iap/iap_test.py
@@ -44,14 +44,19 @@ def test_main(capsys):
         pytest.skip('Only passing on Kokoro.')
     # JWTs are obtained by IAP-protected applications whenever an
     # end-user makes a request.  We've set up an app that echoes back
-    # the JWT in order to expose it to this test.  Thus, this test
+    # the IAP JWT in order to expose it to this test.  Thus, this test
     # exercises both make_iap_request and validate_jwt.
-    iap_jwt = make_iap_request.make_iap_request(
+    resp = make_iap_request.make_iap_request(
         'https://{}/'.format(REFLECT_SERVICE_HOSTNAME),
         IAP_CLIENT_ID)
-    iap_jwt = iap_jwt.split(': ').pop()
-    jwt_validation_result = validate_jwt.validate_iap_jwt_from_app_engine(
-        iap_jwt, IAP_PROJECT_NUMBER, IAP_APP_ID)
+    iap_jwt = resp.split(': ').pop()
+
+    # App Engine JWT audience format below
+    expected_audience = '/projects/{}/apps/{}'.format(
+        IAP_PROJECT_NUMBER, IAP_APP_ID)
+
+    jwt_validation_result = validate_jwt.validate_iap_jwt(
+        iap_jwt, expected_audience)
 
     assert jwt_validation_result[0]
     assert jwt_validation_result[1]
diff --git a/iap/make_iap_request.py b/iap/make_iap_request.py
@@ -41,15 +41,15 @@ def make_iap_request(url, client_id, method='GET', **kwargs):
 
     # Obtain an OpenID Connect (OIDC) token from metadata server or using service
     # account.
-    google_open_id_connect_token = id_token.fetch_id_token(Request(), client_id)
+    open_id_connect_token = id_token.fetch_id_token(Request(), client_id)
 
     # Fetch the Identity-Aware Proxy-protected URL, including an
     # Authorization header containing "Bearer " followed by a
     # Google-issued OpenID Connect token for the service account.
     resp = requests.request(
         method, url,
         headers={'Authorization': 'Bearer {}'.format(
-            google_open_id_connect_token)}, **kwargs)
+            open_id_connect_token)}, **kwargs)
     if resp.status_code == 403:
         raise Exception('Service account does not have permission to '
                         'access the IAP-protected application.')
diff --git a/iap/validate_jwt.py b/iap/validate_jwt.py
@@ -16,102 +16,33 @@
 
 This code should be used by applications in Google Compute Engine-based
 environments (such as Google App Engine flexible environment, Google
-Compute Engine, or Google Container Engine) to provide an extra layer
-of assurance that a request was authorized by IAP.
-
-For applications running in the App Engine standard environment, use
-App Engine's Users API instead.
+Compute Engine, Google Container Engine, Google App Engine) to provide
+an extra layer of assurance that a request was authorized by IAP.
 """
 # [START iap_validate_jwt]
-from google.auth import jwt
-import requests
-
-
-def validate_iap_jwt_from_app_engine(iap_jwt, cloud_project_number,
-                                     cloud_project_id):
-    """Validate a JWT passed to your App Engine app by Identity-Aware Proxy.
-
-    Args:
-      iap_jwt: The contents of the X-Goog-IAP-JWT-Assertion header.
-      cloud_project_number: The project *number* for your Google Cloud project.
-          This is returned by 'gcloud projects describe $PROJECT_ID', or
-          in the Project Info card in Cloud Console.
-      cloud_project_id: The project *ID* for your Google Cloud project.
-
-    Returns:
-      (user_id, user_email, error_str).
-    """
-    expected_audience = '/projects/{}/apps/{}'.format(
-        cloud_project_number, cloud_project_id)
-    return _validate_iap_jwt(iap_jwt, expected_audience)
+from google.auth.transport import requests
+from google.oauth2 import id_token
 
 
-def validate_iap_jwt_from_compute_engine(iap_jwt, cloud_project_number,
-                                         backend_service_id):
-    """Validate an IAP JWT for your (Compute|Container) Engine service.
+def validate_iap_jwt(iap_jwt, expected_audience):
+    """Validate an IAP JWT.
 
     Args:
       iap_jwt: The contents of the X-Goog-IAP-JWT-Assertion header.
-      cloud_project_number: The project *number* for your Google Cloud project.
-          This is returned by 'gcloud projects describe $PROJECT_ID', or
-          in the Project Info card in Cloud Console.
-      backend_service_id: The ID of the backend service used to access the
-          application. See
+      expected_audience: The Signed Header JWT audience. See
           https://cloud.google.com/iap/docs/signed-headers-howto
           for details on how to get this value.
 
     Returns:
       (user_id, user_email, error_str).
     """
-    expected_audience = '/projects/{}/global/backendServices/{}'.format(
-        cloud_project_number, backend_service_id)
-    return _validate_iap_jwt(iap_jwt, expected_audience)
-
 
-def _validate_iap_jwt(iap_jwt, expected_audience):
     try:
-        # Retrieve public key for token signature verification.
-        key_id = jwt.decode_header(iap_jwt).get('kid')
-        if not key_id:
-            return (None, None, '**ERROR: no key ID**')
-        key = get_iap_key(key_id)
-
-        # Verify token signature, expiry and audience.
-        decoded_jwt = jwt.decode(iap_jwt, certs=key, audience=expected_audience)
-
-        # Verify token issuer.
-        if decoded_jwt.get('iss') != 'https://cloud.google.com/iap':
-            return (None, None, '**ERROR: invalid issuer**')
-
+        decoded_jwt = id_token.verify_token(
+            iap_jwt, requests.Request(), audience=expected_audience,
+            certs_url='https://www.gstatic.com/iap/verify/public_key')
         return (decoded_jwt['sub'], decoded_jwt['email'], '')
-    except (ValueError, requests.exceptions.RequestException) as e:
+    except Exception as e:
         return (None, None, '**ERROR: JWT validation error {}**'.format(e))
 
-
-def get_iap_key(key_id):
-    """Retrieves a public key from the list published by Identity-Aware Proxy,
-    re-fetching the key file if necessary.
-    """
-    key_cache = get_iap_key.key_cache
-    key = key_cache.get(key_id)
-    if not key:
-        # Re-fetch the key file.
-        resp = requests.get(
-            'https://www.gstatic.com/iap/verify/public_key')
-        if resp.status_code != 200:
-            raise Exception(
-                'Unable to fetch IAP keys: {} / {} / {}'.format(
-                    resp.status_code, resp.headers, resp.text))
-        key_cache = resp.json()
-        get_iap_key.key_cache = key_cache
-        key = key_cache.get(key_id)
-        if not key:
-            raise Exception('Key {!r} not found'.format(key_id))
-    return key
-
-
-# Used to cache the Identity-Aware Proxy public keys.  This code only
-# refetches the file when a JWT is signed with a key not present in
-# this cache.
-get_iap_key.key_cache = {}
 # [END iap_validate_jwt]
