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):