From 392e1e23e58e893f95552fe5293919bb15da126c Mon Sep 17 00:00:00 2001
From: Rubin Xu <rubinxu@google.com>
Date: Tue, 22 Dec 2020 22:37:43 +0000
Subject: [PATCH] Allow KeyChain.bindAsUser() to be called on the main thread

KeyChain.bindAsUser() couldn't be called on the main thread because
it was using the main thread to handle service connection callback.
Add an overload of KeyChain.bindAsUser() that accepts an alternative
handler to process the connection callback, which makes it possible
to call KeyChain from the main UI thread directly.

Bug: 165641221
Test: atest KeyChainTests
Test: m RunKeyChainRoboTests
Change-Id: I4290bccf5ae04de0d84c7091729e86704b937295
---
 .../android/security/IKeyChainService.aidl    |  1 -
 keystore/java/android/security/KeyChain.java  | 33 +++++++++++++++++--
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/keystore/java/android/security/IKeyChainService.aidl b/keystore/java/android/security/IKeyChainService.aidl
index a9d4094d5245e..f708298a2cbdc 100644
--- a/keystore/java/android/security/IKeyChainService.aidl
+++ b/keystore/java/android/security/IKeyChainService.aidl
@@ -60,7 +60,6 @@ interface IKeyChainService {
     byte[] getEncodedCaCertificate(String alias, boolean includeDeletedSystem);
     List<String> getCaCertificateChainAliases(String rootAlias, boolean includeDeletedSystem);
     void setCredentialManagementApp(String packageName, in AppUriAuthenticationPolicy policy);
-    void updateCredentialManagementAppPolicy(in AppUriAuthenticationPolicy policy);
     boolean hasCredentialManagementApp();
     String getCredentialManagementAppPackageName();
     AppUriAuthenticationPolicy getCredentialManagementAppPolicy();
diff --git a/keystore/java/android/security/KeyChain.java b/keystore/java/android/security/KeyChain.java
index c6e72b0e9f6ea..2f444b34ce819 100644
--- a/keystore/java/android/security/KeyChain.java
+++ b/keystore/java/android/security/KeyChain.java
@@ -31,6 +31,7 @@ import android.content.Intent;
 import android.content.ServiceConnection;
 import android.net.Uri;
 import android.os.Binder;
+import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.Process;
@@ -854,10 +855,26 @@ public final class KeyChain {
     @WorkerThread
     public static KeyChainConnection bindAsUser(@NonNull Context context, UserHandle user)
             throws InterruptedException {
+        return bindAsUser(context, null, user);
+    }
+
+    /**
+     * Bind to KeyChainService in the target user.
+     * Caller should call unbindService on the result when finished.
+     *
+     * @throws InterruptedException if interrupted during binding.
+     * @throws AssertionError if unable to bind to KeyChainService.
+     * @hide
+     */
+    public static KeyChainConnection bindAsUser(@NonNull Context context, @Nullable Handler handler,
+            UserHandle user) throws InterruptedException {
+
         if (context == null) {
             throw new NullPointerException("context == null");
         }
-        ensureNotOnMainThread(context);
+        if (handler == null) {
+            ensureNotOnMainThread(context);
+        }
         if (!UserManager.get(context).isUserUnlocked(user)) {
             throw new IllegalStateException("User must be unlocked");
         }
@@ -884,9 +901,19 @@ public final class KeyChain {
         };
         Intent intent = new Intent(IKeyChainService.class.getName());
         ComponentName comp = intent.resolveSystemService(context.getPackageManager(), 0);
+        if (comp == null) {
+            throw new AssertionError("could not resolve KeyChainService");
+        }
         intent.setComponent(comp);
-        if (comp == null || !context.bindServiceAsUser(
-                intent, keyChainServiceConnection, Context.BIND_AUTO_CREATE, user)) {
+        final boolean bindSucceed;
+        if (handler != null) {
+            bindSucceed = context.bindServiceAsUser(
+                    intent, keyChainServiceConnection, Context.BIND_AUTO_CREATE, handler, user);
+        } else {
+            bindSucceed = context.bindServiceAsUser(
+                    intent, keyChainServiceConnection, Context.BIND_AUTO_CREATE, user);
+        }
+        if (!bindSucceed) {
             throw new AssertionError("could not bind to KeyChainService");
         }
         countDownLatch.await();