Merge Android 12

Bug: 202323961
Merged-In: Id666e2793d42ad6028265fa340ecf4e75ed5a5b5
Change-Id: I6c8ae46410c19cec8a48c353f4de8a0ec4669a7c
diff --git a/Android.bp b/Android.bp
index 9d99db2..6f8bd07 100644
--- a/Android.bp
+++ b/Android.bp
@@ -20,6 +20,7 @@
     name: "KeyChain",
     static_libs: [
         "bouncycastle-unbundled",
+        "com.google.android.material_material",
     ],
     srcs: ["src/**/*.java"],
     platform_apis: true,
diff --git a/AndroidManifest.xml b/AndroidManifest.xml
index 9d1601f..fae32b6 100755
--- a/AndroidManifest.xml
+++ b/AndroidManifest.xml
@@ -11,13 +11,16 @@
             android:allowBackup="false"
             android:usesCleartextTraffic="false"
             android:theme="@android:style/Theme.DeviceDefault.DayNight">
-        <service android:name="com.android.keychain.KeyChainService">
+        <service android:name="com.android.keychain.KeyChainService"
+            android:exported="true">
             <intent-filter>
                 <action android:name="android.security.IKeyChainService"/>
             </intent-filter>
         </service>
         <activity android:name="com.android.keychain.KeyChainActivity"
+                  android:exported="true"
                   android:theme="@style/KeyChainTransparent"
+                  android:launchMode="singleTop"
 		  android:excludeFromRecents="true">
 	    <intent-filter>
 	        <action android:name="com.android.keychain.CHOOSER"/>
diff --git a/res/layout/keychain_activity.xml b/res/layout/keychain_activity.xml
new file mode 100644
index 0000000..b96c504
--- /dev/null
+++ b/res/layout/keychain_activity.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2021 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<RelativeLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="vertical">
+
+    <androidx.coordinatorlayout.widget.CoordinatorLayout
+        android:id="@+id/container"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:layout_alignParentBottom="true"
+        android:layout_centerHorizontal="true" />
+
+</RelativeLayout>
diff --git a/res/values-iw/strings.xml b/res/values-iw/strings.xml
index ea04f37..b28b5ea 100644
--- a/res/values-iw/strings.xml
+++ b/res/values-iw/strings.xml
@@ -19,7 +19,7 @@
     <string name="app_name" msgid="170210454004696382">"Key Chain"</string>
     <string name="title_select_cert" msgid="3588447616418041699">"בחירת אישור"</string>
     <string name="requesting_application" msgid="1589142627467598421">"‏האפליקציה %s ביקשה אישור. בחירת אישור תאפשר לאפליקציה להשתמש בזהות זו מול שרתים כעת ובעתיד."</string>
-    <string name="requesting_server" msgid="5832565605998634370">"‏האפליקציה זיהתה את השרת המבקש כ-%s, אך עליך לתת לאפליקציה גישה לאישור רק אם יש לך אמון באפליקציה."</string>
+    <string name="requesting_server" msgid="5832565605998634370">"‏האפליקציה זיהתה את השרת המבקש כ-%s, אך עליך לתת לאפליקציה גישה לאישור רק אם אתה סומך עליה."</string>
     <string name="install_new_cert_message" msgid="4451971501142085495">"‏ניתן להתקין אישורים מקובץ PKCS#12 עם סיומת %1$s או %2$s הממוקם באחסון חיצוני."</string>
     <string name="allow_button" msgid="3030990695030371561">"בחירה"</string>
     <string name="deny_button" msgid="3766539809121892584">"דחייה"</string>
diff --git a/res/values-kk/strings.xml b/res/values-kk/strings.xml
index 01aa8e4..1c745c1 100644
--- a/res/values-kk/strings.xml
+++ b/res/values-kk/strings.xml
@@ -22,6 +22,6 @@
     <string name="requesting_server" msgid="5832565605998634370">"Бұл қолданба өтінген серверді %s ретінде анықтады, бірақ қолданбаға сенген жағдайда ғана сертификатқа кіруге рұқсат беруге болады."</string>
     <string name="install_new_cert_message" msgid="4451971501142085495">"Сертификаттарды сыртқы жадта орналасқан %1$s немесе %2$s жалғаулы PKCS#12 файлынан орнатуға болады."</string>
     <string name="allow_button" msgid="3030990695030371561">"Таңдау"</string>
-    <string name="deny_button" msgid="3766539809121892584">"Тыйым салу"</string>
+    <string name="deny_button" msgid="3766539809121892584">"Бас тарту"</string>
     <string name="loading_certs_message" msgid="814752048905775439">"Сертификаттар жүктеп алынуда..."</string>
 </resources>
diff --git a/res/values-kn/strings.xml b/res/values-kn/strings.xml
index b10decf..f5c5b06 100644
--- a/res/values-kn/strings.xml
+++ b/res/values-kn/strings.xml
@@ -20,7 +20,7 @@
     <string name="title_select_cert" msgid="3588447616418041699">"ಪ್ರಮಾಣಪತ್ರವನ್ನು ಆಯ್ಕೆಮಾಡಿ"</string>
     <string name="requesting_application" msgid="1589142627467598421">"%s ಅಪ್ಲಿಕೇಶನ್‌‌ ಪ್ರಮಾಣಪತ್ರವೊಂದನ್ನು ವಿನಂತಿಸಿದೆ. ಪ್ರಮಾಣಪತ್ರವನ್ನು ಆಯ್ಕೆ ಮಾಡುವುದರಿಂದ ಸರ್ವರ್‌ಗಳೊಂದಿಗೆ ಈಗ ಮತ್ತು ಭವಿಷ್ಯದಲ್ಲಿ ಈ ಗುರುತನ್ನು ಬಳಸಲು ಈ ಅಪ್ಲಿಕೇಶನ್‌‌ಗೆ ಅನುಮತಿಸುತ್ತದೆ."</string>
     <string name="requesting_server" msgid="5832565605998634370">"ಅಪ್ಲಿಕೇಶನ್‌ ವಿನಂತಿಸಿದ ಸರ್ವರ್‌ ಅನ್ನು %s ರಂತೆ ಗುರುತಿಸಿದೆ, ಆದರೆ ನೀವು ಅಪ್ಲಿಕೇಶನ್‌‌ ಅನ್ನು ನಂಬಿದರೆ ಮಾತ್ರ ಪ್ರಮಾಣಪತ್ರಕ್ಕೆ ಅಪ್ಲಿಕೇಶನ್‌ ಪ್ರವೇಶವನ್ನು ನೀಡಬೇಕು."</string>
