Tenant-aware ID token verification support (#432)

hiranya911 · web-flow · commit 19628b74c48d · 2020-03-26T14:48:14.000-07:00
* Tenant-aware ID token verification support

* Extended InvalidArgumentError in TenantIdMismatchError

* Fixing lint errors
diff --git a/firebase_admin/_auth_utils.py b/firebase_admin/_auth_utils.py
@@ -300,6 +300,13 @@ def __init__(self, message, cause=None, http_response=None):
         exceptions.NotFoundError.__init__(self, message, cause, http_response)
 
 
+class TenantIdMismatchError(exceptions.InvalidArgumentError):
+    """Missing or invalid tenant ID field in the given JWT."""
+
+    def __init__(self, message):
+        exceptions.InvalidArgumentError.__init__(self, message)
+
+
 _CODE_TO_EXC_TYPE = {
     'DUPLICATE_EMAIL': EmailAlreadyExistsError,
     'DUPLICATE_LOCAL_ID': UidAlreadyExistsError,
diff --git a/firebase_admin/auth.py b/firebase_admin/auth.py
@@ -559,12 +559,19 @@ def create_custom_token(self, uid, developer_claims=None):
         return self._token_generator.create_custom_token(uid, developer_claims)
 
     def verify_id_token(self, id_token, check_revoked=False):
+        """Verifies the signature and data for the provided ID token."""
         if not isinstance(check_revoked, bool):
             # guard against accidental wrong assignment.
             raise ValueError('Illegal check_revoked argument. Argument must be of type '
                              ' bool, but given "{0}".'.format(type(check_revoked)))
 
         verified_claims = self._token_verifier.verify_id_token(id_token)
+        if self.tenant_id:
+            token_tenant_id = verified_claims.get('firebase', {}).get('tenant')
+            if self.tenant_id != token_tenant_id:
+                raise _auth_utils.TenantIdMismatchError(
+                    'Invalid tenant ID: {0}'.format(token_tenant_id))
+
         if check_revoked:
             self._check_jwt_revoked(verified_claims, RevokedIdTokenError, 'ID token')
         return verified_claims
diff --git a/firebase_admin/tenant_mgt.py b/firebase_admin/tenant_mgt.py
@@ -36,6 +36,7 @@
 __all__ = [
     'ListTenantsPage',
     'Tenant',
+    'TenantIdMismatchError',
     'TenantNotFoundError',
 
     'auth_for_tenant',
@@ -46,6 +47,8 @@
     'update_tenant',
 ]
 
+
+TenantIdMismatchError = _auth_utils.TenantIdMismatchError
 TenantNotFoundError = _auth_utils.TenantNotFoundError
 
 
diff --git a/tests/test_tenant_mgt.py b/tests/test_tenant_mgt.py
@@ -24,6 +24,7 @@
 from firebase_admin import exceptions
 from firebase_admin import tenant_mgt
 from tests import testutils
+from tests import test_token_gen
 
 
 GET_TENANT_RESPONSE = """{
@@ -696,7 +697,6 @@ def test_tenant_not_found(self, tenant_mgt_app):
         assert excinfo.value.http_response is not None
         assert excinfo.value.cause is not None
 
-
     def _assert_request(self, recorder, want_url, want_body):
         assert len(recorder) == 1
         req = recorder[0]
@@ -706,6 +706,31 @@ def _assert_request(self, recorder, want_url, want_body):
         assert body == want_body
 
 
+class TestVerifyIdToken:
+
+    def test_valid_token(self, tenant_mgt_app):
+        client = tenant_mgt.auth_for_tenant('test-tenant', app=tenant_mgt_app)
+        client._token_verifier.request = test_token_gen.MOCK_REQUEST
+
+        claims = client.verify_id_token(test_token_gen.TEST_ID_TOKEN_WITH_TENANT)
+
+        assert claims['admin'] is True
+        assert claims['uid'] == claims['sub']
+        assert claims['firebase']['tenant'] == 'test-tenant'
+
+    def test_invalid_tenant_id(self, tenant_mgt_app):
+        client = tenant_mgt.auth_for_tenant('other-tenant', app=tenant_mgt_app)
+        client._token_verifier.request = test_token_gen.MOCK_REQUEST
+
+        with pytest.raises(tenant_mgt.TenantIdMismatchError) as excinfo:
+            client.verify_id_token(test_token_gen.TEST_ID_TOKEN_WITH_TENANT)
+
+        assert 'Invalid tenant ID: test-tenant' in str(excinfo.value)
+        assert isinstance(excinfo.value, exceptions.InvalidArgumentError)
+        assert excinfo.value.cause is None
+        assert excinfo.value.http_response is None
+
+
 def _assert_tenant(tenant, tenant_id='tenant-id'):
     assert isinstance(tenant, tenant_mgt.Tenant)
     assert tenant.tenant_id == tenant_id
diff --git a/tests/test_token_gen.py b/tests/test_token_gen.py
@@ -94,6 +94,9 @@ def _get_id_token(payload_overrides=None, header_overrides=None):
         'exp': int(time.time()) + 3600,
         'sub': '1234567890',
         'admin': True,
+        'firebase': {
+            'sign_in_provider': 'provider',
+        },
     }
     if header_overrides:
         headers = _merge_jwt_claims(headers, header_overrides)
@@ -346,6 +349,11 @@ def test_unexpected_response(self, user_mgt_app):
 
 MOCK_GET_USER_RESPONSE = testutils.resource('get_user.json')
 TEST_ID_TOKEN = _get_id_token()
+TEST_ID_TOKEN_WITH_TENANT = _get_id_token({
+    'firebase': {
+        'tenant': 'test-tenant',
+    }
+})
 TEST_SESSION_COOKIE = _get_session_cookie()
 
 
@@ -380,6 +388,14 @@ def test_valid_token(self, user_mgt_app, id_token):
         claims = auth.verify_id_token(id_token, app=user_mgt_app)
         assert claims['admin'] is True
         assert claims['uid'] == claims['sub']
+        assert claims['firebase']['sign_in_provider'] == 'provider'
+
+    def test_valid_token_with_tenant(self, user_mgt_app):
+        _overwrite_cert_request(user_mgt_app, MOCK_REQUEST)
+        claims = auth.verify_id_token(TEST_ID_TOKEN_WITH_TENANT, app=user_mgt_app)
+        assert claims['admin'] is True
+        assert claims['uid'] == claims['sub']
+        assert claims['firebase']['tenant'] == 'test-tenant'
 
     @pytest.mark.parametrize('id_token', valid_tokens.values(), ids=list(valid_tokens))
     def test_valid_token_check_revoked(self, user_mgt_app, id_token):
