Keystore 2.0 SPI: Public key operation workaround.

Test: Keystore cts tests.
Change-Id: I316fdb8beae018ac91c172dede735e6b0759368a

9 files changed, 293 insertions, 55 deletions

diff --git a/keystore/java/android/security/keystore2/AndroidKeyStore3DESCipherSpi.java b/keystore/java/android/security/keystore2/AndroidKeyStore3DESCipherSpi.java
index 0775a1a99886..83ff84d372f3 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStore3DESCipherSpi.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStore3DESCipherSpi.java
@@ -41,7 +41,7 @@ import javax.crypto.spec.IvParameterSpec;
  *
  * @hide
  */
-public class AndroidKeyStore3DESCipherSpi extends AndroidKeyStoreCipherSpiBase {
+public abstract class AndroidKeyStore3DESCipherSpi extends AndroidKeyStoreCipherSpiBase {
 
     private static final int BLOCK_SIZE_BYTES = 8;
 
@@ -73,12 +73,22 @@ public class AndroidKeyStore3DESCipherSpi extends AndroidKeyStoreCipherSpiBase {
             public NoPadding() {
                 super(KeymasterDefs.KM_PAD_NONE);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "DESede/ECB/NoPadding";
+            }
         }
 
         public static class PKCS7Padding extends ECB {
             public PKCS7Padding() {
                 super(KeymasterDefs.KM_PAD_PKCS7);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "DESede/ECB/PKCS7Padding";
+            }
         }
     }
 
@@ -91,12 +101,23 @@ public class AndroidKeyStore3DESCipherSpi extends AndroidKeyStoreCipherSpiBase {
             public NoPadding() {
                 super(KeymasterDefs.KM_PAD_NONE);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "DESede/CBC/NoPadding";
+            }
+
         }
 
         public static class PKCS7Padding extends CBC {
             public PKCS7Padding() {
                 super(KeymasterDefs.KM_PAD_PKCS7);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "DESede/CBC/PKCS7Padding";
+            }
         }
     }
 
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreAuthenticatedAESCipherSpi.java b/keystore/java/android/security/keystore2/AndroidKeyStoreAuthenticatedAESCipherSpi.java
index bc56f015f3bd..aab84e390c73 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreAuthenticatedAESCipherSpi.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreAuthenticatedAESCipherSpi.java
@@ -64,6 +64,11 @@ abstract class AndroidKeyStoreAuthenticatedAESCipherSpi extends AndroidKeyStoreC
         }
 
         @Override
