Expose AES GCM backed by Android Keystore.

Bug: 18088752
Bug: 21786749
Change-Id: Ica90491037d2920f7635195894ba18882fc4406d

7 files changed, 617 insertions, 11 deletions

diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreAuthenticatedAESCipherSpi.java b/keystore/java/android/security/keystore/AndroidKeyStoreAuthenticatedAESCipherSpi.java
new file mode 100644
index 0000000..f412743
--- /dev/null
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreAuthenticatedAESCipherSpi.java
@@ -0,0 +1,442 @@
+/*
+ * 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.keystore;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.IBinder;
+import android.security.KeyStore;
+import android.security.KeyStoreException;
+import android.security.keymaster.KeymasterArguments;
+import android.security.keymaster.KeymasterDefs;
+import android.security.keymaster.OperationResult;
+import android.security.keystore.KeyStoreCryptoOperationChunkedStreamer.Stream;
+
+import libcore.util.EmptyArray;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.security.AlgorithmParameters;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.InvalidKeyException;
+import java.security.Key;
+import java.security.NoSuchAlgorithmException;
+import java.security.ProviderException;
+import java.security.spec.AlgorithmParameterSpec;
+import java.security.spec.InvalidParameterSpecException;
+import java.util.Arrays;
+
+import javax.crypto.CipherSpi;
+import javax.crypto.spec.GCMParameterSpec;
+
+/**
+ * Base class for Android Keystore authenticated AES {@link CipherSpi} implementations.
+ *
+ * @hide
+ */
+abstract class AndroidKeyStoreAuthenticatedAESCipherSpi extends AndroidKeyStoreCipherSpiBase {
+
+    abstract static class GCM extends AndroidKeyStoreAuthenticatedAESCipherSpi {
+        private static final int MIN_SUPPORTED_TAG_LENGTH_BITS = 96;
+        private static final int MAX_SUPPORTED_TAG_LENGTH_BITS = 128;
+        private static final int DEFAULT_TAG_LENGTH_BITS = 128;
+        private static final int IV_LENGTH_BYTES = 12;
+
+        private int mTagLengthBits = DEFAULT_TAG_LENGTH_BITS;
+
+        GCM(int keymasterPadding) {
+            super(KeymasterDefs.KM_MODE_GCM, keymasterPadding);
+        }
+
+        @Override
+        protected final void resetAll() {
+            mTagLengthBits = DEFAULT_TAG_LENGTH_BITS;
+            super.resetAll();
+        }
+
+        @Override
+        protected final void resetWhilePreservingInitState() {
+            super.resetWhilePreservingInitState();
+        }
+
+        @Override
+        protected final void initAlgorithmSpecificParameters() throws InvalidKeyException {
+            if (!isEncrypting()) {
+                throw new InvalidKeyException("IV required when decrypting"
+                        + ". Use IvParameterSpec or AlgorithmParameters to provide it.");
+            }
+        }
+
+        @Override
+        protected final void initAlgorithmSpecificParameters(AlgorithmParameterSpec params)
+                throws InvalidAlgorithmParameterException {
+            // IV is used
+            if (params == null) {
+                if (!isEncrypting()) {
+                    // IV must be provided by the caller
+                    throw new InvalidAlgorithmParameterException(
+                            "GCMParameterSpec must be provided when decrypting");
+                }
+                return;
+            }
+            if (!(params instanceof GCMParameterSpec)) {
+                throw new InvalidAlgorithmParameterException("Only GCMParameterSpec supported");
+            }
+            GCMParameterSpec spec = (GCMParameterSpec) params;
+            byte[] iv = spec.getIV();
+            if (iv == null) {
+                throw new InvalidAlgorithmParameterException("Null IV in GCMParameterSpec");
+            } else if (iv.length != IV_LENGTH_BYTES) {
+                throw new InvalidAlgorithmParameterException("Unsupported IV length: "
+                        + iv.length + " bytes. Only " + IV_LENGTH_BYTES
+                        + " bytes long IV supported");
+            }
+            int tagLengthBits = spec.getTLen();
+            if ((tagLengthBits < MIN_SUPPORTED_TAG_LENGTH_BITS)
+                    || (tagLengthBits > MAX_SUPPORTED_TAG_LENGTH_BITS)
+                    || ((tagLengthBits % 8) != 0)) {
+                throw new InvalidAlgorithmParameterException(
+                        "Unsupported tag length: " + tagLengthBits + " bits"
+                        + ". Supported lengths: 96, 104, 112, 120, 128");
+            }
+            setIv(iv);
+            mTagLengthBits = tagLengthBits;
+        }
+
+        @Override
+        protected final void initAlgorithmSpecificParameters(AlgorithmParameters params)
+                throws InvalidAlgorithmParameterException {
+            if (params == null) {
+                if (!isEncrypting()) {
+                    // IV must be provided by the caller
+                    throw new InvalidAlgorithmParameterException("IV required when decrypting"
+                            + ". Use GCMParameterSpec or GCM AlgorithmParameters to provide it.");
+                }
+                return;
+            }
+
+            GCMParameterSpec spec;
+            try {
+                spec = params.getParameterSpec(GCMParameterSpec.class);
+            } catch (InvalidParameterSpecException e) {
+                if (!isEncrypting()) {
+                    // IV must be provided by the caller
+                    throw new InvalidAlgorithmParameterException("IV and tag length required when"
+                            + " decrypting, but not found in parameters: " + params, e);
+                }
+                setIv(null);
+                return;
+            }
+            initAlgorithmSpecificParameters(spec);
+        }
+
+        @Nullable
+        @Override
+        protected final AlgorithmParameters engineGetParameters() {
+            byte[] iv = getIv();
+            if ((iv != null) && (iv.length > 0)) {
+                try {
+                    AlgorithmParameters params = AlgorithmParameters.getInstance("GCM");
+                    params.init(new GCMParameterSpec(mTagLengthBits, iv));
+                    return params;
+                } catch (NoSuchAlgorithmException e) {
+                    throw new ProviderException(
+                            "Failed to obtain GCM AlgorithmParameters", e);
+                } catch (InvalidParameterSpecException e) {
+                    throw new ProviderException(
+                            "Failed to initialize GCM AlgorithmParameters", e);
+                }
+            }
+            return null;
+        }
+
+        @NonNull
+        @Override
+        protected KeyStoreCryptoOperationStreamer createMainDataStreamer(
+                KeyStore keyStore, IBinder operationToken) {
+            KeyStoreCryptoOperationStreamer streamer = new KeyStoreCryptoOperationChunkedStreamer(
+                    new KeyStoreCryptoOperationChunkedStreamer.MainDataStream(
+                            keyStore, operationToken));
+            if (isEncrypting()) {
+                return streamer;
+            } else {
+                // When decrypting, to avoid leaking unauthenticated plaintext, do not return any
+                // plaintext before ciphertext is authenticated by KeyStore.finish.
+                return new BufferAllOutputUntilDoFinalStreamer(streamer);
+            }
+        }
+
+        @NonNull
+        @Override
+        protected final KeyStoreCryptoOperationStreamer createAdditionalAuthenticationDataStreamer(
+                KeyStore keyStore, IBinder operationToken) {
+            return new KeyStoreCryptoOperationChunkedStreamer(
+                    new AdditionalAuthenticationDataStream(keyStore, operationToken));
+        }
+
+        @Override
+        protected final int getAdditionalEntropyAmountForBegin() {
+            if ((getIv() == null) && (isEncrypting())) {
+                // IV will need to be generated
+                return IV_LENGTH_BYTES;
+            }
+
+            return 0;
+        }
+
+        @Override
+        protected final int getAdditionalEntropyAmountForFinish() {
+            return 0;
+        }
+
+        @Override
+        protected final void addAlgorithmSpecificParametersToBegin(
+                @NonNull KeymasterArguments keymasterArgs) {
+            super.addAlgorithmSpecificParametersToBegin(keymasterArgs);
+            keymasterArgs.addInt(KeymasterDefs.KM_TAG_MAC_LENGTH, mTagLengthBits);
+        }
+
+        protected final int getTagLengthBits() {
+            return mTagLengthBits;
+        }
+
+        public static final class NoPadding extends GCM {
+            public NoPadding() {
+                super(KeymasterDefs.KM_PAD_NONE);
+            }
+
+            @Override
+            protected final int engineGetOutputSize(int inputLen) {
+                int tagLengthBytes = (getTagLengthBits() + 7) / 8;
+                long result;
+                if (isEncrypting()) {
+                    result = getConsumedInputSizeBytes() - getProducedOutputSizeBytes() + inputLen
+                            + tagLengthBytes;
+                } else {
+                    result = getConsumedInputSizeBytes() - getProducedOutputSizeBytes() + inputLen
+                            - tagLengthBytes;
+                }
+                if (result < 0) {
+                    return 0;
+                } else if (result > Integer.MAX_VALUE) {
+                    return Integer.MAX_VALUE;
+                }
+                return (int) result;
+            }
+        }
+    }
+
+    private static final int BLOCK_SIZE_BYTES = 16;
+
+    private final int mKeymasterBlockMode;
+    private final int mKeymasterPadding;
+
+    private byte[] mIv;
+
+    /** Whether the current {@code #mIv} has been used by the underlying crypto operation. */
+    private boolean mIvHasBeenUsed;
+
+    AndroidKeyStoreAuthenticatedAESCipherSpi(
+            int keymasterBlockMode,
+            int keymasterPadding) {
+        mKeymasterBlockMode = keymasterBlockMode;
+        mKeymasterPadding = keymasterPadding;
+    }
+
+    @Override
+    protected void resetAll() {
+        mIv = null;
+        mIvHasBeenUsed = false;
+        super.resetAll();
+    }
+
+    @Override
+    protected final void initKey(int opmode, Key key) throws InvalidKeyException {
+        if (!(key instanceof AndroidKeyStoreSecretKey)) {
+            throw new InvalidKeyException(
+                    "Unsupported key: " + ((key != null) ? key.getClass().getName() : "null"));
+        }
+        if (!KeyProperties.KEY_ALGORITHM_AES.equalsIgnoreCase(key.getAlgorithm())) {
+            throw new InvalidKeyException(
+                    "Unsupported key algorithm: " + key.getAlgorithm() + ". Only " +
+                    KeyProperties.KEY_ALGORITHM_AES + " supported");
+        }
+        setKey((AndroidKeyStoreSecretKey) key);
+    }
+
+    @Override
+    protected void addAlgorithmSpecificParametersToBegin(
+            @NonNull KeymasterArguments keymasterArgs) {
+        if ((isEncrypting()) && (mIvHasBeenUsed)) {
+            // IV is being reused for encryption: this violates security best practices.
+            throw new IllegalStateException(
+                    "IV has already been used. Reusing IV in encryption mode violates security best"
+                    + " practices.");
+        }
+
+        keymasterArgs.addInt(KeymasterDefs.KM_TAG_ALGORITHM, KeymasterDefs.KM_ALGORITHM_AES);
+        keymasterArgs.addInt(KeymasterDefs.KM_TAG_BLOCK_MODE, mKeymasterBlockMode);
+        keymasterArgs.addInt(KeymasterDefs.KM_TAG_PADDING, mKeymasterPadding);
+        if (mIv != null) {
+            keymasterArgs.addBlob(KeymasterDefs.KM_TAG_NONCE, mIv);
+        }
+    }
+
+    @Override
+    protected final void loadAlgorithmSpecificParametersFromBeginResult(
+            @NonNull KeymasterArguments keymasterArgs) {
+        mIvHasBeenUsed = true;
+
+        // NOTE: Keymaster doesn't always return an IV, even if it's used.
+        byte[] returnedIv = keymasterArgs.getBlob(KeymasterDefs.KM_TAG_NONCE, null);
+        if ((returnedIv != null) && (returnedIv.length == 0)) {
+            returnedIv = null;
+        }
+
+        if (mIv == null) {
+            mIv = returnedIv;
+        } else if ((returnedIv != null) && (!Arrays.equals(returnedIv, mIv))) {
+            throw new ProviderException("IV in use differs from provided IV");
+        }
+    }
+
+    @Override
+    protected final int engineGetBlockSize() {
+        return BLOCK_SIZE_BYTES;
+    }
+
+    @Override
+    protected final byte[] engineGetIV() {
+        return ArrayUtils.cloneIfNotEmpty(mIv);
+    }
+
+    protected void setIv(byte[] iv) {
+        mIv = iv;
+    }
+
+    protected byte[] getIv() {
+        return mIv;
+    }
+
+    /**
+     * {@link KeyStoreCryptoOperationStreamer} which buffers all output until {@code doFinal} from
+     * which it returns all output in one go, provided {@code doFinal} succeeds.
+     */
+    private static class BufferAllOutputUntilDoFinalStreamer
+        implements KeyStoreCryptoOperationStreamer {
+
+        private final KeyStoreCryptoOperationStreamer mDelegate;
+        private ByteArrayOutputStream mBufferedOutput = new ByteArrayOutputStream();
+        private long mProducedOutputSizeBytes;
+
+        private BufferAllOutputUntilDoFinalStreamer(KeyStoreCryptoOperationStreamer delegate) {
+            mDelegate = delegate;
+        }
+
+        @Override
+        public byte[] update(byte[] input, int inputOffset, int inputLength)
+                throws KeyStoreException {
+            byte[] output = mDelegate.update(input, inputOffset, inputLength);
+            if (output != null) {
+                try {
+                    mBufferedOutput.write(output);
+                } catch (IOException e) {
+                    throw new ProviderException("Failed to buffer output", e);
+                }
+            }
+            return EmptyArray.BYTE;
+        }
+
+        @Override
+        public byte[] doFinal(byte[] input, int inputOffset, int inputLength,
+                byte[] additionalEntropy) throws KeyStoreException {
+            byte[] output = mDelegate.doFinal(input, inputOffset, inputLength, additionalEntropy);
+            if (output != null) {
+                try {
+                    mBufferedOutput.write(output);
+                } catch (IOException e) {
+                    throw new ProviderException("Failed to buffer output", e);
+                }
+            }
+            byte[] result = mBufferedOutput.toByteArray();
+            mBufferedOutput.reset();
+            mProducedOutputSizeBytes += result.length;
+            return result;
+        }
+
+        @Override
+        public long getConsumedInputSizeBytes() {
+            return mDelegate.getConsumedInputSizeBytes();
+        }
+
+        @Override
+        public long getProducedOutputSizeBytes() {
+            return mProducedOutputSizeBytes;
+        }
+    }
+
+    /**
+     * Additional Authentication Data (AAD) stream via a KeyStore streaming operation. This stream
+     * sends AAD into the KeyStore.
+     */
+    private static class AdditionalAuthenticationDataStream implements Stream {
+
+        private final KeyStore mKeyStore;
+        private final IBinder mOperationToken;
+
+        private AdditionalAuthenticationDataStream(KeyStore keyStore, IBinder operationToken) {
+            mKeyStore = keyStore;
+            mOperationToken = operationToken;
+        }
+
+        @Override
+        public OperationResult update(byte[] input) {
+            KeymasterArguments keymasterArgs = new KeymasterArguments();
+            keymasterArgs.addBlob(KeymasterDefs.KM_TAG_ASSOCIATED_DATA, input);
+
+            // KeyStore does not reflect AAD in inputConsumed, but users of Stream rely on this
+            // field. We fix this discrepancy here. KeyStore.update contract is that all of AAD
+            // has been consumed if the method succeeds.
+            OperationResult result = mKeyStore.update(mOperationToken, keymasterArgs, null);
+            if (result.resultCode == KeyStore.NO_ERROR) {
+                result = new OperationResult(
+                        result.resultCode,
+                        result.token,
+                        result.operationHandle,
+                        input.length, // inputConsumed
+                        result.output,
+                        result.outParams);
+            }
+            return result;
+        }
+
+        @Override
+        public OperationResult finish(byte[] additionalEntropy) {
+            if ((additionalEntropy != null) && (additionalEntropy.length > 0)) {
+                throw new ProviderException("AAD stream does not support additional entropy");
+            }
+            return new OperationResult(
+                    KeyStore.NO_ERROR,
+                    mOperationToken,
+                    0, // operation handle -- nobody cares about this being returned from finish
+                    0, // inputConsumed
+                    EmptyArray.BYTE, // output
+                    new KeymasterArguments() // additional params returned by finish
+                    );
+        }
+    }
+}
\ No newline at end of file
diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreBCWorkaroundProvider.java b/keystore/java/android/security/keystore/AndroidKeyStoreBCWorkaroundProvider.java
index f37cf07..06b76fa 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreBCWorkaroundProvider.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreBCWorkaroundProvider.java
@@ -93,6 +93,9 @@ class AndroidKeyStoreBCWorkaroundProvider extends Provider {
         putSymmetricCipherImpl("AES/CTR/NoPadding",
                 PACKAGE_NAME + ".AndroidKeyStoreUnauthenticatedAESCipherSpi$CTR$NoPadding");
 
+        putSymmetricCipherImpl("AES/GCM/NoPadding",
+                PACKAGE_NAME + ".AndroidKeyStoreAuthenticatedAESCipherSpi$GCM$NoPadding");
+
         putAsymmetricCipherImpl("RSA/ECB/NoPadding",
                 PACKAGE_NAME + ".AndroidKeyStoreRSACipherSpi$NoPadding");
         put("Alg.Alias.Cipher.RSA/None/NoPadding", "RSA/ECB/NoPadding");
diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreCipherSpiBase.java b/keystore/java/android/security/keystore/AndroidKeyStoreCipherSpiBase.java
index d2d5850..fc53451 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreCipherSpiBase.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreCipherSpiBase.java
@@ -26,6 +26,8 @@ import android.security.keymaster.KeymasterArguments;
 import android.security.keymaster.KeymasterDefs;
 import android.security.keymaster.OperationResult;
 
+import libcore.util.EmptyArray;
+
 import java.nio.ByteBuffer;
 import java.security.AlgorithmParameters;
 import java.security.GeneralSecurityException;
@@ -78,6 +80,8 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
     private IBinder mOperationToken;
     private long mOperationHandle;
     private KeyStoreCryptoOperationStreamer mMainDataStreamer;
+    private KeyStoreCryptoOperationStreamer mAdditionalAuthenticationDataStreamer;
+    private boolean mAdditionalAuthenticationDataStreamerClosed;
 
     /**
      * Encountered exception which could not be immediately thrown because it was encountered inside
@@ -189,6 +193,8 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
         mOperationToken = null;
         mOperationHandle = 0;
         mMainDataStreamer = null;
+        mAdditionalAuthenticationDataStreamer = null;
+        mAdditionalAuthenticationDataStreamerClosed = false;
         mCachedException = null;
     }
 
@@ -209,6 +215,8 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
         mOperationToken = null;
         mOperationHandle = 0;
         mMainDataStreamer = null;
+        mAdditionalAuthenticationDataStreamer = null;
+        mAdditionalAuthenticationDataStreamerClosed = false;
         mCachedException = null;
     }
 
@@ -273,6 +281,9 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
 
         loadAlgorithmSpecificParametersFromBeginResult(opResult.outParams);
         mMainDataStreamer = createMainDataStreamer(mKeyStore, opResult.token);
+        mAdditionalAuthenticationDataStreamer =
+                createAdditionalAuthenticationDataStreamer(mKeyStore, opResult.token);
+        mAdditionalAuthenticationDataStreamerClosed = false;
     }
 
     /**
@@ -289,6 +300,20 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
                         keyStore, operationToken));
     }
 
+    /**
+     * Creates a streamer which sends Additional Authentication Data (AAD) into the KeyStore.
+     *
+     * <p>This implementation returns {@code null}.
+     *
+     * @returns stream or {@code null} if AAD is not supported by this cipher.
+     */
+    @Nullable
+    protected KeyStoreCryptoOperationStreamer createAdditionalAuthenticationDataStreamer(
+            @SuppressWarnings("unused") KeyStore keyStore,
+            @SuppressWarnings("unused") IBinder operationToken) {
+        return null;
+    }
+
     @Override
     protected final byte[] engineUpdate(byte[] input, int inputOffset, int inputLen) {
         if (mCachedException != null) {
@@ -307,6 +332,7 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
 
         byte[] output;
         try {
+            flushAAD();
             output = mMainDataStreamer.update(input, inputOffset, inputLen);
         } catch (KeyStoreException e) {
             mCachedException = e;
@@ -320,6 +346,25 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
         return output;
     }
 
+    private void flushAAD() throws KeyStoreException {
+        if ((mAdditionalAuthenticationDataStreamer != null)
+                && (!mAdditionalAuthenticationDataStreamerClosed)) {
+            byte[] output;
+            try {
+                output = mAdditionalAuthenticationDataStreamer.doFinal(
+                        EmptyArray.BYTE, 0, 0,
+                        null // no additional entropy needed flushing AAD
+                        );
+            } finally {
+                mAdditionalAuthenticationDataStreamerClosed = true;
+            }
+            if ((output != null) && (output.length > 0)) {
+                throw new ProviderException(
+                        "AAD update unexpectedly returned data: " + output.length + " bytes");
+            }
+        }
+    }
+
     @Override
     protected final int engineUpdate(byte[] input, int inputOffset, int inputLen, byte[] output,
             int outputOffset) throws ShortBufferException {
@@ -344,12 +389,64 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
 
     @Override
     protected final void engineUpdateAAD(byte[] input, int inputOffset, int inputLen) {
-        super.engineUpdateAAD(input, inputOffset, inputLen);
+        if (mCachedException != null) {
+            return;
+        }
+
+        try {
+            ensureKeystoreOperationInitialized();
+        } catch (InvalidKeyException | InvalidAlgorithmParameterException e) {
+            mCachedException = e;
+            return;
+        }
+
+        if (mAdditionalAuthenticationDataStreamerClosed) {
+            throw new IllegalStateException(
+                    "AAD can only be provided before Cipher.update is invoked");
+        }
+
+        if (mAdditionalAuthenticationDataStreamer == null) {
+            throw new IllegalStateException("This cipher does not support AAD");
+        }
+
+        byte[] output;
+        try {
+            output = mAdditionalAuthenticationDataStreamer.update(input, inputOffset, inputLen);
+        } catch (KeyStoreException e) {
+            mCachedException = e;
+            return;
+        }
+
+        if ((output != null) && (output.length > 0)) {
+            throw new ProviderException("AAD update unexpectedly produced output: "
+                    + output.length + " bytes");
+        }
     }
 
     @Override
     protected final void engineUpdateAAD(ByteBuffer src) {
-        super.engineUpdateAAD(src);
+        if (src == null) {
+            throw new IllegalArgumentException("src == null");
+        }
+        if (!src.hasRemaining()) {
+            return;
+        }
+
+        byte[] input;
+        int inputOffset;
+        int inputLen;
+        if (src.hasArray()) {
+            input = src.array();
+            inputOffset = src.arrayOffset() + src.position();
+            inputLen = src.remaining();
+            src.position(src.limit());
+        } else {
+            input = new byte[src.remaining()];
+            inputOffset = 0;
+            inputLen = input.length;
+            src.get(input);
+        }
+        super.engineUpdateAAD(input, inputOffset, inputLen);
     }
 
     @Override
@@ -368,6 +465,7 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
 
         byte[] output;
         try {
+            flushAAD();
             byte[] additionalEntropy =
                     KeyStoreCryptoOperationUtils.getRandomBytesToMixIntoKeystoreRng(
                             mRng, getAdditionalEntropyAmountForFinish());
@@ -615,6 +713,20 @@ abstract class AndroidKeyStoreCipherSpiBase extends CipherSpi implements KeyStor
         return mKeyStore;
     }
 
+    protected final long getConsumedInputSizeBytes() {
+        if (mMainDataStreamer == null) {
+            throw new IllegalStateException("Not initialized");
+        }
+        return mMainDataStreamer.getConsumedInputSizeBytes();
+    }
+
+    protected final long getProducedOutputSizeBytes() {
+        if (mMainDataStreamer == null) {
+            throw new IllegalStateException("Not initialized");
+        }
+        return mMainDataStreamer.getProducedOutputSizeBytes();
+    }
+
     // The methods below need to be implemented by subclasses.
 
     /**
diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreRSACipherSpi.java b/keystore/java/android/security/keystore/AndroidKeyStoreRSACipherSpi.java
index 6abdf19..38e216d 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreRSACipherSpi.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreRSACipherSpi.java
@@ -129,6 +129,7 @@ abstract class AndroidKeyStoreRSACipherSpi extends AndroidKeyStoreCipherSpiBase
             private final KeyStoreCryptoOperationStreamer mDelegate;
             private final int mModulusSizeBytes;
             private final ByteArrayOutputStream mInputBuffer = new ByteArrayOutputStream();
+            private long mConsumedInputSizeBytes;
 
             private ZeroPaddingEncryptionStreamer(
                     KeyStoreCryptoOperationStreamer delegate,
@@ -142,6 +143,7 @@ abstract class AndroidKeyStoreRSACipherSpi extends AndroidKeyStoreCipherSpiBase
                     throws KeyStoreException {
                 if (inputLength > 0) {
                     mInputBuffer.write(input, inputOffset, inputLength);
+                    mConsumedInputSizeBytes += inputLength;
                 }
                 return EmptyArray.BYTE;
             }
@@ -151,6 +153,7 @@ abstract class AndroidKeyStoreRSACipherSpi extends AndroidKeyStoreCipherSpiBase
                     byte[] additionalEntropy)
                     throws KeyStoreException {
                 if (inputLength > 0) {
+                    mConsumedInputSizeBytes += inputLength;
                     mInputBuffer.write(input, inputOffset, inputLength);
                 }
                 byte[] bufferedInput = mInputBuffer.toByteArray();
@@ -173,6 +176,16 @@ abstract class AndroidKeyStoreRSACipherSpi extends AndroidKeyStoreCipherSpiBase
                 }
                 return mDelegate.doFinal(paddedInput, 0, paddedInput.length, additionalEntropy);
             }
+
+            @Override
+            public long getConsumedInputSizeBytes() {
+                return mConsumedInputSizeBytes;
+            }
+
+            @Override
+            public long getProducedOutputSizeBytes() {
+                return mDelegate.getProducedOutputSizeBytes();
+            }
         }
     }
 
diff --git a/keystore/java/android/security/keystore/AndroidKeyStoreUnauthenticatedAESCipherSpi.java b/keystore/java/android/security/keystore/AndroidKeyStoreUnauthenticatedAESCipherSpi.java
index 76804a9..6c53c6a 100644
--- a/keystore/java/android/security/keystore/AndroidKeyStoreUnauthenticatedAESCipherSpi.java
+++ b/keystore/java/android/security/keystore/AndroidKeyStoreUnauthenticatedAESCipherSpi.java
@@ -1,3 +1,19 @@
+/*
+ * 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.keystore;
 
 import android.annotation.NonNull;
diff --git a/keystore/java/android/security/keystore/KeyStoreCryptoOperationChunkedStreamer.java b/keystore/java/android/security/keystore/KeyStoreCryptoOperationChunkedStreamer.java
index 9957e79..894d52a 100644
--- a/keystore/java/android/security/keystore/KeyStoreCryptoOperationChunkedStreamer.java
+++ b/keystore/java/android/security/keystore/KeyStoreCryptoOperationChunkedStreamer.java
@@ -73,6 +73,8 @@ class KeyStoreCryptoOperationChunkedStreamer implements KeyStoreCryptoOperationS
     private byte[] mBuffered = EmptyArray.BYTE;
     private int mBufferedOffset;
     private int mBufferedLength;
+    private long mConsumedInputSizeBytes;
+    private long mProducedOutputSizeBytes;
 
     public KeyStoreCryptoOperationChunkedStreamer(Stream operation) {
         this(operation, DEFAULT_MAX_CHUNK_SIZE);
@@ -119,6 +121,7 @@ class KeyStoreCryptoOperationChunkedStreamer implements KeyStoreCryptoOperationS
             // Update input array references to reflect that some of its bytes are now in mBuffered.
             inputOffset += inputBytesInChunk;
             inputLength -= inputBytesInChunk;
+            mConsumedInputSizeBytes += inputBytesInChunk;
 
             OperationResult opResult = mKeyStoreStream.update(chunk);
             if (opResult == null) {
@@ -167,9 +170,10 @@ class KeyStoreCryptoOperationChunkedStreamer implements KeyStoreCryptoOperationS
                     }
                 } else {
                     // No more output will be produced in this loop
+                    byte[] result;
                     if (bufferedOutput == null) {
                         // No previously buffered output
-                        return opResult.output;
+                        result = opResult.output;
                     } else {
                         // There was some previously buffered output
                         try {
@@ -177,18 +181,23 @@ class KeyStoreCryptoOperationChunkedStreamer implements KeyStoreCryptoOperationS
                         } catch (IOException e) {
                             throw new IllegalStateException("Failed to buffer output", e);
                         }
-                        return bufferedOutput.toByteArray();
+                        result = bufferedOutput.toByteArray();
                     }
+                    mProducedOutputSizeBytes += result.length;
+                    return result;
                 }
             }
         }
 
+        byte[] result;
         if (bufferedOutput == null) {
             // No output produced
-            return EmptyArray.BYTE;
+            result = EmptyArray.BYTE;
         } else {
-            return bufferedOutput.toByteArray();
+            result = bufferedOutput.toByteArray();
         }
+        mProducedOutputSizeBytes += result.length;
+        return result;
     }
 
     @Override
@@ -210,14 +219,11 @@ class KeyStoreCryptoOperationChunkedStreamer implements KeyStoreCryptoOperationS
         } else if (opResult.resultCode != KeyStore.NO_ERROR) {
             throw KeyStore.getKeyStoreException(opResult.resultCode);
         }
+        mProducedOutputSizeBytes += opResult.output.length;
 
         return ArrayUtils.concat(output, opResult.output);
     }
 
-    /**
-     * Passes all of buffered input into the the KeyStore operation (via the {@code update}
-     * operation) and returns output.
-     */
     public byte[] flush() throws KeyStoreException {
         if (mBufferedLength <= 0) {
             return EmptyArray.BYTE;
@@ -243,7 +249,19 @@ class KeyStoreCryptoOperationChunkedStreamer implements KeyStoreCryptoOperationS
                     + " . Provided: " + chunk.length + ", consumed: " + opResult.inputConsumed);
         }
 
-        return (opResult.output != null) ? opResult.output : EmptyArray.BYTE;
+        byte[] result = (opResult.output != null) ? opResult.output : EmptyArray.BYTE;
+        mProducedOutputSizeBytes += result.length;
+        return result;
+    }
+
+    @Override
+    public long getConsumedInputSizeBytes() {
+        return mConsumedInputSizeBytes;
+    }
+
+    @Override
+    public long getProducedOutputSizeBytes() {
+        return mProducedOutputSizeBytes;
     }
 
     /**
diff --git a/keystore/java/android/security/keystore/KeyStoreCryptoOperationStreamer.java b/keystore/java/android/security/keystore/KeyStoreCryptoOperationStreamer.java
index 1c6de2d..897bd71 100644
--- a/keystore/java/android/security/keystore/KeyStoreCryptoOperationStreamer.java
+++ b/keystore/java/android/security/keystore/KeyStoreCryptoOperationStreamer.java
@@ -37,4 +37,6 @@ interface KeyStoreCryptoOperationStreamer {
     byte[] update(byte[] input, int inputOffset, int inputLength) throws KeyStoreException;
     byte[] doFinal(byte[] input, int inputOffset, int inputLength, byte[] additionalEntropy)
             throws KeyStoreException;
+    long getConsumedInputSizeBytes();
+    long getProducedOutputSizeBytes();
 }