Add ability to replace chain for PrivateKeyEntry

For the AndroidKeyStore API, allow entries to have their certificate
chain replaced without destroying the underlying PrivateKey. Since
entries are backed by unexportable private keys, requiring them to be
supplied again doesn't make sense and is impossible.

Change-Id: I629ce2a625315c8d8020a082892650ac5eba22ae

4 files changed, 242 insertions, 24 deletions

diff --git a/keystore/java/android/security/AndroidKeyStore.java b/keystore/java/android/security/AndroidKeyStore.java
index e19217f..7b6e540 100644
--- a/keystore/java/android/security/AndroidKeyStore.java
+++ b/keystore/java/android/security/AndroidKeyStore.java
@@ -16,7 +16,9 @@
 
 package android.security;
 
+import org.apache.harmony.xnet.provider.jsse.OpenSSLDSAPrivateKey;
 import org.apache.harmony.xnet.provider.jsse.OpenSSLEngine;
+import org.apache.harmony.xnet.provider.jsse.OpenSSLRSAPrivateKey;
 
 import android.util.Log;
 
@@ -193,17 +195,41 @@ public class AndroidKeyStore extends KeyStoreSpi {
 
     private void setPrivateKeyEntry(String alias, PrivateKey key, Certificate[] chain)
             throws KeyStoreException {
-        // Make sure the PrivateKey format is the one we support.
-        final String keyFormat = key.getFormat();
-        if ((keyFormat == null) || (!"PKCS#8".equals(keyFormat))) {
-            throw new KeyStoreException(
-                    "Only PrivateKeys that can be encoded into PKCS#8 are supported");
+        byte[] keyBytes = null;
+
+        final String pkeyAlias;
+        if (key instanceof OpenSSLRSAPrivateKey) {
+            pkeyAlias = ((OpenSSLRSAPrivateKey) key).getPkeyAlias();
+        } else if (key instanceof OpenSSLDSAPrivateKey) {
+            pkeyAlias = ((OpenSSLDSAPrivateKey) key).getPkeyAlias();
+        } else {
+            pkeyAlias = null;
         }
 
-        // Make sure we can actually encode the key.
-        final byte[] keyBytes = key.getEncoded();
-        if (keyBytes == null) {
-            throw new KeyStoreException("PrivateKey has no encoding");
+        final boolean shouldReplacePrivateKey;
+        if (pkeyAlias != null && pkeyAlias.startsWith(Credentials.USER_PRIVATE_KEY)) {
+            final String keySubalias = pkeyAlias.substring(Credentials.USER_PRIVATE_KEY.length());
+            if (!alias.equals(keySubalias)) {
+                throw new KeyStoreException("Can only replace keys with same alias: " + alias
+                        + " != " + keySubalias);
+            }
+
+            shouldReplacePrivateKey = false;
+        } else {
+            // Make sure the PrivateKey format is the one we support.
+            final String keyFormat = key.getFormat();
+            if ((keyFormat == null) || (!"PKCS#8".equals(keyFormat))) {
+                throw new KeyStoreException(
+                        "Only PrivateKeys that can be encoded into PKCS#8 are supported");
+            }
+
+            // Make sure we can actually encode the key.
+            keyBytes = key.getEncoded();
+            if (keyBytes == null) {
+                throw new KeyStoreException("PrivateKey has no encoding");
+            }
+
+            shouldReplacePrivateKey = true;
         }
 
         // Make sure the chain exists since this is a PrivateKey
@@ -273,17 +299,25 @@ public class AndroidKeyStore extends KeyStoreSpi {
         }
 
         /*
-         * Make sure we clear out all the types we know about before trying to
+         * Make sure we clear out all the appropriate types before trying to
          * write.
          */
-        Credentials.deleteAllTypesForAlias(mKeyStore, alias);
+        if (shouldReplacePrivateKey) {
+            Credentials.deleteAllTypesForAlias(mKeyStore, alias);
+        } else {
+            Credentials.deleteCertificateTypesForAlias(mKeyStore, alias);
+        }
 
-        if (!mKeyStore.importKey(Credentials.USER_PRIVATE_KEY + alias, keyBytes)) {
+        if (shouldReplacePrivateKey
+                && !mKeyStore.importKey(Credentials.USER_PRIVATE_KEY + alias, keyBytes)) {
+            Credentials.deleteAllTypesForAlias(mKeyStore, alias);
             throw new KeyStoreException("Couldn't put private key in keystore");
         } else if (!mKeyStore.put(Credentials.USER_CERTIFICATE + alias, userCertBytes)) {
+            Credentials.deleteAllTypesForAlias(mKeyStore, alias);
             throw new KeyStoreException("Couldn't put certificate #1 in keystore");
         } else if (chainBytes != null
                 && !mKeyStore.put(Credentials.CA_CERTIFICATE + alias, chainBytes)) {
+            Credentials.deleteAllTypesForAlias(mKeyStore, alias);
             throw new KeyStoreException("Couldn't put certificate chain in keystore");
         }
     }
diff --git a/keystore/java/android/security/Credentials.java b/keystore/java/android/security/Credentials.java
index 72332eb..f6bf432 100644
--- a/keystore/java/android/security/Credentials.java
+++ b/keystore/java/android/security/Credentials.java
@@ -197,7 +197,20 @@ public class Credentials {
          * don't use a conditional here.
          */
         return keystore.delKey(Credentials.USER_PRIVATE_KEY + alias)
-                | keystore.delete(Credentials.USER_CERTIFICATE + alias)
+                | deleteCertificateTypesForAlias(keystore, alias);
+    }
+
+    /**
+     * Delete all types (private key, certificate, CA certificate) for a
+     * particular {@code alias}. All three can exist for any given alias.
+     * Returns {@code true} if there was at least one of those types.
+     */
+    static boolean deleteCertificateTypesForAlias(KeyStore keystore, String alias) {
+        /*
+         * Make sure every certificate type is deleted. There can be two types,
+         * so don't use a conditional here.
+         */
+        return keystore.delete(Credentials.USER_CERTIFICATE + alias)
                 | keystore.delete(Credentials.CA_CERTIFICATE + alias);
     }
 }
diff --git a/keystore/tests/Android.mk b/keystore/tests/Android.mk
index 95604c6..61cf640 100644
--- a/keystore/tests/Android.mk
+++ b/keystore/tests/Android.mk
@@ -5,7 +5,7 @@ include $(CLEAR_VARS)
 LOCAL_MODULE_TAGS := tests
 LOCAL_CERTIFICATE := platform
 
-LOCAL_JAVA_LIBRARIES := android.test.runner
+LOCAL_JAVA_LIBRARIES := android.test.runner bouncycastle
 
 # Include all test java files.
 LOCAL_SRC_FILES := $(call all-java-files-under, src)
diff --git a/keystore/tests/src/android/security/AndroidKeyStoreTest.java b/keystore/tests/src/android/security/AndroidKeyStoreTest.java
index bff01b8..49e2f12 100644
--- a/keystore/tests/src/android/security/AndroidKeyStoreTest.java
+++ b/keystore/tests/src/android/security/AndroidKeyStoreTest.java
@@ -16,12 +16,17 @@
 
 package android.security;
 
+import com.android.org.bouncycastle.x509.X509V3CertificateGenerator;
+
+import org.apache.harmony.xnet.provider.jsse.OpenSSLEngine;
+
 import android.test.AndroidTestCase;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
-import java.io.IOException;
 import java.io.OutputStream;
+import java.math.BigInteger;
+import java.security.InvalidKeyException;
 import java.security.Key;
 import java.security.KeyFactory;
 import java.security.KeyStore.Entry;
@@ -30,12 +35,14 @@ import java.security.KeyStore.TrustedCertificateEntry;
 import java.security.KeyStoreException;
 import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
+import java.security.PublicKey;
 import java.security.cert.Certificate;
-import java.security.cert.CertificateException;
 import java.security.cert.CertificateFactory;
+import java.security.cert.X509Certificate;
 import java.security.interfaces.RSAPrivateKey;
 import java.security.spec.InvalidKeySpecException;
 import java.security.spec.PKCS8EncodedKeySpec;
+import java.security.spec.X509EncodedKeySpec;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Date;
@@ -44,6 +51,8 @@ import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Set;
 
+import javax.security.auth.x500.X500Principal;
+
 public class AndroidKeyStoreTest extends AndroidTestCase {
     private android.security.KeyStore mAndroidKeyStore;
 
@@ -55,6 +64,22 @@ public class AndroidKeyStoreTest extends AndroidTestCase {
 
     private static final String TEST_ALIAS_3 = "test3";
 
+    private static final X500Principal TEST_DN_1 = new X500Principal("CN=test1");
+
+    private static final X500Principal TEST_DN_2 = new X500Principal("CN=test2");
+
+    private static final BigInteger TEST_SERIAL_1 = BigInteger.ONE;
+
+    private static final BigInteger TEST_SERIAL_2 = BigInteger.valueOf(2L);
+
+    private static final long NOW_MILLIS = System.currentTimeMillis();
+
+    /* We have to round this off because X509v3 doesn't store milliseconds. */
+    private static final Date NOW = new Date(NOW_MILLIS - (NOW_MILLIS % 1000L));
+
+    @SuppressWarnings("deprecation")
+    private static final Date NOW_PLUS_10_YEARS = new Date(NOW.getYear() + 10, 0, 1);
+
     /*
      * The keys and certificates below are generated with:
      *
@@ -759,17 +784,31 @@ public class AndroidKeyStoreTest extends AndroidTestCase {
         assertPrivateKeyEntryEquals(keyEntry, FAKE_KEY_1, FAKE_USER_1, FAKE_CA_1);
     }
 
+    @SuppressWarnings("unchecked")
     private void assertPrivateKeyEntryEquals(PrivateKeyEntry keyEntry, byte[] key, byte[] cert,
             byte[] ca) throws Exception {
         KeyFactory keyFact = KeyFactory.getInstance("RSA");
         PrivateKey expectedKey = keyFact.generatePrivate(new PKCS8EncodedKeySpec(key));
 
-        assertEquals("Returned PrivateKey should be what we inserted", expectedKey,
-                keyEntry.getPrivateKey());
-
         CertificateFactory certFact = CertificateFactory.getInstance("X.509");
         Certificate expectedCert = certFact.generateCertificate(new ByteArrayInputStream(cert));
 
+        final Collection<Certificate> expectedChain;
+        if (ca != null) {
+            expectedChain = (Collection<Certificate>) certFact
+                    .generateCertificates(new ByteArrayInputStream(ca));
+        } else {
+            expectedChain = null;
+        }
+
+        assertPrivateKeyEntryEquals(keyEntry, expectedKey, expectedCert, expectedChain);
+    }
+
+    private void assertPrivateKeyEntryEquals(PrivateKeyEntry keyEntry, PrivateKey expectedKey,
+            Certificate expectedCert, Collection<Certificate> expectedChain) throws Exception {
+        assertEquals("Returned PrivateKey should be what we inserted", expectedKey,
+                keyEntry.getPrivateKey());
+
         assertEquals("Returned Certificate should be what we inserted", expectedCert,
                 keyEntry.getCertificate());
 
@@ -777,13 +816,9 @@ public class AndroidKeyStoreTest extends AndroidTestCase {
 
         assertEquals("First certificate in chain should be user cert", expectedCert, actualChain[0]);
 
-        if (ca == null) {
+        if (expectedChain == null) {
             assertEquals("Certificate chain should not include CAs", 1, actualChain.length);
         } else {
-            @SuppressWarnings("unchecked")
-            Collection<Certificate> expectedChain = (Collection<Certificate>) certFact
-                    .generateCertificates(new ByteArrayInputStream(ca));
-
             int i = 1;
             final Iterator<Certificate> it = expectedChain.iterator();
             while (it.hasNext()) {
@@ -1306,6 +1341,142 @@ public class AndroidKeyStoreTest extends AndroidTestCase {
         }
     }
 
+    @SuppressWarnings("deprecation")
+    private static X509Certificate generateCertificate(android.security.KeyStore keyStore,
+            String alias, BigInteger serialNumber, X500Principal subjectDN, Date notBefore,
+            Date notAfter) throws Exception {
+        final String privateKeyAlias = Credentials.USER_PRIVATE_KEY + alias;
+
+        final PrivateKey privKey;
+        final OpenSSLEngine engine = OpenSSLEngine.getInstance("keystore");
+        try {
+            privKey = engine.getPrivateKeyById(privateKeyAlias);
+        } catch (InvalidKeyException e) {
+            throw new RuntimeException("Can't get key", e);
+        }
+
+        final byte[] pubKeyBytes = keyStore.getPubkey(privateKeyAlias);
+
+        final PublicKey pubKey;
+        try {
+            final KeyFactory keyFact = KeyFactory.getInstance("RSA");
+            pubKey = keyFact.generatePublic(new X509EncodedKeySpec(pubKeyBytes));
+        } catch (NoSuchAlgorithmException e) {
+            throw new IllegalStateException("Can't instantiate RSA key generator", e);
+        } catch (InvalidKeySpecException e) {
+            throw new IllegalStateException("keystore returned invalid key encoding", e);
+        }
+
+        final X509V3CertificateGenerator certGen = new X509V3CertificateGenerator();
+        certGen.setPublicKey(pubKey);
+        certGen.setSerialNumber(serialNumber);
+        certGen.setSubjectDN(subjectDN);
+        certGen.setIssuerDN(subjectDN);
+        certGen.setNotBefore(notBefore);
+        certGen.setNotAfter(notAfter);
+        certGen.setSignatureAlgorithm("sha1WithRSA");
+
+        final X509Certificate cert = certGen.generate(privKey);
+
+        return cert;
+    }
+
+    public void testKeyStore_SetKeyEntry_ReplacedChain_Success() throws Exception {
+        mKeyStore.load(null, null);
+
+        // Create key #1
+        {
+            final String privateKeyAlias = Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1;
+            assertTrue(mAndroidKeyStore.generate(privateKeyAlias));
+
+            Key key = mKeyStore.getKey(TEST_ALIAS_1, null);
+
+            assertTrue(key instanceof PrivateKey);
+
+            PrivateKey expectedKey = (PrivateKey) key;
+
+            X509Certificate expectedCert = generateCertificate(mAndroidKeyStore, TEST_ALIAS_1,
+                    TEST_SERIAL_1, TEST_DN_1, NOW, NOW_PLUS_10_YEARS);
+
+            assertTrue(mAndroidKeyStore.put(Credentials.USER_CERTIFICATE + TEST_ALIAS_1,
+                    expectedCert.getEncoded()));
+
+            Entry entry = mKeyStore.getEntry(TEST_ALIAS_1, null);
+
+            assertTrue(entry instanceof PrivateKeyEntry);
+
+            PrivateKeyEntry keyEntry = (PrivateKeyEntry) entry;
+
+            assertPrivateKeyEntryEquals(keyEntry, expectedKey, expectedCert, null);
+        }
+
+        // Replace key #1 with new chain
+        {
+            Key key = mKeyStore.getKey(TEST_ALIAS_1, null);
+
+            assertTrue(key instanceof PrivateKey);
+
+            PrivateKey expectedKey = (PrivateKey) key;
+
+            X509Certificate expectedCert = generateCertificate(mAndroidKeyStore, TEST_ALIAS_1,
+                    TEST_SERIAL_2, TEST_DN_2, NOW, NOW_PLUS_10_YEARS);
+
+            mKeyStore.setKeyEntry(TEST_ALIAS_1, expectedKey, null,
+                    new Certificate[] { expectedCert });
+
+            Entry entry = mKeyStore.getEntry(TEST_ALIAS_1, null);
+
+            assertTrue(entry instanceof PrivateKeyEntry);
+
+            PrivateKeyEntry keyEntry = (PrivateKeyEntry) entry;
+
+            assertPrivateKeyEntryEquals(keyEntry, expectedKey, expectedCert, null);
+        }
+    }
+
+    public void testKeyStore_SetKeyEntry_ReplacedChain_DifferentPrivateKey_Failure()
+            throws Exception {
+        mKeyStore.load(null, null);
+
+        // Create key #1
+        {
+            final String privateKeyAlias = Credentials.USER_PRIVATE_KEY + TEST_ALIAS_1;
+            assertTrue(mAndroidKeyStore.generate(privateKeyAlias));
+
+            X509Certificate cert = generateCertificate(mAndroidKeyStore, TEST_ALIAS_1,
+                    TEST_SERIAL_1, TEST_DN_1, NOW, NOW_PLUS_10_YEARS);
+
+            assertTrue(mAndroidKeyStore.put(Credentials.USER_CERTIFICATE + TEST_ALIAS_1,
+                    cert.getEncoded()));
+        }
+
+        // Create key #2
+        {
+            final String privateKeyAlias = Credentials.USER_PRIVATE_KEY + TEST_ALIAS_2;
+            assertTrue(mAndroidKeyStore.generate(privateKeyAlias));
+
+            X509Certificate cert = generateCertificate(mAndroidKeyStore, TEST_ALIAS_2,
+                    TEST_SERIAL_2, TEST_DN_2, NOW, NOW_PLUS_10_YEARS);
+
+            assertTrue(mAndroidKeyStore.put(Credentials.USER_CERTIFICATE + TEST_ALIAS_2,
+                    cert.getEncoded()));
+        }
+
+        // Replace key #1 with key #2
+        {
+            Key key1 = mKeyStore.getKey(TEST_ALIAS_2, null);
+
+            X509Certificate cert = generateCertificate(mAndroidKeyStore, TEST_ALIAS_2,
+                    TEST_SERIAL_2, TEST_DN_2, NOW, NOW_PLUS_10_YEARS);
+
+            try {
+                mKeyStore.setKeyEntry(TEST_ALIAS_1, key1, null, new Certificate[] { cert });
+                fail("Should not allow setting of KeyEntry with wrong PrivaetKey");
+            } catch (KeyStoreException success) {
+            }
+        }
+    }
+
     public void testKeyStore_Size_Success() throws Exception {
         mKeyStore.load(null, null);