KeyChain/Keystore 2.0

Key uses mostly public Keystore API which works the same for Keystore
and Keystore 2.0. The only exception is:
 * The public API does not allow for grants.
In this cases we fall back on hidden API.

Keystore 2.0 and KeyMint do not allow for key attestation outside of
key generation or import, so this patch also removes attestKey from
the KeyChainService implementation.

Test: KeyChain tests and CTS tests.
Bug: 171305387
Merged-In: Ieefaba81e36dc0adc87d0eebde8a0901c1687960
Change-Id: Ieefaba81e36dc0adc87d0eebde8a0901c1687960

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();
     }
 }