AppCheck integration with Firestore (#3027)

* AppCheck integration with Firestore.

* Add unit tests.

* Apply suggestions from code review

Co-authored-by: Sebastian Schmidt <[email protected]>

* Simplify the applyRequestMetadata logic.

* Switch to using Deferred.

Co-authored-by: Sebastian Schmidt <[email protected]>

Author: ehsannas

firebase-firestore/firebase-firestore.gradle

@@ -141,6 +141,7 @@ dependencies {
     implementation project(':protolite-well-known-types')
     implementation project(':firebase-database-collection')
     implementation project(':firebase-components')
+    implementation project(':appcheck:firebase-appcheck-interop')
 
 
     //To provide @Generated annotations

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/AccessHelper.java

@@ -18,6 +18,7 @@
 import com.google.android.gms.tasks.Task;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.firestore.auth.CredentialsProvider;
+import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.model.DatabaseId;
 import com.google.firebase.firestore.util.AsyncQueue;
 
@@ -29,15 +30,17 @@ public static FirebaseFirestore newFirebaseFirestore(
       Context context,
       DatabaseId databaseId,
       String persistenceKey,
-      CredentialsProvider credentialsProvider,
+      CredentialsProvider<User> authProvider,
+      CredentialsProvider<String> appCheckProvider,
       AsyncQueue asyncQueue,
       FirebaseApp firebaseApp,
       FirebaseFirestore.InstanceRegistry instanceRegistry) {
     return new FirebaseFirestore(
         context,
         databaseId,
         persistenceKey,
-        credentialsProvider,
+        authProvider,
+        appCheckProvider,
         asyncQueue,
         firebaseApp,
         instanceRegistry,

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/remote/RemoteStoreTest.java

@@ -45,6 +45,7 @@ public void testRemoteStoreStreamStopsWhenNetworkUnreachable() {
             IntegrationTestUtil.testEnvDatabaseInfo(),
             testQueue,
             null,
+            null,
             ApplicationProvider.getApplicationContext(),
             null);
     Semaphore networkChangeSemaphore = new Semaphore(0);

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/remote/StreamTest.java

@@ -32,6 +32,7 @@
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.MutationResult;
+import com.google.firebase.firestore.testutil.EmptyAppCheckTokenProvider;
 import com.google.firebase.firestore.testutil.EmptyCredentialsProvider;
 import com.google.firebase.firestore.testutil.IntegrationTestUtil;
 import com.google.firebase.firestore.util.AsyncQueue;
@@ -113,6 +114,7 @@ private WriteStream createAndOpenWriteStream(
             IntegrationTestUtil.testEnvDatabaseInfo(),
             testQueue,
             new EmptyCredentialsProvider(),
+            new EmptyAppCheckTokenProvider(),
             ApplicationProvider.getApplicationContext(),
             null);
     final WriteStream writeStream = datastore.createWriteStream(callback);
@@ -138,6 +140,7 @@ public void testWatchStreamStopBeforeHandshake() throws Exception {
             IntegrationTestUtil.testEnvDatabaseInfo(),
             testQueue,
             new EmptyCredentialsProvider(),
+            new EmptyAppCheckTokenProvider(),
             ApplicationProvider.getApplicationContext(),
             mockGrpcProvider);
     StreamStatusCallback streamCallback = new StreamStatusCallback() {};
@@ -160,6 +163,7 @@ public void testWriteStreamStopAfterHandshake() throws Exception {
             IntegrationTestUtil.testEnvDatabaseInfo(),
             testQueue,
             new EmptyCredentialsProvider(),
+            new EmptyAppCheckTokenProvider(),
             ApplicationProvider.getApplicationContext(),
             null);
     final WriteStream[] writeStreamWrapper = new WriteStream[1];
@@ -207,6 +211,7 @@ public void testWriteStreamStopPartial() throws Exception {
             IntegrationTestUtil.testEnvDatabaseInfo(),
             testQueue,
             new EmptyCredentialsProvider(),
+            new EmptyAppCheckTokenProvider(),
             ApplicationProvider.getApplicationContext(),
             null);
     StreamStatusCallback streamCallback = new StreamStatusCallback() {};
@@ -287,6 +292,7 @@ public void testStreamRefreshesTokenUponExpiration() throws Exception {
             IntegrationTestUtil.testEnvDatabaseInfo(),
             testQueue,
             mockCredentialsProvider,
+            new EmptyAppCheckTokenProvider(),
             ApplicationProvider.getApplicationContext(),
             null);
     StreamStatusCallback callback = new StreamStatusCallback();
@@ -316,6 +322,7 @@ public void testTokenIsNotInvalidatedOnceStreamIsHealthy() throws Exception {
             IntegrationTestUtil.testEnvDatabaseInfo(),
             testQueue,
             mockCredentialsProvider,
+            new EmptyAppCheckTokenProvider(),
             ApplicationProvider.getApplicationContext(),
             null);
     StreamStatusCallback callback = new StreamStatusCallback();

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/IntegrationTestUtil.java

@@ -272,6 +272,7 @@ public static FirebaseFirestore testFirestore(
             databaseId,
             persistenceKey,
             MockCredentialsProvider.instance(),
+            new EmptyAppCheckTokenProvider(),
             asyncQueue,
             /*firebaseApp=*/ null,
             /*instanceRegistry=*/ (dbId) -> {});

firebase-firestore/src/main/java/com/google/firebase/firestore/FirebaseFirestore.java

@@ -27,11 +27,14 @@
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
+import com.google.firebase.appcheck.interop.InternalAppCheckTokenProvider;
 import com.google.firebase.auth.internal.InternalAuthProvider;
 import com.google.firebase.emulators.EmulatedServiceSettings;
 import com.google.firebase.firestore.FirebaseFirestoreException.Code;
 import com.google.firebase.firestore.auth.CredentialsProvider;
+import com.google.firebase.firestore.auth.FirebaseAppCheckTokenProvider;
 import com.google.firebase.firestore.auth.FirebaseAuthCredentialsProvider;
+import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.core.ActivityScope;
 import com.google.firebase.firestore.core.AsyncEventListener;
 import com.google.firebase.firestore.core.DatabaseInfo;
@@ -85,7 +88,8 @@ public interface InstanceRegistry {
   // databaseId itself that needs locking; it just saves us creating a separate lock object.
   private final DatabaseId databaseId;
   private final String persistenceKey;
-  private final CredentialsProvider credentialsProvider;
+  private final CredentialsProvider<User> authProvider;
+  private final CredentialsProvider<String> appCheckProvider;
   private final AsyncQueue asyncQueue;
   private final FirebaseApp firebaseApp;
   private final UserDataReader userDataReader;
@@ -124,7 +128,8 @@ private static FirebaseFirestore getInstance(@NonNull FirebaseApp app, @NonNull
   static FirebaseFirestore newInstance(
       @NonNull Context context,
       @NonNull FirebaseApp app,
-      @NonNull Deferred<InternalAuthProvider> authProvider,
+      @NonNull Deferred<InternalAuthProvider> deferredAuthProvider,
+      @NonNull Deferred<InternalAppCheckTokenProvider> deferredAppCheckTokenProvider,
       @NonNull String database,
       @NonNull InstanceRegistry instanceRegistry,
       @Nullable GrpcMetadataProvider metadataProvider) {
@@ -136,7 +141,10 @@ static FirebaseFirestore newInstance(
 
     AsyncQueue queue = new AsyncQueue();
 
-    CredentialsProvider provider = new FirebaseAuthCredentialsProvider(authProvider);
+    CredentialsProvider<User> authProvider =
+        new FirebaseAuthCredentialsProvider(deferredAuthProvider);
+    CredentialsProvider<String> appCheckProvider =
+        new FirebaseAppCheckTokenProvider(deferredAppCheckTokenProvider);
 
     // Firestore uses a different database for each app name. Note that we don't use
     // app.getPersistenceKey() here because it includes the application ID which is related
@@ -149,7 +157,8 @@ static FirebaseFirestore newInstance(
             context,
             databaseId,
             persistenceKey,
-            provider,
+            authProvider,
+            appCheckProvider,
             queue,
             app,
             instanceRegistry,
@@ -162,7 +171,8 @@ static FirebaseFirestore newInstance(
       Context context,
       DatabaseId databaseId,
       String persistenceKey,
-      CredentialsProvider credentialsProvider,
+      CredentialsProvider<User> authProvider,
+      CredentialsProvider<String> appCheckProvider,
       AsyncQueue asyncQueue,
       @Nullable FirebaseApp firebaseApp,
       InstanceRegistry instanceRegistry,
@@ -171,7 +181,8 @@ static FirebaseFirestore newInstance(
     this.databaseId = checkNotNull(checkNotNull(databaseId));
     this.userDataReader = new UserDataReader(databaseId);
     this.persistenceKey = checkNotNull(persistenceKey);
-    this.credentialsProvider = checkNotNull(credentialsProvider);
+    this.authProvider = checkNotNull(authProvider);
+    this.appCheckProvider = checkNotNull(appCheckProvider);
     this.asyncQueue = checkNotNull(asyncQueue);
     // NOTE: We allow firebaseApp to be null in tests only.
     this.firebaseApp = firebaseApp;
@@ -242,7 +253,13 @@ private void ensureClientConfigured() {
 
       client =
           new FirestoreClient(
-              context, databaseInfo, settings, credentialsProvider, asyncQueue, metadataProvider);
+              context,
+              databaseInfo,
+              settings,
+              authProvider,
+              appCheckProvider,
+              asyncQueue,
+              metadataProvider);
     }
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/FirestoreMultiDbComponent.java

@@ -22,6 +22,7 @@
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseAppLifecycleListener;
 import com.google.firebase.FirebaseOptions;
+import com.google.firebase.appcheck.interop.InternalAppCheckTokenProvider;
 import com.google.firebase.auth.internal.InternalAuthProvider;
 import com.google.firebase.firestore.remote.GrpcMetadataProvider;
 import com.google.firebase.inject.Deferred;
@@ -42,16 +43,19 @@ class FirestoreMultiDbComponent
   private final FirebaseApp app;
   private final Context context;
   private final Deferred<InternalAuthProvider> authProvider;
+  private final Deferred<InternalAppCheckTokenProvider> appCheckProvider;
   private final GrpcMetadataProvider metadataProvider;
 
   FirestoreMultiDbComponent(
       @NonNull Context context,
       @NonNull FirebaseApp app,
       @NonNull Deferred<InternalAuthProvider> authProvider,
+      @NonNull Deferred<InternalAppCheckTokenProvider> appCheckProvider,
       @Nullable GrpcMetadataProvider metadataProvider) {
     this.context = context;
     this.app = app;
     this.authProvider = authProvider;
+    this.appCheckProvider = appCheckProvider;
     this.metadataProvider = metadataProvider;
     this.app.addLifecycleEventListener(this);
   }
@@ -63,7 +67,7 @@ synchronized FirebaseFirestore get(@NonNull String databaseId) {
     if (firestore == null) {
       firestore =
           FirebaseFirestore.newInstance(
-              context, app, authProvider, databaseId, this, metadataProvider);
+              context, app, authProvider, appCheckProvider, databaseId, this, metadataProvider);
       instances.put(databaseId, firestore);
     }
     return firestore;

firebase-firestore/src/main/java/com/google/firebase/firestore/FirestoreRegistrar.java

@@ -19,6 +19,7 @@
 import androidx.annotation.RestrictTo;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseOptions;
+import com.google.firebase.appcheck.interop.InternalAppCheckTokenProvider;
 import com.google.firebase.auth.internal.InternalAuthProvider;
 import com.google.firebase.components.Component;
 import com.google.firebase.components.ComponentRegistrar;
@@ -48,13 +49,15 @@ public List<Component<?>> getComponents() {
             .add(Dependency.optionalProvider(HeartBeatInfo.class))
             .add(Dependency.optionalProvider(UserAgentPublisher.class))
             .add(Dependency.deferred(InternalAuthProvider.class))
+            .add(Dependency.deferred(InternalAppCheckTokenProvider.class))
             .add(Dependency.optional(FirebaseOptions.class))
             .factory(
                 c ->
                     new FirestoreMultiDbComponent(
                         c.get(Context.class),
                         c.get(FirebaseApp.class),
                         c.getDeferred(InternalAuthProvider.class),
+                        c.getDeferred(InternalAppCheckTokenProvider.class),
                         new FirebaseClientGrpcMetadataProvider(
                             c.getProvider(UserAgentPublisher.class),
                             c.getProvider(HeartBeatInfo.class),

firebase-firestore/src/main/java/com/google/firebase/firestore/auth/CredentialsProvider.java

@@ -18,7 +18,7 @@
 import com.google.firebase.firestore.util.Listener;
 
 /** A CredentialsProvider has a method to fetch an authorization token. */
-public abstract class CredentialsProvider {
+public abstract class CredentialsProvider<T> {
   /**
    * Requests token for the current user. Use {@link #invalidateToken} to force-refresh the token.
    *
@@ -36,7 +36,7 @@ public abstract class CredentialsProvider {
    * Sets the listener to be notified of credential changes (sign-in / sign-out, token changes). It
    * is immediately called once with the initial user.
    */
-  public abstract void setChangeListener(Listener<User> changeListener);
+  public abstract void setChangeListener(Listener<T> changeListener);
 
   /** Removes the listener set with {@link #setChangeListener}. */
   public abstract void removeChangeListener();