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

2 files changed, 60 insertions, 13 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);
     }
 }