+        protected final String getTransform() {
+            return "AES/GCM/NoPadding";
+        }
+
+        @Override
         protected final void resetAll() {
             mTagLengthBits = DEFAULT_TAG_LENGTH_BITS;
             super.resetAll();
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreBCWorkaroundProvider.java b/keystore/java/android/security/keystore2/AndroidKeyStoreBCWorkaroundProvider.java
index dd943d422e62..9ad6f3adbd33 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreBCWorkaroundProvider.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreBCWorkaroundProvider.java
@@ -254,13 +254,13 @@ class AndroidKeyStoreBCWorkaroundProvider extends Provider {
     private void putAsymmetricCipherImpl(String transformation, String implClass) {
         put("Cipher." + transformation, implClass);
         put("Cipher." + transformation + " SupportedKeyClasses",
-                KEYSTORE_PRIVATE_KEY_CLASS_NAME + "|" + KEYSTORE_PUBLIC_KEY_CLASS_NAME);
+                KEYSTORE_PRIVATE_KEY_CLASS_NAME);
     }
 
     private void putSignatureImpl(String algorithm, String implClass) {
         put("Signature." + algorithm, implClass);
         put("Signature." + algorithm + " SupportedKeyClasses",
-                KEYSTORE_PRIVATE_KEY_CLASS_NAME + "|" + KEYSTORE_PUBLIC_KEY_CLASS_NAME);
+                KEYSTORE_PRIVATE_KEY_CLASS_NAME);
     }
 
     public static String[] getSupportedEcdsaSignatureDigests() {
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreCipherSpiBase.java b/keystore/java/android/security/keystore2/AndroidKeyStoreCipherSpiBase.java
index a3b04abfba3f..2ee952cbc5fb 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreCipherSpiBase.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreCipherSpiBase.java
@@ -43,6 +43,7 @@ import java.security.PublicKey;
 import java.security.SecureRandom;
 import java.security.spec.AlgorithmParameterSpec;
 import java.security.spec.InvalidKeySpecException;
+import java.security.spec.MGF1ParameterSpec;
 import java.security.spec.PKCS8EncodedKeySpec;
 import java.security.spec.X509EncodedKeySpec;
 import java.util.ArrayList;
@@ -57,6 +58,8 @@ import javax.crypto.NoSuchPaddingException;
 import javax.crypto.SecretKey;
 import javax.crypto.SecretKeyFactory;
 import javax.crypto.ShortBufferException;
+import javax.crypto.spec.OAEPParameterSpec;
+import javax.crypto.spec.PSource;
 import javax.crypto.spec.SecretKeySpec;
 
 /**
@@ -99,6 +102,8 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
      */
     private Exception mCachedException;
 
+    private Cipher mCipher;
+
     AndroidKeyStoreCipherSpiBase() {
         mOperation = null;
         mEncrypting = false;
@@ -110,6 +115,7 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
         mAdditionalAuthenticationDataStreamer = null;
         mAdditionalAuthenticationDataStreamerClosed = false;
         mCachedException = null;
+        mCipher = null;
     }
 
     @Override
@@ -117,6 +123,45 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
             throws InvalidKeyException {
         resetAll();
 
+        if (!(key instanceof AndroidKeyStorePrivateKey
+                || key instanceof AndroidKeyStoreSecretKey)) {
+            try {
+                mCipher = Cipher.getInstance(getTransform());
+                String transform = getTransform();
+
+                if ("RSA/ECB/OAEPWithSHA-224AndMGF1Padding".equals(transform)) {
+                    OAEPParameterSpec spec =
+                            new OAEPParameterSpec("SHA-224", "MGF1",
+                                    new MGF1ParameterSpec("SHA1"), PSource.PSpecified.DEFAULT);
+                    mCipher.init(opmode, key, spec, random);
+                } else if ("RSA/ECB/OAEPWithSHA-256AndMGF1Padding".equals(transform)) {
+                    OAEPParameterSpec spec =
+                            new OAEPParameterSpec("SHA-256", "MGF1",
+                                    new MGF1ParameterSpec("SHA1"), PSource.PSpecified.DEFAULT);
+                    mCipher.init(opmode, key, spec, random);
+
+                } else if ("RSA/ECB/OAEPWithSHA-384AndMGF1Padding".equals(transform)) {
+                    OAEPParameterSpec spec =
+                            new OAEPParameterSpec("SHA-384", "MGF1",
+                                    new MGF1ParameterSpec("SHA1"), PSource.PSpecified.DEFAULT);
+                    mCipher.init(opmode, key, spec, random);
+
+                } else if ("RSA/ECB/OAEPWithSHA-512AndMGF1Padding".equals(transform)) {
+                    OAEPParameterSpec spec =
+                            new OAEPParameterSpec("SHA-512", "MGF1",
+                                    new MGF1ParameterSpec("SHA1"), PSource.PSpecified.DEFAULT);
+                    mCipher.init(opmode, key, spec, random);
+                } else {
+                    mCipher.init(opmode, key, random);
+                }
+                return;
+            } catch (NoSuchAlgorithmException
+                    | NoSuchPaddingException
+                    | InvalidAlgorithmParameterException e) {
+                throw new InvalidKeyException(e);
+            }
+        }
+
         boolean success = false;
         try {
             init(opmode, key, random);
@@ -139,6 +184,17 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
             SecureRandom random) throws InvalidKeyException, InvalidAlgorithmParameterException {
         resetAll();
 
+        if (!(key instanceof AndroidKeyStorePrivateKey
+                || key instanceof AndroidKeyStoreSecretKey)) {
+            try {
+                mCipher = Cipher.getInstance(getTransform());
+                mCipher.init(opmode, key, params, random);
+                return;
+            } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
+                throw new InvalidKeyException(e);
+            }
+        }
+
         boolean success = false;
         try {
             init(opmode, key, random);
@@ -157,6 +213,17 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
             SecureRandom random) throws InvalidKeyException, InvalidAlgorithmParameterException {
         resetAll();
 
+        if (!(key instanceof AndroidKeyStorePrivateKey
+                || key instanceof AndroidKeyStoreSecretKey)) {
+            try {
+                mCipher = Cipher.getInstance(getTransform());
+                mCipher.init(opmode, key, params, random);
+                return;
+            } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
+                throw new InvalidKeyException(e);
+            }
+        }
+
         boolean success = false;
         try {
             init(opmode, key, random);
@@ -214,6 +281,7 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
         mAdditionalAuthenticationDataStreamer = null;
         mAdditionalAuthenticationDataStreamerClosed = false;
         mCachedException = null;
+        mCipher = null;
     }
 
     /**
@@ -320,6 +388,10 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
 
     @Override
     protected final byte[] engineUpdate(byte[] input, int inputOffset, int inputLen) {
+        if (mCipher != null) {
+            return mCipher.update(input, inputOffset, inputLen);
+        }
+
         if (mCachedException != null) {
             return null;
         }
@@ -371,6 +443,9 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
     @Override
     protected final int engineUpdate(byte[] input, int inputOffset, int inputLen, byte[] output,
             int outputOffset) throws ShortBufferException {
+        if (mCipher != null) {
+            return mCipher.update(input, inputOffset, inputLen, output);
+        }
         byte[] outputCopy = engineUpdate(input, inputOffset, inputLen);
         if (outputCopy == null) {
             return 0;
@@ -387,6 +462,10 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
     @Override
     protected final int engineUpdate(ByteBuffer input, ByteBuffer output)
             throws ShortBufferException {
+        if (mCipher != null) {
+            return mCipher.update(input, output);
+        }
+
         if (input == null) {
             throw new NullPointerException("input == null");
         }
@@ -423,6 +502,11 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
 
     @Override
     protected final void engineUpdateAAD(byte[] input, int inputOffset, int inputLen) {
+        if (mCipher != null) {
+            mCipher.updateAAD(input, inputOffset, inputLen);
+            return;
+        }
+
         if (mCachedException != null) {
             return;
         }
@@ -459,6 +543,11 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
 
     @Override
     protected final void engineUpdateAAD(ByteBuffer src) {
+        if (mCipher != null) {
+            mCipher.updateAAD(src);
+            return;
+        }
+
         if (src == null) {
             throw new IllegalArgumentException("src == null");
         }
@@ -486,6 +575,10 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
     @Override
     protected final byte[] engineDoFinal(byte[] input, int inputOffset, int inputLen)
             throws IllegalBlockSizeException, BadPaddingException {
+        if (mCipher != null) {
+            return mCipher.doFinal(input, inputOffset, inputLen);
+        }
+
         if (mCachedException != null) {
             throw (IllegalBlockSizeException)
                     new IllegalBlockSizeException().initCause(mCachedException);
@@ -522,6 +615,10 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
     protected final int engineDoFinal(byte[] input, int inputOffset, int inputLen, byte[] output,
             int outputOffset) throws ShortBufferException, IllegalBlockSizeException,
             BadPaddingException {
+        if (mCipher != null) {
+            return mCipher.doFinal(input, inputOffset, inputLen, output);
+        }
+
         byte[] outputCopy = engineDoFinal(input, inputOffset, inputLen);
         if (outputCopy == null) {
             return 0;
@@ -538,6 +635,10 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
     @Override
     protected final int engineDoFinal(ByteBuffer input, ByteBuffer output)
             throws ShortBufferException, IllegalBlockSizeException, BadPaddingException {
+        if (mCipher != null) {
+            return mCipher.doFinal(input, output);
+        }
+
         if (input == null) {
             throw new NullPointerException("input == null");
         }
@@ -575,6 +676,10 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
     @Override
     protected final byte[] engineWrap(Key key)
             throws IllegalBlockSizeException, InvalidKeyException {
+        if (mCipher != null) {
+            return mCipher.wrap(key);
+        }
+
         if (mKey == null) {
             throw new IllegalStateException("Not initilized");
         }
@@ -656,6 +761,10 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
     @Override
     protected final Key engineUnwrap(byte[] wrappedKey, String wrappedKeyAlgorithm,
             int wrappedKeyType) throws InvalidKeyException, NoSuchAlgorithmException {
+        if (mCipher != null) {
+            return mCipher.unwrap(wrappedKey, wrappedKeyAlgorithm, wrappedKeyType);
+        }
+
         if (mKey == null) {
             throw new IllegalStateException("Not initilized");
         }
@@ -902,4 +1011,6 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
      */
     protected abstract void loadAlgorithmSpecificParametersFromBeginResult(
             KeyParameter[] parameters);
+
+    protected abstract String getTransform();
 }
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreECDSASignatureSpi.java b/keystore/java/android/security/keystore2/AndroidKeyStoreECDSASignatureSpi.java
index d1ef1df817e6..8289671de212 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreECDSASignatureSpi.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreECDSASignatureSpi.java
@@ -44,6 +44,11 @@ abstract class AndroidKeyStoreECDSASignatureSpi extends AndroidKeyStoreSignature
         }
 
         @Override
