Cleanup logic for per-op auth keys.

This streamlines the exception throwing logic for per-op auth keys of
AndroidKeyStore.

Change-Id: I7e27c17fd89d5a7f71f5d7578f584189c5236fb8

5 files changed, 173 insertions, 65 deletions

diff --git a/keystore/java/android/security/KeyStoreCipherSpi.java b/keystore/java/android/security/KeyStoreCipherSpi.java
index 917f716..20dd524 100644
--- a/keystore/java/android/security/KeyStoreCipherSpi.java
+++ b/keystore/java/android/security/KeyStoreCipherSpi.java
@@ -22,6 +22,7 @@ import android.security.keymaster.KeymasterDefs;
 import android.security.keymaster.OperationResult;
 
 import java.security.AlgorithmParameters;
+import java.security.GeneralSecurityException;
 import java.security.InvalidAlgorithmParameterException;
 import java.security.InvalidKeyException;
 import java.security.Key;
@@ -298,38 +299,36 @@ public abstract class KeyStoreCipherSpi extends CipherSpi implements KeyStoreCry
         mAdditionalEntropyForBegin = null;
         if (opResult == null) {
             throw new KeyStoreConnectException();
-        } else if ((opResult.resultCode != KeyStore.NO_ERROR)
-                && (opResult.resultCode != KeyStore.OP_AUTH_NEEDED)) {
-            switch (opResult.resultCode) {
-                case KeymasterDefs.KM_ERROR_INVALID_NONCE:
-                    throw new InvalidAlgorithmParameterException("Invalid IV");
+        }
+
+        // Store operation token and handle regardless of the error code returned by KeyStore to
+        // ensure that the operation gets aborted immediately if the code below throws an exception.
+        mOperationToken = opResult.token;
+        mOperationHandle = opResult.operationHandle;
+
+        // If necessary, throw an exception due to KeyStore operation having failed.
+        GeneralSecurityException e = KeyStoreCryptoOperationUtils.getExceptionForCipherInit(
+                mKeyStore, mKey, opResult.resultCode);
+        if (e != null) {
+            if (e instanceof InvalidKeyException) {
+                throw (InvalidKeyException) e;
+            } else if (e instanceof InvalidAlgorithmParameterException) {
+                throw (InvalidAlgorithmParameterException) e;
+            } else {
+                throw new RuntimeException("Unexpected exception type", e);
             }
-            throw mKeyStore.getInvalidKeyException(mKey.getAlias(), opResult.resultCode);
         }
 
-        if (opResult.token == null) {
+        if (mOperationToken == null) {
             throw new IllegalStateException("Keystore returned null operation token");
         }
-        // The operation handle/token is now either valid for use immediately or needs to be
-        // authorized through user authentication (if the error code was OP_AUTH_NEEDED).
-        mOperationToken = opResult.token;
-        mOperationHandle = opResult.operationHandle;
+
         loadAlgorithmSpecificParametersFromBeginResult(keymasterOutputArgs);
         mFirstOperationInitiated = true;
         mIvHasBeenUsed = true;
         mMainDataStreamer = new KeyStoreCryptoOperationChunkedStreamer(
                 new KeyStoreCryptoOperationChunkedStreamer.MainDataStream(
                         mKeyStore, opResult.token));
-
-        if (opResult.resultCode != KeyStore.NO_ERROR) {
-            // The operation requires user authentication. Check whether such authentication is
-            // possible (e.g., the key may have been permanently invalidated).
-            InvalidKeyException e =
-                    mKeyStore.getInvalidKeyException(mKey.getAlias(), opResult.resultCode);
-            if (!(e instanceof UserNotAuthenticatedException)) {
-                throw e;
-            }
-        }
     }
 
     @Override
diff --git a/keystore/java/android/security/KeyStoreCryptoOperationUtils.java b/keystore/java/android/security/KeyStoreCryptoOperationUtils.java
new file mode 100644
index 0000000..313b527
--- /dev/null
+++ b/keystore/java/android/security/KeyStoreCryptoOperationUtils.java
@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2015 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.security;
+
+import android.security.keymaster.KeymasterDefs;
+
+import java.security.GeneralSecurityException;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.InvalidKeyException;
+
+/**
+ * Assorted utility methods for implementing crypto operations on top of KeyStore.
+ *
+ * @hide
+ */
+abstract class KeyStoreCryptoOperationUtils {
+    private KeyStoreCryptoOperationUtils() {}
+
+    /**
+     * Returns the {@link InvalidKeyException} to be thrown by the {@code init} method of
+     * the crypto operation in response to {@code KeyStore.begin} operation or {@code null} if
+     * the {@code init} method should succeed.
+     */
+    static InvalidKeyException getInvalidKeyExceptionForInit(
+            KeyStore keyStore, KeyStoreKey key, int beginOpResultCode) {
+        if (beginOpResultCode == KeyStore.NO_ERROR) {
+            return null;
+        }
+
+        // An error occured. However, some errors should not lead to init throwing an exception.
+        // See below.
+        InvalidKeyException e =
+                keyStore.getInvalidKeyException(key.getAlias(), beginOpResultCode);
+        switch (beginOpResultCode) {
+            case KeyStore.OP_AUTH_NEEDED:
+                // Operation needs to be authorized by authenticating the user. Don't throw an
+                // exception is such authentication is possible for this key
+                // (UserNotAuthenticatedException). An example of when it's not possible is where
+                // the key is permanently invalidated (KeyPermanentlyInvalidatedException).
+                if (e instanceof UserNotAuthenticatedException) {
+                    return null;
+                }
+                break;
+        }
+        return e;
+    }
+
+    /**
+     * Returns the exception to be thrown by the {@code Cipher.init} method of the crypto operation
+     * in response to {@code KeyStore.begin} operation or {@code null} if the {@code init} method
+     * should succeed.
+     */
+    static GeneralSecurityException getExceptionForCipherInit(
+            KeyStore keyStore, KeyStoreKey key, int beginOpResultCode) {
+        if (beginOpResultCode == KeyStore.NO_ERROR) {
+            return null;
+        }
+
+        // Cipher-specific cases
+        switch (beginOpResultCode) {
+            case KeymasterDefs.KM_ERROR_INVALID_NONCE:
+                return new InvalidAlgorithmParameterException("Invalid IV");
+        }
+
+        // General cases
+        return getInvalidKeyExceptionForInit(keyStore, key, beginOpResultCode);
+    }
+}
diff --git a/keystore/java/android/security/KeyStoreHmacSpi.java b/keystore/java/android/security/KeyStoreHmacSpi.java
index 4590b9c..2705304 100644
--- a/keystore/java/android/security/KeyStoreHmacSpi.java
+++ b/keystore/java/android/security/KeyStoreHmacSpi.java
@@ -168,31 +168,27 @@ public abstract class KeyStoreHmacSpi extends MacSpi implements KeyStoreCryptoOp
                 new KeymasterArguments());
         if (opResult == null) {
             throw new KeyStoreConnectException();
-        } else if ((opResult.resultCode != KeyStore.NO_ERROR)
-                && (opResult.resultCode != KeyStore.OP_AUTH_NEEDED)) {
-            throw mKeyStore.getInvalidKeyException(mKey.getAlias(), opResult.resultCode);
         }
 
