1 files changed, 39 insertions, 19 deletions
diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreSignatureSpiBase.java b/keystore/java/android/security/keystore/AndroidKeyStoreSignatureSpiBase.java
index 5cdcc41..76240dd 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreSignatureSpiBase.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreSignatureSpiBase.java
@@ -58,7 +58,7 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
      */
     private IBinder mOperationToken;
     private long mOperationHandle;
-    private KeyStoreCryptoOperationChunkedStreamer mMessageStreamer;
+    private KeyStoreCryptoOperationStreamer mMessageStreamer;
 
     /**
      * Encountered exception which could not be immediately thrown because it was encountered inside
@@ -229,9 +229,20 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
             throw new ProviderException("Keystore returned invalid operation handle");
         }
 
-        mMessageStreamer = new KeyStoreCryptoOperationChunkedStreamer(
+        mMessageStreamer = createMainDataStreamer(mKeyStore, opResult.token);
+    }
+
+    /**
+     * Creates a streamer which sends the message to be signed/verified into the provided KeyStore
+     *
+     * <p>This implementation returns a working streamer.
+     */
+    @NonNull
+    protected KeyStoreCryptoOperationStreamer createMainDataStreamer(
+            KeyStore keyStore, IBinder operationToken) {
+        return new KeyStoreCryptoOperationChunkedStreamer(
                 new KeyStoreCryptoOperationChunkedStreamer.MainDataStream(
-                        mKeyStore, opResult.token));
+                        keyStore, operationToken));
     }
 
     @Override
@@ -314,7 +325,10 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
             byte[] additionalEntropy =
                     KeyStoreCryptoOperationUtils.getRandomBytesToMixIntoKeystoreRng(
                             appRandom, getAdditionalEntropyAmountForSign());
-            signature = mMessageStreamer.doFinal(EmptyArray.BYTE, 0, 0, additionalEntropy);
+            signature = mMessageStreamer.doFinal(
+                    EmptyArray.BYTE, 0, 0,
+                    null, // no signature provided -- it'll be generated by this invocation
+                    additionalEntropy);
         } catch (InvalidKeyException | KeyStoreException e) {
             throw new SignatureException(e);
         }
@@ -329,31 +343,37 @@ abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
             throw new SignatureException(mCachedException);
         }
 
-        boolean result;
         try {
             ensureKeystoreOperationInitialized();
-            mMessageStreamer.flush();
-            OperationResult opResult = mKeyStore.finish(mOperationToken, null, signature);
-            if (opResult == null) {
-                throw new KeyStoreConnectException();
+        } catch (InvalidKeyException e) {
+            throw new SignatureException(e);
+        }
+
+        boolean verified;
+        try {
+            byte[] output = mMessageStreamer.doFinal(
+                    EmptyArray.BYTE, 0, 0,
+                    signature,
+                    null // no additional entropy needed -- verification is deterministic
+                    );
+            if (output.length != 0) {
+                throw new ProviderException(
+                        "Signature verification unexpected produced output: " + output.length
+                        + " bytes");
             }
-            switch (opResult.resultCode) {
-                case KeyStore.NO_ERROR:
-                    result = true;
-                    break;
+            verified = true;
+        } catch (KeyStoreException e) {
+            switch (e.getErrorCode()) {
                 case KeymasterDefs.KM_ERROR_VERIFICATION_FAILED:
-                    result = false;
+                    verified = false;
                     break;
                 default:
-                    throw new SignatureException(
-                            KeyStore.getKeyStoreException(opResult.resultCode));
+                    throw new SignatureException(e);
             }
-        } catch (InvalidKeyException | KeyStoreException e) {
-            throw new SignatureException(e);
         }
 
         resetWhilePreservingInitState();
-        return result;
+        return verified;
     }
 
     @Override