Reliably delete keys if key generation fails.

Bug: 18088752
Change-Id: Iea68f3f96fc872d5628f163a1314ebd080c9d39e

2 files changed, 73 insertions, 57 deletions

diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreKeyGeneratorSpi.java b/keystore/java/android/security/keystore/AndroidKeyStoreKeyGeneratorSpi.java
index dc4c8a3..4d6178f 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreKeyGeneratorSpi.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreKeyGeneratorSpi.java
@@ -296,19 +296,33 @@ public abstract class AndroidKeyStoreKeyGeneratorSpi extends KeyGeneratorSpi {
         int flags = 0;
         String keyAliasInKeystore = Credentials.USER_SECRET_KEY + spec.getKeystoreAlias();
         KeyCharacteristics resultingKeyCharacteristics = new KeyCharacteristics();
-        int errorCode = mKeyStore.generateKey(
-                keyAliasInKeystore, args, additionalEntropy, flags, resultingKeyCharacteristics);
-        if (errorCode != KeyStore.NO_ERROR) {
-            throw new ProviderException(
-                    "Keystore operation failed", KeyStore.getKeyStoreException(errorCode));
-        }
-        @KeyProperties.KeyAlgorithmEnum String keyAlgorithmJCA;
+        boolean success = false;
         try {
-            keyAlgorithmJCA = KeyProperties.KeyAlgorithm.fromKeymasterSecretKeyAlgorithm(
-                    mKeymasterAlgorithm, mKeymasterDigest);
-        } catch (IllegalArgumentException e) {
-            throw new ProviderException("Failed to obtain JCA secret key algorithm name", e);
+            Credentials.deleteAllTypesForAlias(mKeyStore, spec.getKeystoreAlias());
+            int errorCode = mKeyStore.generateKey(
+                    keyAliasInKeystore,
+                    args,
+                    additionalEntropy,
+                    flags,
+                    resultingKeyCharacteristics);
+            if (errorCode != KeyStore.NO_ERROR) {
+                throw new ProviderException(
+                        "Keystore operation failed", KeyStore.getKeyStoreException(errorCode));
+            }
+            @KeyProperties.KeyAlgorithmEnum String keyAlgorithmJCA;
+            try {
+                keyAlgorithmJCA = KeyProperties.KeyAlgorithm.fromKeymasterSecretKeyAlgorithm(
+                        mKeymasterAlgorithm, mKeymasterDigest);
+            } catch (IllegalArgumentException e) {
+                throw new ProviderException("Failed to obtain JCA secret key algorithm name", e);
+            }
+            SecretKey result = new AndroidKeyStoreSecretKey(keyAliasInKeystore, keyAlgorithmJCA);
+            success = true;
+            return result;
+        } finally {
+            if (!success) {
+                Credentials.deleteAllTypesForAlias(mKeyStore, spec.getKeystoreAlias());
+            }
         }
-        return new AndroidKeyStoreSecretKey(keyAliasInKeystore, keyAlgorithmJCA);
     }
 }
diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreKeyPairGeneratorSpi.java b/keystore/java/android/security/keystore/AndroidKeyStoreKeyPairGeneratorSpi.java
index 4b45fd7..7b5ca3a 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreKeyPairGeneratorSpi.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreKeyPairGeneratorSpi.java
@@ -121,7 +121,6 @@ public abstract class AndroidKeyStoreKeyPairGeneratorSpi extends KeyPairGenerato
     public KeyPair generateKeyPair() {
         if (mKeyStore == null || mSpec == null) {
             throw new IllegalStateException("Not initialized");
-
         }
 
         final int flags = (mEncryptionAtRestRequired) ? KeyStore.FLAG_ENCRYPTED : 0;
@@ -134,62 +133,65 @@ public abstract class AndroidKeyStoreKeyPairGeneratorSpi extends KeyPairGenerato
 
         final String alias = mSpec.getKeystoreAlias();
 
-        Credentials.deleteAllTypesForAlias(mKeyStore, alias);
-
         byte[][] args = getArgsForKeyType(mKeyType, mSpec.getAlgorithmParameterSpec());
 
         final String privateKeyAlias = Credentials.USER_PRIVATE_KEY + alias;
 
-        if (!mKeyStore.generate(privateKeyAlias, KeyStore.UID_SELF, mKeyType, mKeySize,
-                flags, args)) {
-            throw new IllegalStateException("could not generate key in keystore");
-        }
+        boolean success = false;
+        try {
+            Credentials.deleteAllTypesForAlias(mKeyStore, alias);
+            if (!mKeyStore.generate(privateKeyAlias, KeyStore.UID_SELF, mKeyType, mKeySize,
+                    flags, args)) {
+                throw new IllegalStateException("could not generate key in keystore");
+            }
 
-        Credentials.deleteSecretKeyTypeForAlias(mKeyStore, 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 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 = mKeyStore.getPubkey(privateKeyAlias);
 
-        final byte[] pubKeyBytes = mKeyStore.getPubkey(privateKeyAlias);
+            final PublicKey pubKey;
+            try {
+                final KeyFactory keyFact = KeyFactory.getInstance(mKeyAlgorithm);
+                pubKey = keyFact.generatePublic(new X509EncodedKeySpec(pubKeyBytes));
+            } catch (NoSuchAlgorithmException e) {
+                throw new IllegalStateException("Can't instantiate key generator", e);
+            } catch (InvalidKeySpecException e) {
+                throw new IllegalStateException("keystore returned invalid key encoding", e);
+            }
 
-        final PublicKey pubKey;
-        try {
-            final KeyFactory keyFact = KeyFactory.getInstance(mKeyAlgorithm);
-            pubKey = keyFact.generatePublic(new X509EncodedKeySpec(pubKeyBytes));
-        } catch (NoSuchAlgorithmException e) {
-            throw new IllegalStateException("Can't instantiate key generator", e);
-        } catch (InvalidKeySpecException e) {
-            throw new IllegalStateException("keystore returned invalid key encoding", e);
-        }
+            final X509Certificate cert;
+            try {
+                cert = generateCertificate(privKey, pubKey);
+            } catch (Exception e) {
+                throw new IllegalStateException("Can't generate certificate", e);
+            }
 
-        final X509Certificate cert;
-        try {
-            cert = generateCertificate(privKey, pubKey);
-        } catch (Exception e) {
-            Credentials.deleteAllTypesForAlias(mKeyStore, alias);
-            throw new IllegalStateException("Can't generate certificate", e);
-        }
+            byte[] certBytes;
+            try {
+                certBytes = cert.getEncoded();
+            } catch (CertificateEncodingException e) {
+                throw new IllegalStateException("Can't get encoding of certificate", e);
+            }
 
-        byte[] certBytes;
-        try {
-            certBytes = cert.getEncoded();
-        } catch (CertificateEncodingException e) {
-            Credentials.deleteAllTypesForAlias(mKeyStore, alias);
-            throw new IllegalStateException("Can't get encoding of certificate", e);
-        }
+            if (!mKeyStore.put(Credentials.USER_CERTIFICATE + alias, certBytes, KeyStore.UID_SELF,
+                    flags)) {
+                throw new IllegalStateException("Can't store certificate in AndroidKeyStore");
+            }
 
-        if (!mKeyStore.put(Credentials.USER_CERTIFICATE + alias, certBytes, KeyStore.UID_SELF,
-                flags)) {
-            Credentials.deleteAllTypesForAlias(mKeyStore, alias);
-            throw new IllegalStateException("Can't store certificate in AndroidKeyStore");
+            KeyPair result = new KeyPair(pubKey, privKey);
+            success = true;
+            return result;
+        } finally {
+            if (!success) {
+                Credentials.deleteAllTypesForAlias(mKeyStore, alias);
+            }
         }
-
-        return new KeyPair(pubKey, privKey);
     }
 
     @SuppressWarnings("deprecation")