Added tenant-scoped custom token support (#433)

* Added tenant-scoped custom token support

* Fixed a lint error; Improved test assertion

Author: hiranya911

firebase_admin/_token_gen.py

@@ -130,7 +130,7 @@ def signing_provider(self):
                     'details on creating custom tokens.'.format(error, url))
         return self._signing_provider
 
-    def create_custom_token(self, uid, developer_claims=None):
+    def create_custom_token(self, uid, developer_claims=None, tenant_id=None):
         """Builds and signs a Firebase custom auth token."""
         if developer_claims is not None:
             if not isinstance(developer_claims, dict):
@@ -161,6 +161,8 @@ def create_custom_token(self, uid, developer_claims=None):
             'iat': now,
             'exp': now + MAX_TOKEN_LIFETIME_SECONDS,
         }
+        if tenant_id:
+            payload['tenant_id'] = tenant_id
 
         if developer_claims is not None:
             payload['claims'] = developer_claims

firebase_admin/auth.py

@@ -556,7 +556,8 @@ def tenant_id(self):
         return self._tenant_id
 
     def create_custom_token(self, uid, developer_claims=None):
-        return self._token_generator.create_custom_token(uid, developer_claims)
+        return self._token_generator.create_custom_token(
+            uid, developer_claims, tenant_id=self.tenant_id)
 
     def verify_id_token(self, id_token, check_revoked=False):
         """Verifies the signature and data for the provided ID token."""

tests/test_tenant_mgt.py

@@ -21,6 +21,7 @@
 
 import firebase_admin
 from firebase_admin import auth
+from firebase_admin import credentials
 from firebase_admin import exceptions
 from firebase_admin import tenant_mgt
 from tests import testutils
@@ -731,6 +732,34 @@ def test_invalid_tenant_id(self, tenant_mgt_app):
         assert excinfo.value.http_response is None
 
 
+@pytest.fixture(scope='module')
+def tenant_aware_custom_token_app():
+    cred = credentials.Certificate(testutils.resource_filename('service_account.json'))
+    app = firebase_admin.initialize_app(cred, name='tenantAwareCustomToken')
+    yield app
+    firebase_admin.delete_app(app)
+
+
+class TestCreateCustomToken:
+
+    def test_custom_token(self, tenant_aware_custom_token_app):
+        client = tenant_mgt.auth_for_tenant('test-tenant', app=tenant_aware_custom_token_app)
+
+        custom_token = client.create_custom_token('user1')
+
+        test_token_gen.verify_custom_token(
+            custom_token, expected_claims=None, tenant_id='test-tenant')
+
+    def test_custom_token_with_claims(self, tenant_aware_custom_token_app):
+        client = tenant_mgt.auth_for_tenant('test-tenant', app=tenant_aware_custom_token_app)
+        claims = {'admin': True}
+
+        custom_token = client.create_custom_token('user1', claims)
+
+        test_token_gen.verify_custom_token(
+            custom_token, expected_claims=claims, tenant_id='test-tenant')
+
+
 def _assert_tenant(tenant, tenant_id='tenant-id'):
     assert isinstance(tenant, tenant_mgt.Tenant)
     assert tenant.tenant_id == tenant_id

tests/test_token_gen.py

@@ -66,7 +66,7 @@ def _merge_jwt_claims(defaults, overrides):
             del defaults[key]
     return defaults
 
-def _verify_custom_token(custom_token, expected_claims):
+def verify_custom_token(custom_token, expected_claims, tenant_id=None):
     assert isinstance(custom_token, bytes)
     token = google.oauth2.id_token.verify_token(
         custom_token,
@@ -75,6 +75,11 @@ def _verify_custom_token(custom_token, expected_claims):
     assert token['uid'] == MOCK_UID
     assert token['iss'] == MOCK_SERVICE_ACCOUNT_EMAIL
     assert token['sub'] == MOCK_SERVICE_ACCOUNT_EMAIL
+    if tenant_id is None:
+        assert 'tenant_id' not in token
+    else:
+        assert token['tenant_id'] == tenant_id
+
     header = jwt.decode_header(custom_token)
     assert header.get('typ') == 'JWT'
     assert header.get('alg') == 'RS256'
@@ -198,7 +203,7 @@ class TestCreateCustomToken:
     def test_valid_params(self, auth_app, values):
         user, claims = values
         custom_token = auth.create_custom_token(user, claims, app=auth_app)
-        _verify_custom_token(custom_token, claims)
+        verify_custom_token(custom_token, claims)
 
     @pytest.mark.parametrize('values', invalid_args.values(), ids=list(invalid_args))
     def test_invalid_params(self, auth_app, values):