-        if (opResult.token == null) {
-            throw new IllegalStateException("Keystore returned null operation token");
-        }
-        // The operation handle/token is now either valid for use immediately or needs to be
-        // authorized through user authentication (if the error code was OP_AUTH_NEEDED).
+        // Store operation token and handle regardless of the error code returned by KeyStore to
+        // ensure that the operation gets aborted immediately if the code below throws an exception.
         mOperationToken = opResult.token;
         mOperationHandle = opResult.operationHandle;
+
+        // If necessary, throw an exception due to KeyStore operation having failed.
+        InvalidKeyException e = KeyStoreCryptoOperationUtils.getInvalidKeyExceptionForInit(
+                mKeyStore, mKey, opResult.resultCode);
+        if (e != null) {
+            throw e;
+        }
+
+        if (mOperationToken == null) {
+            throw new IllegalStateException("Keystore returned null operation token");
+        }
+
         mChunkedStreamer = new KeyStoreCryptoOperationChunkedStreamer(
                 new KeyStoreCryptoOperationChunkedStreamer.MainDataStream(
                         mKeyStore, mOperationToken));
-
-        if (opResult.resultCode != KeyStore.NO_ERROR) {
-            // The operation requires user authentication. Check whether such authentication is
-            // possible (e.g., the key may have been permanently invalidated).
-            InvalidKeyException e =
-                    mKeyStore.getInvalidKeyException(mKey.getAlias(), opResult.resultCode);
-            if (!(e instanceof UserNotAuthenticatedException)) {
-                throw e;
-            }
-        }
     }
 
     @Override
diff --git a/keystore/java/android/security/KeyStoreKey.java b/keystore/java/android/security/KeyStoreKey.java
new file mode 100644
index 0000000..7a34829
--- /dev/null
+++ b/keystore/java/android/security/KeyStoreKey.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2015 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.security;
+
+import java.security.Key;
+
+/**
+ * {@link Key} backed by AndroidKeyStore.
+ *
+ * @hide
+ */
+public class KeyStoreKey implements Key {
+    private final String mAlias;
+    private final String mAlgorithm;
+
+    public KeyStoreKey(String alias, String algorithm) {
+        mAlias = alias;
+        mAlgorithm = algorithm;
+    }
+
+    String getAlias() {
+        return mAlias;
+    }
+
+    @Override
+    public String getAlgorithm() {
+        return mAlgorithm;
+    }
+
+    @Override
+    public String getFormat() {
+        // This key does not export its key material
+        return null;
+    }
+
+    @Override
+    public byte[] getEncoded() {
+        // This key does not export its key material
+        return null;
+    }
+}
diff --git a/keystore/java/android/security/KeyStoreSecretKey.java b/keystore/java/android/security/KeyStoreSecretKey.java
index 7f0e3d3..ee25465 100644
--- a/keystore/java/android/security/KeyStoreSecretKey.java
+++ b/keystore/java/android/security/KeyStoreSecretKey.java
@@ -23,33 +23,9 @@ import javax.crypto.SecretKey;
  *
  * @hide
  */
-public class KeyStoreSecretKey implements SecretKey {
-    private final String mAlias;
-    private final String mAlgorithm;
+public class KeyStoreSecretKey extends KeyStoreKey implements SecretKey {
 
     public KeyStoreSecretKey(String alias, String algorithm) {
-        mAlias = alias;
-        mAlgorithm = algorithm;
-    }
-
-    String getAlias() {
-        return mAlias;
-    }
-
-    @Override
-    public String getAlgorithm() {
-        return mAlgorithm;
-    }
-
-    @Override
-    public String getFormat() {
-        // This key does not export its key material
-        return null;
-    }
-
-    @Override
-    public byte[] getEncoded() {
-        // This key does not export its key material
-        return null;
+        super(alias, algorithm);
     }
 }