-    <string name="install_new_cert_message" msgid="4451971501142085495">"ಬಾಹ್ಯ ಸಂಗ್ರಹಣೆಯಲ್ಲಿ ಸಂಗ್ರಹವಾಗಿರುವ %1$s ಅಥವಾ %2$s ವಿಸ್ತರಣೆಯೊಂದಿಗೆ PKCS#12 ಫೈಲ್‌ನಿಂದ ನೀವು ಪ್ರಮಾಣಪತ್ರಗಳನ್ನು ಇನ್‌ಸ್ಟಾಲ್ ಮಾಡಬಹುದು."</string>
+    <string name="install_new_cert_message" msgid="4451971501142085495">"ಬಾಹ್ಯ ಸಂಗ್ರಹಣೆಯಲ್ಲಿ ಸಂಗ್ರಹವಾಗಿರುವ %1$s ಅಥವಾ %2$s ವಿಸ್ತರಣೆಯೊಂದಿಗೆ PKCS#12 ಫೈಲ್‌ನಿಂದ ನೀವು ಪ್ರಮಾಣಪತ್ರಗಳನ್ನು ಸ್ಥಾಪಿಸಬಹುದು."</string>
     <string name="allow_button" msgid="3030990695030371561">"ಆಯ್ಕೆಮಾಡಿ"</string>
     <string name="deny_button" msgid="3766539809121892584">"ನಿರಾಕರಿಸಿ"</string>
     <string name="loading_certs_message" msgid="814752048905775439">"ಪ್ರಮಾಣಪತ್ರವನ್ನು ಲೋಡ್ ಮಾಡಲಾಗುತ್ತಿದೆ..."</string>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 2da7102..7cc4b11 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -35,5 +35,5 @@
     <string name="deny_button">Deny</string>
 
     <!-- Text to show while the KeyChain activity is loading certificates. -->
-    <string name="loading_certs_message">Loading certificates...</string>
+    <string name="loading_certs_message">%s is checking for a certificate&#8230;</string>
 </resources>
diff --git a/res/values/styles.xml b/res/values/styles.xml
index 2e1081b..6135255 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -18,9 +18,10 @@
 -->
 
 <resources xmlns:android="http://schemas.android.com/apk/res/android">
-    <style name="KeyChainTransparent" parent="@android:style/Theme.DeviceDefault.DayNight">
+    <style name="KeyChainTransparent" parent="@style/Theme.AppCompat.DayNight.NoActionBar">
         <item name="android:windowBackground">@android:color/transparent</item>
         <item name="android:windowNoTitle">true</item>
         <item name="android:windowIsFloating">true</item>
+        <item name="android:backgroundDimEnabled">false</item>
     </style>
 </resources>
diff --git a/robotests/src/com/android/keychain/AliasLoaderTest.java b/robotests/src/com/android/keychain/AliasLoaderTest.java
index 24372dd..130edde 100644
--- a/robotests/src/com/android/keychain/AliasLoaderTest.java
+++ b/robotests/src/com/android/keychain/AliasLoaderTest.java
@@ -19,26 +19,12 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-import android.security.Credentials;
 import android.util.Base64;
+
 import com.android.keychain.internal.KeyInfoProvider;
 
 import com.google.common.collect.ImmutableList;
 
-import java.io.ByteArrayInputStream;
-import java.security.KeyStore;
-import java.security.KeyStoreException;
-import java.security.cert.Certificate;
-import java.security.cert.CertificateException;
-import java.security.cert.CertificateFactory;
-import java.security.cert.X509Certificate;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.concurrent.CancellationException;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
-import javax.security.auth.x500.X500Principal;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -48,6 +34,22 @@
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.shadows.ShadowApplication;
 