+        protected String getAlgorithm() {
+            return "NONEwithECDSA";
+        }
+
+        @Override
         protected KeyStoreCryptoOperationStreamer createMainDataStreamer(
                 KeyStoreOperation operation) {
             return new TruncateToFieldSizeMessageStreamer(
@@ -113,30 +118,50 @@ abstract class AndroidKeyStoreECDSASignatureSpi extends AndroidKeyStoreSignature
         public SHA1() {
             super(KeymasterDefs.KM_DIGEST_SHA1);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA1withECDSA";
+        }
     }
 
     public final static class SHA224 extends AndroidKeyStoreECDSASignatureSpi {
         public SHA224() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_224);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA224withECDSA";
+        }
     }
 
     public final static class SHA256 extends AndroidKeyStoreECDSASignatureSpi {
         public SHA256() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_256);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA256withECDSA";
+        }
     }
 
     public final static class SHA384 extends AndroidKeyStoreECDSASignatureSpi {
         public SHA384() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_384);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA384withECDSA";
+        }
     }
 
     public final static class SHA512 extends AndroidKeyStoreECDSASignatureSpi {
         public SHA512() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_512);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA512withECDSA";
+        }
     }
 
     private final int mKeymasterDigest;
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreRSACipherSpi.java b/keystore/java/android/security/keystore2/AndroidKeyStoreRSACipherSpi.java
index 6546bb6654e2..6ff9432905ed 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreRSACipherSpi.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreRSACipherSpi.java
@@ -316,6 +316,25 @@ abstract class AndroidKeyStoreRSACipherSpi extends AndroidKeyStoreCipherSpiBase
         protected final int getAdditionalEntropyAmountForFinish() {
             return (isEncrypting()) ? mDigestOutputSizeBytes : 0;
         }
