Merge "KeyChain/Keystore 2.0" am: 4d9cd24e64

Original change: https://android-review.googlesource.com/c/platform/packages/apps/KeyChain/+/1560128

MUST ONLY BE SUBMITTED BY AUTOMERGER

Change-Id: I08dd82c6edc2418148378e638cf0abd221d9e4df

7 files changed, 564 insertions, 373 deletions

diff --git a/robotests/src/com/android/keychain/AliasLoaderTest.java b/robotests/src/com/android/keychain/AliasLoaderTest.java
index 78c6c85..24372dd 100644
--- a/robotests/src/com/android/keychain/AliasLoaderTest.java
+++ b/robotests/src/com/android/keychain/AliasLoaderTest.java
@@ -20,10 +20,20 @@ import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import android.security.Credentials;
-import android.security.KeyStore;
 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;
@@ -126,6 +136,11 @@ public final class AliasLoaderTest {
     private byte[] mECCertTwo;
     private ArrayList<byte[]> mIssuers;
 
+    private Certificate toCertificate(byte[] bytes) throws CertificateException {
+        CertificateFactory cf = CertificateFactory.getInstance("X.509");
+        return cf.generateCertificate(new ByteArrayInputStream(bytes));
+    }
+
     @Before
     public void setUp() {
         mRSACertOne = Base64.decode(SELF_SIGNED_RSA_CERT_1_B64, Base64.DEFAULT);
@@ -148,9 +163,10 @@ public final class AliasLoaderTest {
     @Test
     public void testAliasLoader_loadsAllAliases()
             throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException {
+                    TimeoutException, KeyStoreException {
         KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.list(Credentials.USER_PRIVATE_KEY)).thenReturn(new String[] {"b", "c", "a"});
+        when(keyStore.aliases()).thenReturn(
+                Collections.enumeration(ImmutableList.of("b", "c", "a")));
 
         KeyChainActivity.AliasLoader loader =
                 new KeyChainActivity.AliasLoader(
@@ -172,9 +188,9 @@ public final class AliasLoaderTest {
     @Test
     public void testAliasLoader_copesWithNoAliases()
             throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException {
+                    TimeoutException, KeyStoreException {
         KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.list(Credentials.USER_PRIVATE_KEY)).thenReturn(null);
+        when(keyStore.aliases()).thenReturn(null);
 
         KeyChainActivity.AliasLoader loader =
                 new KeyChainActivity.AliasLoader(
@@ -193,9 +209,10 @@ public final class AliasLoaderTest {
     @Test
     public void testAliasLoader_filtersNonUserSelectableAliases()
             throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException {
+                    TimeoutException, KeyStoreException {
         KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.list(Credentials.USER_PRIVATE_KEY)).thenReturn(new String[] {"a", "b", "c"});
+        when(keyStore.aliases()).thenReturn(
+                Collections.enumeration(ImmutableList.of("a", "b", "c")));
 
         KeyInfoProvider infoProvider = mock(KeyInfoProvider.class);
         when(infoProvider.isUserSelectable("a")).thenReturn(false);
@@ -217,10 +234,10 @@ public final class AliasLoaderTest {
     @Test
     public void testAliasLoader_filtersAliasesWithNonConformingParameters()
             throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException {
+                    TimeoutException, KeyStoreException {
         KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.list(Credentials.USER_PRIVATE_KEY))
-                .thenReturn(new String[] {"a", "b", "c", "d"});
+        when(keyStore.aliases()).thenReturn(
+                Collections.enumeration(ImmutableList.of("a", "b", "c", "d")));
 
         KeyInfoProvider infoProvider = mock(KeyInfoProvider.class);
         when(infoProvider.isUserSelectable("a")).thenReturn(true);
@@ -251,18 +268,20 @@ public final class AliasLoaderTest {
         Assert.assertEquals("a", result.getItem(0));
     }
 
-    private KeyStore prepareKeyStoreWithCertificates() {
+    private KeyStore prepareKeyStoreWithCertificates()
+            throws CertificateException, KeyStoreException {
         KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.get(Credentials.USER_CERTIFICATE + "rsa1")).thenReturn(mRSACertOne);
-        when(keyStore.get(Credentials.USER_CERTIFICATE + "ec1")).thenReturn(mECCertOne);
-        when(keyStore.get(Credentials.USER_CERTIFICATE + "rsa2")).thenReturn(mRSACertTwo);
-        when(keyStore.get(Credentials.USER_CERTIFICATE + "ec2")).thenReturn(mECCertTwo);
+        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));
 
         return keyStore;
     }
 
     @Test
-    public void testCertificateParametersFilter_filtersByKey() throws CancellationException {
+    public void testCertificateParametersFilter_filtersByKey()
+            throws CancellationException, KeyStoreException, CertificateException {
         KeyStore keyStore = prepareKeyStoreWithCertificates();
 
         KeyChainActivity.CertificateParametersFilter ec_checker =
@@ -279,7 +298,8 @@ public final class AliasLoaderTest {
     }
 
     @Test
-    public void testCertificateParametersFilter_filtersByIssuer() throws CancellationException {
+    public void testCertificateParametersFilter_filtersByIssuer()
+            throws CancellationException, KeyStoreException, CertificateException {
         KeyStore keyStore = prepareKeyStoreWithCertificates();
 
         KeyChainActivity.CertificateParametersFilter issuer_checker =
@@ -294,7 +314,7 @@ public final class AliasLoaderTest {
     @Test
     public void testCertificateParametersFilter_filtersByIssuerAndKey()
             throws InterruptedException, ExecutionException, CancellationException,
-                    TimeoutException {
+                    TimeoutException, KeyStoreException, CertificateException {
         KeyStore keyStore = prepareKeyStoreWithCertificates();
 
         KeyChainActivity.CertificateParametersFilter issuer_checker =
diff --git a/src/com/android/keychain/KeyChainActivity.java b/src/com/android/keychain/KeyChainActivity.java
index 15a1cad..62401fd 100644
--- a/src/com/android/keychain/KeyChainActivity.java
+++ b/src/com/android/keychain/KeyChainActivity.java
@@ -36,10 +36,8 @@ import android.os.Looper;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.UserManager;
-import android.security.Credentials;
 import android.security.IKeyChainAliasCallback;
 import android.security.KeyChain;
-import android.security.KeyStore;
 import android.stats.devicepolicy.DevicePolicyEnums;
 import android.util.Log;
 import android.view.ContextThemeWrapper;
@@ -54,17 +52,20 @@ import android.widget.TextView;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.keychain.internal.KeyInfoProvider;
+
 import org.bouncycastle.asn1.x509.X509Name;
 
-import java.io.ByteArrayInputStream;
 import java.io.IOException;
-import java.io.InputStream;
+import java.security.KeyStore;
+import java.security.KeyStoreException;
+import java.security.NoSuchAlgorithmException;
+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.Arrays;
 import java.util.Collections;
+import java.util.Enumeration;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -85,7 +86,19 @@ public class KeyChainActivity extends Activity {
     // get do file I/O in the remote keystore process and while they
     // do not cause StrictMode violations, they logically should not
     // be done on the UI thread.
-    private KeyStore mKeyStore = KeyStore.getInstance();
+    private final KeyStore mKeyStore = getKeyStore();
+
+    private static KeyStore getKeyStore() {
+        try {
+            final KeyStore keystore = KeyStore.getInstance("AndroidKeyStore");
+            keystore.load(null);
+            return keystore;
+        } catch (KeyStoreException | IOException | NoSuchAlgorithmException
+                | CertificateException e) {
+            Log.e(TAG, "Error opening AndroidKeyStore.", e);
+            throw new RuntimeException("Error opening AndroidKeyStore.", e);
+        }
+    }
 
     // A dialog to show the user while the KeyChain Activity is loading the
     // certificates.
@@ -375,8 +388,8 @@ public class KeyChainActivity extends Activity {
         private final KeyInfoProvider mInfoProvider;
         private final CertificateParametersFilter mCertificateFilter;
 
-        public AliasLoader(KeyStore keyStore, Context context, KeyInfoProvider infoProvider,
-                CertificateParametersFilter certificateFilter) {
+        public AliasLoader(KeyStore keyStore, Context context,
+                KeyInfoProvider infoProvider, CertificateParametersFilter certificateFilter) {
           mKeyStore = keyStore;
           mContext = context;
           mInfoProvider = infoProvider;
@@ -384,10 +397,19 @@ public class KeyChainActivity extends Activity {
         }
 
         @Override protected CertificateAdapter doInBackground(Void... params) {
-            String[] aliasArray = mKeyStore.list(Credentials.USER_PRIVATE_KEY);
-            List<String> rawAliasList = ((aliasArray == null)
-                                      ? Collections.<String>emptyList()
-                                      : Arrays.asList(aliasArray));
+            final List<String> rawAliasList = new ArrayList<>();
+            try {
+                final Enumeration<String> aliases = mKeyStore.aliases();
+                while (aliases.hasMoreElements()) {
+                    final String alias = aliases.nextElement();
+                    if (mKeyStore.isKeyEntry(alias)) {
+                        rawAliasList.add(alias);
+                    }
+                }
+            } catch (KeyStoreException e) {
+                Log.e(TAG, "Error while loading entries from keystore. "
+                        + "List may be empty or incomplete.");
+            }
 
             return new CertificateAdapter(mKeyStore, mContext,
                     rawAliasList.stream().filter(mInfoProvider::isUserSelectable)
@@ -685,33 +707,51 @@ public class KeyChainActivity extends Activity {
     }
 
     private static X509Certificate loadCertificate(KeyStore keyStore, String alias) {
-        byte[] bytes = keyStore.get(Credentials.USER_CERTIFICATE + alias);
-        if (bytes == null) {
-            Log.i(TAG, String.format("Missing user certificate for key alias %s", alias));
-            return null;
-        }
-        InputStream in = new ByteArrayInputStream(bytes);
+        final Certificate cert;
         try {
-            CertificateFactory cf = CertificateFactory.getInstance("X.509");
-            return (X509Certificate)cf.generateCertificate(in);
-        } catch (CertificateException ignored) {
-            Log.w(TAG, "Error generating certificate", ignored);
+            if (keyStore.isCertificateEntry(alias)) {
+                return null;
+            }
+            cert = keyStore.getCertificate(alias);
+        } catch (KeyStoreException e) {
+            Log.e(TAG, String.format("Error trying to retrieve certificate for \"%s\".", alias), e);
             return null;
         }
-    }
-
-    private static List<X509Certificate> loadCertificateChain(KeyStore keyStore, String alias) {
-        byte[] chainBytes = keyStore.get(Credentials.CA_CERTIFICATE + alias);
-        if (chainBytes == null) {
-            Log.i(TAG, String.format("Missing certificate chain for key alias %s", alias));
-            return Collections.emptyList();
+        if (cert != null) {
+            if (cert instanceof X509Certificate) {
+                return (X509Certificate) cert;
+            } else {
+                Log.w(TAG, String.format("Certificate associated with alias \"%s\" is not X509.",
+                        alias));
+            }
         }
+        return null;
+    }
 
+    private static List<X509Certificate> loadCertificateChain(KeyStore keyStore,
+            String alias) {
+        final Certificate[] certs;
+        final boolean isCertificateEntry;
         try {
-            return Credentials.convertFromPem(chainBytes);
-        } catch (IOException | CertificateException e) {
-            Log.w(TAG, String.format("Error parsing certificate chain for alias %s", alias), e);
+            isCertificateEntry = keyStore.isCertificateEntry(alias);
+            certs = keyStore.getCertificateChain(alias);
+        } catch (KeyStoreException e) {
+            Log.e(TAG, String.format("Error trying to retrieve certificate chain for \"%s\".",
+                    alias), e);
             return Collections.emptyList();
         }
+        final List<X509Certificate> result = new ArrayList<>();
+        // If this is a certificate entry we return the single certificate. Otherwise we trim the
+        // leaf and return only the rest of the chain.
+        for (int i = isCertificateEntry ? 0 : 1; i < certs.length; ++i) {
+            if (certs[i] instanceof X509Certificate) {
+                result.add((X509Certificate) certs[i]);
+            } else {
+                Log.w(TAG,"A certificate in the chain of alias \""
+                        + alias + "\" is not X509.");
+                return Collections.emptyList();
+            }
+        }
+        return result;
     }
 }
diff --git a/src/com/android/keychain/KeyChainService.java b/src/com/android/keychain/KeyChainService.java
index f96efa8..91b6c31 100644
--- a/src/com/android/keychain/KeyChainService.java
+++ b/src/com/android/keychain/KeyChainService.java
@@ -31,6 +31,7 @@ import android.content.Intent;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.StringParceledListSlice;
+import android.hardware.security.keymint.ErrorCode;
 import android.net.Uri;
 import android.os.Binder;
 import android.os.Build;
@@ -42,14 +43,16 @@ import android.security.CredentialManagementApp;
 import android.security.Credentials;
 import android.security.IKeyChainService;
 import android.security.KeyChain;
-import android.security.KeyStore;
-import android.security.keymaster.KeymasterArguments;
-import android.security.keymaster.KeymasterCertificateChain;
-import android.security.keystore.AttestationUtils;
-import android.security.keystore.DeviceIdAttestationException;
+import android.security.KeyStore2;
+import android.security.keystore.AndroidKeyStoreProvider;
 import android.security.keystore.KeyGenParameterSpec;
+import android.security.keystore.KeyProperties;
 import android.security.keystore.ParcelableKeyGenParameterSpec;
 import android.security.keystore.StrongBoxUnavailableException;
+import android.security.keystore2.AndroidKeyStoreLoadStoreParameter;
+import android.system.keystore2.Domain;
+import android.system.keystore2.KeyDescriptor;
+import android.system.keystore2.KeyPermission;
 import android.text.TextUtils;
 import android.util.Base64;
 import android.util.Log;
@@ -62,20 +65,32 @@ import com.android.keychain.internal.GrantsDatabase;
 import com.android.org.conscrypt.TrustedCertificateStore;
 
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.security.InvalidAlgorithmParameterException;
+import java.security.Key;
+import java.security.KeyFactory;
 import java.security.KeyPair;
 import java.security.KeyPairGenerator;
+import java.security.KeyStore;
+import java.security.KeyStoreException;
 import java.security.NoSuchAlgorithmException;
 import java.security.NoSuchProviderException;
+import java.security.PrivateKey;
+import java.security.ProviderException;
+import java.security.UnrecoverableKeyException;
 import java.security.cert.Certificate;
 import java.security.cert.CertificateEncodingException;
 import java.security.cert.CertificateException;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
+import java.security.spec.InvalidKeySpecException;
+import java.security.spec.PKCS8EncodedKeySpec;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.Enumeration;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -88,12 +103,12 @@ public class KeyChainService extends IntentService {
     private static final String TAG = "KeyChain";
     private static final String CERT_INSTALLER_PACKAGE = "com.android.certinstaller";
     private final Set<Integer> ALLOWED_UIDS = Collections.unmodifiableSet(
-            new HashSet(Arrays.asList(KeyStore.UID_SELF, Process.WIFI_UID)));
+            new HashSet(Arrays.asList(android.security.KeyStore.UID_SELF, Process.WIFI_UID)));
 
     /** created in onCreate(), closed in onDestroy() */
     private GrantsDatabase mGrantsDb;
     private Injector mInjector;
-    private final KeyStore mKeyStore = KeyStore.getInstance();
+    private final KeyStore mKeyStore = getKeyStore();
     private KeyChainStateStorage mStateStorage;
 
     private Object mCredentialManagementAppLock = new Object();
@@ -106,6 +121,40 @@ public class KeyChainService extends IntentService {
         mInjector = new Injector();
     }
 
+    private static KeyStore getKeyStore() {
+        try {
+            final KeyStore keystore = KeyStore.getInstance("AndroidKeyStore");
+            keystore.load(null);
+            return keystore;
+        } catch (KeyStoreException | IOException | NoSuchAlgorithmException
+                | CertificateException e) {
+            Log.e(TAG, "Error opening AndroidKeyStore.", e);
+            throw new RuntimeException("Error opening AndroidKeyStore.", e);
+        }
+    }
+
+    private KeyStore getKeyStore(boolean useWifiNamespace) {
+        if (!useWifiNamespace) {
+            return mKeyStore;
+        }
+        try {
+            KeyStore keystore = null;
+            if (AndroidKeyStoreProvider.isKeystore2Enabled()) {
+                keystore = KeyStore.getInstance("AndroidKeyStore");
+                keystore.load(
+                        new AndroidKeyStoreLoadStoreParameter(
+                                KeyProperties.NAMESPACE_WIFI));
+            } else {
+                keystore = AndroidKeyStoreProvider.getKeyStoreForUid(Process.WIFI_UID);
+            }
+            return keystore;
+        } catch (IOException | CertificateException | KeyStoreException
+                | NoSuchProviderException | NoSuchAlgorithmException e) {
+            Log.e(TAG, "Failed to open AndroidKeyStore for WI-FI namespace.", e);
+            return null;
+        }
+    }
+
     @Override public void onCreate() {
         super.onCreate();
         mGrantsDb = new GrantsDatabase(this, new KeyStoreAliasesProvider(mKeyStore));
@@ -132,23 +181,35 @@ public class KeyChainService extends IntentService {
 
         @Override
         public List<String> getExistingKeyAliases() {
-            List<String> aliases = new ArrayList<String>();
-            String[] keyStoreAliases = mKeyStore.list(Credentials.USER_PRIVATE_KEY);
-            if (keyStoreAliases == null) {
-                return aliases;
-            }
-
-            for (String alias: keyStoreAliases) {
-                Log.w(TAG, "Got Alias from KeyStore: " + alias);
-                String unPrefixedAlias = alias.replaceFirst("^" + Credentials.USER_PRIVATE_KEY, "");
-                if (!unPrefixedAlias.startsWith(LockPatternUtils.SYNTHETIC_PASSWORD_KEY_PREFIX)) {
-                    aliases.add(unPrefixedAlias);
+            final List<String> keyStoreAliases = new ArrayList<>();
+            try {
+                final Enumeration<String> aliases = mKeyStore.aliases();
+                while (aliases.hasMoreElements()) {
+                    final String alias = aliases.nextElement();
+                    if (!alias.startsWith(LockPatternUtils.SYNTHETIC_PASSWORD_KEY_PREFIX)) {
+                        if (mKeyStore.isKeyEntry(alias)) {
+                            keyStoreAliases.add(alias);
+                        }
+                    }
                 }
+            } catch (KeyStoreException e) {
+                Log.e(TAG, "Error while loading entries from keystore. "
+                        + "List may be empty or incomplete.");
             }
-            return aliases;
+
+            return keyStoreAliases;
         }
     }
 
+    private KeyDescriptor makeKeyDescriptor(String alias) {
+        final KeyDescriptor key = new KeyDescriptor();
+        key.domain = Domain.APP;
+        key.nspace = KeyProperties.NAMESPACE_APPLICATION;
+        key.alias = alias;
+        key.blob = null;
+        return key;
+    }
+
     private final IKeyChainService.Stub mIKeyChainService = new IKeyChainService.Stub() {
         private final TrustedCertificateStore mTrustedCertificateStore
                 = new TrustedCertificateStore();
@@ -160,24 +221,65 @@ public class KeyChainService extends IntentService {
                 return null;
             }
 
-            final String keystoreAlias = Credentials.USER_PRIVATE_KEY + alias;
-            final int uid = mInjector.getCallingUid();
-            Log.i(TAG, String.format("UID %d will be granted access to %s", uid, keystoreAlias));
-            return mKeyStore.grant(keystoreAlias, uid);
+            final int granteeUid = mInjector.getCallingUid();
+
+            if (AndroidKeyStoreProvider.isKeystore2Enabled()) {
+                KeyStore2 keyStore2 = KeyStore2.getInstance();
+                try {
+                    KeyDescriptor grant = keyStore2.grant(makeKeyDescriptor(alias), granteeUid,
+                            KeyPermission.USE | KeyPermission.GET_INFO);
+                    return KeyChain.getGrantString(grant);
+                } catch (android.security.KeyStoreException e) {
+                    Log.e(TAG, "Failed to grant " + alias + " to uid: " + granteeUid, e);
+                    return null;
+                }
+            } else {
+                android.security.KeyStore keyStore = android.security.KeyStore.getInstance();
+                final String keystoreAlias = Credentials.USER_PRIVATE_KEY + alias;
+                Log.i(TAG, String.format("UID %d will be granted access to %s", granteeUid,
+                        keystoreAlias));
+                return keyStore.grant(keystoreAlias, granteeUid);
+            }
         }
 
         @Override public byte[] getCertificate(String alias) {
-            if (!hasGrant(alias)) {
+            if (!hasGrant(alias) && !isCallerWithSystemUid()) {
+                return null;
+            }
+            try {
+                if (!mKeyStore.isCertificateEntry(alias)) {
+                    final Certificate cert = mKeyStore.getCertificate(alias);
+                    if (cert == null) return null;
+                    return cert.getEncoded();
+                } else {
+                    return null;
+                }
+            } catch (KeyStoreException | CertificateEncodingException e) {
+                Log.e(TAG, "Failed to retrieve certificate.", e);
                 return null;
             }
-            return mKeyStore.get(Credentials.USER_CERTIFICATE + alias);
         }
 
         @Override public byte[] getCaCertificates(String alias) {
-            if (!hasGrant(alias)) {
+            if (!hasGrant(alias) && !isCallerWithSystemUid()) {
+                return null;
+            }
+            try {
+                if (mKeyStore.isCertificateEntry(alias)) {
+                    return mKeyStore.getCertificate(alias).getEncoded();
+                } else {
+                    final Certificate[] certs = mKeyStore.getCertificateChain(alias);
+                    if (certs == null || certs.length <= 1) return null;
+                    ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+                    for (int i = 1; i < certs.length; ++i) {
+                        outputStream.write(certs[i].getEncoded());
+                    }
+                    return outputStream.toByteArray();
+                }
+            } catch (KeyStoreException | CertificateEncodingException | IOException e) {
+                Log.e(TAG, "Failed to retrieve certificate(s) from AndroidKeyStore.", e);
                 return null;
             }
-            return mKeyStore.get(Credentials.CA_CERTIFICATE + alias);
         }
 
         @Override public boolean isUserSelectable(String alias) {
@@ -208,23 +310,11 @@ public class KeyChainService extends IntentService {
             }
             // Validate the alias here to avoid relying on KeyGenParameterSpec c'tor preventing
             // the creation of a KeyGenParameterSpec instance with a non-empty alias.
-            if (TextUtils.isEmpty(alias) || spec.getUid() != KeyStore.UID_SELF) {
+            if (TextUtils.isEmpty(alias) || spec.getUid() != android.security.KeyStore.UID_SELF) {
                 Log.e(TAG, "Cannot generate key pair with empty alias or specified uid.");
                 return KeyChain.KEY_GEN_MISSING_ALIAS;
             }
 
-            if (spec.getAttestationChallenge() != null) {
-                Log.e(TAG, "Key generation request should not include an Attestation challenge.");
-                return KeyChain.KEY_GEN_SUPERFLUOUS_ATTESTATION_CHALLENGE;
-            }
-
-            if (!removeKeyPair(alias)) {
-                Log.e(TAG, "Failed to remove previously-installed alias " + alias);
-                //TODO: Introduce a different error code in R to distinguish the failure to remove
-                // old keys from other failures.
-                return KeyChain.KEY_GEN_FAILURE;
-            }
-
             try {
                 KeyPairGenerator generator = KeyPairGenerator.getInstance(
                         algorithm, "AndroidKeyStore");
@@ -250,75 +340,60 @@ public class KeyChainService extends IntentService {
             } catch (StrongBoxUnavailableException e) {
                 Log.e(TAG, "StrongBox unavailable.", e);
                 return KeyChain.KEY_GEN_STRONGBOX_UNAVAILABLE;
+            } catch (ProviderException e) {
+                Throwable t = e.getCause();
+                if (t instanceof android.security.KeyStoreException) {
+                    if (((android.security.KeyStoreException) t).getErrorCode()
+                            == ErrorCode.CANNOT_ATTEST_IDS) {
+                        return KeyChain.KEY_ATTESTATION_CANNOT_ATTEST_IDS;
+                    }
+                }
+                Log.e(TAG, "KeyStore error.", e);
+                return KeyChain.KEY_GEN_FAILURE;
             }
         }
 
-        @Override public int attestKey(
-                String alias, byte[] attestationChallenge,
-                int[] idAttestationFlags,
-                KeymasterCertificateChain attestationChain) {
+        @Override public boolean setKeyPairCertificate(String alias, byte[] userCertificate,
+                byte[] userCertificateChain) {
             checkSystemCaller();
-            validateAlias(alias);
 
-            if (attestationChallenge == null) {
-                Log.e(TAG, String.format("Missing attestation challenge for alias %s", alias));
-                return KeyChain.KEY_ATTESTATION_MISSING_CHALLENGE;
-            }
-
-            if (Log.isLoggable(TAG, Log.DEBUG)) {
-                Log.d(TAG, String.format("About to attest key alias %s, challenge %s, flags %s",
-                        alias, Base64.encodeToString(attestationChallenge, Base64.NO_WRAP),
-                        Arrays.toString(idAttestationFlags)));
+            final PrivateKey privateKey;
+            try {
+                final Key key = mKeyStore.getKey(alias, null);
+                if (! (key instanceof PrivateKey)) {
+                    return false;
+                }
+                privateKey = (PrivateKey) key;
+            } catch (KeyStoreException | NoSuchAlgorithmException | UnrecoverableKeyException e) {
+                Log.e(TAG, "Failed to get private key entry.", e);
+                return false;
             }
 
-            final KeymasterArguments attestArgs;
+            final ArrayList<Certificate> certs = new ArrayList<>();
             try {
-                attestArgs = AttestationUtils.prepareAttestationArguments(
-                        mContext, idAttestationFlags, attestationChallenge);
-            } catch (DeviceIdAttestationException e) {
-                Log.e(TAG, "Failed collecting attestation data", e);
-                return KeyChain.KEY_ATTESTATION_CANNOT_COLLECT_DATA;
-            }
-            final String keystoreAlias = Credentials.USER_PRIVATE_KEY + alias;
-            final int errorCode = mKeyStore.attestKey(keystoreAlias, attestArgs, attestationChain);
-            if (errorCode != KeyStore.NO_ERROR) {
-                Log.e(TAG, String.format("Failure attesting for key %s: %d", alias, errorCode));
-                if (errorCode == KeyStore.CANNOT_ATTEST_IDS) {
-                    return KeyChain.KEY_ATTESTATION_CANNOT_ATTEST_IDS;
-                } else {
-                    // General failure, cannot discern which.
-                    return KeyChain.KEY_ATTESTATION_FAILURE;
+                if (userCertificate != null) {
+                    certs.add(parseCertificate(userCertificate));
+                }
+                if (userCertificateChain != null) {
+                    certs.addAll(parseCertificates(userCertificateChain));
                 }
+            } catch (CertificateException e) {
+                Log.e(TAG, "Failed to parse user certificate.", e);
+                return false;
             }
 
-            return KeyChain.KEY_ATTESTATION_SUCCESS;
-        }
+            final Certificate[] certsArray = certs.toArray(new Certificate[0]);
 
-        @Override public boolean setKeyPairCertificate(String alias, byte[] userCertificate,
-                byte[] userCertificateChain) {
-            checkSystemCaller();
-            if (!mKeyStore.put(Credentials.USER_CERTIFICATE + alias, userCertificate,
-                        KeyStore.UID_SELF, KeyStore.FLAG_NONE)) {
-                Log.e(TAG, "Failed to import user certificate " + userCertificate);
+            try {
+                // setKeyEntry with a private key loaded from AndroidKeyStore replaces
+                // the certificate components without touching the private key if
+                // the alias is the same as that of the private key.
+                mKeyStore.setKeyEntry(alias, privateKey, null, certsArray);
+            } catch (KeyStoreException e) {
+                Log.e(TAG, "Failed update key certificates.", e);
                 return false;
             }
 
-            if (userCertificateChain != null && userCertificateChain.length > 0) {
-                if (!mKeyStore.put(Credentials.CA_CERTIFICATE + alias, userCertificateChain,
-                            KeyStore.UID_SELF, KeyStore.FLAG_NONE)) {
-                    Log.e(TAG, "Failed to import certificate chain" + userCertificateChain);
-                    if (!mKeyStore.delete(Credentials.USER_CERTIFICATE + alias)) {
-                        Log.e(TAG, "Failed to clean up key chain after certificate chain"
-                                + " importing failed");
-                    }
-                    return false;
-                }
-            } else {
-                if (!mKeyStore.delete(Credentials.CA_CERTIFICATE + alias)) {
-                    Log.e(TAG, "Failed to remove CA certificate chain for alias " + alias);
-                }
-            }
-
             if (Log.isLoggable(TAG, Log.DEBUG)) {
                 Log.d(TAG, String.format("Set certificate for key alias %s : user %s CA chain: %s",
                         alias, emptyOrBase64Encoded(userCertificate),
@@ -354,8 +429,9 @@ public class KeyChainService extends IntentService {
             final String alias;
             String subject = null;
             final boolean isSecurityLoggingEnabled = mInjector.isSecurityLoggingEnabled();
+            final X509Certificate cert;
             try {
-                final X509Certificate cert = parseCertificate(caCertificate);
+                cert = parseCertificate(caCertificate);
 
                 final boolean isDebugLoggable = Log.isLoggable(TAG, Log.DEBUG);
                 subject = cert.getSubjectX500Principal().getName(X500Principal.CANONICAL);
@@ -388,13 +464,13 @@ public class KeyChainService extends IntentService {
             // anchor separately and independently of CA certificates, at which point this code
             // should be removed.
             if (CERT_INSTALLER_PACKAGE.equals(callingPackage())) {
-                final boolean result = mKeyStore.put(
-                        String.format("%s%s %s", Credentials.CA_CERTIFICATE, subject, alias),
-                        caCertificate, Process.SYSTEM_UID,
-                        KeyStore.FLAG_NONE);
-                Log.d(TAG, String.format(
-                        "Attempted installing %s (subject: %s) to KeyStore. Result: %b", alias,
-                        subject, result));
+                try {
+                    mKeyStore.setCertificateEntry(String.format("%s %s", subject, alias), cert);
+                } catch(KeyStoreException e) {
+                    Log.e(TAG, String.format(
+                            "Attempted installing %s (subject: %s) to KeyStore. Failed", alias,
+                            subject), e);
+                }
             }
 
             broadcastLegacyStorageChange();
@@ -451,34 +527,60 @@ public class KeyChainService extends IntentService {
                                 emptyOrBase64Encoded(userCertificateChain)));
             }
 
-            if (!removeKeyPair(alias)) {
+            final ArrayList<Certificate> certs = new ArrayList<>();
+            try {
+                if (userCertificate != null) {
+                    certs.add(parseCertificate(userCertificate));
+                }
+                if (userCertificateChain != null) {
+                    certs.addAll(parseCertificates(userCertificateChain));
+                }
+            } catch (CertificateException e) {
+                Log.e(TAG, "Failed to parse user certificate.", e);
                 return false;
             }
-            if (privateKey != null && !mKeyStore.importKey(
-                    Credentials.USER_PRIVATE_KEY + alias, privateKey, uid, KeyStore.FLAG_NONE)) {
-                Log.e(TAG, "Failed to import private key " + alias);
+
+            if (certs.isEmpty()) {
+                Log.e(TAG, "Cannot install private key without public certificate.");
                 return false;
             }
-            if (userCertificate != null &&
-                    !mKeyStore.put(Credentials.USER_CERTIFICATE + alias, userCertificate,
-                        uid, KeyStore.FLAG_NONE)) {
-                Log.e(TAG, "Failed to import user certificate " + userCertificate);
-                if (!mKeyStore.delete(Credentials.USER_PRIVATE_KEY + alias)) {
-                    Log.e(TAG, "Failed to delete private key after certificate importing failed");
+
+            final Certificate[] certificates = certs.toArray(new Certificate[0]);
+
+            final PrivateKey privateKey1;
+            try {
+                if (privateKey != null) {
+                    final KeyFactory keyFactory =
+                            KeyFactory.getInstance(certificates[0].getPublicKey().getAlgorithm());
+                    privateKey1 = keyFactory.generatePrivate(new PKCS8EncodedKeySpec(privateKey));
+                } else {
+                    privateKey1 = null;
                 }
+            } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
+                Log.e(TAG, "Failed to parse private key.", e);
+                return false;
+            }
+
+            KeyStore keystore = getKeyStore(uid == Process.WIFI_UID);
+            if (keystore == null) {
                 return false;
             }
-            if (userCertificateChain != null && userCertificateChain.length > 0) {
-                if (!mKeyStore.put(Credentials.CA_CERTIFICATE + alias, userCertificateChain, uid,
-                        KeyStore.FLAG_NONE)) {
-                    Log.e(TAG, "Failed to import certificate chain" + userCertificateChain);
-                    if (!removeKeyPair(alias)) {
-                        Log.e(TAG, "Failed to clean up key chain after certificate chain"
-                                + " importing failed");
+
+            try {
+                if (privateKey != null) {
+                    keystore.setKeyEntry(alias, privateKey1, null, certificates);
+                } else {
+                    if (certificates.length > 1) {
+                        Log.e(TAG,
+                                "Cannot install key certificate chain without private key.");
+                        return false;
                     }
-                    return false;
+                    keystore.setCertificateEntry(alias, certificates[0]);
                 }
+            } catch (KeyStoreException e) {
+                Log.e(TAG, "Failed to install key pair.", e);
             }
+
             broadcastKeychainChange();
             broadcastLegacyStorageChange();
             return true;
@@ -486,7 +588,11 @@ public class KeyChainService extends IntentService {
 
         @Override public boolean removeKeyPair(String alias) {
             checkCertInstallerOrSystemCallerOrHasPermission();
-            if (!Credentials.deleteAllTypesForAlias(mKeyStore, alias)) {
+            try {
+                mKeyStore.deleteEntry(alias);
+            } catch (KeyStoreException e) {
+                Log.e(TAG, String.format(
+                        "Failed not remove keystore entry with alias %s", alias));
                 return false;
             }
             Log.w(TAG, String.format(
@@ -499,14 +605,25 @@ public class KeyChainService extends IntentService {
 
         @Override public boolean containsKeyPair(String alias) {
             checkSystemCaller();
-            return mKeyStore.contains(Credentials.USER_PRIVATE_KEY + alias) &&
-                    mKeyStore.contains(Credentials.USER_CERTIFICATE + alias);
+            try {
+                final Key key = mKeyStore.getKey(alias, null);
+                return key instanceof PrivateKey;
+            } catch (UnrecoverableKeyException | NoSuchAlgorithmException | KeyStoreException e) {
+                Log.w("Error while trying to check for key presence.", e);
+                return false;
+            }
         }
 
         private X509Certificate parseCertificate(byte[] bytes) throws CertificateException {
             CertificateFactory cf = CertificateFactory.getInstance("X.509");
             return (X509Certificate) cf.generateCertificate(new ByteArrayInputStream(bytes));
         }
+        private Collection<X509Certificate> parseCertificates(byte[] bytes)
+                throws CertificateException {
+            final CertificateFactory cf = CertificateFactory.getInstance("X.509");
+            return (Collection<X509Certificate>)
+                    cf.generateCertificates(new ByteArrayInputStream(bytes));
+        }
 
         @Override public boolean reset() {
             // only Settings should be able to reset
@@ -624,10 +741,14 @@ public class KeyChainService extends IntentService {
 
         @Override public int[] getGrants(String alias) {
             checkSystemCaller();
-            if (!mKeyStore.contains(Credentials.USER_PRIVATE_KEY + alias)) {
-                throw new IllegalArgumentException("Alias not found: " + alias);
+            try {
+                if (mKeyStore.isKeyEntry(alias)) {
+                    return mGrantsDb.getGrants(alias);
+                }
+            } catch (KeyStoreException e) {
+                Log.w(TAG, "Error while checking if key exists.", e);
             }
-            return mGrantsDb.getGrants(alias);
+            throw new IllegalArgumentException("Alias not found: " + alias);
         }
 
         @Override
diff --git a/support/src/com/android/keychain/tests/support/IKeyChainServiceTestSupport.aidl b/support/src/com/android/keychain/tests/support/IKeyChainServiceTestSupport.aidl
index c62c971..c8ecc76 100644
--- a/support/src/com/android/keychain/tests/support/IKeyChainServiceTestSupport.aidl
+++ b/support/src/com/android/keychain/tests/support/IKeyChainServiceTestSupport.aidl
@@ -31,16 +31,11 @@ import android.security.keystore.ParcelableKeyGenParameterSpec;
  * @hide
  */
 interface IKeyChainServiceTestSupport {
-    boolean keystoreReset();
-    boolean keystoreSetPassword(String password);
-    boolean keystorePut(String key, in byte[] value);
-    boolean keystoreImportKey(String key, in byte[] value);
     void revokeAppPermission(int uid, String alias);
     void grantAppPermission(int uid, String alias);
     boolean installKeyPair(in byte[] privateKey, in byte[] userCert, in byte[] certChain, String alias);
     boolean removeKeyPair(String alias);
     void setUserSelectable(String alias, boolean isUserSelectable);
     int generateKeyPair(in String algorithm, in ParcelableKeyGenParameterSpec spec);
-    int attestKey(in String alias, in byte[] challenge, in int[] idAttestationFlags);
     boolean setKeyPairCertificate(String alias, in byte[] userCert, in byte[] certChain);
 }
diff --git a/support/src/com/android/keychain/tests/support/KeyChainServiceTestSupport.java b/support/src/com/android/keychain/tests/support/KeyChainServiceTestSupport.java
index 5f688f8..1e99770 100644
--- a/support/src/com/android/keychain/tests/support/KeyChainServiceTestSupport.java
+++ b/support/src/com/android/keychain/tests/support/KeyChainServiceTestSupport.java
@@ -23,39 +23,14 @@ import android.os.RemoteException;
 import android.security.IKeyChainService;
 import android.security.KeyChain;
 import android.security.KeyStore;
-import android.security.keymaster.KeymasterCertificateChain;
 import android.security.keystore.ParcelableKeyGenParameterSpec;
 import android.util.Log;
 
 public class KeyChainServiceTestSupport extends Service {
     private static final String TAG = "KeyChainServiceTest";
 
-    private final KeyStore mKeyStore = KeyStore.getInstance();
-
     private final IKeyChainServiceTestSupport.Stub mIKeyChainServiceTestSupport
             = new IKeyChainServiceTestSupport.Stub() {
-        @Override public boolean keystoreReset() {
-            Log.d(TAG, "keystoreReset");
-            for (String key : mKeyStore.list("")) {
-                if (!mKeyStore.delete(key, KeyStore.UID_SELF)) {
-                    return false;
-                }
-            }
-            return true;
-        }
-        @Override public boolean keystoreSetPassword(String password) {
-            Log.d(TAG, "keystoreSetPassword");
-            return mKeyStore.onUserPasswordChanged(password);
-        }
-        @Override public boolean keystorePut(String key, byte[] value) {
-            Log.d(TAG, "keystorePut");
-            return mKeyStore.put(key, value, KeyStore.UID_SELF, KeyStore.FLAG_NONE);
-        }
-        @Override public boolean keystoreImportKey(String key, byte[] value) {
-            Log.d(TAG, "keystoreImport");
-            return mKeyStore.importKey(key, value, KeyStore.UID_SELF, KeyStore.FLAG_NONE);
-        }
-
         @Override public void revokeAppPermission(final int uid, final String alias)
                 throws RemoteException {
             Log.d(TAG, "revokeAppPermission");
@@ -102,16 +77,6 @@ public class KeyChainServiceTestSupport extends Service {
             });
         }
 
-        @Override public int attestKey(
-                String alias, byte[] attestationChallenge,
-                int[] idAttestationFlags) throws RemoteException {
-            KeymasterCertificateChain attestationChain = new KeymasterCertificateChain();
-            return performBlockingKeyChainCall(keyChainService -> {
-                return keyChainService.attestKey(alias, attestationChallenge, idAttestationFlags,
-                        attestationChain);
-            });
-        }
-
         @Override public boolean setKeyPairCertificate(String alias, byte[] userCertificate,
                 byte[] userCertificateChain) throws RemoteException {
             return performBlockingKeyChainCall(keyChainService -> {
diff --git a/tests/src/com/android/keychain/tests/BasicKeyChainServiceTest.java b/tests/src/com/android/keychain/tests/BasicKeyChainServiceTest.java
index 0315a51..f3f20db 100644
--- a/tests/src/com/android/keychain/tests/BasicKeyChainServiceTest.java
+++ b/tests/src/com/android/keychain/tests/BasicKeyChainServiceTest.java
@@ -37,6 +37,7 @@ import android.security.KeyChain;
 import android.security.keystore.KeyGenParameterSpec;
 import android.security.keystore.KeyProperties;
 import android.security.keystore.ParcelableKeyGenParameterSpec;
+import android.util.Base64;
 import android.util.Log;
 import androidx.test.InstrumentationRegistry;
 import androidx.test.runner.AndroidJUnit4;
@@ -138,8 +139,8 @@ public class BasicKeyChainServiceTest {
     public void testCanAccessKeyAfterGettingGrant()
             throws RemoteException, IOException, CertificateException {
         Log.d(TAG, "Testing access to imported key after getting grant.");
-        assertThat(mTestSupportService.keystoreReset()).isTrue();
-        installFirstKey();
+        removeKeyPair(ALIAS_1);
+        generateRsaKey(ALIAS_1);
         assertThat(mKeyChainService.requestPrivateKey(ALIAS_1)).isNull();
         mTestSupportService.grantAppPermission(Process.myUid(), ALIAS_1);
         assertThat(mKeyChainService.requestPrivateKey(ALIAS_1)).isNotNull();
@@ -149,7 +150,7 @@ public class BasicKeyChainServiceTest {
     public void testInstallAndRemoveKeyPair()
             throws RemoteException, IOException, CertificateException {
         Log.d(TAG, "Testing importing key.");
-        assertThat(mTestSupportService.keystoreReset()).isTrue();
+        removeKeyPair(ALIAS_IMPORTED);
         // No key installed, all should fail.
         assertThat(mKeyChainService.requestPrivateKey(ALIAS_IMPORTED)).isNull();
         assertThat(mKeyChainService.getCertificate(ALIAS_IMPORTED)).isNull();
@@ -162,10 +163,11 @@ public class BasicKeyChainServiceTest {
                     Credentials.convertToPem(privateKeyEntry.getCertificateChain()),
                     ALIAS_IMPORTED)).isTrue();
 
-        // No grant, all should still fail.
+        // No grant, Private key access should still fail.
         assertThat(mKeyChainService.requestPrivateKey(ALIAS_IMPORTED)).isNull();
-        assertThat(mKeyChainService.getCertificate(ALIAS_IMPORTED)).isNull();
-        assertThat(mKeyChainService.getCaCertificates(ALIAS_IMPORTED)).isNull();
+        // Certificate access succeeds because this test runs as AID_SYSTEM.
+        assertThat(mKeyChainService.getCertificate(ALIAS_IMPORTED)).isNotNull();
+        assertThat(mKeyChainService.getCaCertificates(ALIAS_IMPORTED)).isNotNull();
         // Grant access
         mTestSupportService.grantAppPermission(Process.myUid(), ALIAS_IMPORTED);
         // Has grant, all should succeed.
@@ -179,7 +181,7 @@ public class BasicKeyChainServiceTest {
     @Test
     public void testUserSelectability() throws RemoteException, IOException, CertificateException {
         Log.d(TAG, "Testing user-selectability of a key.");
-        assertThat(mTestSupportService.keystoreReset()).isTrue();
+        removeKeyPair(ALIAS_IMPORTED);
         PrivateKeyEntry privateKeyEntry =
                 TestKeyStore.getClientCertificate().getPrivateKey("RSA", "RSA");
         assertThat(mTestSupportService.installKeyPair(privateKeyEntry.getPrivateKey().getEncoded(),
@@ -210,18 +212,6 @@ public class BasicKeyChainServiceTest {
     }
 
     @Test
-    public void testGenerateKeyPairErrorsOnSuperflousAttestationChallenge() throws RemoteException {
-        KeyGenParameterSpec specWithChallenge =
-                new KeyGenParameterSpec.Builder(buildRsaKeySpec(ALIAS_GENERATED))
-                        .setAttestationChallenge(DUMMY_CHALLENGE)
-                        .build();
-        ParcelableKeyGenParameterSpec parcelableSpec =
-                new ParcelableKeyGenParameterSpec(specWithChallenge);
-        assertThat(mTestSupportService.generateKeyPair("RSA", parcelableSpec)).isEqualTo(
-                KeyChain.KEY_GEN_SUPERFLUOUS_ATTESTATION_CHALLENGE);
-    }
-
-    @Test
     public void testGenerateKeyPairErrorsOnInvalidAlgorithm() throws RemoteException {
         ParcelableKeyGenParameterSpec parcelableSpec = new ParcelableKeyGenParameterSpec(
                 buildRsaKeySpec(ALIAS_GENERATED));
@@ -250,31 +240,69 @@ public class BasicKeyChainServiceTest {
         assertThat(mKeyChainService.requestPrivateKey(ALIAS_GENERATED)).isNotNull();
     }
 
-    @Test
-    public void testAttestKeyFailsOnMissingChallenge() throws RemoteException {
-        generateRsaKey(ALIAS_GENERATED);
-        assertThat(mTestSupportService.attestKey(ALIAS_GENERATED, null, new int[]{}
-                )).isEqualTo(KeyChain.KEY_ATTESTATION_MISSING_CHALLENGE);
-    }
-
-    @Test
-    public void testAttestKeyFailsOnNonExistentKey() throws RemoteException {
-        assertThat(mTestSupportService.attestKey(ALIAS_NON_EXISTING, DUMMY_CHALLENGE, new int[]{}
-                )).isEqualTo(KeyChain.KEY_ATTESTATION_FAILURE);
-    }
-
-    @Test
-    public void testAttestKeySucceedsOnGeneratedKey() throws RemoteException {
-        generateRsaKey(ALIAS_GENERATED);
-        assertThat(mTestSupportService.attestKey(ALIAS_GENERATED, DUMMY_CHALLENGE,
-                null)).isEqualTo(KeyChain.KEY_ATTESTATION_SUCCESS);
-    }
+    private static final String USER_CERT1 =
+            "MIIC+DCCAeACAQEwDQYJKoZIhvcNAQELBQAwRTELMAkGA1UEBhMCVVMxEzARBgNV" +
+            "BAgMClNvbWUtU3RhdGUxITAfBgNVBAoMGEludGVybmV0IFdpZGdpdHMgUHR5IEx0" +
+            "ZDAeFw0yMTAyMTAwMDEyMDNaFw00ODA2MjgwMDEyMDNaMD8xCzAJBgNVBAYTAlVT" +
+            "MQswCQYDVQQIDAJDQTEUMBIGA1UEBwwLTWVkaW9jcmV0b24xDTALBgNVBAoMBENv" +
+            "Q28wggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQC2LieaIDILjZPtNvZI" +
+            "I6VSvKIPIdY5STadoWjmt9TqjyvszJ+oEceBPqwevOdtXpJCiHhP91B314XBT33z" +
+            "8re1b0A5k31YKcy0PU+3OMh4XG6O3/Z5/9GfsekfQZK3jagbn3uqJ2emyj0JK+HY" +
+            "ipD6iwyO21DerUYavPpC0uo8PKAxc9l6XjILg9qoi68yCi8P3tkLLAFWCsc7GUkA" +
+            "v97zUjGq8iz/gIrwmzJBB3O//7e1nuHO5AswgJiwWa9aY6uHKWm97xP0Kw4pShFI" +
+            "eEsdNKj82FURpWIvrWsztYPaEM/+nG8cNiJ1XRT3DfnlKzXVVQi7pE82HkUpFNLQ" +
+            "4wCbAgMBAAEwDQYJKoZIhvcNAQELBQADggEBAJwq05ykEDK5etwC5jDYsX7hyiYI" +
+            "TAephAfhMqT7CtwScVlurSWt4ivv3IBx2huibQqdwq9zyCG9UzQKACMrEAaINerr" +
+            "VPwq+MpfgOUJRlZOCchr1KaN+quoOstStsGqemHZSTyUow6OrG2a+DIGQZy84yza" +
+            "y4+NZYtTjbVKFAmQkji344cVd6qFjCG5Dgo+1u4+YOWV4QTVWwbjgjCGiDNyZKeO" +
+            "Y8n8pdsv4ygat4VCotch22bqiUJLgY/abiprqKz3jfkko9urHkWsrdqzc3mS29AC" +
+            "tWVhomOQq/51A0wfQRbSw2MY+j5LaF22Sz5rnPk0u6Vcm2eLUI5gSJ/HZOU=";
+
+    private static final String USER_CERT2 =
+            "MIIDETCCAfkCFDEBBiFSwkBDAZMyhMjrr7P4wHklMA0GCSqGSIb3DQEBCwUAMEUx" +
+            "CzAJBgNVBAYTAlVTMRMwEQYDVQQIDApTb21lLVN0YXRlMSEwHwYDVQQKDBhJbnRl" +
+            "cm5ldCBXaWRnaXRzIFB0eSBMdGQwHhcNMjEwMjIyMTgwNzE2WhcNNDgwNzEwMTgw" +
+            "NzE2WjBFMQswCQYDVQQGEwJVUzETMBEGA1UECAwKU29tZS1TdGF0ZTEhMB8GA1UE" +
+            "CgwYSW50ZXJuZXQgV2lkZ2l0cyBQdHkgTHRkMIIBIjANBgkqhkiG9w0BAQEFAAOC" +
+            "AQ8AMIIBCgKCAQEA8XLlpvOvG22fJ4ZNSEdc/QIzkbjmtP5rmqX8dgfZdZ7xdUNO" +
+            "24Coyv9f+GiXJdDlL2+KIK3ZTpcuVdd5oO+GdrxB1Dnvhqt2YgOb/+VyeMOdsEre" +
+            "Bqw4Y+v8A92zpWlYlh9xyjVE53ismdQg2RT8kIxQt2ydMyZRx8KzNmZRsXouiVSi" +
+            "ngWlXdtkpXNiVt5CkSBwscfVNMkF4EfcKqYGLJzFTXcSsRVZoknGNtgJZIsHaOJh" +
+            "etmrfxATcbvNYdDhm9xs2ud/WwaRDqbMme9KVOrk9g5NBZIn7SpWDUyBk36W2CgQ" +
+            "U/3OPnDOUS6oW+YKE1xvJ3i2FhLD2ufJyNh4WwIDAQABMA0GCSqGSIb3DQEBCwUA" +
+            "A4IBAQDkl+8ZlYgKwSHjqxHjJsoBXBlwWUw0FQDaFz0tlPv/f5w80NsTXImVgRcP" +
+            "MBWxLrVa/JKCm7GPZgIotvYsRKL0/DRZDUuva86C2hi9C/E7OtgFkDO0d/t9QIQM" +
+            "gSNgdEgda/+lixc7XrsjJKAufFlYhzk/q5uyiD3SbqFzdlADumWtY9Xu6fqU2JB5" +
+            "TFOTxqpU5c2b+sXL6uc+dA2pSP94wL+7g+uKVhdJGsimbXIq1jl9r+C2ykrUjuz+" +
+            "b7uBJ5Qzq81tdmNCZ4pLtqmatTAMj2LYsKiXUe9Fh0nd/2aZek6I1YHaIBGg7jns" +
+            "ZYEWPj4Rd0KiE3L/ymeQ4VQx6SW2";
+
+    private static final String CA_CERT =
+            "MIIDazCCAlOgAwIBAgIUVA2nyBJMc/OcO0C/yPt/E1TwREIwDQYJKoZIhvcNAQEL" +
+            "BQAwRTELMAkGA1UEBhMCVVMxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM" +
+            "GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDAeFw0yMTAyMDkyMzAzMzdaFw00ODA2" +
+            "MjcyMzAzMzdaMEUxCzAJBgNVBAYTAlVTMRMwEQYDVQQIDApTb21lLVN0YXRlMSEw" +
+            "HwYDVQQKDBhJbnRlcm5ldCBXaWRnaXRzIFB0eSBMdGQwggEiMA0GCSqGSIb3DQEB" +
+            "AQUAA4IBDwAwggEKAoIBAQDxcuWm868bbZ8nhk1IR1z9AjORuOa0/muapfx2B9l1" +
+            "nvF1Q07bgKjK/1/4aJcl0OUvb4ogrdlOly5V13mg74Z2vEHUOe+Gq3ZiA5v/5XJ4" +
+            "w52wSt4GrDhj6/wD3bOlaViWH3HKNUTneKyZ1CDZFPyQjFC3bJ0zJlHHwrM2ZlGx" +
+            "ei6JVKKeBaVd22Slc2JW3kKRIHCxx9U0yQXgR9wqpgYsnMVNdxKxFVmiScY22Alk" +
+            "iwdo4mF62at/EBNxu81h0OGb3Gza539bBpEOpsyZ70pU6uT2Dk0FkiftKlYNTIGT" +
+            "fpbYKBBT/c4+cM5RLqhb5goTXG8neLYWEsPa58nI2HhbAgMBAAGjUzBRMB0GA1Ud" +
+            "DgQWBBSsg7KE2+ypUJrrk9pBIVYpgjZM3TAfBgNVHSMEGDAWgBSsg7KE2+ypUJrr" +
+            "k9pBIVYpgjZM3TAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQBx" +
+            "pUWwTL8vsrq3hNEDOui1Y6kYPiaAs2Fd//tBjPhgc9tlKS8KLkGU0yW8kEyPovvK" +
+            "M6/IUDZhrHzli8Y9iJxiZ43o/fni1DOLHLU+CCOaKGrmBDBhgDlz1m2nXsCOb8uF" +
+            "ohz4Yp1SLMy1YpQPmltiaILXNTw33B8xFa0d9ChZbBudyiiNs0vijdFMMYetqAGM" +
+            "aXCjPXdM6hGjHxc6vf2ECJOrfVg3D22VOJ0WzrrJDM9pRxJgJ/IFG0eirKvE7G8s" +
+            "zk1Og79JAFr3qMKEnUS7nuym7J69HSQlFHc7JEMzeS78YR7EHOlOlK/bEX/8cHF9" +
+            "ZlGCOS1Ds0rmVe3CoQIp";
 
     @Test
     public void testSetKeyPairCertificate() throws RemoteException {
         generateRsaKey(ALIAS_GENERATED);
-        final byte[] userCert = new byte[] {'a', 'b', 'c'};
-        final byte[] certChain = new byte[] {'d', 'e', 'f'};
+        final byte[] userCert = Base64.decode(USER_CERT1, Base64.DEFAULT);
+        final byte[] certChain = Base64.decode(CA_CERT, Base64.DEFAULT);
 
         assertThat(mTestSupportService.setKeyPairCertificate(ALIAS_GENERATED, userCert,
                 certChain)).isTrue();
@@ -283,7 +311,8 @@ public class BasicKeyChainServiceTest {
         assertThat(mKeyChainService.getCertificate(ALIAS_GENERATED)).isEqualTo(userCert);
         assertThat(mKeyChainService.getCaCertificates(ALIAS_GENERATED)).isEqualTo(certChain);
 
-        final byte[] newUserCert = new byte[] {'x', 'y', 'z'};
+        final byte[] newUserCert = Base64.decode(USER_CERT2, Base64.DEFAULT);
+
         assertThat(mTestSupportService.setKeyPairCertificate(ALIAS_GENERATED, newUserCert,
                 null)).isTrue();
         assertThat(mKeyChainService.getCertificate(ALIAS_GENERATED)).isEqualTo(newUserCert);
@@ -367,37 +396,6 @@ public class BasicKeyChainServiceTest {
         }
     }
 
-    void installFirstKey() throws RemoteException, IOException, CertificateException {
-        String intermediate = "-intermediate";
-        String root = "-root";
-
-        String alias1PrivateKey = Credentials.USER_PRIVATE_KEY + ALIAS_1;
-        String alias1ClientCert = Credentials.USER_CERTIFICATE + ALIAS_1;
-        String alias1IntermediateCert = (Credentials.CA_CERTIFICATE + ALIAS_1 + intermediate);
-        String alias1RootCert = (Credentials.CA_CERTIFICATE + ALIAS_1 + root);
-        PrivateKeyEntry privateKeyEntry =
-                TestKeyStore.getClientCertificate().getPrivateKey("RSA", "RSA");
-        Certificate intermediate1 = privateKeyEntry.getCertificateChain()[1];
-        Certificate root1 = TestKeyStore.getClientCertificate().getRootCertificate("RSA");
-
-        assertThat(
-                mTestSupportService.keystoreImportKey(
-                    alias1PrivateKey, privateKeyEntry.getPrivateKey().getEncoded()))
-            .isTrue();
-        assertThat(
-                mTestSupportService.keystorePut(
-                    alias1ClientCert,
-                    Credentials.convertToPem(privateKeyEntry.getCertificate())))
-            .isTrue();
-        assertThat(
-                mTestSupportService.keystorePut(
-                    alias1IntermediateCert, Credentials.convertToPem(intermediate1)))
-            .isTrue();
-        assertThat(
-                mTestSupportService.keystorePut(alias1RootCert, Credentials.convertToPem(root1)))
-            .isTrue();
-    }
-
     @Test
     public void testContainsKeyPair_NonExisting() throws RemoteException {
         assertThat(mKeyChainService.containsKeyPair(ALIAS_NON_EXISTING)).isFalse();
diff --git a/tests/src/com/android/keychain/tests/KeyChainActivityTest.java b/tests/src/com/android/keychain/tests/KeyChainActivityTest.java
index 73c3867..9acb6de 100644
--- a/tests/src/com/android/keychain/tests/KeyChainActivityTest.java
+++ b/tests/src/com/android/keychain/tests/KeyChainActivityTest.java
@@ -17,39 +17,40 @@
 package com.android.keychain.tests;
 
 import static com.android.keychain.KeyChainActivity.CertificateParametersFilter;
+
 import static com.google.common.truth.Truth.assertThat;
+
 import static org.junit.Assert.fail;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 import android.platform.test.annotations.LargeTest;
-import android.security.Credentials;
-import android.security.KeyStore;
 import android.util.Base64;
+
 import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.OutputStreamWriter;
-import java.io.Writer;
-import java.nio.charset.StandardCharsets;
+import java.security.KeyFactory;
+import java.security.KeyStore;
+import java.security.KeyStoreException;
+import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
 import java.security.cert.Certificate;
-import java.security.cert.CertificateEncodingException;
 import java.security.cert.CertificateException;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
+import java.security.spec.InvalidKeySpecException;
+import java.security.spec.PKCS8EncodedKeySpec;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.concurrent.CancellationException;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeoutException;
+
 import javax.security.auth.x500.X500Principal;
-import org.bouncycastle.util.io.pem.PemObject;
-import org.bouncycastle.util.io.pem.PemWriter;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
 
 @LargeTest
 @RunWith(AndroidJUnit4.class)
@@ -58,22 +59,25 @@ public final class KeyChainActivityTest {
     // openssl req -newkey rsa:2048 -nodes -keyout key.pem -x509 -days 3650
     //   -out root_ca_certificate.pem
     private static final String ROOT_CA_CERT_RSA =
-            "MIIDazCCAlOgAwIBAgIUWQjj+9olDNtdjcSLzK2RpxI9j6UwDQYJKoZIhvcNAQELBQAwRTELMAkG" +
-            "A1UEBhMCVUsxCzAJBgNVBAgMAk5BMQ8wDQYDVQQHDAZMb25kb24xGDAWBgNVBAoMD0FuZHJvaWQg" +
-            "VGVzdCBDQTAeFw0yMDAzMTIxMTUxNDJaFw0zMDAzMTAxMTUxNDJaMEUxCzAJBgNVBAYTAlVLMQsw" +
-            "CQYDVQQIDAJOQTEPMA0GA1UEBwwGTG9uZG9uMRgwFgYDVQQKDA9BbmRyb2lkIFRlc3QgQ0EwggEi" +
-            "MA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDGbsQO+LqPMr5Nr4Lq0B7C0th93phohSY6hb2w" +
-            "MmZs3MRwamlw8FS64KEgmszX5lnyLNqRs91FOFNuq4f2A+TYQhawi9D2bHB7z2ishDM3SxNAqwQl" +
-            "LzVNBJw7DAtimajy3VvXoprescFbsOZx8wPGGb2xMKqAXg4Yw9F6te4Y4BSIiwCWtammiSR8Ev0B" +
-            "lcnMBrWmSZ4yYF+UgNgNiD/TVrTtRmzQlRhBo5n4F61SGeAxb5p0NRRGmAXKtx358HiLANzZSCiM" +
-            "UE5IrgDvW8AKPn5InuYS1G1K2wG5ar1eanQahimtaIEugQxhqG0+/OiKKq2LGRiBpwV1OomXHNFr" +
-            "AgMBAAGjUzBRMB0GA1UdDgQWBBRrxYWzKZpCHDi/NK4keXIU5iGukzAfBgNVHSMEGDAWgBRrxYWz" +
-            "KZpCHDi/NK4keXIU5iGukzAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQBJTjUY" +
-            "rfi3adJFpF2IFkSnavPRxi+NX6wxfKgQvcScO7sV18gAMG7r2NhjsVeScZ48mxsNkj99lHaVoNm6" +
-            "c+sWmcb3LO7WCqmAgfJcHeQ5VuluwNoJWo+SGuKbh6/yRejNeQFf++uaEwXP3yNydwKJyQDyDwoG" +
-            "vx0jvy8glkVl3fr6u0lGQqmubGU5Q1X6QyA0zJ/sSWBVorLCgk6KvPABQJhjoij+g/GOB1h4g6fb" +
-            "bQ3xnek6TGwjQ1bB7rQlqBF7iP/9iUtuuDf0cR8LwMr1Z2OUEMDjRHQCZQJ3APc0kW1ewJ8nqQ+m" +
-            "NsUKFRuThYtE/OFsV/TfXwMXbc2rMug+";
+            "MIIDazCCAlOgAwIBAgIUKNVj8g/LG+6tqsK8HZqT6EuktpkwDQYJKoZIhvcNAQEL" +
+            "BQAwRTELMAkGA1UEBhMCVUsxCzAJBgNVBAgMAk5BMQ8wDQYDVQQHDAZMb25kb24x" +
+            "GDAWBgNVBAoMD0FuZHJvaWQgVGVzdCBDQTAeFw0yMTAyMjMwMDU0MzdaFw0zMTAy" +
+            "MjEwMDU0MzdaMEUxCzAJBgNVBAYTAlVLMQswCQYDVQQIDAJOQTEPMA0GA1UEBwwG" +
+            "TG9uZG9uMRgwFgYDVQQKDA9BbmRyb2lkIFRlc3QgQ0EwggEiMA0GCSqGSIb3DQEB" +
+            "AQUAA4IBDwAwggEKAoIBAQDLpcaoJijYhS3QUDgG8kVGrwTxaVTS0TE156fJa5za" +
+            "s3RI7TYKHzYwn1KMJJUwoc+cOkv+rFC7j5MQ+6SMK2GpDoCoGn4FV9dDPVnxIgjj" +
+            "/66kuf+we1ur3gz7m/8tFFdZhLFoFMRzcNg+F35jSur8y0dnc8O83gMwuf+91pU7" +
+            "HyNahHzDyMM5sR7u1K91R1MKiOnVqJNTHWVK+rl3G0m0rbDTz7/xbq3/FPBvw764" +
+            "QUJgkSEG15i5CFNn2ww5IYnF30Wbke3kUfdd/Q32MOyfkcp3El/TPJj/3mevGHed" +
+            "vTi0j6ovIXMrDnoJvmeI40p97EMIlaZ3x39i+krE02RfAgMBAAGjUzBRMB0GA1Ud" +
+            "DgQWBBRDpNR2iik8NGDmY5IvgiUy6dMRMjAfBgNVHSMEGDAWgBRDpNR2iik8NGDm" +
+            "Y5IvgiUy6dMRMjAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQDG" +
+            "OIZdoopwX2PlkAnWMR6PIl4GJ9eiPg1cbQHHcqsB8oNYiiiZydJ8TniZYFptEIjQ" +
+            "LIuerSr/a+c353JNOOK76w0vciJld8uQDCTnMChUahJsQjJJgDCyhjkQqV9Srstu" +
+            "IClm5IP/LEYVSzfR+LqhwX5JYUU8zsDnbiAH3CIENhBw5ApXgfDHP54bHOBw5eHE" +
+            "ETATnib8uPZVHsyHQXXULDqMITY9pmjc9/9x6CqcMGEiSVvhDSujerVlnw3I17Te" +
+            "HWT8bf8mJyQFR8kr59NPQMNN4oUIfFzj2VgzjL21mKGR+hBGqoIMJ1ZdPGCJsw0W" +
+            "5WYM4TtQtL7yNVDtJzOL";
 
     // Generated with:
     // openssl genrsa  -out intermediate_key.pem 2048
@@ -81,42 +85,74 @@ public final class KeyChainActivityTest {
     // openssl x509 -req -days 3650 -in intermediate_ca.csr -CA root_ca_certificate.pem
     //   -CAkey key.pem -CAcreateserial -out intermediate_ca.pem
     private static final String INTERMEDIATE_CA_CERT_RSA =
-            "MIIDHjCCAgYCFHVrA0CEKcoc/C8BqKap3a1m0x8CMA0GCSqGSIb3DQEBCwUAMEUxCzAJBgNVBAYT" +
-            "AlVLMQswCQYDVQQIDAJOQTEPMA0GA1UEBwwGTG9uZG9uMRgwFgYDVQQKDA9BbmRyb2lkIFRlc3Qg" +
-            "Q0EwHhcNMjAwMzEyMTE1NzUwWhcNMzAwMzEwMTE1NzUwWjBSMQswCQYDVQQGEwJVSzELMAkGA1UE" +
-            "CAwCTkExDzANBgNVBAcMBkxvbmRvbjElMCMGA1UECgwcQW5kcm9pZCBJbnRlcm1lZGlhdGUgVGVz" +
-            "dCBDQTCCASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAPlEIzApeTyKzQWTvv25z/KVEbsr" +
-            "alrcto7mX56HV1VQ53cGqi4I7dso3cvpg0CYfcZ+mZh6Evkd+njkcc7Dh/nI0KBJIzGZuo2LB+0r" +
-            "qT0RfvI/Xv7CqLO9KOjWJ3+HK3EhSXGvnLvSTsQD1LnE9HXKVdhdOUgLFjbcZrzH62mvTRAO6nhg" +
-            "agWTzprTXOX8okaMJtJl9QGMG63Z/m5DONPPrgASW6X6wksGjyorEaakQTUGuPimapP5mk+Y31Se" +
-            "pLDDumqRavLT2CpfjHfFq0iDmnnJjG5nz6oKlirhg9JjxHwuKm5jIdsO4dIgi1fJ8Goz2ODG6R6E" +
-            "6CjitsQjoMMCAwEAATANBgkqhkiG9w0BAQsFAAOCAQEAdud6PahLSmDLOTr9t0Jqq1HKpeRsjSqn" +
-            "JTpd/GkNUrxcctKLTzpAruMq6en8OfcSEa2s3HUMJ1LVMoO9pp4aQaadH7626Q4uqHbNGHWngVNX" +
-            "lfBioxrH2QJ2wuKBjUipEGWaM3LY0wqNuBFd5qAVuBwQtZ1x/XH7/Y4l38Y5EGGEi4jSw8eCqiiQ" +
-            "2UKItmK8byl3/T5SVgAMbFYz1WJN37EgETEcEgPlosSQ4pha6fVB3Oz6mSfzjGXqHKpHBPUn/N2d" +
-            "Z2kxJG8IuwhhhyemhqJdCfOxT2WpemgLQQCCgqtM9O89peWL8AJUVT9cF9KySOvh/P9lTtSf5bJf" +
-            "sAzfSg==";
+            "MIIDHjCCAgYCFHgZBbZMuJTvvm1wlBBoPE7peS4jMA0GCSqGSIb3DQEBCwUAMEUx" +
+            "CzAJBgNVBAYTAlVLMQswCQYDVQQIDAJOQTEPMA0GA1UEBwwGTG9uZG9uMRgwFgYD" +
+            "VQQKDA9BbmRyb2lkIFRlc3QgQ0EwHhcNMjEwMjIzMDA1ODQ5WhcNMzEwMjIxMDA1" +
+            "ODQ5WjBSMQswCQYDVQQGEwJVSzELMAkGA1UECAwCTkExDzANBgNVBAcMBkxvbmRv" +
+            "bjElMCMGA1UECgwcQW5kcm9pZCBJbnRlcm1lZGlhdGUgVGVzdCBDQTCCASIwDQYJ" +
+            "KoZIhvcNAQEBBQADggEPADCCAQoCggEBAMZ5tLvS8qiZGK8fBUPzZYWRG2kmf3Hi" +
+            "o6db8qjTvzkJ0l43zmI92v0pEy8a4/XlC7wGSaJuX/TqpAR43+a9kYqGWczkASpX" +
+            "7w1W5tjjDlLANZdRP2R7Wb3XWTTMxzDWOqsMGTHw+H92oFm4bGO9+PaBRRzifLmQ" +
+            "OI/Whw6/kHZlXyI7J68EwRbyaZXNJ6iqk8dF7B4iwZ9yNgW1H0m8uxdDAykEA2WX" +
+            "wbzITwD1zdMsQV7/eLe9RIMFN5VMeHtIFUi2AcioG0i4ZLZNFnVrFQKu7u3XQyxk" +
+            "u7ZzgeGtpluM71sDxnqhZv9NaZIq3mV3JrKHPsw6+uJjN4U5AVfzIYUCAwEAATAN" +
+            "BgkqhkiG9w0BAQsFAAOCAQEAXrRQUFlLyS3QlmwkGocLQISY9B8fF8LTH7sl6HFA" +
+            "VSVuhPDuNsmqVhsMH1981MY2rSVfM3fkUMz1WEH7ZbhooYPirax/AlW+oRdRB/xX" +
+            "WEAJRGgybK98PXogI4tqEvicVn2kfcyNzmfMn8yRClxD5GuZ0oOA50lpUwUmeJjo" +
+            "jb3DY8NF+bcA0lW5h7p86ezqjhB836XZRL47jZJj+jgKoiSsdcex7rzikW6bzdlV" +
+            "f9DCuBJMpM1y31AP15Gvg5Jhh9Wc4y8LLipTn7wGdJvvhclUe1U3roerhAEB+rU/" +
+            "Eu7h+Nqjogg3IzmfHlhDe4N8o/XLdc2FnhCZGZklDDDMvA==";
 
     // Generated with:
     // openssl req -new -newkey rsa:2048 -nodes -keyout server.key -out server.csr
     // openssl x509 -req -days 3650 -in server.csr -CA intermediate_ca.pem
     //   -CAkey intermediate_key.pem -CAcreateserial -out server.pem
     private static final String LEAF_SERVER_CERT_RSA =
-            "MIIDIjCCAgoCFGNstHCN7uzPtFlxnbu2FQ3+sc7rMA0GCSqGSIb3DQEBCwUAMFIxCzAJBgNVBAYT" +
-            "AlVLMQswCQYDVQQIDAJOQTEPMA0GA1UEBwwGTG9uZG9uMSUwIwYDVQQKDBxBbmRyb2lkIEludGVy" +
-            "bWVkaWF0ZSBUZXN0IENBMB4XDTIwMDMxMjEyMTQwMloXDTMwMDMxMDEyMTQwMlowSTELMAkGA1UE" +
-            "BhMCVUsxCzAJBgNVBAgMAk5BMQ8wDQYDVQQHDAZMb25kb24xHDAaBgNVBAoME0FuZHJvaWQgU2Vy" +
-            "dmVyIFRlc3QwggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDCBdt0uV01nG2ULMThfnEx" +
-            "CDoVvOSJJLcqKSMsoRUcEYvvqA70x7O2BR6+FZLilr09qAyK1ygyTh+Q6NcBxiB137khrygd/R3S" +
-            "U9nNWI6Xj882EospCiPaewR3j49F6F45PvviAacx7v0mZpU1dVMP7ZvhJWb1AB675/379EFztYyU" +
-            "ghM5ub8zuAgBuvH/O/dJpnmKUmO43n/qSHDM+Q+oDmLAkWD4gd3SOBZKcLgTyO/pD0pghT8m2t0t" +
-            "9X/3KX+tQhif1pAemCOqvxr/HHjlLWxY+QWmcRIAMzTg6+h7sNSKAwwfulMomzMwzkV5sJAYUJOm" +
-            "suhhyAXCNX5rWuqXAgMBAAEwDQYJKoZIhvcNAQELBQADggEBACdFD/uV4iT4jg3/rEtszPHkyP4b" +
-            "zHaYJYpExoWNbJIz7djhycptdM7wjSoesWQMfpJz95aNWoVHrW85DSdGT+7HwZEsW1zUx3KkXURA" +
-            "CdlbVBn1CS4vm0Xk7Rr3LayfhqdFALQVItvBr+LkJPiG/R1jQySp1qaw+NrwukFoepukeZxHH1bF" +
-            "9zjGCLwOcfuRB9g8Gm45wRgoSUTDKbD1SMSLuPyllKeHLKE+chhYjm51Evy2xR0DLm0yaFmeyPPM" +
-            "KQaZJmtkeAzk2uYYb1HsVDlnEoDoXpTKKN1j39qckpCHxF0X0KqbN7D0grDWIDIee5mnSrg3+Xq5" +
-            "aCEDTUn6uU4=";
+            "MIIDIjCCAgoCFDV+Rg1WxBy4ThLxvu1lRJl1YUzwMA0GCSqGSIb3DQEBCwUAMFIx" +
+            "CzAJBgNVBAYTAlVLMQswCQYDVQQIDAJOQTEPMA0GA1UEBwwGTG9uZG9uMSUwIwYD" +
+            "VQQKDBxBbmRyb2lkIEludGVybWVkaWF0ZSBUZXN0IENBMB4XDTIxMDIyMzAxMDEw" +
+            "OFoXDTMxMDIyMTAxMDEwOFowSTELMAkGA1UEBhMCVUsxCzAJBgNVBAgMAk5BMQ8w" +
+            "DQYDVQQHDAZMb25kb24xHDAaBgNVBAoME0FuZHJvaWQgU2VydmVyIFRlc3QwggEi" +
+            "MA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQC7267U6iCagBJFiYMMUkteQliO" +
+            "ljOfnuSZJTG1xxNXBggLDdHmchjOPQgEICINpxz7Hhg+PLME3DEqrwUHo9k/bR3e" +
+            "Tglt1o6qxGirIEKROtdzNyi3medRex6FATXq1g4W/1U0tl8EbMv9kJPZP52Uj0Rq" +
+            "XYZ9Y27IGtWDcudpyJij/nBbV/kfufti2pFNHhnXytyBrQXz6AziZjHnNt316z64" +
+            "Tfqchr0jxqIF3Sup9AVKnGooGymwT8ez5C6VO7WoRZnp40pH78GzSALnRJC9w26V" +
+            "1IVlqjYPWvFwJ0ENb0Nuc8Jr3tW/0gf1UzTsuRsMU9vl+tMjweXS8HI8M4uRAgMB" +
+            "AAEwDQYJKoZIhvcNAQELBQADggEBAIcOLAdu9HU6+Bk/SkQW2qW2mY5+WFSYYRG5" +
+            "KfDiMQr0kUeddBJYk+bLN3Qqi6KUAZ+ITxyarVrjZjxaTr1JV6m9fVxdZ8elAx+7" +
+            "ci9ghBkiUPKFtVz3Y1F31Em4tLRr0LHF49Mjvr62+mQuZlAXZ3TuMdxrwc9AePhN" +
+            "btY8YwUsPPJ2vrIQB14NJ6EIaMaIyTowBPX9eo5K47ISbfnCUKhFYAEK4v5s4Hkt" +
+            "Us/209E0FNdFKOZDKDwclhZSFbyA1tknBRXsP7QCFuj/hPFxcRaj9R8CoHQAqEFr" +
+            "mV8JscA9dV40m+kRkj5TZeHjIw2xI39btv5aeRW2fBNB0MEoNPQ=";
+
+    private static final String PRIVATE_SERVER_KEY =
+            "MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC7267U6iCagBJF" +
+            "iYMMUkteQliOljOfnuSZJTG1xxNXBggLDdHmchjOPQgEICINpxz7Hhg+PLME3DEq" +
+            "rwUHo9k/bR3eTglt1o6qxGirIEKROtdzNyi3medRex6FATXq1g4W/1U0tl8EbMv9" +
+            "kJPZP52Uj0RqXYZ9Y27IGtWDcudpyJij/nBbV/kfufti2pFNHhnXytyBrQXz6Azi" +
+            "ZjHnNt316z64Tfqchr0jxqIF3Sup9AVKnGooGymwT8ez5C6VO7WoRZnp40pH78Gz" +
+            "SALnRJC9w26V1IVlqjYPWvFwJ0ENb0Nuc8Jr3tW/0gf1UzTsuRsMU9vl+tMjweXS" +
+            "8HI8M4uRAgMBAAECggEARWNUhYJhPpAVr6emRxPSkONysGAce1YGW+bYIKuCoj8x" +
+            "E1wsbrEwJmV2o4d27JIQa1TnYX2sJhxq8Lgq5HKJ2Rql0KoEY5S/p6Xaf3LwA5K3" +
+            "Z/A00vQ+8+LFGB2lW7NrCuWPBGRkXk8NXgBcC/+qZegxPhSDi6cBkVoQCXiUr4Zs" +
+            "4wacrmpQbl/FekvpuWQxnVAm95knZhJ87r7izwO6e3VP1VZseG7ld6PbxkLnuTP7" +
+            "Y9Q+viAPwkk1SedvYy6RtIRbyyOKzTVbh9SXirPsLM6N+a8k3J7LY+8HnZhte0O4" +
+            "BFoZwwhXt/kPHilro6gt69Bh2bas+lpP62c8e7Q3HQKBgQDwY6XJzLXV1UQNbwLL" +
+            "FHTDBdf3Afe6W5jcKpsOUaNP/424JeLpc5a2ccaytCPs2p3RCTZ3SHaxe+LNrKbU" +
+            "iKdloGO/XTqcAHbvw/uy5Zzsv6XqUjREmCvCatF7UU7PH9Pztjhb91IXT7taB1Ee" +
+            "yYyPPJU9ioZDEV45/PIUpGwjnwKBgQDIDrtHx5fBH0od8XWSBzuQ3QjLl3WZcEAv" +
+            "Bf2GcdYUV7migWOXm28k51l2VbQQRCVOulJjf45TZgZgzKInZZJ+rZ69FLrOx7q3" +
+            "7X+8rTyWKa2u/IlZ/EBtNIRuqH7A+OuY5qroFngtYa/qsaX58/eUAe64I6HslenG" +
+            "ktvmwdSCzwKBgAtQ1YQLU9/t+xcay6ndm6V2h/UDrbKjDy4F/2iMJUDlybkKZ4UP" +
+            "wN9zuaO94RcML3OgmGTDD3tJVqLR5sSIbkDVbPycGd8wEmk084s3TczDNL80AWvd" +
+            "Meoj9xpz+F69o8+MG1kQ6ldYlHwnbgUh/bDcbDYKaEmN7r6SDp80IjcHAoGAOrqQ" +
+            "Yf8G3qu3z1h94jN7Wgh5N4MsA7I/NU614Uzzwp8KINmJCg2YMCY2ThXUuV238gei" +
+            "fhEJEBSIVMxd4eDgg42mZu159ZAOkUYIVLQqcA6mLRN3otH5e9WJ9w5Bv5aTWxyE" +
+            "GYPXHcNqqCQkjF8BVBLJKIdVVqWfriqYoYJPR2MCgYAWWIrS8XN1c2iLParb+xOJ" +
+            "q1WKe8q5wFucqJCVFbJXjtpGgZFxZLAFlT8VpaiBwDLQBPC+CYhzLVwAvzW2h46W" +
+            "KONqO7zoxiuwOEj466iH4YrgviNw6lGtgSPB6wx91c7se/1lcZhviBMB5rUjtxxj" +
+            "qKIC9Y+gz77w1M3pwMlhDA==";
 
     private static final X500Principal LEAF_SUBJECT =
             new X500Principal("O=Android Server Test, L=London, ST=NA, C=UK");
@@ -127,44 +163,63 @@ public final class KeyChainActivityTest {
     private static final X500Principal ROOT_SUBJECT =
             new X500Principal("O=Android Test CA, L=London, ST=NA, C=UK");
 
+    private byte[] mPrivateKey;
     private byte[] mLeafRsaCertificate;
     private byte[] mIntermediateRsaCertificate;
     private byte[] mRootRsaCertificate;
 
     @Before
     public void setUp() {
+        mPrivateKey = Base64.decode(PRIVATE_SERVER_KEY, Base64.DEFAULT);
         mLeafRsaCertificate = Base64.decode(LEAF_SERVER_CERT_RSA, Base64.DEFAULT);
         mIntermediateRsaCertificate = Base64.decode(INTERMEDIATE_CA_CERT_RSA, Base64.DEFAULT);
         mRootRsaCertificate = Base64.decode(ROOT_CA_CERT_RSA, Base64.DEFAULT);
     }
 
+    @After
+    public void tearDown() {
+        KeyStore keyStore = null;
+        try {
+            keyStore = KeyStore.getInstance("AndroidKeyStore");
+            keyStore.load(null);
+            keyStore.deleteEntry("testCertificateParametersFilter_client");
+        } catch (KeyStoreException | CertificateException | NoSuchAlgorithmException
+                | IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     @Test
     public void testCertificateParametersFilter_filtersByIntermediateIssuer()
-            throws InterruptedException, ExecutionException, CancellationException,
-            TimeoutException, IOException, CertificateEncodingException {
+            throws CancellationException, IOException, CertificateException,
+            KeyStoreException, NoSuchAlgorithmException {
         KeyStore keyStore = prepareKeyStoreWithLongChainCertificates();
 
         assertThat(createCheckerForIssuer(keyStore, ROOT_SUBJECT)
-                .shouldPresentCertificate("client")).isTrue();
+                .shouldPresentCertificate("testCertificateParametersFilter_client")).isTrue();
 
         assertThat(createCheckerForIssuer(keyStore, INTERMEDIATE_SUBJECT)
-                .shouldPresentCertificate("client")).isTrue();
+                .shouldPresentCertificate("testCertificateParametersFilter_client")).isTrue();
 
         assertThat(createCheckerForIssuer(keyStore, LEAF_SUBJECT)
-                .shouldPresentCertificate("client")).isFalse();
+                .shouldPresentCertificate("testCertificateParametersFilter_client")).isFalse();
     }
 
     // Return a KeyStore instance that has both a client certificate as well as a certificate
     // chain associated with it.
     private KeyStore prepareKeyStoreWithLongChainCertificates()
-            throws IOException, CertificateEncodingException {
-        KeyStore keyStore = mock(KeyStore.class);
-        when(keyStore.get(Credentials.USER_CERTIFICATE + "client")).thenReturn(mLeafRsaCertificate);
-        Certificate[] intermediates = new Certificate[] {
-            parseCertificate(mRootRsaCertificate), parseCertificate(mIntermediateRsaCertificate)};
-        byte[] intermediatesPem = convertToPem(intermediates);
-        assertThat(intermediatesPem).isNotNull();
-        when(keyStore.get(Credentials.CA_CERTIFICATE + "client")).thenReturn(intermediatesPem);
+            throws IOException, CertificateException, KeyStoreException, NoSuchAlgorithmException {
+
+        KeyStore keyStore = KeyStore.getInstance("AndroidKeyStore");
+        keyStore.load(null);
+
+        Certificate[] certs = new Certificate[3];
+        certs[0] = parseCertificate(mLeafRsaCertificate);
+        certs[1] = parseCertificate(mIntermediateRsaCertificate);
+        certs[2] = parseCertificate(mRootRsaCertificate);
+
+        keyStore.setKeyEntry("testCertificateParametersFilter_client", parseKey(mPrivateKey),
+                null, certs);
 
         return keyStore;
     }
@@ -189,16 +244,13 @@ public final class KeyChainActivityTest {
         }
     }
 
-    // Copied from android.security.Credentials, as that is a framework class.
-    public static byte[] convertToPem(Certificate... objects)
-            throws IOException, CertificateEncodingException {
-        ByteArrayOutputStream bao = new ByteArrayOutputStream();
-        Writer writer = new OutputStreamWriter(bao, StandardCharsets.US_ASCII);
-        PemWriter pw = new PemWriter(writer);
-        for (Certificate o : objects) {
-            pw.writeObject(new PemObject("CERTIFICATE", o.getEncoded()));
+    private static PrivateKey parseKey(byte[] key) {
+        try {
+            KeyFactory kf = KeyFactory.getInstance("RSA");
+            return kf.generatePrivate(new PKCS8EncodedKeySpec(key));
+        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
+            fail(String.format("Could not parse private key: %s", e));
+            return null;
         }
-        pw.close();
-        return bao.toByteArray();
     }
 }