Allow createCustomToken() to work with tenant-aware auth

rsgowman · rsgowman · commit 8b02f7ad9a88 · 2019-11-28T11:30:38.000-05:00
diff --git a/src/auth/auth.ts b/src/auth/auth.ts
@@ -626,18 +626,16 @@ export class TenantAwareAuth extends BaseAuth<TenantAwareAuthRequestHandler> {
 
   /**
    * Creates a new custom token that can be sent back to a client to use with
-   * signInWithCustomToken().
+   * signInWithCustomToken(). The tenant id will be embedded in the token and
+   * will be verified during the call to signInWithCustomToken().
    *
    * @param {string} uid The uid to use as the JWT subject.
    * @param {object=} developerClaims Optional additional claims to include in the JWT payload.
    *
    * @return {Promise<string>} A JWT for the provided payload.
    */
   public createCustomToken(uid: string, developerClaims?: object): Promise<string> {
-    // This is not yet supported by the Auth server. It is also not yet determined how this will be
-    // supported.
-    return Promise.reject(
-        new FirebaseAuthError(AuthClientErrorCode.UNSUPPORTED_TENANT_OPERATION));
+    return this.tokenGenerator.createCustomTokenWithTenantId(uid, this.tenantId, developerClaims);
   }
 
   /**
diff --git a/src/auth/token-generator.ts b/src/auth/token-generator.ts
@@ -74,6 +74,7 @@ interface JWTBody {
   exp: number;
   iss: string;
   sub: string;
+  tenant_id?: string;
 }
 
 /**
@@ -254,13 +255,38 @@ export class FirebaseTokenGenerator {
    *                           service account key and containing the provided payload.
    */
   public createCustomToken(uid: string, developerClaims?: {[key: string]: any}): Promise<string> {
+    return this.createCustomTokenInternal(uid, null, developerClaims);
+  }
+
+  /**
+   * Creates a new Firebase Auth Custom token.
+   *
+   * @param {string} uid The user ID to use for the generated Firebase Auth Custom token.
+   * @param {string} tenantId The tenant ID to use for the generated Firebase Auth Custom token.
+   * @param {object} [developerClaims] Optional developer claims to include in the generated Firebase
+   *                 Auth Custom token.
+   * @return {Promise<string>} A Promise fulfilled with a Firebase Auth Custom token signed with a
+   *                           service account key and containing the provided payload.
+   */
+  public createCustomTokenWithTenantId(
+      uid: string, tenantId: string, developerClaims?: {[key: string]: any}): Promise<string> {
+    if (!validator.isNonEmptyString(tenantId)) {
+      throw new FirebaseAuthError(
+        AuthClientErrorCode.INVALID_ARGUMENT,
+        '`tenantId` argument must be a non-empty string uid.');
+    }
+    return this.createCustomTokenInternal(uid, tenantId, developerClaims);
+  }
+
+  private createCustomTokenInternal(
+      uid: string, tenantId: string | null, developerClaims?: {[key: string]: any}): Promise<string> {
     let errorMessage: string;
-    if (typeof uid !== 'string' || uid === '') {
-      errorMessage = 'First argument to createCustomToken() must be a non-empty string uid.';
+    if (!validator.isNonEmptyString(uid)) {
+      errorMessage = '`uid` argument must be a non-empty string uid.';
     } else if (uid.length > 128) {
-      errorMessage = 'First argument to createCustomToken() must a uid with less than or equal to 128 characters.';
+      errorMessage = '`uid` argument must a uid with less than or equal to 128 characters.';
     } else if (!this.isDeveloperClaimsValid_(developerClaims)) {
-      errorMessage = 'Second argument to createCustomToken() must be an object containing the developer claims.';
+      errorMessage = '`developerClaims` argument must be a valid, non-null object containing the developer claims.';
     }
 
     if (typeof errorMessage !== 'undefined') {
@@ -296,6 +322,9 @@ export class FirebaseTokenGenerator {
         sub: account,
         uid,
       };
+      if (tenantId) {
+        body.tenant_id = tenantId;
+      }
       if (Object.keys(claims).length > 0) {
         body.claims = claims;
       }
diff --git a/test/integration/auth.spec.ts b/test/integration/auth.spec.ts
@@ -710,6 +710,49 @@ describe('admin.auth', () => {
             expect(userRecord.uid).to.equal(createdUserUid);
           });
       });
+
+      it('createCustomToken() mints a JWT that can be used to sign in tenant users', () => {
+        return tenantAwareAuth.createCustomToken('uid1')
+          .then((customToken) => {
+            firebase.auth().tenantId = createdTenantId;
+            return firebase.auth().signInWithCustomToken(customToken);
+          })
+          .then(({user}) => {
+            return user.getIdToken();
+          })
+          .then((idToken) => {
+            return tenantAwareAuth.verifyIdToken(idToken);
+          })
+          .then((token) => {
+            expect(token.uid).to.equal('uid1');
+            expect(token.firebase.tenant).to.equal(createdTenantId);
+          });
+      });
+
+      it('createCustomToken() should fail if tenantIds mismatch', async () => {
+        const aDifferentTenant = await admin.auth().tenantManager().createTenant({displayName: 'A-Different-Tenant'});
+        try {
+          const customToken = await tenantAwareAuth.createCustomToken('uid1');
+
+          firebase.auth().tenantId = aDifferentTenant.tenantId;
+
+          await expect(firebase.auth().signInWithCustomToken(customToken))
+            .to.be.rejectedWith('Specified tenant ID does not match the custom token.');
+
+          // TODO(rsgowman): Currently, the above error has a code of
+          // 'auth/internal-error', i.e. you could add the following chain onto
+          // it:
+          //
+          //     .and.eventually.have.property('code', 'auth/internal-error');
+          //
+          // However, this doesn't strike me as an internal error, implying
+          // that the client sdk is slightly "wrong". Fix the client sdk to
+          // return a better error code, and then add the above statement to
+          // the expect statement (with the new error code.)
+        } finally {
+          admin.auth().tenantManager().deleteTenant(aDifferentTenant.tenantId);
+        }
+      });
     });
 
     // Sanity check OIDC/SAML config management API.