+
+        @Override
+        protected final String getTransform() {
+            switch (mKeymasterDigest) {
+                case KeymasterDefs.KM_DIGEST_SHA1:
+                    return "RSA/ECB/OAEPWithSHA-1AndMGF1Padding";
+                case KeymasterDefs.KM_DIGEST_SHA_2_224:
+                    return "RSA/ECB/OAEPWithSHA-224AndMGF1Padding";
+                case KeymasterDefs.KM_DIGEST_SHA_2_256:
+                    return "RSA/ECB/OAEPWithSHA-256AndMGF1Padding";
+                case KeymasterDefs.KM_DIGEST_SHA_2_384:
+                    return "RSA/ECB/OAEPWithSHA-384AndMGF1Padding";
+                case KeymasterDefs.KM_DIGEST_SHA_2_512:
+                    return "RSA/ECB/OAEPWithSHA-512AndMGF1Padding";
+                default:
+                    return "RSA/ECB/OAEPPadding";
+            }
+        }
+
     }
 
     public static class OAEPWithSHA1AndMGF1Padding extends OAEPWithMGF1Padding {
@@ -358,6 +377,11 @@ abstract class AndroidKeyStoreRSACipherSpi extends AndroidKeyStoreCipherSpiBase
     }
 
     @Override
+    protected String getTransform() {
+        return "RSA/ECB/" + KeyProperties.EncryptionPadding.fromKeymaster(mKeymasterPadding);
+    }
+
+    @Override
     protected final void initKey(int opmode, Key key) throws InvalidKeyException {
         if (key == null) {
             throw new InvalidKeyException("Unsupported key: null");
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreRSASignatureSpi.java b/keystore/java/android/security/keystore2/AndroidKeyStoreRSASignatureSpi.java
index ab7559116a41..931c2f864eba 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreRSASignatureSpi.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreRSASignatureSpi.java
@@ -48,42 +48,70 @@ abstract class AndroidKeyStoreRSASignatureSpi extends AndroidKeyStoreSignatureSp
         public NONEWithPKCS1Padding() {
             super(KeymasterDefs.KM_DIGEST_NONE);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "NONEwithRSA";
+        }
     }
 
     public static final class MD5WithPKCS1Padding extends PKCS1Padding {
         public MD5WithPKCS1Padding() {
             super(KeymasterDefs.KM_DIGEST_MD5);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "MD5withRSA";
+        }
     }
 
     public static final class SHA1WithPKCS1Padding extends PKCS1Padding {
         public SHA1WithPKCS1Padding() {
             super(KeymasterDefs.KM_DIGEST_SHA1);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA1withRSA";
+        }
     }
 
     public static final class SHA224WithPKCS1Padding extends PKCS1Padding {
         public SHA224WithPKCS1Padding() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_224);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA224withRSA";
+        }
     }
 
     public static final class SHA256WithPKCS1Padding extends PKCS1Padding {
         public SHA256WithPKCS1Padding() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_256);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA256withRSA";
