Merge "Make KeyStoreCryptoOperationChunkedStreamer lazy."

am: c13f012cc5

Change-Id: I8387f3e35250fd814f010c8b649caa06e4259551

4 files changed, 115 insertions, 204 deletions

diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreAuthenticatedAESCipherSpi.java b/keystore/java/android/security/keystore/AndroidKeyStoreAuthenticatedAESCipherSpi.java
index c6515efd2c61..feb6101cb080 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreAuthenticatedAESCipherSpi.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreAuthenticatedAESCipherSpi.java
@@ -176,7 +176,7 @@ abstract class AndroidKeyStoreAuthenticatedAESCipherSpi extends AndroidKeyStoreC
                 KeyStore keyStore, IBinder operationToken) {
             KeyStoreCryptoOperationStreamer streamer = new KeyStoreCryptoOperationChunkedStreamer(
                     new KeyStoreCryptoOperationChunkedStreamer.MainDataStream(
-                            keyStore, operationToken));
+                            keyStore, operationToken), 0);
             if (isEncrypting()) {
                 return streamer;
             } else {
@@ -191,7 +191,7 @@ abstract class AndroidKeyStoreAuthenticatedAESCipherSpi extends AndroidKeyStoreC
         protected final KeyStoreCryptoOperationStreamer createAdditionalAuthenticationDataStreamer(
                 KeyStore keyStore, IBinder operationToken) {
             return new KeyStoreCryptoOperationChunkedStreamer(
-                    new AdditionalAuthenticationDataStream(keyStore, operationToken));
+                    new AdditionalAuthenticationDataStream(keyStore, operationToken), 0);
         }
 
         @Override
diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreCipherSpiBase.java b/keystore/java/android/security/keystore/AndroidKeyStoreCipherSpiBase.java
index 5bcb34a67a54..ccc3153749fb 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreCipherSpiBase.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreCipherSpiBase.java
@@ -299,7 +299,7 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
             KeyStore keyStore, IBinder operationToken) {
         return new KeyStoreCryptoOperationChunkedStreamer(
                 new KeyStoreCryptoOperationChunkedStreamer.MainDataStream(
-                        keyStore, operationToken));
+                        keyStore, operationToken), 0);
     }
 
     /**
diff --git a/keystore/java/android/security/keystore/ArrayUtils.java b/keystore/java/android/security/keystore/ArrayUtils.java
index 26172d276b8c..f519c7cdd3d2 100644
--- a/keystore/java/android/security/keystore/ArrayUtils.java
+++ b/keystore/java/android/security/keystore/ArrayUtils.java
@@ -55,6 +55,34 @@ public abstract class ArrayUtils {
         }
     }
 
+    /**
+     * Copies a subset of the source array to the destination array.
+     * Length will be limited to the bounds of source and destination arrays.
+     * The length actually copied is returned, which will be <= length argument.
+     * @param src is the source array
+     * @param srcOffset is the offset in the source array.
+     * @param dst is the destination array.
+     * @param dstOffset is the offset in the destination array.
+     * @param length is the length to be copied from source to destination array.
+     * @return The length actually copied from source array.
+     */
+    public static int copy(byte[] src, int srcOffset, byte[] dst, int dstOffset, int length) {
+        if (dst == null || src == null) {
+            return 0;
+        }
+        if (length > dst.length - dstOffset) {
+            length = dst.length - dstOffset;
+        }
+        if (length > src.length - srcOffset) {
+            length = src.length - srcOffset;
+        }
+        if (length <= 0) {
+            return 0;
+        }
+        System.arraycopy(src, srcOffset, dst, dstOffset, length);
+        return length;
+    }
+
     public static byte[] subarray(byte[] arr, int offset, int len) {
         if (len == 0) {
             return EmptyArray.BYTE;
diff --git a/keystore/java/android/security/keystore/KeyStoreCryptoOperationChunkedStreamer.java b/keystore/java/android/security/keystore/KeyStoreCryptoOperationChunkedStreamer.java
index 75bea26aecef..2c0f40d528d2 100644
--- a/keystore/java/android/security/keystore/KeyStoreCryptoOperationChunkedStreamer.java
+++ b/keystore/java/android/security/keystore/KeyStoreCryptoOperationChunkedStreamer.java
@@ -24,19 +24,20 @@ import android.security.keymaster.OperationResult;
 
 import libcore.util.EmptyArray;
 
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.security.ProviderException;
-
 /**
  * Helper for streaming a crypto operation's input and output via {@link KeyStore} service's
  * {@code update} and {@code finish} operations.
  *
- * <p>The helper abstracts away to issues that need to be solved in most code that uses KeyStore's
+ * <p>The helper abstracts away issues that need to be solved in most code that uses KeyStore's
  * update and finish operations. Firstly, KeyStore's update operation can consume only a limited
  * amount of data in one go because the operations are marshalled via Binder. Secondly, the update
  * operation may consume less data than provided, in which case the caller has to buffer the
- * remainder for next time. The helper exposes {@link #update(byte[], int, int) update} and
+ * remainder for next time. Thirdly, when the input is smaller than a threshold, skipping update
+ * and passing input data directly to final improves performance. This threshold is configurable;
+ * using a threshold <= 1 causes the helper act eagerly, which may be required for some types of
+ * operations (e.g. ciphers).
+ *
+ * <p>The helper exposes {@link #update(byte[], int, int) update} and
  * {@link #doFinal(byte[], int, int, byte[], byte[]) doFinal} operations which can be used to
  * conveniently implement various JCA crypto primitives.
  *
@@ -67,240 +68,122 @@ class KeyStoreCryptoOperationChunkedStreamer implements KeyStoreCryptoOperationS
 
     // Binder buffer is about 1MB, but it's shared between all active transactions of the process.
     // Thus, it's safer to use a much smaller upper bound.
-    private static final int DEFAULT_MAX_CHUNK_SIZE = 64 * 1024;
+    private static final int DEFAULT_CHUNK_SIZE_MAX = 64 * 1024;
+    // The chunk buffer will be sent to update until its size under this threshold.
+    // This threshold should be <= the max input allowed for finish.
+    // Setting this threshold <= 1 will effectivley disable buffering between updates.
+    private static final int DEFAULT_CHUNK_SIZE_THRESHOLD = 2 * 1024;
 
     private final Stream mKeyStoreStream;
-    private final int mMaxChunkSize;
-
-    private byte[] mBuffered = EmptyArray.BYTE;
-    private int mBufferedOffset;
-    private int mBufferedLength;
+    private final int mChunkSizeMax;
+    private final int mChunkSizeThreshold;
+    private final byte[] mChunk;
+    private int mChunkLength = 0;
     private long mConsumedInputSizeBytes;
     private long mProducedOutputSizeBytes;
 
-    public KeyStoreCryptoOperationChunkedStreamer(Stream operation) {
-        this(operation, DEFAULT_MAX_CHUNK_SIZE);
+    KeyStoreCryptoOperationChunkedStreamer(Stream operation) {
+        this(operation, DEFAULT_CHUNK_SIZE_THRESHOLD, DEFAULT_CHUNK_SIZE_MAX);
     }
 
-    public KeyStoreCryptoOperationChunkedStreamer(Stream operation, int maxChunkSize) {
+    KeyStoreCryptoOperationChunkedStreamer(Stream operation, int chunkSizeThreshold) {
+        this(operation, chunkSizeThreshold, DEFAULT_CHUNK_SIZE_MAX);
+    }
+
+    KeyStoreCryptoOperationChunkedStreamer(Stream operation, int chunkSizeThreshold,
+            int chunkSizeMax) {
         mKeyStoreStream = operation;
-        mMaxChunkSize = maxChunkSize;
+        mChunkSizeMax = chunkSizeMax;
+        if (chunkSizeThreshold <= 0) {
+            mChunkSizeThreshold = 1;
+        } else if (chunkSizeThreshold > chunkSizeMax) {
+            mChunkSizeThreshold = chunkSizeMax;
+        } else {
+            mChunkSizeThreshold = chunkSizeThreshold;
+        }
+        mChunk = new byte[mChunkSizeMax];
     }
 
-    @Override
     public byte[] update(byte[] input, int inputOffset, int inputLength) throws KeyStoreException {
-        if (inputLength == 0) {
+        if (inputLength == 0 || input == null) {
             // No input provided
             return EmptyArray.BYTE;
         }
+        if (inputLength < 0 || inputOffset < 0 || (inputOffset + inputLength) > input.length) {
+            throw new KeyStoreException(KeymasterDefs.KM_ERROR_UNKNOWN_ERROR,
+                "Input offset and length out of bounds of input array");
+        }
 
-        ByteArrayOutputStream bufferedOutput = null;
+        byte[] output = EmptyArray.BYTE;
 
-        while (inputLength > 0) {
-            byte[] chunk;
-            int inputBytesInChunk;
-            if ((mBufferedLength + inputLength) > mMaxChunkSize) {
-                // Too much input for one chunk -- extract one max-sized chunk and feed it into the
-                // update operation.
-                inputBytesInChunk = mMaxChunkSize - mBufferedLength;
-                chunk = ArrayUtils.concat(mBuffered, mBufferedOffset, mBufferedLength,
-                        input, inputOffset, inputBytesInChunk);
-            } else {
-                // All of available input fits into one chunk.
-                if ((mBufferedLength == 0) && (inputOffset == 0)
-                        && (inputLength == input.length)) {
-                    // Nothing buffered and all of input array needs to be fed into the update
-                    // operation.
-                    chunk = input;
-                    inputBytesInChunk = input.length;
-                } else {
-                    // Need to combine buffered data with input data into one array.
-                    inputBytesInChunk = inputLength;
-                    chunk = ArrayUtils.concat(mBuffered, mBufferedOffset, mBufferedLength,
-                            input, inputOffset, inputBytesInChunk);
-                }
-            }
-            // Update input array references to reflect that some of its bytes are now in mBuffered.
-            inputOffset += inputBytesInChunk;
-            inputLength -= inputBytesInChunk;
-            mConsumedInputSizeBytes += inputBytesInChunk;
+        while (inputLength > 0 || mChunkLength >= mChunkSizeThreshold) {
+            int inputConsumed = ArrayUtils.copy(input, inputOffset, mChunk, mChunkLength,
+                    inputLength);
+            inputLength -= inputConsumed;
+            inputOffset += inputConsumed;
+            mChunkLength += inputConsumed;
+            mConsumedInputSizeBytes += inputConsumed;
 
-            OperationResult opResult = mKeyStoreStream.update(chunk);
-            if (opResult == null) {
-                throw new KeyStoreConnectException();
-            } else if (opResult.resultCode != KeyStore.NO_ERROR) {
-                throw KeyStore.getKeyStoreException(opResult.resultCode);
+            if (mChunkLength > mChunkSizeMax) {
+                throw new KeyStoreException(KeymasterDefs.KM_ERROR_INVALID_INPUT_LENGTH,
+                    "Chunk size exceeded max chunk size. Max: " + mChunkSizeMax
+                    + " Actual: " + mChunkLength);
             }
 
-            if (opResult.inputConsumed == chunk.length) {
-                // The whole chunk was consumed
-                mBuffered = EmptyArray.BYTE;
-                mBufferedOffset = 0;
-                mBufferedLength = 0;
-            } else if (opResult.inputConsumed <= 0) {
-                // Nothing was consumed. More input needed.
-                if (inputLength > 0) {
-                    // More input is available, but it wasn't included into the previous chunk
-                    // because the chunk reached its maximum permitted size.
-                    // Shouldn't have happened.
+            if (mChunkLength >= mChunkSizeThreshold) {
+                OperationResult opResult = mKeyStoreStream.update(
+                        ArrayUtils.subarray(mChunk, 0, mChunkLength));
+
+                if (opResult == null) {
+                    throw new KeyStoreConnectException();
+                } else if (opResult.resultCode != KeyStore.NO_ERROR) {
+                    throw KeyStore.getKeyStoreException(opResult.resultCode);
+                }
+                if (opResult.inputConsumed <= 0) {
+                    throw new KeyStoreException(KeymasterDefs.KM_ERROR_INVALID_INPUT_LENGTH,
+                        "Keystore consumed 0 of " + mChunkLength + " bytes provided.");
+                } else if (opResult.inputConsumed > mChunkLength) {
                     throw new KeyStoreException(KeymasterDefs.KM_ERROR_UNKNOWN_ERROR,
-                            "Keystore consumed nothing from max-sized chunk: " + chunk.length
-                                    + " bytes");
+                        "Keystore consumed more input than provided. Provided: "
+                            + mChunkLength + ", consumed: " + opResult.inputConsumed);
                 }
-                mBuffered = chunk;
-                mBufferedOffset = 0;
-                mBufferedLength = chunk.length;
-            } else if (opResult.inputConsumed < chunk.length) {
-                // The chunk was consumed only partially -- buffer the rest of the chunk
-                mBuffered = chunk;
-                mBufferedOffset = opResult.inputConsumed;
-                mBufferedLength = chunk.length - opResult.inputConsumed;
-            } else {
-                throw new KeyStoreException(KeymasterDefs.KM_ERROR_UNKNOWN_ERROR,
-                        "Keystore consumed more input than provided. Provided: " + chunk.length
-                                + ", consumed: " + opResult.inputConsumed);
-            }
+                mChunkLength -= opResult.inputConsumed;
 
-            if ((opResult.output != null) && (opResult.output.length > 0)) {
-                if (inputLength + mBufferedLength > 0) {
-                    // More output might be produced in this loop -- buffer the current output
-                    if (bufferedOutput == null) {
-                        bufferedOutput = new ByteArrayOutputStream();
-                    }
-                    try {
-                        bufferedOutput.write(opResult.output);
-                    } catch (IOException e) {
-                        throw new ProviderException("Failed to buffer output", e);
-                    }
-                } else {
-                    // No more output will be produced in this loop
-                    byte[] result;
-                    if (bufferedOutput == null) {
-                        // No previously buffered output
-                        result = opResult.output;
-                    } else {
-                        // There was some previously buffered output
-                        try {
-                            bufferedOutput.write(opResult.output);
-                        } catch (IOException e) {
-                            throw new ProviderException("Failed to buffer output", e);
-                        }
-                        result = bufferedOutput.toByteArray();
-                    }
-                    mProducedOutputSizeBytes += result.length;
-                    return result;
+                if (mChunkLength > 0) {
+                    // Partialy consumed, shift chunk contents
+                    ArrayUtils.copy(mChunk, opResult.inputConsumed, mChunk, 0, mChunkLength);
                 }
-            }
-        }
 
-        byte[] result;
-        if (bufferedOutput == null) {
-            // No output produced
-            result = EmptyArray.BYTE;
-        } else {
-            result = bufferedOutput.toByteArray();
+                if ((opResult.output != null) && (opResult.output.length > 0)) {
+                    // Output was produced
+                    mProducedOutputSizeBytes += opResult.output.length;
+                    output = ArrayUtils.concat(output, opResult.output);
+                }
+            }
         }
-        mProducedOutputSizeBytes += result.length;
-        return result;
+        return output;
     }
 
-    @Override
     public byte[] doFinal(byte[] input, int inputOffset, int inputLength,
             byte[] signature, byte[] additionalEntropy) throws KeyStoreException {
-        if (inputLength == 0) {
-            // No input provided -- simplify the rest of the code
-            input = EmptyArray.BYTE;
-            inputOffset = 0;
-        }
-
-        // Flush all buffered input and provided input into keystore/keymaster.
         byte[] output = update(input, inputOffset, inputLength);
-        output = ArrayUtils.concat(output, flush());
+        byte[] finalChunk = ArrayUtils.subarray(mChunk, 0, mChunkLength);
+        OperationResult opResult = mKeyStoreStream.finish(finalChunk, signature, additionalEntropy);
 
-        OperationResult opResult = mKeyStoreStream.finish(EmptyArray.BYTE, signature,
-                                                          additionalEntropy);
         if (opResult == null) {
             throw new KeyStoreConnectException();
         } else if (opResult.resultCode != KeyStore.NO_ERROR) {
             throw KeyStore.getKeyStoreException(opResult.resultCode);
         }
-        mProducedOutputSizeBytes += opResult.output.length;
-
-        return ArrayUtils.concat(output, opResult.output);
-    }
-
-    public byte[] flush() throws KeyStoreException {
-        if (mBufferedLength <= 0) {
-            return EmptyArray.BYTE;
-        }
-
-        // Keep invoking the update operation with remaining buffered data until either all of the
-        // buffered data is consumed or until update fails to consume anything.
-        ByteArrayOutputStream bufferedOutput = null;
-        while (mBufferedLength > 0) {
-            byte[] chunk = ArrayUtils.subarray(mBuffered, mBufferedOffset, mBufferedLength);
-            OperationResult opResult = mKeyStoreStream.update(chunk);
-            if (opResult == null) {
-                throw new KeyStoreConnectException();
-            } else if (opResult.resultCode != KeyStore.NO_ERROR) {
-                throw KeyStore.getKeyStoreException(opResult.resultCode);
-            }
-
-            if (opResult.inputConsumed <= 0) {
-                // Nothing was consumed. Break out of the loop to avoid an infinite loop.
-                break;
-            }
-
-            if (opResult.inputConsumed >= chunk.length) {
-                // All of the input was consumed
-                mBuffered = EmptyArray.BYTE;
-                mBufferedOffset = 0;
-                mBufferedLength = 0;
-            } else {
-                // Some of the input was not consumed
-                mBuffered = chunk;
-                mBufferedOffset = opResult.inputConsumed;
-                mBufferedLength = chunk.length - opResult.inputConsumed;
-            }
-
-            if (opResult.inputConsumed > chunk.length) {
-                throw new KeyStoreException(KeymasterDefs.KM_ERROR_UNKNOWN_ERROR,
-                        "Keystore consumed more input than provided. Provided: "
-                                + chunk.length + ", consumed: " + opResult.inputConsumed);
-            }
-
-            if ((opResult.output != null) && (opResult.output.length > 0)) {
-                // Some output was produced by this update operation
-                if (bufferedOutput == null) {
-                    // No output buffered yet.
-                    if (mBufferedLength == 0) {
-                        // No more output will be produced by this flush operation
-                        mProducedOutputSizeBytes += opResult.output.length;
-                        return opResult.output;
-                    } else {
-                        // More output might be produced by this flush operation -- buffer output.
-                        bufferedOutput = new ByteArrayOutputStream();
-                    }
-                }
-                // Buffer the output from this update operation
-                try {
-                    bufferedOutput.write(opResult.output);
-                } catch (IOException e) {
-                    throw new ProviderException("Failed to buffer output", e);
-                }
-            }
-        }
+        // If no error, assume all input consumed
+        mConsumedInputSizeBytes += finalChunk.length;
 
-        if (mBufferedLength > 0) {
-            throw new KeyStoreException(KeymasterDefs.KM_ERROR_INVALID_INPUT_LENGTH,
-                    "Keystore failed to consume last "
-                            + ((mBufferedLength != 1) ? (mBufferedLength + " bytes") : "byte")
-                            + " of input");
+        if ((opResult.output != null) && (opResult.output.length > 0)) {
+            mProducedOutputSizeBytes += opResult.output.length;
+            output = ArrayUtils.concat(output, opResult.output);
         }
 
-        byte[] result = (bufferedOutput != null) ? bufferedOutput.toByteArray() : EmptyArray.BYTE;
-        mProducedOutputSizeBytes += result.length;
-        return result;
+        return output;
     }
 
     @Override