Tomandersen/handshake rewrite

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

@@ -15,6 +15,7 @@
 package com.google.firebase.firestore;
 
 import android.content.Context;
+import androidx.core.util.Supplier;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.firestore.auth.CredentialsProvider;
 import com.google.firebase.firestore.auth.User;
@@ -31,17 +32,17 @@ public static FirebaseFirestore newFirebaseFirestore(
       Context context,
       DatabaseId databaseId,
       String persistenceKey,
-      CredentialsProvider<User> authProvider,
-      CredentialsProvider<String> appCheckProvider,
+      Supplier<CredentialsProvider<User>> authProviderFactory,
+      Supplier<CredentialsProvider<String>> appCheckTokenProviderFactory,
       Function<FirebaseFirestoreSettings, ComponentProvider> componentProviderFactory,
       FirebaseApp firebaseApp,
       FirebaseFirestore.InstanceRegistry instanceRegistry) {
     return new FirebaseFirestore(
         context,
         databaseId,
         persistenceKey,
-        authProvider,
-        appCheckProvider,
+        authProviderFactory,
+        appCheckTokenProviderFactory,
         componentProviderFactory,
         firebaseApp,
         instanceRegistry,

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

@@ -0,0 +1,11 @@
+package com.google.firebase.firestore;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+public class FirestoreClientProviderTest {
+
+  // TODO(requires backend/emulator support)
+
+}

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

@@ -38,12 +38,17 @@
 import com.google.firebase.firestore.testutil.IntegrationTestUtil;
 import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.AsyncQueue.TimerId;
+import com.google.firestore.v1.InitResponse;
+import com.google.protobuf.ByteString;
 import io.grpc.Status;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.Timeout;
 import org.junit.runner.RunWith;
 
 class MockCredentialsProvider extends EmptyCredentialsProvider {
@@ -68,6 +73,9 @@ public List<String> observedStates() {
 
 @RunWith(AndroidJUnit4.class)
 public class StreamTest {
+
+  @Rule public Timeout timeout = new Timeout(10, TimeUnit.SECONDS);
+
   /** Single mutation to send to the write stream. */
   private static final List<Mutation> mutations =
       Collections.singletonList(setMutation("foo/bar", map()));
@@ -96,7 +104,7 @@ public void onClose(Status status) {
     }
 
     @Override
-    public void onHandshakeComplete() {
+    public void onHandshake(InitResponse initResponse) {
       handshakeSemaphore.release();
     }
 
@@ -131,7 +139,7 @@ private void waitForWriteStreamOpen(
       AsyncQueue testQueue, WriteStream writeStream, StreamStatusCallback callback) {
     testQueue.enqueueAndForget(writeStream::start);
     waitFor(callback.openSemaphore);
-    testQueue.enqueueAndForget(writeStream::writeHandshake);
+    testQueue.enqueueAndForget(() -> writeStream.sendHandshake(ByteString.EMPTY));
     waitFor(callback.handshakeSemaphore);
   }
 
@@ -180,9 +188,9 @@ public void testWriteStreamStopAfterHandshake() throws Exception {
     StreamStatusCallback streamCallback =
         new StreamStatusCallback() {
           @Override
-          public void onHandshakeComplete() {
+          public void onHandshake(InitResponse initResponse) {
             assertThat(writeStreamWrapper[0].getLastStreamToken()).isNotEmpty();
-            super.onHandshakeComplete();
+            super.onHandshake(initResponse);
           }
 
           @Override
@@ -202,7 +210,7 @@ public void onWriteResponse(
         () -> assertThrows(Throwable.class, () -> writeStream.writeMutations(mutations)));
 
     // Handshake should always be called
-    testQueue.enqueueAndForget(writeStream::writeHandshake);
+    testQueue.enqueueAndForget(() -> writeStream.sendHandshake(ByteString.EMPTY));
     waitFor(streamCallback.handshakeSemaphore);
 
     // Now writes should succeed

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

@@ -47,7 +47,6 @@
 import com.google.firebase.firestore.core.DatabaseInfo;
 import com.google.firebase.firestore.model.DatabaseId;
 import com.google.firebase.firestore.testutil.provider.FirestoreProvider;
-import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.Listener;
 import com.google.firebase.firestore.util.Logger;
 import com.google.firebase.firestore.util.Logger.Level;
@@ -311,15 +310,13 @@ public static FirebaseFirestore testFirestore(
 
     ensureStrictMode();
 
-    AsyncQueue asyncQueue = new AsyncQueue();
-
     FirebaseFirestore firestore =
         AccessHelper.newFirebaseFirestore(
             context,
             databaseId,
             persistenceKey,
-            MockCredentialsProvider.instance(),
-            new EmptyAppCheckTokenProvider(),
+            MockCredentialsProvider::instance,
+            EmptyAppCheckTokenProvider::new,
             ComponentProvider::defaultFactory,
             /* firebaseApp= */ null,
             /* instanceRegistry= */ (dbId) -> {});

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

@@ -20,10 +20,12 @@
 import android.annotation.SuppressLint;
 import android.app.Activity;
 import android.content.Context;
+import androidx.annotation.GuardedBy;
 import androidx.annotation.Keep;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
+import androidx.core.util.Supplier;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
@@ -56,6 +58,7 @@
 import com.google.firebase.firestore.util.Logger.Level;
 import com.google.firebase.firestore.util.Preconditions;
 import com.google.firebase.inject.Deferred;
+import com.google.protobuf.ByteString;
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
@@ -76,6 +79,10 @@
 public class FirebaseFirestore {
 
   private final Function<FirebaseFirestoreSettings, ComponentProvider> componentProviderFactory;
+  private volatile ByteString sessionToken;
+
+  @GuardedBy("clientProvider")
+  private boolean networkEnabled;
 
   /**
    * Provides a registry management interface for {@code FirebaseFirestore} instances.
@@ -93,8 +100,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<User> authProvider;
-  private final CredentialsProvider<String> appCheckProvider;
+  private final Supplier<CredentialsProvider<User>> authProviderFactory;
+  private final Supplier<CredentialsProvider<String>> appCheckTokenProviderFactory;
   private final FirebaseApp firebaseApp;
   private final UserDataReader userDataReader;
   // When user requests to terminate, use this to notify `FirestoreMultiDbComponent` to deregister
@@ -105,6 +112,8 @@ public interface InstanceRegistry {
   final FirestoreClientProvider clientProvider;
   private final GrpcMetadataProvider metadataProvider;
 
+  @VisibleForTesting Function<Executor, Task<Void>> clearPersistenceMethod;
+
   @Nullable private PersistentCacheIndexManager persistentCacheIndexManager;
 
   @NonNull
@@ -193,46 +202,44 @@ static FirebaseFirestore newInstance(
     }
     DatabaseId databaseId = DatabaseId.forDatabase(projectId, database);
 
-    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
     // to the project ID. We already include the project ID when resolving the database,
     // so there is no need to include it in the persistence key.
     String persistenceKey = app.getName();
 
-    return new FirebaseFirestore(
-        context,
-        databaseId,
-        persistenceKey,
-        authProvider,
-        appCheckProvider,
-        ComponentProvider::defaultFactory,
-        app,
-        instanceRegistry,
-        metadataProvider);
+    FirebaseFirestore firestore =
+        new FirebaseFirestore(
+            context,
+            databaseId,
+            persistenceKey,
+            () -> new FirebaseAuthCredentialsProvider(deferredAuthProvider),
+            () -> new FirebaseAppCheckTokenProvider(deferredAppCheckTokenProvider),
+            ComponentProvider::defaultFactory,
+            app,
+            instanceRegistry,
+            metadataProvider);
+    return firestore;
   }
 
   @VisibleForTesting
   FirebaseFirestore(
       Context context,
       DatabaseId databaseId,
-      String persistenceKey,
-      CredentialsProvider<User> authProvider,
-      CredentialsProvider<String> appCheckProvider,
+      @NonNull String persistenceKey,
+      @NonNull Supplier<CredentialsProvider<User>> authProviderFactory,
+      @NonNull Supplier<CredentialsProvider<String>> appCheckTokenProviderFactory,
       @NonNull Function<FirebaseFirestoreSettings, ComponentProvider> componentProviderFactory,
       @Nullable FirebaseApp firebaseApp,
       InstanceRegistry instanceRegistry,
       @Nullable GrpcMetadataProvider metadataProvider) {
+    this.networkEnabled = true;
     this.context = checkNotNull(context);
     this.databaseId = checkNotNull(checkNotNull(databaseId));
     this.userDataReader = new UserDataReader(databaseId);
     this.persistenceKey = checkNotNull(persistenceKey);
-    this.authProvider = checkNotNull(authProvider);
-    this.appCheckProvider = checkNotNull(appCheckProvider);
+    this.authProviderFactory = checkNotNull(authProviderFactory);
+    this.appCheckTokenProviderFactory = checkNotNull(appCheckTokenProviderFactory);
     this.componentProviderFactory = checkNotNull(componentProviderFactory);
     this.clientProvider = new FirestoreClientProvider(this::newClient);
     // NOTE: We allow firebaseApp to be null in tests only.
@@ -241,6 +248,21 @@ static FirebaseFirestore newInstance(
     this.metadataProvider = metadataProvider;
 
     this.settings = new FirebaseFirestoreSettings.Builder().build();
+
+    this.clearPersistenceMethod =
+        executor -> {
+          final TaskCompletionSource<Void> source = new TaskCompletionSource<>();
+          executor.execute(
+              () -> {
+                try {
+                  SQLitePersistence.clearPersistence(context, databaseId, persistenceKey);
+                  source.setResult(null);
+                } catch (FirebaseFirestoreException e) {
+                  source.setException(e);
+                }
+              });
+          return source.getTask();
+        };
   }
 
   /** Returns the settings used by this {@code FirebaseFirestore} object. */
@@ -296,14 +318,36 @@ private FirestoreClient newClient(AsyncQueue asyncQueue) {
       DatabaseInfo databaseInfo =
           new DatabaseInfo(databaseId, persistenceKey, settings.getHost(), settings.isSslEnabled());
 
-      return new FirestoreClient(
-          context,
-          databaseInfo,
-          authProvider,
-          appCheckProvider,
-          asyncQueue,
-          metadataProvider,
-          componentProviderFactory.apply(settings));
+      FirestoreClient client =
+          new FirestoreClient(
+              context,
+              databaseInfo,
+              authProviderFactory.get(),
+              appCheckTokenProviderFactory.get(),
+              asyncQueue,
+              metadataProvider,
+              componentProviderFactory.apply(settings));
+
+      client.setClearPersistenceCallback(
+          sessionToken -> {
+            synchronized (clientProvider) {
+              if (client.isTerminated()) return;
+              this.sessionToken = sessionToken;
+              clearPersistence();
+            }
+          });
+
+      // Session token must be set before we enable network, since it is part of stream handshake.
+      if (sessionToken != null) {
+        client.setSessionToken(sessionToken);
+        sessionToken = null;
+      }
+
+      if (networkEnabled) {
+        client.enableNetwork();
+      }
+
+      return client;
     }
   }
 
@@ -624,7 +668,14 @@ public Task<Void> waitForPendingWrites() {
    */
   @NonNull
   public Task<Void> enableNetwork() {
-    return clientProvider.call(FirestoreClient::enableNetwork);
+    synchronized (clientProvider) {
+      networkEnabled = true;
+      if (clientProvider.isConfigured()) {
+        return clientProvider.call(FirestoreClient::enableNetwork);
+      } else {
+        return Tasks.forResult(null);
+      }
+    }
   }
 
   /**
@@ -636,7 +687,14 @@ public Task<Void> enableNetwork() {
    */
   @NonNull
   public Task<Void> disableNetwork() {
-    return clientProvider.call(FirestoreClient::disableNetwork);
+    synchronized (clientProvider) {
+      networkEnabled = false;
+      if (clientProvider.isConfigured()) {
+        return clientProvider.call(FirestoreClient::disableNetwork);
+      } else {
+        return Tasks.forResult(null);
+      }
+    }
   }
 
   /** Globally enables / disables Cloud Firestore logging for the SDK. */
@@ -669,30 +727,14 @@ public static void setLoggingEnabled(boolean loggingEnabled) {
   @NonNull
   public Task<Void> clearPersistence() {
     return clientProvider.executeIfShutdown(
-        this::clearPersistence,
+        clearPersistenceMethod,
         executor ->
             Tasks.forException(
                 new FirebaseFirestoreException(
                     "Persistence cannot be cleared while the firestore instance is running.",
                     FirebaseFirestoreException.Code.FAILED_PRECONDITION)));
   }
 
-  @NonNull
-  private Task<Void> clearPersistence(Executor executor) {
-    final TaskCompletionSource<Void> source = new TaskCompletionSource<>();
-    executor.execute(
-        () -> {
-          try {
-            SQLitePersistence.clearPersistence(context, databaseId, persistenceKey);
-            source.setResult(null);
-          } catch (FirebaseFirestoreException e) {
-            source.setException(e);
-          }
-        });
-    return source.getTask();
-  }
-  ;
-
   /**
    * Attaches a listener for a snapshots-in-sync event. The snapshots-in-sync event indicates that
    * all listeners affected by a given change have fired, even if a single server-generated change