Resolve StrictMode violation in App Check. (#4085)

* Resolve StrictMode violation in App Check.

* Attempt to fix some tests.

* Fix unit tests.

* Make `retrieveStoredAppCheckTokenInBackground` private instead of package-private.

* Move listener invocations back to the main thread while keeping disk write on background thread.

* Refactor to use lambda syntax.

Author: rosalyntan

appcheck/firebase-appcheck/src/main/java/com/google/firebase/appcheck/internal/DefaultFirebaseAppCheck.java

@@ -21,6 +21,7 @@
 import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.tasks.Continuation;
 import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseException;
@@ -35,6 +36,8 @@
 import com.google.firebase.inject.Provider;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 
 public class DefaultFirebaseAppCheck extends FirebaseAppCheck {
 
@@ -46,6 +49,8 @@ public class DefaultFirebaseAppCheck extends FirebaseAppCheck {
   private final List<AppCheckListener> appCheckListenerList;
   private final StorageHelper storageHelper;
   private final TokenRefreshManager tokenRefreshManager;
+  private final ExecutorService backgroundExecutor;
+  private final Task<Void> retrieveStoredTokenTask;
   private final Clock clock;
 
   private AppCheckProviderFactory appCheckProviderFactory;
@@ -55,6 +60,17 @@ public class DefaultFirebaseAppCheck extends FirebaseAppCheck {
   public DefaultFirebaseAppCheck(
       @NonNull FirebaseApp firebaseApp,
       @NonNull Provider<HeartBeatController> heartBeatController) {
+    this(
+        checkNotNull(firebaseApp),
+        checkNotNull(heartBeatController),
+        Executors.newCachedThreadPool());
+  }
+
+  @VisibleForTesting
+  DefaultFirebaseAppCheck(
+      @NonNull FirebaseApp firebaseApp,
+      @NonNull Provider<HeartBeatController> heartBeatController,
+      @NonNull ExecutorService backgroundExecutor) {
     checkNotNull(firebaseApp);
     checkNotNull(heartBeatController);
     this.firebaseApp = firebaseApp;
@@ -65,8 +81,22 @@ public DefaultFirebaseAppCheck(
         new StorageHelper(firebaseApp.getApplicationContext(), firebaseApp.getPersistenceKey());
     this.tokenRefreshManager =
         new TokenRefreshManager(firebaseApp.getApplicationContext(), /* firebaseAppCheck= */ this);
+    this.backgroundExecutor = backgroundExecutor;
+    this.retrieveStoredTokenTask = retrieveStoredAppCheckTokenInBackground(backgroundExecutor);
     this.clock = new Clock.DefaultClock();
-    setCachedToken(storageHelper.retrieveAppCheckToken());
+  }
+
+  private Task<Void> retrieveStoredAppCheckTokenInBackground(@NonNull ExecutorService executor) {
+    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
+    executor.execute(
+        () -> {
+          AppCheckToken token = storageHelper.retrieveAppCheckToken();
+          if (token != null) {
+            setCachedToken(token);
+          }
+          taskCompletionSource.setResult(null);
+        });
+    return taskCompletionSource.getTask();
   }
 
   @Override
@@ -146,44 +176,50 @@ public void removeAppCheckListener(@NonNull AppCheckListener listener) {
   @NonNull
   @Override
   public Task<AppCheckTokenResult> getToken(boolean forceRefresh) {
-    if (!forceRefresh && hasValidToken()) {
-      return Tasks.forResult(DefaultAppCheckTokenResult.constructFromAppCheckToken(cachedToken));
-    }
-    if (appCheckProvider == null) {
-      return Tasks.forResult(
-          DefaultAppCheckTokenResult.constructFromError(
-              new FirebaseException("No AppCheckProvider installed.")));
-    }
-    // TODO: Cache the in-flight task.
-    return fetchTokenFromProvider()
-        .continueWithTask(
-            new Continuation<AppCheckToken, Task<AppCheckTokenResult>>() {
-              @Override
-              public Task<AppCheckTokenResult> then(@NonNull Task<AppCheckToken> task) {
-                if (task.isSuccessful()) {
-                  return Tasks.forResult(
-                      DefaultAppCheckTokenResult.constructFromAppCheckToken(task.getResult()));
-                }
-                // If the token exchange failed, return a dummy token for integrators to attach in
-                // their headers.
-                return Tasks.forResult(
-                    DefaultAppCheckTokenResult.constructFromError(
-                        new FirebaseException(
-                            task.getException().getMessage(), task.getException())));
-              }
-            });
+    return retrieveStoredTokenTask.continueWithTask(
+        unused -> {
+          if (!forceRefresh && hasValidToken()) {
+            return Tasks.forResult(
+                DefaultAppCheckTokenResult.constructFromAppCheckToken(cachedToken));
+          }
+          if (appCheckProvider == null) {
+            return Tasks.forResult(
+                DefaultAppCheckTokenResult.constructFromError(
+                    new FirebaseException("No AppCheckProvider installed.")));
+          }
+          // TODO: Cache the in-flight task.
+          return fetchTokenFromProvider()
+              .continueWithTask(
+                  appCheckTokenTask -> {
+                    if (appCheckTokenTask.isSuccessful()) {
+                      return Tasks.forResult(
+                          DefaultAppCheckTokenResult.constructFromAppCheckToken(
+                              appCheckTokenTask.getResult()));
+                    }
+                    // If the token exchange failed, return a dummy token for integrators to attach
+                    // in their headers.
+                    return Tasks.forResult(
+                        DefaultAppCheckTokenResult.constructFromError(
+                            new FirebaseException(
+                                appCheckTokenTask.getException().getMessage(),
+                                appCheckTokenTask.getException())));
+                  });
+        });
   }
 
   @NonNull
   @Override
   public Task<AppCheckToken> getAppCheckToken(boolean forceRefresh) {
-    if (!forceRefresh && hasValidToken()) {
-      return Tasks.forResult(cachedToken);
-    }
-    if (appCheckProvider == null) {
-      return Tasks.forException(new FirebaseException("No AppCheckProvider installed."));
-    }
-    return fetchTokenFromProvider();
+    return retrieveStoredTokenTask.continueWithTask(
+        unused -> {
+          if (!forceRefresh && hasValidToken()) {
+            return Tasks.forResult(cachedToken);
+          }
+          if (appCheckProvider == null) {
+            return Tasks.forException(new FirebaseException("No AppCheckProvider installed."));
+          }
+          return fetchTokenFromProvider();
+        });
   }
 
   /** Fetches an {@link AppCheckToken} via the installed {@link AppCheckProvider}. */
@@ -227,7 +263,7 @@ void setCachedToken(@NonNull AppCheckToken token) {
    * well as the in-memory cached {@link AppCheckToken}.
    */
   private void updateStoredToken(@NonNull AppCheckToken token) {
-    storageHelper.saveAppCheckToken(token);
+    backgroundExecutor.execute(() -> storageHelper.saveAppCheckToken(token));
     setCachedToken(token);
 
     tokenRefreshManager.maybeScheduleTokenRefresh(token);

appcheck/firebase-appcheck/src/test/java/com/google/firebase/appcheck/internal/DefaultFirebaseAppCheckTest.java

@@ -24,6 +24,7 @@
 import androidx.test.core.app.ApplicationProvider;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.Tasks;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.appcheck.AppCheckProvider;
 import com.google.firebase.appcheck.AppCheckProviderFactory;
@@ -74,7 +75,10 @@ public void setup() {
     when(mockAppCheckProvider.getToken()).thenReturn(Tasks.forResult(validDefaultAppCheckToken));
 
     defaultFirebaseAppCheck =
-        new DefaultFirebaseAppCheck(mockFirebaseApp, () -> mockHeartBeatController);
+        new DefaultFirebaseAppCheck(
+            mockFirebaseApp,
+            () -> mockHeartBeatController,
+            MoreExecutors.newDirectExecutorService());
   }
 
   @Test