+import java.io.ByteArrayInputStream;
+import java.security.KeyStore;
+import java.security.KeyStoreException;
+import java.security.KeyStoreSpi;
+import java.security.cert.Certificate;
+import java.security.cert.CertificateException;
+import java.security.cert.CertificateFactory;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import javax.security.auth.x500.X500Principal;
+
 @RunWith(RobolectricTestRunner.class)
 public final class AliasLoaderTest {
     // Generated using:
@@ -135,14 +137,22 @@
     private byte[] mECCertOne;
     private byte[] mECCertTwo;
     private ArrayList<byte[]> mIssuers;
+    private KeyStoreSpi mKeyStoreSpi;
+    private KeyStore mKeyStore;
+    private KeyInfoProvider mInfoProvider;
 
     private Certificate toCertificate(byte[] bytes) throws CertificateException {
         CertificateFactory cf = CertificateFactory.getInstance("X.509");
         return cf.generateCertificate(new ByteArrayInputStream(bytes));
     }
 
+    private Certificate[] toCertificateChain(byte[] bytes) throws CertificateException {
+        CertificateFactory cf = CertificateFactory.getInstance("X.509");
+        return new Certificate[]{cf.generateCertificate(new ByteArrayInputStream(bytes))};
+    }
+
     @Before
-    public void setUp() {
+    public void setUp() throws Exception {
         mRSACertOne = Base64.decode(SELF_SIGNED_RSA_CERT_1_B64, Base64.DEFAULT);
         mRSACertTwo = Base64.decode(SELF_SIGNED_RSA_CERT_2_B64, Base64.DEFAULT);
         mECCertOne = Base64.decode(SELF_SIGNED_EC_CERT_1_B64, Base64.DEFAULT);
@@ -158,19 +168,22 @@
 
         mDummyChecker = mock(KeyChainActivity.CertificateParametersFilter.class);
         when(mDummyChecker.shouldPresentCertificate(Mockito.anyString())).thenReturn(true);
+
+        mKeyStoreSpi = mock(KeyStoreSpi.class);
+        mKeyStore = new KeyStore(mKeyStoreSpi, null, "test") {};
+        mKeyStore.load(null);
+        mInfoProvider = mock(KeyInfoProvider.class);
     }
 
     @Test
     public void testAliasLoader_loadsAllAliases()
             throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException, KeyStoreException {
-        KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.aliases()).thenReturn(
-                Collections.enumeration(ImmutableList.of("b", "c", "a")));
+            TimeoutException {
+        prepareKeyStoreWithAliases(ImmutableList.of("b", "c", "a"));
 
         KeyChainActivity.AliasLoader loader =
                 new KeyChainActivity.AliasLoader(
-                        keyStore,
+                        mKeyStore,
                         RuntimeEnvironment.application,
                         mDummyInfoProvider,
                         mDummyChecker);
@@ -188,13 +201,12 @@
     @Test
     public void testAliasLoader_copesWithNoAliases()
             throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException, KeyStoreException {
-        KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.aliases()).thenReturn(null);
+                    TimeoutException {
+        when(mKeyStoreSpi.engineAliases()).thenReturn(Collections.enumeration(ImmutableList.of()));
 
         KeyChainActivity.AliasLoader loader =
                 new KeyChainActivity.AliasLoader(
-                        keyStore,
+                        mKeyStore,
                         RuntimeEnvironment.application,
                         mDummyInfoProvider,
                         mDummyChecker);
@@ -209,19 +221,15 @@
     @Test
     public void testAliasLoader_filtersNonUserSelectableAliases()
             throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException, KeyStoreException {
-        KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.aliases()).thenReturn(
-                Collections.enumeration(ImmutableList.of("a", "b", "c")));
-
-        KeyInfoProvider infoProvider = mock(KeyInfoProvider.class);
-        when(infoProvider.isUserSelectable("a")).thenReturn(false);
-        when(infoProvider.isUserSelectable("b")).thenReturn(true);
-        when(infoProvider.isUserSelectable("c")).thenReturn(false);
+                    TimeoutException {
+        prepareKeyStoreWithAliases(ImmutableList.of("b", "c", "a"));
+        when(mInfoProvider.isUserSelectable("a")).thenReturn(false);
+        when(mInfoProvider.isUserSelectable("b")).thenReturn(true);
+        when(mInfoProvider.isUserSelectable("c")).thenReturn(false);
 
         KeyChainActivity.AliasLoader loader =
                 new KeyChainActivity.AliasLoader(
-                        keyStore, RuntimeEnvironment.application, infoProvider, mDummyChecker);
+                        mKeyStore, RuntimeEnvironment.application, mInfoProvider, mDummyChecker);
         loader.execute();
 
         ShadowApplication.runBackgroundTasks();
@@ -235,15 +243,11 @@
     public void testAliasLoader_filtersAliasesWithNonConformingParameters()
             throws InterruptedException, ExecutionException, CancellationException,
                     TimeoutException, KeyStoreException {
-        KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.aliases()).thenReturn(
-                Collections.enumeration(ImmutableList.of("a", "b", "c", "d")));
-
-        KeyInfoProvider infoProvider = mock(KeyInfoProvider.class);
-        when(infoProvider.isUserSelectable("a")).thenReturn(true);
-        when(infoProvider.isUserSelectable("b")).thenReturn(true);
-        when(infoProvider.isUserSelectable("c")).thenReturn(false);
-        when(infoProvider.isUserSelectable("d")).thenReturn(false);
+        prepareKeyStoreWithAliases(ImmutableList.of("a", "b", "c", "d"));
+        when(mInfoProvider.isUserSelectable("a")).thenReturn(true);
+        when(mInfoProvider.isUserSelectable("b")).thenReturn(true);
+        when(mInfoProvider.isUserSelectable("c")).thenReturn(false);
+        when(mInfoProvider.isUserSelectable("d")).thenReturn(false);
 
         KeyChainActivity.CertificateParametersFilter checker =
                 mock(KeyChainActivity.CertificateParametersFilter.class);
@@ -258,7 +262,7 @@
 
         KeyChainActivity.AliasLoader loader =
                 new KeyChainActivity.AliasLoader(
-                        keyStore, RuntimeEnvironment.application, infoProvider, checker);
+                        mKeyStore, RuntimeEnvironment.application, mInfoProvider, checker);
         loader.execute();
 
         ShadowApplication.runBackgroundTasks();
@@ -268,43 +272,55 @@
         Assert.assertEquals("a", result.getItem(0));
     }
 
-    private KeyStore prepareKeyStoreWithCertificates()
-            throws CertificateException, KeyStoreException {
-        KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.getCertificate("rsa1")).thenReturn(toCertificate(mRSACertOne));
-        when(keyStore.getCertificate("ec1")).thenReturn(toCertificate(mECCertOne));
-        when(keyStore.getCertificate("rsa2")).thenReturn(toCertificate(mRSACertTwo));
-        when(keyStore.getCertificate("ec2")).thenReturn(toCertificate(mECCertTwo));
+    private void prepareKeyStoreWithAliases(ImmutableList<String> aliases) {
+        when(mKeyStoreSpi.engineAliases()).thenReturn(Collections.enumeration(aliases));
+        for (int i = 0; i < aliases.size(); i++) {
+            when(mKeyStoreSpi.engineIsKeyEntry(aliases.get(i))).thenReturn(true);
+        }
+    }
 
-        return keyStore;
+    private void prepareKeyStoreWithCertificates() throws CertificateException {
+        when(mKeyStoreSpi.engineGetCertificate("rsa1")).thenReturn(toCertificate(mRSACertOne));
+        when(mKeyStoreSpi.engineGetCertificate("ec1")).thenReturn(toCertificate(mECCertOne));
+        when(mKeyStoreSpi.engineGetCertificate("rsa2")).thenReturn(toCertificate(mRSACertTwo));
+        when(mKeyStoreSpi.engineGetCertificate("ec2")).thenReturn(toCertificate(mECCertTwo));
+
+        when(mKeyStoreSpi.engineGetCertificateChain("rsa1"))
+                .thenReturn(toCertificateChain(mRSACertOne));
+        when(mKeyStoreSpi.engineGetCertificateChain("ec1"))
+                .thenReturn(toCertificateChain(mECCertOne));
+        when(mKeyStoreSpi.engineGetCertificateChain("rsa2"))
+                .thenReturn(toCertificateChain(mRSACertTwo));
+        when(mKeyStoreSpi.engineGetCertificateChain("ec2"))
+                .thenReturn(toCertificateChain(mECCertTwo));
     }
 
     @Test
     public void testCertificateParametersFilter_filtersByKey()
-            throws CancellationException, KeyStoreException, CertificateException {
-        KeyStore keyStore = prepareKeyStoreWithCertificates();
+            throws CancellationException, CertificateException {
+        prepareKeyStoreWithCertificates();
 
         KeyChainActivity.CertificateParametersFilter ec_checker =
                 new KeyChainActivity.CertificateParametersFilter(
-                        keyStore, new String[] {"EC"}, new ArrayList<byte[]>());
+                        mKeyStore, new String[] {"EC"}, new ArrayList<byte[]>());
         Assert.assertFalse(ec_checker.shouldPresentCertificate("rsa1"));
         Assert.assertTrue(ec_checker.shouldPresentCertificate("ec1"));
 
         KeyChainActivity.CertificateParametersFilter rsa_and_ec_checker =
                 new KeyChainActivity.CertificateParametersFilter(
-                        keyStore, new String[] {"EC", "RSA"}, new ArrayList<byte[]>());
+                        mKeyStore, new String[] {"EC", "RSA"}, new ArrayList<byte[]>());
         Assert.assertTrue(rsa_and_ec_checker.shouldPresentCertificate("rsa1"));
         Assert.assertTrue(rsa_and_ec_checker.shouldPresentCertificate("ec1"));
     }
 
     @Test
     public void testCertificateParametersFilter_filtersByIssuer()
-            throws CancellationException, KeyStoreException, CertificateException {
-        KeyStore keyStore = prepareKeyStoreWithCertificates();
+            throws CancellationException, CertificateException {
+        prepareKeyStoreWithCertificates();
 
         KeyChainActivity.CertificateParametersFilter issuer_checker =
                 new KeyChainActivity.CertificateParametersFilter(
-                        keyStore, new String[] {}, mIssuers);
+                        mKeyStore, new String[] {}, mIssuers);
         Assert.assertTrue(issuer_checker.shouldPresentCertificate("rsa1"));
         Assert.assertTrue(issuer_checker.shouldPresentCertificate("ec1"));
         Assert.assertFalse(issuer_checker.shouldPresentCertificate("rsa2"));
@@ -313,13 +329,12 @@
 
     @Test
     public void testCertificateParametersFilter_filtersByIssuerAndKey()
-            throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException, KeyStoreException, CertificateException {
-        KeyStore keyStore = prepareKeyStoreWithCertificates();
+            throws CancellationException, CertificateException {
+        prepareKeyStoreWithCertificates();
 
         KeyChainActivity.CertificateParametersFilter issuer_checker =
                 new KeyChainActivity.CertificateParametersFilter(
-                        keyStore, new String[] {"EC"}, mIssuers);
+                        mKeyStore, new String[] {"EC"}, mIssuers);
         Assert.assertFalse(issuer_checker.shouldPresentCertificate("rsa1"));
         Assert.assertTrue(issuer_checker.shouldPresentCertificate("ec1"));
         Assert.assertFalse(issuer_checker.shouldPresentCertificate("rsa2"));
diff --git a/robotests/src/com/android/keychain/KeyChainServiceRoboTest.java b/robotests/src/com/android/keychain/KeyChainServiceRoboTest.java
index 3016af4..6392af4 100644
--- a/robotests/src/com/android/keychain/KeyChainServiceRoboTest.java
+++ b/robotests/src/com/android/keychain/KeyChainServiceRoboTest.java
@@ -18,8 +18,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
@@ -39,7 +39,6 @@
 import android.security.IKeyChainService;
 
 import com.android.org.conscrypt.TrustedCertificateStore;
-import com.android.keychain.ShadowKeyStore;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -55,6 +54,7 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.security.KeyStore;
 import java.security.cert.CertificateException;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
@@ -64,9 +64,11 @@
 @RunWith(RobolectricTestRunner.class)
 @Config(shadows = {
     ShadowTrustedCertificateStore.class,
-    ShadowKeyStore.class
 })
 public final class KeyChainServiceRoboTest {
+
+    private static final String DEFAULT_KEYSTORE_TYPE = "BKS";
+
     private IKeyChainService.Stub mKeyChain;
 
     @Mock
@@ -128,9 +130,13 @@
         mShadowPackageManager = shadowOf(packageManager);
 
         final ServiceController<KeyChainService> serviceController =
-                Robolectric.buildService(KeyChainService.class).create().bind();
+                Robolectric.buildService(KeyChainService.class);
         final KeyChainService service = serviceController.get();
         service.setInjector(mockInjector);
+        doReturn(KeyStore.getInstance(DEFAULT_KEYSTORE_TYPE))
+                .when(mockInjector).getKeyStoreInstance();
+        serviceController.create().bind();
+
         final Intent intent = new Intent(IKeyChainService.class.getName());
         mKeyChain = (IKeyChainService.Stub) service.onBind(intent);
     }
diff --git a/robotests/src/com/android/keychain/ShadowKeyStore.java b/robotests/src/com/android/keychain/ShadowKeyStore.java
deleted file mode 100644
index 1a7a2b3..0000000
--- a/robotests/src/com/android/keychain/ShadowKeyStore.java
+++ /dev/null
@@ -1,32 +0,0 @@
-/*
- * Copyright (C) 2019 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.keychain;
-
-import android.security.KeyStore;
-
-import org.robolectric.annotation.Implementation;
-import org.robolectric.annotation.Implements;
-
-import java.util.List;
-
-@Implements(KeyStore.class)
-public class ShadowKeyStore {
-    @Implementation
-    public String[] list(String prefix) {
-        return new String[0];
-    }
-}
diff --git a/src/com/android/keychain/KeyChainActivity.java b/src/com/android/keychain/KeyChainActivity.java
index 62401fd..105dc6a 100644
--- a/src/com/android/keychain/KeyChainActivity.java
+++ b/src/com/android/keychain/KeyChainActivity.java
@@ -17,7 +17,6 @@
 package com.android.keychain;
 
 import android.annotation.NonNull;
-import android.app.Activity;
 import android.app.AlertDialog;
 import android.app.PendingIntent;
 import android.app.admin.DevicePolicyEventLogger;
@@ -31,6 +30,7 @@
 import android.content.res.Resources;
 import android.net.Uri;
 import android.os.AsyncTask;
+import android.os.Bundle;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.RemoteException;
@@ -40,7 +40,6 @@
 import android.security.KeyChain;
 import android.stats.devicepolicy.DevicePolicyEnums;
 import android.util.Log;
-import android.view.ContextThemeWrapper;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
@@ -50,9 +49,13 @@
 import android.widget.RadioButton;
 import android.widget.TextView;
 
+import androidx.appcompat.app.AppCompatActivity;
+
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.keychain.internal.KeyInfoProvider;
 
+import com.google.android.material.snackbar.Snackbar;
+
 import org.bouncycastle.asn1.x509.X509Name;
 
 import java.io.IOException;
@@ -74,9 +77,16 @@
 
 import javax.security.auth.x500.X500Principal;
 
-public class KeyChainActivity extends Activity {
+public class KeyChainActivity extends AppCompatActivity {
     private static final String TAG = "KeyChain";
 
+    // The amount of time to delay showing a snackbar. If the alias is received before the snackbar
+    // is shown, the activity will finish. If the certificate selection dialog is shown before the
+    // snackbar, no snackbar will be shown.
+    private static final long SNACKBAR_DELAY_TIME = 2000;
+    // The minimum amount of time to display a snackbar while loading certificates.
+    private static final long SNACKBAR_MIN_TIME = 1000;
+
     private int mSenderUid;
     private String mSenderPackageName;
 
@@ -100,12 +110,25 @@
         }
     }
 
-    // A dialog to show the user while the KeyChain Activity is loading the
-    // certificates.
-    AlertDialog mLoadingDialog;
+    // A snackbar to show the user while the KeyChain Activity is loading the certificates.
+    private Snackbar mSnackbar;
+
+    // A remote service may call {@link android.security.KeyChain#choosePrivateKeyAlias} multiple
+    // times, which will result in multiple intents being sent to KeyChainActivity. The time of the
+    // first received intent is recorded in order to ensure the snackbar is displayed for a
+    // minimum amount of time after receiving the first intent.
+    private long mFirstIntentReceivedTimeMillis = 0L;
 
     private ExecutorService executor = Executors.newSingleThreadExecutor();
     private Handler handler = new Handler(Looper.getMainLooper());
+    private final Runnable mFinishActivity = KeyChainActivity.this::finish;
+    private final Runnable mShowSnackBar = this::showSnackBar;
+
+    @Override
+    protected void onCreate(Bundle savedState) {
+        super.onCreate(savedState);
+        setContentView(R.layout.keychain_activity);
+    }
 
     @Override public void onResume() {
         super.onResume();
@@ -117,6 +140,9 @@
             return;
         }
         try {
+            // getTargetPackage guarantees that the returned string is
+            // supplied by the system, so that an application can not
+            // spoof its package.
             mSenderPackageName = mSender.getIntentSender().getTargetPackage();
             mSenderUid = getPackageManager().getPackageInfo(
                     mSenderPackageName, 0).applicationInfo.uid;
@@ -130,14 +156,27 @@
         chooseCertificate();
     }
 
-    private void showLoadingDialog() {
-        final Context themedContext = new ContextThemeWrapper(
-                this, com.android.internal.R.style.Theme_Translucent_NoTitleBar);
-        mLoadingDialog = new AlertDialog.Builder(themedContext)
-                .setTitle(R.string.app_name)
-                .setMessage(R.string.loading_certs_message)
-                .create();
-        mLoadingDialog.show();
+    @Override
+    protected void onNewIntent(Intent intent) {
+        super.onNewIntent(intent);
+        handler.removeCallbacks(mFinishActivity);
+    }
+
+    private void showSnackBar() {
+        mFirstIntentReceivedTimeMillis = System.currentTimeMillis();
+        mSnackbar = Snackbar.make(findViewById(R.id.container),
+                String.format(getResources().getString(R.string.loading_certs_message),
+                        getApplicationLabel()), Snackbar.LENGTH_INDEFINITE);
+        mSnackbar.show();
+    }
+
+    private void finishSnackBar() {
+        if (mSnackbar != null) {
+            mSnackbar.dismiss();
+            mSnackbar = null;
+        } else {
+            handler.removeCallbacks(mShowSnackBar);
+        }
     }
 
     private void chooseCertificate() {
@@ -209,20 +248,17 @@
                     finish(null);
                     return;
                 }
-                runOnUiThread(new Runnable() {
-                    @Override public void run() {
-                        if (mLoadingDialog != null) {
-                            mLoadingDialog.dismiss();
-                            mLoadingDialog = null;
-                        }
-                        displayCertChooserDialog(certAdapter);
-                    }
+                runOnUiThread(() -> {
+                    finishSnackBar();
+                    displayCertChooserDialog(certAdapter);
                 });
             }
         };
 
-        // Show a dialog to the user to indicate long-running task.
-        showLoadingDialog();
+        // Show a snackbar to the user to indicate long-running task.
+        if (mSnackbar == null) {
+            handler.postDelayed(mShowSnackBar, SNACKBAR_DELAY_TIME);
+        }
         Uri uri = getIntent().getParcelableExtra(KeyChain.EXTRA_URI);
         String alias = getIntent().getStringExtra(KeyChain.EXTRA_ALIAS);
 
@@ -489,20 +525,8 @@
             }
         });
 
-        // getTargetPackage guarantees that the returned string is
-        // supplied by the system, so that an application can not
-        // spoof its package.
-        String pkg = mSender.getIntentSender().getTargetPackage();
-        PackageManager pm = getPackageManager();
-        CharSequence applicationLabel;
-        try {
-            applicationLabel = pm.getApplicationLabel(pm.getApplicationInfo(pkg, 0)).toString();
-        } catch (PackageManager.NameNotFoundException e) {
-            applicationLabel = pkg;
-        }
-        String appMessage = String.format(res.getString(R.string.requesting_application),
-                                          applicationLabel);
-        String contextMessage = appMessage;
+        String contextMessage = String.format(res.getString(R.string.requesting_application),
+                getApplicationLabel());
         Uri uri = getIntent().getParcelableExtra(KeyChain.EXTRA_URI);
         if (uri != null) {
             String hostMessage = String.format(res.getString(R.string.requesting_server),
@@ -531,6 +555,15 @@
         dialog.show();
     }
 
+    private String getApplicationLabel() {
+        PackageManager pm = getPackageManager();
+        try {
+            return pm.getApplicationLabel(pm.getApplicationInfo(mSenderPackageName, 0)).toString();
+        } catch (PackageManager.NameNotFoundException e) {
+            return mSenderPackageName;
+        }
+    }
+
     @VisibleForTesting
     static class CertificateAdapter extends BaseAdapter {
         private final List<String> mAliases;
@@ -632,10 +665,6 @@
     }
 
     private void finish(String alias, boolean isAliasFromPolicy) {
-        if (mLoadingDialog != null) {
-            mLoadingDialog.dismiss();
-            mLoadingDialog = null;
-        }
         if (alias == null || alias.equals(KeyChain.KEY_ALIAS_SELECTION_DENIED)) {
             alias = null;
             setResult(RESULT_CANCELED);
@@ -651,7 +680,7 @@
             new ResponseSender(keyChainAliasResponse, alias, isAliasFromPolicy).execute();
             return;
         }
-        finish();
+        finishActivity();
     }
 
     private class ResponseSender extends AsyncTask<Void, Void, Void> {
@@ -698,7 +727,20 @@
             return null;
         }
         @Override protected void onPostExecute(Void unused) {
+            finishActivity();
+        }
+    }
+
+    private void finishActivity() {
+        long timeElapsedSinceFirstIntent =
+                System.currentTimeMillis() - mFirstIntentReceivedTimeMillis;
+        if (mFirstIntentReceivedTimeMillis == 0L
+                || timeElapsedSinceFirstIntent > SNACKBAR_MIN_TIME) {
+            finishSnackBar();
             finish();
+        } else {
+            long remainingTimeToShowSnackBar = SNACKBAR_MIN_TIME - timeElapsedSinceFirstIntent;
+            handler.postDelayed(mFinishActivity, remainingTimeToShowSnackBar);
         }
     }
 
diff --git a/src/com/android/keychain/KeyChainService.java b/src/com/android/keychain/KeyChainService.java
index 886a55f..6686542 100644
--- a/src/com/android/keychain/KeyChainService.java
+++ b/src/com/android/keychain/KeyChainService.java
@@ -58,6 +58,7 @@
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.util.Preconditions;
 import com.android.internal.widget.LockPatternUtils;
 import com.android.keychain.internal.ExistingKeysProvider;
 import com.android.keychain.internal.GrantsDatabase;
@@ -104,10 +105,16 @@
     private final Set<Integer> ALLOWED_UIDS = Collections.unmodifiableSet(
             new HashSet(Arrays.asList(UID_SELF, Process.WIFI_UID)));
 
+    private static final String MSG_NOT_SYSTEM = "Not system package";
+    private static final String MSG_NOT_SYSTEM_OR_CERT_INSTALLER =
+            "Not system or cert installer package";
+    private static final String MSG_NOT_SYSTEM_OR_CRED_MNG_APP =
+            "Not system or credential management app package";
+
     /** created in onCreate(), closed in onDestroy() */
     private GrantsDatabase mGrantsDb;
     private Injector mInjector;
-    private final KeyStore mKeyStore = getKeyStore();
+    private KeyStore mKeyStore;
     private KeyChainStateStorage mStateStorage;
 
     private Object mCredentialManagementAppLock = new Object();
@@ -120,9 +127,9 @@
         mInjector = new Injector();
     }
 
-    private static KeyStore getKeyStore() {
+    private KeyStore getKeyStore() {
         try {
-            final KeyStore keystore = KeyStore.getInstance("AndroidKeyStore");
+            final KeyStore keystore = mInjector.getKeyStoreInstance();
             keystore.load(null);
             return keystore;
         } catch (KeyStoreException | IOException | NoSuchAlgorithmException
@@ -137,7 +144,7 @@
             return mKeyStore;
         }
         try {
-            final KeyStore keystore = KeyStore.getInstance("AndroidKeyStore");
+            final KeyStore keystore = mInjector.getKeyStoreInstance();
             keystore.load(
                     new AndroidKeyStoreLoadStoreParameter(
                             KeyProperties.NAMESPACE_WIFI));
@@ -151,6 +158,7 @@
 
     @Override public void onCreate() {
         super.onCreate();
+        mKeyStore = getKeyStore();
         mGrantsDb = new GrantsDatabase(this, new KeyStoreAliasesProvider(mKeyStore));
         mStateStorage = new KeyChainStateStorage(getDataDir());
 
@@ -211,11 +219,12 @@
 
         @Override
         public String requestPrivateKey(String alias) {
-            if (!hasGrant(alias)) {
+            final CallerIdentity caller = getCaller();
+            if (!hasGrant(alias, caller)) {
                 return null;
             }
 
-            final int granteeUid = mInjector.getCallingUid();
+            final int granteeUid = caller.mUid;
 
             try {
                 final KeyStore2 keyStore2 = KeyStore2.getInstance();
@@ -228,8 +237,28 @@
             }
         }
 
+        @Override
+        public String getWifiKeyGrantAsUser(String alias) {
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
+
+            if (!hasGrant(alias, Process.WIFI_UID)) {
+                return null;
+            }
+
+            KeyStore2 keyStore2 = KeyStore2.getInstance();
+            try {
+                KeyDescriptor grant = keyStore2.grant(makeKeyDescriptor(alias),
+                        Process.WIFI_UID, KeyPermission.USE | KeyPermission.GET_INFO);
+                return KeyStore2.makeKeystoreEngineGrantString(grant.nspace);
+            } catch (android.security.KeyStoreException e) {
+                Log.e(TAG, "Failed to grant " + alias + " to uid: " + Process.WIFI_UID, e);
+                return null;
+            }
+        }
+
         @Override public byte[] getCertificate(String alias) {
-            if (!hasGrant(alias) && !isCallerWithSystemUid()) {
+            final CallerIdentity caller = getCaller();
+            if (!hasGrant(alias, caller) && !isSystemUid(caller)) {
                 return null;
             }
             try {
@@ -247,7 +276,8 @@
         }
 
         @Override public byte[] getCaCertificates(String alias) {
-            if (!hasGrant(alias) && !isCallerWithSystemUid()) {
+            final CallerIdentity caller = getCaller();
+            if (!hasGrant(alias, caller) && !isSystemUid(caller)) {
                 return null;
             }
             try {
@@ -275,7 +305,7 @@
 
         @Override public void setUserSelectable(String alias, boolean isUserSelectable) {
             validateAlias(alias);
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             Log.i(TAG, String.format("Marking certificate %s as user-selectable: %b", alias,
                     isUserSelectable));
             mGrantsDb.setIsUserSelectable(alias, isUserSelectable);
@@ -283,7 +313,7 @@
 
         @Override public int generateKeyPair(
                 String algorithm, ParcelableKeyGenParameterSpec parcelableSpec) {
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             final KeyGenParameterSpec spec = parcelableSpec.getSpec();
             final String alias = spec.getKeystoreAlias();
 
@@ -341,7 +371,7 @@
 
         @Override public boolean setKeyPairCertificate(String alias, byte[] userCertificate,
                 byte[] userCertificateChain) {
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
 
             final PrivateKey privateKey;
             try {
@@ -396,14 +426,17 @@
             }
         }
 
-        private boolean hasGrant(String alias) {
+        private boolean hasGrant(String alias, CallerIdentity caller) {
+            return hasGrant(alias, caller.mUid);
+        }
+
+        private boolean hasGrant(String alias, int targetUid) {
             validateAlias(alias);
 
-            final int callingUid = mInjector.getCallingUid();
-            if (!mGrantsDb.hasGrant(callingUid, alias)) {
+            if (!mGrantsDb.hasGrant(targetUid, alias)) {
                 Log.w(TAG, String.format(
                         "uid %d doesn't have permission to access the requested alias %s",
-                        callingUid, alias));
+                        targetUid, alias));
                 return false;
             }
 
@@ -411,7 +444,9 @@
         }
 
         @Override public String installCaCertificate(byte[] caCertificate) {
-            checkCertInstallerOrSystemCaller();
+            final CallerIdentity caller = getCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(caller) || isCertInstaller(caller),
+                    MSG_NOT_SYSTEM_OR_CERT_INSTALLER);
             final String alias;
             String subject = null;
             final boolean isSecurityLoggingEnabled = mInjector.isSecurityLoggingEnabled();
@@ -449,7 +484,7 @@
             // anchors. Ultimately, the user should explicitly choose to install the VPN trust
             // anchor separately and independently of CA certificates, at which point this code
             // should be removed.
-            if (CERT_INSTALLER_PACKAGE.equals(callingPackage())) {
+            if (CERT_INSTALLER_PACKAGE.equals(caller.mPackageName)) {
                 try {
                     mKeyStore.setCertificateEntry(String.format("%s %s", subject, alias), cert);
                 } catch(KeyStoreException e) {
@@ -482,7 +517,9 @@
         @Override public boolean installKeyPair(@Nullable byte[] privateKey,
                 @Nullable byte[] userCertificate, @Nullable byte[] userCertificateChain,
                 String alias, int uid) {
-            checkCertInstallerOrSystemCaller();
+            final CallerIdentity caller = getCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(caller) || isCertInstaller(caller),
+                    MSG_NOT_SYSTEM_OR_CERT_INSTALLER);
             if (KeyChain.KEY_ALIAS_SELECTION_DENIED.equals(alias)) {
                 throw new IllegalArgumentException("The alias specified for the key denotes "
                         + "a reserved value and cannot be used to name a key");
@@ -573,7 +610,13 @@
         }
 
         @Override public boolean removeKeyPair(String alias) {
-            checkCertInstallerOrSystemCallerOrHasPermission();
+            final CallerIdentity caller = getCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(caller) || isCertInstaller(caller),
+                    MSG_NOT_SYSTEM_OR_CERT_INSTALLER);
+            return removeKeyPairInternal(alias);
+        }
+
+        private boolean removeKeyPairInternal(String alias) {
             try {
                 mKeyStore.deleteEntry(alias);
             } catch (KeyStoreException e) {
@@ -590,7 +633,7 @@
         }
 
         @Override public boolean containsKeyPair(String alias) {
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             try {
                 final Key key = mKeyStore.getKey(alias, null);
                 return key instanceof PrivateKey;
@@ -613,7 +656,7 @@
 
         @Override public boolean reset() {
             // only Settings should be able to reset
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             mGrantsDb.removeAllAliasesInformation();
             boolean ok = true;
             synchronized (mTrustedCertificateStore) {
@@ -634,7 +677,7 @@
 
         @Override public boolean deleteCaCertificate(String alias) {
             // only Settings should be able to delete
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             boolean ok = true;
             Log.i(TAG, String.format("Deleting CA certificate %s", alias));
             synchronized (mTrustedCertificateStore) {
@@ -674,59 +717,50 @@
             }
         }
 
-        private void checkCertInstallerOrSystemCaller() {
-            final String caller = callingPackage();
-            if (!isCallerWithSystemUid() && !CERT_INSTALLER_PACKAGE.equals(caller)) {
-                throw new SecurityException("Not system or cert installer package: " + caller);
+        private boolean hasManageCredentialManagementAppPermission(CallerIdentity caller) {
+            return mContext.checkPermission(Manifest.permission.MANAGE_CREDENTIAL_MANAGEMENT_APP,
+                    caller.mPid, caller.mUid) == PackageManager.PERMISSION_GRANTED;
+        }
+
+        private boolean isCertInstaller(CallerIdentity caller) {
+            return caller.mPackageName != null
+                    && CERT_INSTALLER_PACKAGE.equals(caller.mPackageName);
+        }
+
+        private boolean isCredentialManagementApp(CallerIdentity caller) {
+            synchronized (mCredentialManagementAppLock) {
+                return mCredentialManagementApp != null && caller.mPackageName != null
+                        && caller.mPackageName.equals(mCredentialManagementApp.getPackageName());
             }
         }
 
-        private void checkSystemCaller() {
-            if (!isCallerWithSystemUid()) {
-                throw new SecurityException("Not system package: " + callingPackage());
-            }
-        }
-
-        private void checkCertInstallerOrSystemCallerOrHasPermission() {
-            if (!hasManageCredentialManagementAppPermission()) {
-                checkCertInstallerOrSystemCaller();
-            }
-        }
-
-        private void checkSystemCallerOrHasPermission() {
-            if (!hasManageCredentialManagementAppPermission()) {
-                checkSystemCaller();
-            }
-        }
-
-        private boolean hasManageCredentialManagementAppPermission() {
-            return mContext.checkCallingPermission(
-                    Manifest.permission.MANAGE_CREDENTIAL_MANAGEMENT_APP)
-                    == PackageManager.PERMISSION_GRANTED;
-        }
-
-        private boolean isCallerWithSystemUid() {
-            return UserHandle.isSameApp(mInjector.getCallingUid(), Process.SYSTEM_UID);
-        }
-
-        private String callingPackage() {
-            return getPackageManager().getNameForUid(mInjector.getCallingUid());
+        private boolean isSystemUid(CallerIdentity caller) {
+            return UserHandle.isSameApp(caller.mUid, Process.SYSTEM_UID);
         }
 
         @Override public boolean hasGrant(int uid, String alias) {
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             return mGrantsDb.hasGrant(uid, alias);
         }
 
-        @Override public void setGrant(int uid, String alias, boolean value) {
-            checkSystemCaller();
-            mGrantsDb.setGrant(uid, alias, value);
-            broadcastPermissionChange(uid, alias, value);
+        @Override public boolean setGrant(int uid, String alias, boolean granted) {
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
+            mGrantsDb.setGrant(uid, alias, granted);
+            if (!granted) {
+                try {
+                    KeyStore2.getInstance().ungrant(makeKeyDescriptor(alias), uid);
+                } catch (android.security.KeyStoreException e) {
+                    Log.e(TAG, "Failed to ungrant " + alias + " to uid: " + uid, e);
+                    return false;
+                }
+            }
+            broadcastPermissionChange(uid, alias, granted);
             broadcastLegacyStorageChange();
+            return true;
         }
 
         @Override public int[] getGrants(String alias) {
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             try {
                 if (mKeyStore.isKeyEntry(alias)) {
                     return mGrantsDb.getGrants(alias);
@@ -805,7 +839,9 @@
         @Override
         public void setCredentialManagementApp(@NonNull String packageName,
                 @NonNull AppUriAuthenticationPolicy authenticationPolicy) {
-            checkSystemCallerOrHasPermission();
+            final CallerIdentity caller = getCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(caller)
+                    || hasManageCredentialManagementAppPermission(caller), MSG_NOT_SYSTEM);
             checkValidAuthenticationPolicy(authenticationPolicy);
 
             synchronized (mCredentialManagementAppLock) {
@@ -848,7 +884,7 @@
             // authentication policy
             for (String existingAlias : existingAliases) {
                 if (!newAliases.contains(existingAlias)) {
-                    removeKeyPair(existingAlias);
+                    removeKeyPairInternal(existingAlias);
                 }
             }
         }
@@ -870,7 +906,7 @@
 
         @Override
         public boolean hasCredentialManagementApp() {
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             synchronized (mCredentialManagementAppLock) {
                 return mCredentialManagementApp != null;
             }
@@ -879,7 +915,7 @@
         @Nullable
         @Override
         public String getCredentialManagementAppPackageName() {
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             synchronized (mCredentialManagementAppLock) {
                 return mCredentialManagementApp != null
                         ? mCredentialManagementApp.getPackageName()
@@ -890,7 +926,9 @@
         @Nullable
         @Override
         public AppUriAuthenticationPolicy getCredentialManagementAppPolicy() {
-            checkSystemCaller();
+            final CallerIdentity caller = getCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(caller)
+                            || isCredentialManagementApp(caller), MSG_NOT_SYSTEM_OR_CRED_MNG_APP);
             synchronized (mCredentialManagementAppLock) {
                 return mCredentialManagementApp != null
                         ? mCredentialManagementApp.getAuthenticationPolicy()
@@ -902,20 +940,24 @@
         @Override
         public String getPredefinedAliasForPackageAndUri(@NonNull String packageName,
                 @Nullable Uri uri) {
-            checkSystemCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(getCaller()), MSG_NOT_SYSTEM);
             synchronized (mCredentialManagementAppLock) {
                 if (mCredentialManagementApp == null || uri == null) {
                     return null;
                 }
                 Map<Uri, String> urisToAliases = mCredentialManagementApp.getAuthenticationPolicy()
                         .getAppAndUriMappings().get(packageName);
-                return urisToAliases.get(uri);
+                return urisToAliases != null ? urisToAliases.get(uri) : null;
             }
         }
 
         @Override
         public void removeCredentialManagementApp() {
-            checkSystemCallerOrHasPermission();
+            final CallerIdentity caller = getCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(caller)
+                            || isCredentialManagementApp(caller)
+                            || hasManageCredentialManagementAppPermission(caller),
+                    MSG_NOT_SYSTEM_OR_CRED_MNG_APP);
             synchronized (mCredentialManagementAppLock) {
                 if (mCredentialManagementApp != null) {
                     // Remove all certificates
@@ -926,6 +968,16 @@
                 mStateStorage.saveCredentialManagementApp(mCredentialManagementApp);
             }
         }
+
+        @Override
+        public boolean isCredentialManagementApp(@NonNull String packageName) {
+            final CallerIdentity caller = getCaller();
+            Preconditions.checkCallAuthorization(isSystemUid(caller)
+                    || isCredentialManagementApp(caller), MSG_NOT_SYSTEM_OR_CRED_MNG_APP);
+            synchronized (mCredentialManagementAppLock) {
+                return packageName.equals(mCredentialManagementApp.getPackageName());
+            }
+        }
     };
 
     @Override public IBinder onBind(Intent intent) {
@@ -982,6 +1034,23 @@
         return Base64.encodeToString(cert, Base64.NO_WRAP);
     }
 
+    private final class CallerIdentity {
+
+        final int mUid;
+        final int mPid;
+        final String mPackageName;
+
+        CallerIdentity() {
+            mUid = mInjector.getCallingUid();
+            mPid = Binder.getCallingPid();
+            mPackageName = getPackageManager().getNameForUid(mUid);
+        }
+    }
+
+    private CallerIdentity getCaller() {
+        return new CallerIdentity();
+    }
+
     @VisibleForTesting
     void setInjector(Injector injector) {
         mInjector = injector;
@@ -1003,5 +1072,9 @@
         public int getCallingUid() {
             return Binder.getCallingUid();
         }
+
+        public KeyStore getKeyStoreInstance() throws KeyStoreException {
+            return KeyStore.getInstance("AndroidKeyStore");
+        }
     }
 }
diff --git a/support/Android.bp b/support/Android.bp
index 196a3ba..3859d9b 100644
--- a/support/Android.bp
+++ b/support/Android.bp
@@ -30,7 +30,4 @@
         "junit",
     ],
     certificate: "platform",
-    test_suites: [
-        "general-tests",
-    ]
 }
diff --git a/support/AndroidManifest.xml b/support/AndroidManifest.xml
index 934660a..11e7f55 100644
--- a/support/AndroidManifest.xml
+++ b/support/AndroidManifest.xml
@@ -18,7 +18,8 @@
           package="com.android.keychain.tests.support"
 	  android:sharedUserId="android.uid.system">
     <application android:process="system">
-        <service android:name="com.android.keychain.tests.support.KeyChainServiceTestSupport">
+        <service android:name="com.android.keychain.tests.support.KeyChainServiceTestSupport"
+            android:exported="true">
             <intent-filter>
                 <action android:name="com.android.keychain.tests.support.IKeyChainServiceTestSupport"/>
             </intent-filter>
diff --git a/tests/Android.bp b/tests/Android.bp
index 5293e8e..8d2415a 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -36,9 +36,9 @@
     libs: [
         "android.test.base",
     ],
-    test_suites: ["general-tests"],
-    required: [
-        "KeyChainTestsSupport",
+    test_suites: ["device-tests"],
+    data: [
+        ":KeyChainTestsSupport",
     ],
 
     instrumentation_for: "KeyChain",
diff --git a/tests/AndroidManifest.xml b/tests/AndroidManifest.xml
index c1f5021..754b76b 100644
--- a/tests/AndroidManifest.xml
+++ b/tests/AndroidManifest.xml
@@ -40,12 +40,14 @@
     -->
     <application>
         <uses-library android:name="android.test.runner" />
-        <activity android:name="com.android.keychain.tests.KeyChainTestActivity">
+        <activity android:name="com.android.keychain.tests.KeyChainTestActivity"
+            android:exported="true">
             <intent-filter>
                 <action android:name="com.android.keychain.tests.KeyChainTestActivity"/>
             </intent-filter>
         </activity>
-        <activity android:name="com.android.keychain.tests.KeyChainSocketTestActivity">
+        <activity android:name="com.android.keychain.tests.KeyChainSocketTestActivity"
+            android:exported="true">
             <intent-filter>
                 <action android:name="com.android.keychain.tests.KeyChainSocketTestActivity"/>
             </intent-filter>