Make fetching of sync vault keys async on Android

TrustedVaultClient::FetchKeys() has an asynchronous signature, but the
Android implementation (and the JNI bridge) adopted a synchronous
implementation prior to this patch.

In this proposal, the operation becomes non-blocking, based on:
1. The introduction of a dedicated Java->C++ API call, invoked when the
   operation completes.
2. The adoption of base.Promise in TrustedVaultClient.Backend, which is
   an idiomatic way to express that keys will eventually be read.

Bug: 1012659
Change-Id: I04b9b32549e5a880d8cbe3da1fffb60ddaa1ae37
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1968991Reviewed-by: Boris Sazonov <bsazonov@chromium.org>
Commit-Queue: Mikel Astiz <mastiz@chromium.org>
Cr-Commit-Position: refs/heads/master@{#726428}

chrome/android/java/src/org/chromium/chrome/browser/sync/TrustedVaultClient.java

@@ -10,12 +10,16 @@ import android.content.Intent;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
+import org.chromium.base.Promise;
 import org.chromium.base.annotations.CalledByNative;
+import org.chromium.base.annotations.NativeMethods;
 import org.chromium.chrome.browser.AppHooks;
 import org.chromium.chrome.browser.util.IntentUtils;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
+import java.util.TreeSet;
 
 /**
  * Client used to communicate with GmsCore about sync encryption keys.
@@ -29,9 +33,9 @@ public class TrustedVaultClient {
          * Reads and returns available encryption keys without involving any user action.
          *
          * @param gaiaId String representation of the Gaia ID.
-         * @return the list of known keys, if any, where the last one is the most recent.
+         * @return a promise with known keys, if any, where the last one is the most recent.
          */
-        List<byte[]> fetchKeys(String gaiaId);
+        Promise<List<byte[]>> fetchKeys(String gaiaId);
 
         /**
          * Gets an Intent that can be used to display a UI that allows the user to reauthenticate
@@ -48,8 +52,8 @@ public class TrustedVaultClient {
      */
     public static class EmptyBackend implements Backend {
         @Override
-        public List<byte[]> fetchKeys(String gaiaId) {
-            return Collections.emptyList();
+        public Promise<List<byte[]>> fetchKeys(String gaiaId) {
+            return Promise.fulfilled(Collections.emptyList());
         }
 
         @Override
@@ -62,6 +66,9 @@ public class TrustedVaultClient {
 
     private final Backend mBackend;
 
+    // Registered native TrustedVaultClientAndroid instances. Usually exactly one.
+    private final Set<Long> mNativeTrustedVaultClientAndroidSet = new TreeSet<Long>();
+
     @VisibleForTesting
     public TrustedVaultClient(Backend backend) {
         assert backend != null;
@@ -107,11 +114,56 @@ public class TrustedVaultClient {
     }
 
     /**
-     * Forwards calls to Backend.fetchKeys().
+     * Registers a C++ client, which is a prerequisite before interacting with Java.
      */
     @CalledByNative
-    private static byte[][] fetchKeys(String gaiaId) {
-        List<byte[]> keys = get().mBackend.fetchKeys(gaiaId);
-        return keys.toArray(new byte[0][]);
+    private static void registerNative(long nativeTrustedVaultClientAndroid) {
+        assert !isNativeRegistered(nativeTrustedVaultClientAndroid);
+        get().mNativeTrustedVaultClientAndroidSet.add(nativeTrustedVaultClientAndroid);
+    }
+
+    /**
+     * Unregisters a previously-registered client, canceling any in-flight requests.
+     */
+    @CalledByNative
+    private static void unregisterNative(long nativeTrustedVaultClientAndroid) {
+        assert isNativeRegistered(nativeTrustedVaultClientAndroid);
+        get().mNativeTrustedVaultClientAndroidSet.remove(nativeTrustedVaultClientAndroid);
+    }
+
+    /**
+     * Convenience function to check if a native client has been registered.
+     */
+    private static boolean isNativeRegistered(long nativeTrustedVaultClientAndroid) {
+        return get().mNativeTrustedVaultClientAndroidSet.contains(nativeTrustedVaultClientAndroid);
+    }
+
+    /**
+     * Forwards calls to Backend.fetchKeys() and upon completion invokes native method
+     * fetchKeysCompleted().
+     */
+    @CalledByNative
+    private static void fetchKeys(long nativeTrustedVaultClientAndroid, String gaiaId) {
+        assert isNativeRegistered(nativeTrustedVaultClientAndroid);
+        get().mBackend.fetchKeys(gaiaId).then(
+                (keys)
+                        -> {
+                    if (isNativeRegistered(nativeTrustedVaultClientAndroid)) {
+                        TrustedVaultClientJni.get().fetchKeysCompleted(
+                                nativeTrustedVaultClientAndroid, gaiaId,
+                                keys.toArray(new byte[0][]));
+                    }
+                },
+                (exception) -> {
+                    if (isNativeRegistered(nativeTrustedVaultClientAndroid)) {
+                        TrustedVaultClientJni.get().fetchKeysCompleted(
+                                nativeTrustedVaultClientAndroid, gaiaId, new byte[0][]);
+                    }
+                });
+    }
+
+    @NativeMethods
+    interface Natives {
+        void fetchKeysCompleted(long nativeTrustedVaultClientAndroid, String gaiaId, byte[][] keys);
     }
 }

chrome/browser/sync/trusted_vault_client_android.cc

@@ -4,36 +4,74 @@
 
 #include "chrome/browser/sync/trusted_vault_client_android.h"
 
+#include <utility>
+
 #include "base/android/jni_android.h"
-#include "base/android/jni_array.h"
-#include "base/android/jni_string.h"
-#include "base/android/scoped_java_ref.h"
-#include "base/callback.h"
 #include "base/logging.h"
 #include "chrome/android/chrome_jni_headers/TrustedVaultClient_jni.h"
+#include "content/public/browser/browser_thread.h"
+
+TrustedVaultClientAndroid::OngoingFetchKeys::OngoingFetchKeys(
+    const std::string& gaia_id,
+    base::OnceCallback<void(const std::vector<std::string>&)> callback)
+    : gaia_id(gaia_id), callback(std::move(callback)) {}
 
-TrustedVaultClientAndroid::TrustedVaultClientAndroid() = default;
+TrustedVaultClientAndroid::OngoingFetchKeys::~OngoingFetchKeys() = default;
 
-TrustedVaultClientAndroid::~TrustedVaultClientAndroid() = default;
+TrustedVaultClientAndroid::TrustedVaultClientAndroid() {
+  JNIEnv* const env = base::android::AttachCurrentThread();
+  Java_TrustedVaultClient_registerNative(env, reinterpret_cast<intptr_t>(this));
+}
+
+TrustedVaultClientAndroid::~TrustedVaultClientAndroid() {
+  JNIEnv* const env = base::android::AttachCurrentThread();
+  Java_TrustedVaultClient_unregisterNative(env,
+                                           reinterpret_cast<intptr_t>(this));
+}
+
+void TrustedVaultClientAndroid::FetchKeysCompleted(
+    JNIEnv* env,
+    const base::android::JavaParamRef<jstring>& gaia_id,
+    const base::android::JavaParamRef<jobjectArray>& keys) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+  DCHECK(ongoing_fetch_keys_) << "No ongoing FetchKeys() request";
+  DCHECK_EQ(ongoing_fetch_keys_->gaia_id,
+            base::android::ConvertJavaStringToUTF8(env, gaia_id))
+      << "User mismatch in FetchKeys() response";
+
+  // Make a copy of the callback and reset |ongoing_fetch_keys_| before invoking
+  // the callback, in case it has side effects.
+  base::OnceCallback<void(const std::vector<std::string>&)> cb =
+      std::move(ongoing_fetch_keys_->callback);
+  ongoing_fetch_keys_.reset();
+
+  // Convert from Java bytes[][] to the C++ equivalent, in this case
+  // std::vector<std::string>.
+  // TODO(crbug.com/1027676): Avoid std::string for binary keys.
+  std::vector<std::string> converted_keys;
+  JavaArrayOfByteArrayToStringVector(env, keys, &converted_keys);
+  std::move(cb).Run(converted_keys);
+}
 
 void TrustedVaultClientAndroid::FetchKeys(
     const std::string& gaia_id,
     base::OnceCallback<void(const std::vector<std::string>&)> cb) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::UI);
+  DCHECK(!ongoing_fetch_keys_)
+      << "Only one FetchKeys() request is allowed at any time";
+
+  // Store for later completion when Java invokes FetchKeysCompleted().
+  ongoing_fetch_keys_ =
+      std::make_unique<OngoingFetchKeys>(gaia_id, std::move(cb));
+
   JNIEnv* const env = base::android::AttachCurrentThread();
   const base::android::ScopedJavaLocalRef<jstring> java_gaia_id =
       base::android::ConvertUTF8ToJavaString(env, gaia_id);
 
-  // Fetch keys from the implementation in Java.
-  const base::android::ScopedJavaLocalRef<jobjectArray> java_keys =
-      Java_TrustedVaultClient_fetchKeys(env, java_gaia_id);
-  DCHECK(java_keys);
-
-  // Convert from Java bytes[][] to the C++ equivalent, in this case
-  // std::vector<std::string>.
-  // TODO(crbug.com/1027676): Avoid std::string for binary keys.
-  std::vector<std::string> keys;
-  JavaArrayOfByteArrayToStringVector(env, java_keys, &keys);
-  std::move(cb).Run(keys);
+  // Trigger the fetching keys from the implementation in Java, which will
+  // eventually call FetchKeysCompleted().
+  Java_TrustedVaultClient_fetchKeys(env, reinterpret_cast<intptr_t>(this),
+                                    java_gaia_id);
 }
 
 void TrustedVaultClientAndroid::StoreKeys(

chrome/browser/sync/trusted_vault_client_android.h

@@ -5,10 +5,14 @@
 #ifndef CHROME_BROWSER_SYNC_TRUSTED_VAULT_CLIENT_ANDROID_H_
 #define CHROME_BROWSER_SYNC_TRUSTED_VAULT_CLIENT_ANDROID_H_
 
+#include <memory>
 #include <string>
 #include <vector>
 
-#include "base/callback_forward.h"
+#include "base/android/jni_array.h"
+#include "base/android/jni_string.h"
+#include "base/android/scoped_java_ref.h"
+#include "base/callback.h"
 #include "components/sync/driver/trusted_vault_client.h"
 
 // JNI bridge for a Java implementation of the TrustedVaultClient interface,
@@ -24,12 +28,35 @@ class TrustedVaultClientAndroid : public syncer::TrustedVaultClient {
   TrustedVaultClientAndroid& operator=(const TrustedVaultClientAndroid&) =
       delete;
 
+  // Called from Java to notify the completion of a FetchKeys() operation
+  // previously initiated from C++. This must correspond to an ongoing
+  // FetchKeys() request, and |gaia_id| must match the user's ID.
+  void FetchKeysCompleted(
+      JNIEnv* env,
+      const base::android::JavaParamRef<jstring>& gaia_id,
+      const base::android::JavaParamRef<jobjectArray>& keys);
+
   // TrustedVaultClient implementation.
   void FetchKeys(
       const std::string& gaia_id,
       base::OnceCallback<void(const std::vector<std::string>&)> cb) override;
   void StoreKeys(const std::string& gaia_id,
                  const std::vector<std::string>& keys) override;
+
+ private:
+  // Struct representing an in-flight FetchKeys() call invoked from C++.
+  struct OngoingFetchKeys {
+    OngoingFetchKeys(
+        const std::string& gaia_id,
+        base::OnceCallback<void(const std::vector<std::string>&)> callback);
+    ~OngoingFetchKeys();
+
+    const std::string gaia_id;
+    base::OnceCallback<void(const std::vector<std::string>&)> callback;
+  };
+
+  // Null if no in-flight FetchKeys().
+  std::unique_ptr<OngoingFetchKeys> ongoing_fetch_keys_;
 };
 
 #endif  // CHROME_BROWSER_SYNC_TRUSTED_VAULT_CLIENT_ANDROID_H_