Reset AndroidKeyStore Mac and Cipher state when init fails.

For consistency, this also switches AndroidKeyStore's Mac init/reset
implementation to the same approach as used in Cipher.

Bug: 18088752
Change-Id: Id34caf7cfe04f2058e22d8632890f762927bb31c

2 files changed, 79 insertions, 22 deletions

diff --git a/keystore/java/android/security/KeyStoreCipherSpi.java b/keystore/java/android/security/KeyStoreCipherSpi.java
index 37e00b28cd83..1f8d8ec4ff43 100644
--- a/keystore/java/android/security/KeyStoreCipherSpi.java
+++ b/keystore/java/android/security/KeyStoreCipherSpi.java
@@ -152,29 +152,58 @@ public abstract class KeyStoreCipherSpi extends CipherSpi implements KeyStoreCry
 
     @Override
     protected void engineInit(int opmode, Key key, SecureRandom random) throws InvalidKeyException {
-        init(opmode, key, random);
-        initAlgorithmSpecificParameters();
-        ensureKeystoreOperationInitialized();
+        resetAll();
+
+        boolean success = false;
+        try {
+            init(opmode, key, random);
+            initAlgorithmSpecificParameters();
+            ensureKeystoreOperationInitialized();
+            success = true;
+        } finally {
+            if (!success) {
+                resetAll();
+            }
+        }
     }
 
     @Override
     protected void engineInit(int opmode, Key key, AlgorithmParameters params, SecureRandom random)
             throws InvalidKeyException, InvalidAlgorithmParameterException {
-        init(opmode, key, random);
-        initAlgorithmSpecificParameters(params);
-        ensureKeystoreOperationInitialized();
+        resetAll();
+
+        boolean success = false;
+        try {
+            init(opmode, key, random);
+            initAlgorithmSpecificParameters(params);
+            ensureKeystoreOperationInitialized();
+            success = true;
+        } finally {
+            if (!success) {
+                resetAll();
+            }
+        }
     }
 
     @Override
     protected void engineInit(int opmode, Key key, AlgorithmParameterSpec params,
             SecureRandom random) throws InvalidKeyException, InvalidAlgorithmParameterException {
-        init(opmode, key, random);
-        initAlgorithmSpecificParameters(params);
-        ensureKeystoreOperationInitialized();
+        resetAll();
+
+        boolean success = false;
+        try {
+            init(opmode, key, random);
+            initAlgorithmSpecificParameters(params);
+            ensureKeystoreOperationInitialized();
+            success = true;
+        } finally {
+            if (!success) {
+                resetAll();
+            }
+        }
     }
 
     private void init(int opmode, Key key, SecureRandom random) throws InvalidKeyException {
-        resetAll();
         if (!(key instanceof KeyStoreSecretKey)) {
             throw new InvalidKeyException(
                     "Unsupported key: " + ((key != null) ? key.getClass().getName() : "null"));
diff --git a/keystore/java/android/security/KeyStoreHmacSpi.java b/keystore/java/android/security/KeyStoreHmacSpi.java
index a19bbda3b9ec..f8b6fef2a90f 100644
--- a/keystore/java/android/security/KeyStoreHmacSpi.java
+++ b/keystore/java/android/security/KeyStoreHmacSpi.java
@@ -69,9 +69,10 @@ public abstract class KeyStoreHmacSpi extends MacSpi implements KeyStoreCryptoOp
     private final int mKeymasterDigest;
     private final int mMacSizeBytes;
 
-    private String mKeyAliasInKeyStore;
+    // Fields below are populated by engineInit and should be preserved after engineDoFinal.
+    private KeyStoreSecretKey mKey;
 
-    // The fields below are reset by the engineReset operation.
+    // Fields below are reset when engineDoFinal succeeds.
     private KeyStoreCryptoOperationChunkedStreamer mChunkedStreamer;
     private IBinder mOperationToken;
     private Long mOperationHandle;
@@ -89,28 +90,49 @@ public abstract class KeyStoreHmacSpi extends MacSpi implements KeyStoreCryptoOp
     @Override
     protected void engineInit(Key key, AlgorithmParameterSpec params) throws InvalidKeyException,
             InvalidAlgorithmParameterException {
+        resetAll();
+
+        boolean success = false;
+        try {
+            init(key, params);
+            ensureKeystoreOperationInitialized();
+            success = true;
+        } finally {
+            if (!success) {
+                resetAll();
+            }
+        }
+    }
+
+    private void init(Key key, AlgorithmParameterSpec params) throws InvalidKeyException,
+        InvalidAlgorithmParameterException {
         if (key == null) {
             throw new InvalidKeyException("key == null");
         } else if (!(key instanceof KeyStoreSecretKey)) {
             throw new InvalidKeyException(
                     "Only Android KeyStore secret keys supported. Key: " + key);
         }
+        mKey = (KeyStoreSecretKey) key;
 
         if (params != null) {
             throw new InvalidAlgorithmParameterException(
                     "Unsupported algorithm parameters: " + params);
         }
 
-        mKeyAliasInKeyStore = ((KeyStoreSecretKey) key).getAlias();
-        if (mKeyAliasInKeyStore == null) {
-            throw new InvalidKeyException("Key's KeyStore alias not known");
+    }
+
+    private void resetAll() {
+        mKey = null;
+        IBinder operationToken = mOperationToken;
+        if (operationToken != null) {
+            mOperationToken = null;
+            mKeyStore.abort(operationToken);
         }
-        engineReset();
-        ensureKeystoreOperationInitialized();
+        mOperationHandle = null;
+        mChunkedStreamer = null;
     }
 
-    @Override
-    protected void engineReset() {
+    private void resetWhilePreservingInitState() {
         IBinder operationToken = mOperationToken;
         if (operationToken != null) {
             mOperationToken = null;
@@ -120,11 +142,16 @@ public abstract class KeyStoreHmacSpi extends MacSpi implements KeyStoreCryptoOp
         mChunkedStreamer = null;
     }
 
+    @Override
+    protected void engineReset() {
+        resetWhilePreservingInitState();
+    }
+
     private void ensureKeystoreOperationInitialized() {
         if (mChunkedStreamer != null) {
             return;
         }
-        if (mKeyAliasInKeyStore == null) {
+        if (mKey == null) {
             throw new IllegalStateException("Not initialized");
         }
 
@@ -132,7 +159,8 @@ public abstract class KeyStoreHmacSpi extends MacSpi implements KeyStoreCryptoOp
         keymasterArgs.addInt(KeymasterDefs.KM_TAG_ALGORITHM, KeymasterDefs.KM_ALGORITHM_HMAC);
         keymasterArgs.addInt(KeymasterDefs.KM_TAG_DIGEST, mKeymasterDigest);
 
-        OperationResult opResult = mKeyStore.begin(mKeyAliasInKeyStore,
+        OperationResult opResult = mKeyStore.begin(
+                mKey.getAlias(),
                 KeymasterDefs.KM_PURPOSE_SIGN,
                 true,
                 keymasterArgs,
@@ -184,7 +212,7 @@ public abstract class KeyStoreHmacSpi extends MacSpi implements KeyStoreCryptoOp
             throw KeyStore.getCryptoOperationException(e);
         }
 
-        engineReset();
+        resetWhilePreservingInitState();
         return result;
     }