Add tenant-aware token generation and verification.

src/main/java/com/google/firebase/auth/AbstractFirebaseAuth.java

@@ -1092,7 +1092,7 @@ final void destroy() {
       destroyed.set(true);
     }
   }
-  
+
   /** Performs any additional required clean up. */
   protected abstract void doDestroy();
 

src/main/java/com/google/firebase/auth/FirebaseToken.java

@@ -44,7 +44,11 @@ public String getUid() {
 
   /** Returns the tenant ID for the this token. */
   public String getTenantId() {
-    return (String) claims.get("tenant_id");
+    Map<String, Object> firebase = (Map<String, Object>) claims.get("firebase");
+    if (firebase == null) {
+      return null;
+    }
+    return (String) firebase.get("tenant");
   }
 
   /** Returns the Issuer for the this token. */

src/main/java/com/google/firebase/auth/FirebaseTokenUtils.java

@@ -30,6 +30,7 @@
 import com.google.firebase.ImplFirebaseTrampolines;
 import com.google.firebase.auth.internal.CryptoSigners;
 import com.google.firebase.auth.internal.FirebaseTokenFactory;
+import com.google.firebase.internal.Nullable;
 
 import java.io.IOException;
 
@@ -52,11 +53,17 @@ final class FirebaseTokenUtils {
   private FirebaseTokenUtils() { }
 
   static FirebaseTokenFactory createTokenFactory(FirebaseApp firebaseApp, Clock clock) {
+    return createTokenFactory(firebaseApp, clock, null);
+  }
+
+  static FirebaseTokenFactory createTokenFactory(
+      FirebaseApp firebaseApp, Clock clock, @Nullable String tenantId) {
     try {
       return new FirebaseTokenFactory(
           firebaseApp.getOptions().getJsonFactory(),
           clock,
-          CryptoSigners.getCryptoSigner(firebaseApp));
+          CryptoSigners.getCryptoSigner(firebaseApp),
+          tenantId);
     } catch (IOException e) {
       throw new IllegalStateException(
           "Failed to initialize FirebaseTokenFactory. Make sure to initialize the SDK "
@@ -68,6 +75,11 @@ static FirebaseTokenFactory createTokenFactory(FirebaseApp firebaseApp, Clock cl
   }
 
   static FirebaseTokenVerifierImpl createIdTokenVerifier(FirebaseApp app, Clock clock) {
+    return createIdTokenVerifier(app, clock, null);
+  }
+
+  static FirebaseTokenVerifierImpl createIdTokenVerifier(
+      FirebaseApp app, Clock clock, @Nullable String tenantId) {
     String projectId = ImplFirebaseTrampolines.getProjectId(app);
     checkState(!Strings.isNullOrEmpty(projectId),
         "Must initialize FirebaseApp with a project ID to call verifyIdToken()");
@@ -82,6 +94,7 @@ static FirebaseTokenVerifierImpl createIdTokenVerifier(FirebaseApp app, Clock cl
         .setJsonFactory(app.getOptions().getJsonFactory())
         .setPublicKeysManager(publicKeysManager)
         .setIdTokenVerifier(idTokenVerifier)
+        .setTenantId(tenantId)
         .build();
   }
 

src/main/java/com/google/firebase/auth/FirebaseTokenVerifierImpl.java

@@ -28,6 +28,7 @@
 import com.google.api.client.util.ArrayMap;
 import com.google.common.base.Joiner;
 import com.google.common.base.Strings;
+import com.google.firebase.internal.Nullable;
 import java.io.IOException;
 import java.math.BigDecimal;
 import java.security.GeneralSecurityException;
@@ -53,6 +54,7 @@ final class FirebaseTokenVerifierImpl implements FirebaseTokenVerifier {
   private final String shortName;
   private final String articledShortName;
   private final String docUrl;
+  private final String tenantId;
 
   private FirebaseTokenVerifierImpl(Builder builder) {
     this.jsonFactory = checkNotNull(builder.jsonFactory);
@@ -65,6 +67,7 @@ private FirebaseTokenVerifierImpl(Builder builder) {
     this.shortName = builder.shortName;
     this.articledShortName = prefixWithIndefiniteArticle(this.shortName);
     this.docUrl = builder.docUrl;
+    this.tenantId = Strings.nullToEmpty(builder.tenantId);
   }
 
   /**
@@ -90,7 +93,9 @@ public FirebaseToken verifyToken(String token) throws FirebaseAuthException {
     IdToken idToken = parse(token);
     checkContents(idToken);
     checkSignature(idToken);
-    return new FirebaseToken(idToken.getPayload());
+    FirebaseToken firebaseToken = new FirebaseToken(idToken.getPayload());
+    checkTenantId(firebaseToken);
+    return firebaseToken;
   }
 
   GooglePublicKeysManager getPublicKeysManager() {
@@ -278,6 +283,18 @@ private boolean containsLegacyUidField(IdToken.Payload payload) {
     return false;
   }
 
+  private void checkTenantId(final FirebaseToken firebaseToken) throws FirebaseAuthException {
+    String tokenTenantId = Strings.nullToEmpty(firebaseToken.getTenantId());
+    if (!this.tenantId.equals(tokenTenantId)) {
+      throw new FirebaseAuthException(
+          FirebaseUserManager.TENANT_ID_MISMATCH_ERROR,
+          String.format(
+            "The tenant ID ('%s') of the token did not match the expected ('%s') value",
+            tokenTenantId,
+            tenantId));
+    }
+  }
+
   static Builder builder() {
     return new Builder();
   }
@@ -290,6 +307,7 @@ static final class Builder {
     private String shortName;
     private IdTokenVerifier idTokenVerifier;
     private String docUrl;
+    private String tenantId;
 
     private Builder() { }
 
@@ -323,6 +341,11 @@ Builder setDocUrl(String docUrl) {
       return this;
     }
 
+    Builder setTenantId(@Nullable String tenantId) {
+      this.tenantId = tenantId;
+      return this;
+    }
+
     FirebaseTokenVerifierImpl build() {
       return new FirebaseTokenVerifierImpl(this);
     }

src/main/java/com/google/firebase/auth/FirebaseUserManager.java

@@ -65,6 +65,7 @@
  */
 class FirebaseUserManager {
 
+  static final String TENANT_ID_MISMATCH_ERROR = "tenant-id-mismatch";
   static final String TENANT_NOT_FOUND_ERROR = "tenant-not-found";
   static final String USER_NOT_FOUND_ERROR = "user-not-found";
   static final String INTERNAL_ERROR = "internal-error";
@@ -84,6 +85,7 @@ class FirebaseUserManager {
       .put("INVALID_PHONE_NUMBER", "invalid-phone-number")
       .put("PHONE_NUMBER_EXISTS", "phone-number-already-exists")
       .put("PROJECT_NOT_FOUND", "project-not-found")
+      .put("TENANT_ID_MISMATCH", TENANT_ID_MISMATCH_ERROR)
       .put("TENANT_NOT_FOUND", TENANT_NOT_FOUND_ERROR)
       .put("USER_NOT_FOUND", USER_NOT_FOUND_ERROR)
       .put("WEAK_PASSWORD", "invalid-password")

src/main/java/com/google/firebase/auth/TenantAwareFirebaseAuth.java

@@ -41,22 +41,20 @@ public class TenantAwareFirebaseAuth extends AbstractFirebaseAuth {
   }
 
   private static Builder builderFromAppAndTenantId(final FirebaseApp app, final String tenantId) {
-    // TODO(micahstairs): Incorporate tenant ID into token generation as well as ID token and
-    // session cookie verification.
     return AbstractFirebaseAuth.builder()
         .setFirebaseApp(app)
         .setTokenFactory(
             new Supplier<FirebaseTokenFactory>() {
               @Override
               public FirebaseTokenFactory get() {
-                return FirebaseTokenUtils.createTokenFactory(app, Clock.SYSTEM);
+                return FirebaseTokenUtils.createTokenFactory(app, Clock.SYSTEM, tenantId);
               }
             })
         .setIdTokenVerifier(
             new Supplier<FirebaseTokenVerifier>() {
               @Override
               public FirebaseTokenVerifier get() {
-                return FirebaseTokenUtils.createIdTokenVerifier(app, Clock.SYSTEM);
+                return FirebaseTokenUtils.createIdTokenVerifier(app, Clock.SYSTEM, tenantId);
               }
             })
         .setCookieVerifier(

src/main/java/com/google/firebase/auth/internal/FirebaseCustomAuthToken.java

@@ -22,6 +22,7 @@
 import com.google.api.client.json.webtoken.JsonWebSignature;
 import com.google.api.client.util.Key;
 import com.google.firebase.auth.FirebaseToken;
+import com.google.firebase.internal.Nullable;
 
 import java.io.IOException;
 
@@ -77,6 +78,9 @@ public static class Payload extends IdToken.Payload {
     @Key("claims")
     private GenericJson developerClaims;
 
+    @Key("tenant_id")
+    private String tenantId;
+
     public final String getUid() {
       return uid;
     }
@@ -95,6 +99,15 @@ public Payload setDeveloperClaims(GenericJson developerClaims) {
       return this;
     }
 
+    public final String getTenantId() {
+      return tenantId;
+    }
+
+    public Payload setTenantId(String tenantId) {
+      this.tenantId = tenantId;
+      return this;
+    }
+
     @Override
     public Payload setIssuer(String issuer) {
       return (Payload) super.setIssuer(issuer);

src/main/java/com/google/firebase/auth/internal/FirebaseTokenFactory.java

@@ -25,8 +25,9 @@
 import com.google.api.client.util.Base64;
 import com.google.api.client.util.Clock;
 import com.google.api.client.util.StringUtils;
-
 import com.google.common.base.Strings;
+import com.google.firebase.internal.Nullable;
+
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Map;
@@ -41,11 +42,18 @@ public class FirebaseTokenFactory {
   private final JsonFactory jsonFactory;
   private final Clock clock;
   private final CryptoSigner signer;
+  private final String tenantId;
 
   public FirebaseTokenFactory(JsonFactory jsonFactory, Clock clock, CryptoSigner signer) {
+    this(jsonFactory, clock, signer, null);
+  }
+
+  public FirebaseTokenFactory(
+      JsonFactory jsonFactory, Clock clock, CryptoSigner signer, @Nullable String tenantId) {
     this.jsonFactory = checkNotNull(jsonFactory);
     this.clock = checkNotNull(clock);
     this.signer = checkNotNull(signer);
+    this.tenantId = tenantId;
   }
 
   String createSignedCustomAuthTokenForUser(String uid) throws IOException {
@@ -68,6 +76,9 @@ public String createSignedCustomAuthTokenForUser(
             .setAudience(FirebaseCustomAuthToken.FIREBASE_AUDIENCE)
             .setIssuedAtTimeSeconds(issuedAt)
             .setExpirationTimeSeconds(issuedAt + FirebaseCustomAuthToken.TOKEN_DURATION_SECONDS);
+    if (!Strings.isNullOrEmpty(tenantId)) {
+      payload.setTenantId(tenantId);
+    }
 
     if (developerClaims != null) {
       Collection<String> reservedNames = payload.getClassInfo().getNames();

src/test/java/com/google/firebase/auth/FirebaseAuthIT.java

@@ -47,6 +47,7 @@
 import com.google.firebase.FirebaseOptions;
 import com.google.firebase.ImplFirebaseTrampolines;
 import com.google.firebase.auth.hash.Scrypt;
+import com.google.firebase.internal.Nullable;
 import com.google.firebase.testing.IntegrationTestUtils;
 import java.io.IOException;
 import java.net.URLDecoder;
@@ -705,6 +706,58 @@ public void testCustomTokenWithIAM() throws Exception {
     }
   }
 
+  @Test
+  public void testTenantAwareCustomToken() throws Exception {
+    // Create tenant to use.
+    TenantManager tenantManager = auth.getTenantManager();
+    Tenant.CreateRequest tenantCreateRequest =
+        new Tenant.CreateRequest().setDisplayName("DisplayName");
+    String tenantId = tenantManager.createTenant(tenantCreateRequest).getTenantId();
+
+    try {
+      // Create and decode a token with a tenant-aware client.
+      TenantAwareFirebaseAuth tenantAwareAuth = auth.getTenantManager().getAuthForTenant(tenantId);
+      String customToken = tenantAwareAuth.createCustomTokenAsync("user1").get();
+      String idToken = signInWithCustomToken(customToken, tenantId);
+      FirebaseToken decoded = tenantAwareAuth.verifyIdTokenAsync(idToken).get();
+      assertEquals("user1", decoded.getUid());
+      assertEquals(tenantId, decoded.getTenantId());
+    } finally {
+      // Delete tenant.
+      tenantManager.deleteTenantAsync(tenantId).get();
+    }
+  }
+
+  @Test
+  public void testVerifyTokenWithWrongTenantAwareClient() throws Exception {
+    // Create tenant to use.
+    TenantManager tenantManager = auth.getTenantManager();
+    Tenant.CreateRequest tenantCreateRequest =
+        new Tenant.CreateRequest().setDisplayName("DisplayName");
+    String tenantId = tenantManager.createTenant(tenantCreateRequest).getTenantId();
+
+    // Create tenant-aware clients.
+    TenantAwareFirebaseAuth tenantAwareAuth1 = auth.getTenantManager().getAuthForTenant(tenantId);
+    TenantAwareFirebaseAuth tenantAwareAuth2 = auth.getTenantManager().getAuthForTenant("OTHER");
+
+    try {
+      // Create a token with one client and decode with the other.
+      String customToken = tenantAwareAuth1.createCustomTokenAsync("user").get();
+      String idToken = signInWithCustomToken(customToken, tenantId);
+      try {
+        tenantAwareAuth2.verifyIdTokenAsync(idToken).get();
+        fail("No error thrown for verifying a token with the wrong tenant-aware client");
+      } catch (ExecutionException e) {
+        assertTrue(e.getCause() instanceof FirebaseAuthException);
+        assertEquals(FirebaseUserManager.TENANT_ID_MISMATCH_ERROR,
+            ((FirebaseAuthException) e.getCause()).getErrorCode());
+      }
+    } finally {
+      // Delete tenant.
+      tenantManager.deleteTenantAsync(tenantId).get();
+    }
+  }
+
   @Test
   public void testVerifyIdToken() throws Exception {
     String customToken = auth.createCustomTokenAsync("user2").get();
@@ -931,12 +984,21 @@ private String randomPhoneNumber() {
   }
 
   private String signInWithCustomToken(String customToken) throws IOException {
-    GenericUrl url = new GenericUrl(VERIFY_CUSTOM_TOKEN_URL + "?key="
+    return signInWithCustomToken(customToken, null);
+  }
+
+  private String signInWithCustomToken(
+      String customToken, @Nullable String tenantId) throws IOException {
+    final GenericUrl url = new GenericUrl(VERIFY_CUSTOM_TOKEN_URL + "?key="
         + IntegrationTestUtils.getApiKey());
-    Map<String, Object> content = ImmutableMap.<String, Object>of(
-        "token", customToken, "returnSecureToken", true);
+    ImmutableMap.Builder<String, Object> content = ImmutableMap.<String, Object>builder();
+    content.put("token", customToken);
+    content.put("returnSecureToken", true);
+    if (tenantId != null) {
+      content.put("tenantId", tenantId);
+    }
     HttpRequest request = transport.createRequestFactory().buildPostRequest(url,
-        new JsonHttpContent(jsonFactory, content));
+        new JsonHttpContent(jsonFactory, content.build()));
     request.setParser(new JsonObjectParser(jsonFactory));
     HttpResponse response = request.execute();
     try {