diff --git a/test/unit/auth/auth.spec.ts b/test/unit/auth/auth.spec.ts
@@ -16,6 +16,7 @@
 
 'use strict';
 
+import * as jwt from 'jsonwebtoken';
 import * as _ from 'lodash';
 import * as chai from 'chai';
 import * as sinon from 'sinon';
@@ -28,7 +29,6 @@ import * as mocks from '../../resources/mocks';
 import {Auth, TenantAwareAuth, BaseAuth, DecodedIdToken} from '../../../src/auth/auth';
 import {UserRecord} from '../../../src/auth/user-record';
 import {FirebaseApp} from '../../../src/firebase-app';
-import {FirebaseTokenGenerator} from '../../../src/auth/token-generator';
 import {
   AuthRequestHandler, TenantAwareAuthRequestHandler, AbstractAuthRequestHandler,
 } from '../../../src/auth/auth-api-request';
@@ -328,63 +328,49 @@ AUTH_CONFIGS.forEach((testConfig) => {
     }
 
     describe('createCustomToken()', () => {
-      let spy: sinon.SinonSpy;
-      beforeEach(() => {
-        spy = sinon.spy(FirebaseTokenGenerator.prototype, 'createCustomToken');
-      });
-
-      afterEach(() => {
-        spy.restore();
+      it('should return a jwt', async () => {
+        const token = await auth.createCustomToken('uid1');
+        const decodedToken = jwt.decode(token, {complete: true});
+        expect(decodedToken).to.have.property('header').that.has.property('typ', 'JWT');
       });
 
       if (testConfig.Auth === TenantAwareAuth) {
-        it('should reject with an unsupported tenant operation error', () => {
-          const expectedError = new FirebaseAuthError(AuthClientErrorCode.UNSUPPORTED_TENANT_OPERATION);
-          return auth.createCustomToken(mocks.uid)
-            .then(() => {
-              throw new Error('Unexpected success');
-            })
-            .catch((error) => {
-              expect(error).to.deep.equal(expectedError);
-            });
+        it('should contain tenant_id', async () => {
+          const token = await auth.createCustomToken('uid1');
+          expect(jwt.decode(token)).to.have.property('tenant_id', TENANT_ID);
         });
       } else {
-        it('should throw if a cert credential is not specified', () => {
-          const mockCredentialAuth = testConfig.init(mocks.mockCredentialApp());
-
-          expect(() => {
-            mockCredentialAuth.createCustomToken(mocks.uid, mocks.developerClaims);
-          }).not.to.throw;
+        it('should not contain tenant_id', async () => {
+          const token = await auth.createCustomToken('uid1');
+          expect(jwt.decode(token)).to.not.have.property('tenant_id');
         });
+      }
 
-        it('should forward on the call to the token generator\'s createCustomToken() method', () => {
-          const developerClaimsCopy = deepCopy(mocks.developerClaims);
-          return auth.createCustomToken(mocks.uid, mocks.developerClaims)
-            .then(() => {
-              expect(spy)
-                .to.have.been.calledOnce
-                .and.calledWith(mocks.uid, developerClaimsCopy);
-            });
-        });
+      it('should throw if a cert credential is not specified', () => {
+        const mockCredentialAuth = testConfig.init(mocks.mockCredentialApp());
 
-        it('should be fulfilled given an app which returns null access tokens', () => {
-          // createCustomToken() does not rely on an access token and therefore works in this scenario.
-          return nullAccessTokenAuth.createCustomToken(mocks.uid, mocks.developerClaims)
-            .should.eventually.be.fulfilled;
-        });
+        expect(() => {
+          mockCredentialAuth.createCustomToken(mocks.uid, mocks.developerClaims);
+        }).not.to.throw;
+      });
 
-        it('should be fulfilled given an app which returns invalid access tokens', () => {
-          // createCustomToken() does not rely on an access token and therefore works in this scenario.
-          return malformedAccessTokenAuth.createCustomToken(mocks.uid, mocks.developerClaims)
-            .should.eventually.be.fulfilled;
-        });
+      it('should be fulfilled given an app which returns null access tokens', () => {
+        // createCustomToken() does not rely on an access token and therefore works in this scenario.
+        return nullAccessTokenAuth.createCustomToken(mocks.uid, mocks.developerClaims)
+          .should.eventually.be.fulfilled;
+      });
 
-        it('should be fulfilled given an app which fails to generate access tokens', () => {
-          // createCustomToken() does not rely on an access token and therefore works in this scenario.
-          return rejectedPromiseAccessTokenAuth.createCustomToken(mocks.uid, mocks.developerClaims)
-            .should.eventually.be.fulfilled;
-        });
-      }
+      it('should be fulfilled given an app which returns invalid access tokens', () => {
+        // createCustomToken() does not rely on an access token and therefore works in this scenario.
+        return malformedAccessTokenAuth.createCustomToken(mocks.uid, mocks.developerClaims)
+          .should.eventually.be.fulfilled;
+      });
+
+      it('should be fulfilled given an app which fails to generate access tokens', () => {
+        // createCustomToken() does not rely on an access token and therefore works in this scenario.
+        return rejectedPromiseAccessTokenAuth.createCustomToken(mocks.uid, mocks.developerClaims)
+          .should.eventually.be.fulfilled;
+      });
     });
 
     it('verifyIdToken() should throw when project ID is not specified', () => {
diff --git a/test/unit/auth/token-generator.spec.ts b/test/unit/auth/token-generator.spec.ts