+        }
     }
 
     public static final class SHA384WithPKCS1Padding extends PKCS1Padding {
         public SHA384WithPKCS1Padding() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_384);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA384withRSA";
+        }
     }
 
     public static final class SHA512WithPKCS1Padding extends PKCS1Padding {
         public SHA512WithPKCS1Padding() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_512);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA512withRSA";
+        }
     }
 
     abstract static class PSSPadding extends AndroidKeyStoreRSASignatureSpi {
@@ -103,30 +131,50 @@ abstract class AndroidKeyStoreRSASignatureSpi extends AndroidKeyStoreSignatureSp
         public SHA1WithPSSPadding() {
             super(KeymasterDefs.KM_DIGEST_SHA1);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA1withRSA/PSS";
+        }
     }
 
     public static final class SHA224WithPSSPadding extends PSSPadding {
         public SHA224WithPSSPadding() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_224);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA224withRSA/PSS";
+        }
     }
 
     public static final class SHA256WithPSSPadding extends PSSPadding {
         public SHA256WithPSSPadding() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_256);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA256withRSA/PSS";
+        }
     }
 
     public static final class SHA384WithPSSPadding extends PSSPadding {
         public SHA384WithPSSPadding() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_384);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA384withRSA/PSS";
+        }
     }
 
     public static final class SHA512WithPSSPadding extends PSSPadding {
         public SHA512WithPSSPadding() {
             super(KeymasterDefs.KM_DIGEST_SHA_2_512);
         }
+        @Override
+        protected String getAlgorithm() {
+            return "SHA512withRSA/PSS";
+        }
     }
 
     private final int mKeymasterDigest;
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreSignatureSpiBase.java b/keystore/java/android/security/keystore2/AndroidKeyStoreSignatureSpiBase.java
index 9b4f01e744f7..96da1e00ade8 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreSignatureSpiBase.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreSignatureSpiBase.java
@@ -30,10 +30,12 @@ import libcore.util.EmptyArray;
 import java.nio.ByteBuffer;
 import java.security.InvalidKeyException;
 import java.security.InvalidParameterException;
+import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
 import java.security.ProviderException;
 import java.security.PublicKey;
 import java.security.SecureRandom;
+import java.security.Signature;
 import java.security.SignatureException;
 import java.security.SignatureSpi;
 import java.util.ArrayList;
@@ -76,6 +78,13 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
      */
     private Exception mCachedException;
 
+    /**
+     * This signature object is used for public key operations, i.e, signatrue verification.
+     * The Android Keystore backend does not perform public key operations and defers to the
+     * Highest priority provider.
+     */
+    private Signature mSignature;
+
     AndroidKeyStoreSignatureSpiBase() {
         mOperation = null;
         mOperationChallenge = 0;
@@ -84,6 +93,7 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
         appRandom = null;
         mMessageStreamer = null;
         mCachedException = null;
+        mSignature = null;
     }
 
     @Override
