Allow createCustomToken() to work with tenant-aware auth (#708)

rsgowman · web-flow · commit 99d321437271 · 2019-12-16T15:51:36.000-05:00
diff --git a/src/auth/auth.ts b/src/auth/auth.ts
@@ -98,13 +98,21 @@ export class BaseAuth<T extends AbstractAuthRequestHandler> {
   /**
    * The BaseAuth class constructor.
    *
-   * @param {T} authRequestHandler The RPC request handler
-   *     for this instance.
+   * @param app The FirebaseApp to associate with this Auth instance.
+   * @param authRequestHandler The RPC request handler for this instance.
+   * @param tokenGenerator Optional token generator. If not specified, a
+   *     (non-tenant-aware) instance will be created. Use this paramter to
+   *     specify a tenant-aware tokenGenerator.
    * @constructor
    */
-  constructor(app: FirebaseApp, protected readonly authRequestHandler: T) {
-    const cryptoSigner = cryptoSignerFromApp(app);
-    this.tokenGenerator = new FirebaseTokenGenerator(cryptoSigner);
+  constructor(app: FirebaseApp, protected readonly authRequestHandler: T, tokenGenerator?: FirebaseTokenGenerator) {
+    if (tokenGenerator) {
+      this.tokenGenerator = tokenGenerator;
+    } else {
+      const cryptoSigner = cryptoSignerFromApp(app);
+      this.tokenGenerator = new FirebaseTokenGenerator(cryptoSigner);
+    }
+
     this.sessionCookieVerifier = createSessionCookieVerifier(app);
     this.idTokenVerifier = createIdTokenVerifier(app);
   }
@@ -613,26 +621,12 @@ export class TenantAwareAuth extends BaseAuth<TenantAwareAuthRequestHandler> {
    * @constructor
    */
   constructor(app: FirebaseApp, tenantId: string) {
-    super(app, new TenantAwareAuthRequestHandler(app, tenantId));
+    const cryptoSigner = cryptoSignerFromApp(app);
+    const tokenGenerator = new FirebaseTokenGenerator(cryptoSigner, tenantId);
+    super(app, new TenantAwareAuthRequestHandler(app, tenantId), tokenGenerator);
     utils.addReadonlyGetter(this, 'tenantId', tenantId);
   }
 
-  /**
-   * Creates a new custom token that can be sent back to a client to use with
-   * 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));
-  }
-
   /**
    * Verifies a JWT auth token. Returns a Promise with the tokens claims. Rejects
    * the promise if the token could not be verified. If checkRevoked is set to true,
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;
 }
 
 /**
@@ -247,33 +248,43 @@ export class FirebaseTokenGenerator {
 
   private readonly signer: CryptoSigner;
 
-  constructor(signer: CryptoSigner) {
+  /**
+   * @param tenantId The tenant ID to use for the generated Firebase Auth
+   *     Custom token. If absent, then no tenant ID claim will be set in the
+   *     resulting JWT.
+   */
+  constructor(signer: CryptoSigner, public readonly tenantId?: string) {
     if (!validator.isNonNullObject(signer)) {
       throw new FirebaseAuthError(
         AuthClientErrorCode.INVALID_CREDENTIAL,
         'INTERNAL ASSERT: Must provide a CryptoSigner to use FirebaseTokenGenerator.',
       );
     }
+    if (typeof tenantId !== 'undefined' && !validator.isNonEmptyString(tenantId)) {
+      throw new FirebaseAuthError(
+        AuthClientErrorCode.INVALID_ARGUMENT,
+        '`tenantId` argument must be a non-empty string.');
+    }
     this.signer = signer;
   }
 
   /**
    * Creates a new Firebase Auth Custom token.
    *
-   * @param {string} uid The user 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.
+   * @param uid The user ID to use for the generated Firebase Auth Custom token.
+   * @param developerClaims Optional developer claims to include in the generated Firebase
+   *     Auth Custom token.
+   * @return A Promise fulfilled with a Firebase Auth Custom token signed with a
+   *     service account key and containing the provided payload.
    */
   public createCustomToken(uid: string, developerClaims?: {[key: string]: any}): Promise<string> {
     let errorMessage: string | undefined;
-    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 (errorMessage) {
@@ -309,6 +320,9 @@ export class FirebaseTokenGenerator {
         sub: account,
         uid,
       };
+      if (this.tenantId) {
+        body.tenant_id = this.tenantId;
+      }
       if (Object.keys(claims).length > 0) {
         body.claims = claims;
       }
diff --git a/src/index.d.ts b/src/index.d.ts
@@ -1445,7 +1445,8 @@ declare namespace admin.auth {
     /**
      * Creates a new Firebase custom token (JWT) that can be sent back to a client
      * device to use to sign in with the client SDKs' `signInWithCustomToken()`
-     * methods.
+     * methods. (Tenant-aware instances will also embed the tenant ID in the
+     * token.)
      *
      * See [Create Custom Tokens](/docs/auth/admin/create-custom-tokens) for code
      * samples and detailed documentation.
diff --git a/test/integration/auth.spec.ts b/test/integration/auth.spec.ts
@@ -30,7 +30,7 @@ import url = require('url');
 import * as mocks from '../resources/mocks';
 import { AuthProviderConfig } from '../../src/auth/auth-config';
 import { deepExtend, deepCopy } from '../../src/utils/deep-copy';
-import { User } from '@firebase/auth-types';
+import { User, FirebaseAuth } from '@firebase/auth-types';
 
 /* tslint:disable:no-var-requires */
 const chalk = require('chalk');
@@ -91,6 +91,10 @@ function randomOidcProviderId(): string {
   return 'oidc.' + generateRandomString(10, false).toLowerCase();
 }
 
+function clientAuth(): FirebaseAuth {
+  expect(firebase.auth).to.be.ok;
+  return firebase.auth!();
+}
 
 describe('admin.auth', () => {
 
@@ -213,7 +217,7 @@ describe('admin.auth', () => {
     let currentIdToken: string;
     let currentUser: User;
     // Sign in with an email and password account.
-    return firebase.auth!().signInWithEmailAndPassword(mockUserData.email, mockUserData.password)
+    return clientAuth().signInWithEmailAndPassword(mockUserData.email, mockUserData.password)
       .then(({user}) => {
         expect(user).to.exist;
         currentUser = user!;
@@ -248,7 +252,7 @@ describe('admin.auth', () => {
       })
       .then(() => {
         // New sign-in should succeed.
-        return firebase.auth!().signInWithEmailAndPassword(
+        return clientAuth().signInWithEmailAndPassword(
             mockUserData.email, mockUserData.password);
       })
       .then(({user}) => {
@@ -273,7 +277,7 @@ describe('admin.auth', () => {
         // Confirm custom claims set on the UserRecord.
         expect(userRecord.customClaims).to.deep.equal(customClaims);
         expect(userRecord.email).to.exist;
-        return firebase.auth!().signInWithEmailAndPassword(
+        return clientAuth().signInWithEmailAndPassword(
           userRecord.email!, mockUserData.password);
       })
       .then(({user}) => {
@@ -302,8 +306,8 @@ describe('admin.auth', () => {
         // Custom claims should be cleared.
         expect(userRecord.customClaims).to.deep.equal({});
         // Force token refresh. All claims should be cleared.
-        expect(firebase.auth!().currentUser).to.exist;
-        return firebase.auth!().currentUser!.getIdToken(true);
+        expect(clientAuth().currentUser).to.exist;
+        return clientAuth().currentUser!.getIdToken(true);
       })
       .then((idToken) => {
         // Verify ID token contents.
@@ -368,7 +372,7 @@ describe('admin.auth', () => {
       isAdmin: true,
     })
     .then((customToken) => {
-      return firebase.auth!().signInWithCustomToken(customToken);
+      return clientAuth().signInWithCustomToken(customToken);
     })
     .then(({user}) => {
       expect(user).to.exist;
@@ -388,7 +392,7 @@ describe('admin.auth', () => {
       isAdmin: true,
     })
     .then((customToken) => {
-      return firebase.auth!().signInWithCustomToken(customToken);
+      return clientAuth().signInWithCustomToken(customToken);
     })
     .then(({user}) => {
       expect(user).to.exist;
@@ -426,7 +430,7 @@ describe('admin.auth', () => {
 
     // Sign out after each test.
     afterEach(() => {
-      return firebase.auth!().signOut();
+      return clientAuth().signOut();
     });
 
     // Delete test user at the end of test suite.
@@ -443,10 +447,10 @@ describe('admin.auth', () => {
         .then((link) => {
           const code = getActionCode(link);
           expect(getContinueUrl(link)).equal(actionCodeSettings.url);
-          return firebase.auth!().confirmPasswordReset(code, newPassword);
+          return clientAuth().confirmPasswordReset(code, newPassword);
         })
         .then(() => {
-          return firebase.auth!().signInWithEmailAndPassword(email, newPassword);
+          return clientAuth().signInWithEmailAndPassword(email, newPassword);
         })
         .then((result) => {
           expect(result.user).to.exist;
@@ -466,10 +470,10 @@ describe('admin.auth', () => {
         .then((link) => {
           const code = getActionCode(link);
           expect(getContinueUrl(link)).equal(actionCodeSettings.url);
-          return firebase.auth!().applyActionCode(code);
+          return clientAuth().applyActionCode(code);
         })
         .then(() => {
-          return firebase.auth!().signInWithEmailAndPassword(email, userData.password);
+          return clientAuth().signInWithEmailAndPassword(email, userData.password);
         })
         .then((result) => {
           expect(result.user).to.exist;
@@ -482,7 +486,7 @@ describe('admin.auth', () => {
       return admin.auth().generateSignInWithEmailLink(email, actionCodeSettings)
         .then((link) => {
           expect(getContinueUrl(link)).equal(actionCodeSettings.url);
-          return firebase.auth!().signInWithEmailLink(email, link);
+          return clientAuth().signInWithEmailLink(email, link);
         })
         .then((result) => {
           expect(result.user).to.exist;
@@ -722,6 +726,23 @@ describe('admin.auth', () => {
             expect(userRecord.uid).to.equal(createdUserUid);
           });
       });
+
+      it('createCustomToken() mints a JWT that can be used to sign in tenant users', async () => {
+        try {
+          clientAuth().tenantId = createdTenantId;
+
+          const customToken = await tenantAwareAuth.createCustomToken('uid1');
+          const {user} = await clientAuth().signInWithCustomToken(customToken);
+          expect(user).to.not.be.null;
+          const idToken = await user!.getIdToken();
+          const token = await tenantAwareAuth.verifyIdToken(idToken);
+
+          expect(token.uid).to.equal('uid1');
+          expect(token.firebase.tenant).to.equal(createdTenantId);
+        } finally {
+          clientAuth().tenantId = null;
+        }
+      });
     });
 
     // Sanity check OIDC/SAML config management API.
@@ -1203,7 +1224,7 @@ describe('admin.auth', () => {
 
     it('creates a valid Firebase session cookie', () => {
       return admin.auth().createCustomToken(uid, {admin: true, groupId: '1234'})
-        .then((customToken) => firebase.auth!().signInWithCustomToken(customToken))
+        .then((customToken) => clientAuth().signInWithCustomToken(customToken))
         .then(({user}) => {
           expect(user).to.exist;
           return user!.getIdToken();
@@ -1239,7 +1260,7 @@ describe('admin.auth', () => {
     it('creates a revocable session cookie', () => {
       let currentSessionCookie: string;
       return admin.auth().createCustomToken(uid2)
-        .then((customToken) => firebase.auth!().signInWithCustomToken(customToken))
+        .then((customToken) => clientAuth().signInWithCustomToken(customToken))
         .then(({user}) => {
           expect(user).to.exist;
           return user!.getIdToken();
@@ -1266,7 +1287,7 @@ describe('admin.auth', () => {
 
     it('fails when called with a revoked ID token', () => {
       return admin.auth().createCustomToken(uid3, {admin: true, groupId: '1234'})
-        .then((customToken) => firebase.auth!().signInWithCustomToken(customToken))
+        .then((customToken) => clientAuth().signInWithCustomToken(customToken))
         .then(({user}) => {
           expect(user).to.exist;
           return user!.getIdToken();
@@ -1294,7 +1315,7 @@ describe('admin.auth', () => {
 
     it('fails when called with a Firebase ID token', () => {
       return admin.auth().createCustomToken(uid)
-        .then((customToken) => firebase.auth!().signInWithCustomToken(customToken))
+        .then((customToken) => clientAuth().signInWithCustomToken(customToken))
         .then(({user}) => {
           expect(user).to.exist;
           return user!.getIdToken();
@@ -1580,7 +1601,7 @@ function testImportAndSignInUser(
       expect(result.successCount).to.equal(1);
       expect(result.errors.length).to.equal(0);
       // Sign in with an email and password to the imported account.
-      return firebase.auth!().signInWithEmailAndPassword(users[0].email, rawPassword);
+      return clientAuth().signInWithEmailAndPassword(users[0].email, rawPassword);
     })
     .then(({user}) => {
       // Confirm successful sign-in.
diff --git a/test/unit/auth/auth.spec.ts b/test/unit/auth/auth.spec.ts
diff --git a/test/unit/auth/token-generator.spec.ts b/test/unit/auth/token-generator.spec.ts