@@ -123,27 +133,13 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
     protected final void engineInitVerify(PublicKey publicKey) throws InvalidKeyException {
         resetAll();
 
-        boolean success = false;
         try {
-            if (publicKey == null) {
-                throw new InvalidKeyException("Unsupported key: null");
-            }
-            AndroidKeyStoreKey keystoreKey;
-            if (publicKey instanceof AndroidKeyStorePublicKey) {
-                keystoreKey = (AndroidKeyStorePublicKey) publicKey;
-            } else {
-                throw new InvalidKeyException("Unsupported public key type: " + publicKey);
-            }
-            mSigning = false;
-            initKey(keystoreKey);
-            appRandom = null;
-            ensureKeystoreOperationInitialized();
-            success = true;
-        } finally {
-            if (!success) {
-                resetAll();
-            }
+            mSignature = Signature.getInstance(getAlgorithm());
+        } catch (NoSuchAlgorithmException e) {
+            throw new InvalidKeyException(e);
         }
+
+        mSignature.initVerify(publicKey);
     }
 
     /**
@@ -251,6 +247,11 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
 
     @Override
     protected final void engineUpdate(byte[] b, int off, int len) throws SignatureException {
+        if (mSignature != null) {
+            mSignature.update(b, off, len);
+            return;
+        }
+
         if (mCachedException != null) {
             throw new SignatureException(mCachedException);
         }
@@ -337,39 +338,10 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
 
     @Override
     protected final boolean engineVerify(byte[] signature) throws SignatureException {
-        if (mCachedException != null) {
-            throw new SignatureException(mCachedException);
-        }
-
-        try {
-            ensureKeystoreOperationInitialized();
-        } catch (InvalidKeyException e) {
-            throw new SignatureException(e);
-        }
-
-        boolean verified;
-        try {
-            byte[] output = mMessageStreamer.doFinal(
-                    EmptyArray.BYTE, 0, 0,
-                    signature);
-            if (output.length != 0) {
-                throw new ProviderException(
-                        "Signature verification unexpected produced output: " + output.length
-                        + " bytes");
-            }
-            verified = true;
-        } catch (KeyStoreException e) {
-            switch (e.getErrorCode()) {
-                case KeymasterDefs.KM_ERROR_VERIFICATION_FAILED:
-                    verified = false;
-                    break;
-                default:
-                    throw new SignatureException(e);
-            }
+        if (mSignature != null) {
+            return mSignature.verify(signature);
         }
-
-        resetWhilePreservingInitState();
-        return verified;
+        throw new IllegalStateException("Not initialised.");
     }
 
     @Override
@@ -392,6 +364,13 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
     }
 
     /**
+     * Implementations need to report the algorithm they implement so that we can delegate to the
+     * highest priority provider.
+     * @return Algorithm string.
+     */
+    protected abstract String getAlgorithm();
+
+    /**
      * Returns {@code true} if this signature is initialized for signing, {@code false} if this
      * signature is initialized for verification.
      */
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreUnauthenticatedAESCipherSpi.java b/keystore/java/android/security/keystore2/AndroidKeyStoreUnauthenticatedAESCipherSpi.java
index 4d4b0d8f183b..fd3d28976b2e 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreUnauthenticatedAESCipherSpi.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreUnauthenticatedAESCipherSpi.java
@@ -42,7 +42,7 @@ import javax.crypto.spec.IvParameterSpec;
  *
  * @hide
  */
-class AndroidKeyStoreUnauthenticatedAESCipherSpi extends AndroidKeyStoreCipherSpiBase {
+abstract class AndroidKeyStoreUnauthenticatedAESCipherSpi extends AndroidKeyStoreCipherSpiBase {
 
     abstract static class ECB extends AndroidKeyStoreUnauthenticatedAESCipherSpi {
         protected ECB(int keymasterPadding) {
@@ -53,12 +53,22 @@ class AndroidKeyStoreUnauthenticatedAESCipherSpi extends AndroidKeyStoreCipherSp
             public NoPadding() {
                 super(KeymasterDefs.KM_PAD_NONE);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "AES/ECB/NoPadding";
+            }
         }
 
         public static class PKCS7Padding extends ECB {
             public PKCS7Padding() {
                 super(KeymasterDefs.KM_PAD_PKCS7);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "AES/ECB/PKCS7Padding";
+            }
         }
     }
 
@@ -71,12 +81,22 @@ class AndroidKeyStoreUnauthenticatedAESCipherSpi extends AndroidKeyStoreCipherSp
             public NoPadding() {
                 super(KeymasterDefs.KM_PAD_NONE);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "AES/CBC/NoPadding";
+            }
         }
 
         public static class PKCS7Padding extends CBC {
             public PKCS7Padding() {
                 super(KeymasterDefs.KM_PAD_PKCS7);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "AES/CBC/PKCS7Padding";
+            }
         }
     }
 
@@ -89,6 +109,11 @@ class AndroidKeyStoreUnauthenticatedAESCipherSpi extends AndroidKeyStoreCipherSp
             public NoPadding() {
                 super(KeymasterDefs.KM_PAD_NONE);
             }
+
+            @Override
+            protected final String getTransform() {
+                return "AES/CTR/NoPadding";
+            }
         }
     }