resolve merge conflicts of 4aa5650699 to master.

Change-Id: Ie9469cedffd3f974f02c14582c7d48224a90c667

15 files changed, 177 insertions, 65 deletions

diff --git a/aes_operation.cpp b/aes_operation.cpp
index bbaa0ad..4c5c88c 100644
--- a/aes_operation.cpp
+++ b/aes_operation.cpp
@@ -211,17 +211,18 @@ inline bool is_bad_decrypt(unsigned long error) {
             ERR_GET_REASON(error) == CIPHER_R_BAD_DECRYPT);
 }
 
-keymaster_error_t AesEvpOperation::Finish(const AuthorizationSet& /* additional_params */,
-                                          const Buffer& /* signature */,
-                                          AuthorizationSet* /* output_params */, Buffer* output) {
+keymaster_error_t AesEvpOperation::Finish(const AuthorizationSet& additional_params,
+                                          const Buffer& input, const Buffer& /* signature */,
+                                          AuthorizationSet* output_params, Buffer* output) {
+    keymaster_error_t error;
+    if (!UpdateForFinish(additional_params, input, output_params, output, &error))
+        return error;
+
     if (!output->reserve(AES_BLOCK_SIZE))
         return KM_ERROR_MEMORY_ALLOCATION_FAILED;
 
-    keymaster_error_t error;
-    if (block_mode_ == KM_MODE_GCM && aad_block_buf_length_ > 0 &&
-        !ProcessBufferedAadBlock(&error)) {
+    if (block_mode_ == KM_MODE_GCM && aad_block_buf_length_ > 0 && !ProcessBufferedAadBlock(&error))
         return error;
-    }
 
     int output_written = -1;
     if (!EVP_CipherFinal_ex(&ctx_, output->peek_write(), &output_written)) {
@@ -464,6 +465,23 @@ bool AesEvpOperation::InternalUpdate(const uint8_t* input, size_t input_length,
     return output->advance_write(output_written);
 }
 
+bool AesEvpOperation::UpdateForFinish(const AuthorizationSet& additional_params,
+                                      const Buffer& input, AuthorizationSet* output_params,
+                                      Buffer* output, keymaster_error_t* error) {
+    if (input.available_read() || !additional_params.empty()) {
+        size_t input_consumed;
+        *error = Update(additional_params, input, output_params, output, &input_consumed);
+        if (*error != KM_ERROR_OK)
+            return false;
+        if (input_consumed != input.available_read()) {
+            *error = KM_ERROR_INVALID_INPUT_LENGTH;
+            return false;
+        }
+    }
+
+    return true;
+}
+
 keymaster_error_t AesEvpEncryptOperation::Begin(const AuthorizationSet& input_params,
                                                 AuthorizationSet* output_params) {
     if (!output_params)
@@ -488,18 +506,18 @@ keymaster_error_t AesEvpEncryptOperation::Begin(const AuthorizationSet& input_pa
 }
 
 keymaster_error_t AesEvpEncryptOperation::Finish(const AuthorizationSet& additional_params,
-                                                 const Buffer& signature,
+                                                 const Buffer& input, const Buffer& signature,
                                                  AuthorizationSet* output_params, Buffer* output) {
-    if (!output->reserve(AES_BLOCK_SIZE + tag_length_))
+    if (!output->reserve(input.available_read() + AES_BLOCK_SIZE + tag_length_))
         return KM_ERROR_MEMORY_ALLOCATION_FAILED;
 
     keymaster_error_t error =
-        AesEvpOperation::Finish(additional_params, signature, output_params, output);
+        AesEvpOperation::Finish(additional_params, input, signature, output_params, output);
     if (error != KM_ERROR_OK)
         return error;
 
     if (tag_length_ > 0) {
-        if (!output->reserve(output->available_read() + tag_length_))
+        if (!output->reserve(tag_length_))
             return KM_ERROR_MEMORY_ALLOCATION_FAILED;
 
         if (!EVP_CIPHER_CTX_ctrl(&ctx_, EVP_CTRL_GCM_GET_TAG, tag_length_, output->peek_write()))
@@ -609,15 +627,21 @@ void AesEvpDecryptOperation::BufferCandidateTagData(const uint8_t* data, size_t
 }
 
 keymaster_error_t AesEvpDecryptOperation::Finish(const AuthorizationSet& additional_params,
-                                                 const Buffer& signature,
+                                                 const Buffer& input, const Buffer& signature,
                                                  AuthorizationSet* output_params, Buffer* output) {
+    keymaster_error_t error;
+    if (!UpdateForFinish(additional_params, input, output_params, output, &error))
+        return error;
+
     if (tag_buf_length_ < tag_length_)
         return KM_ERROR_INVALID_INPUT_LENGTH;
     else if (tag_length_ > 0 &&
              !EVP_CIPHER_CTX_ctrl(&ctx_, EVP_CTRL_GCM_SET_TAG, tag_length_, tag_buf_.get()))
         return TranslateLastOpenSslError();
 
-    return AesEvpOperation::Finish(additional_params, signature, output_params, output);
+    AuthorizationSet empty_params;
+    Buffer empty_input;
+    return AesEvpOperation::Finish(empty_params, empty_input, signature, output_params, output);
 }
 
 keymaster_error_t AesEvpOperation::Abort() {
diff --git a/aes_operation.h b/aes_operation.h
index 7c0bd72..1a296bb 100644
--- a/aes_operation.h
+++ b/aes_operation.h
@@ -68,8 +68,9 @@ class AesEvpOperation : public Operation {
     keymaster_error_t Update(const AuthorizationSet& additional_params, const Buffer& input,
                              AuthorizationSet* output_params, Buffer* output,
                              size_t* input_consumed) override;
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
     keymaster_error_t Abort() override;
 
     virtual int evp_encrypt_mode() = 0;
@@ -85,6 +86,8 @@ class AesEvpOperation : public Operation {
     bool ProcessBufferedAadBlock(keymaster_error_t* error);
     bool InternalUpdate(const uint8_t* input, size_t input_length, Buffer* output,
                         keymaster_error_t* error);
+    bool UpdateForFinish(const AuthorizationSet& additional_params, const Buffer& input,
+                         AuthorizationSet* output_params, Buffer* output, keymaster_error_t* error);
 
     const keymaster_block_mode_t block_mode_;
     EVP_CIPHER_CTX ctx_;
@@ -111,8 +114,9 @@ class AesEvpEncryptOperation : public AesEvpOperation {
 
     keymaster_error_t Begin(const AuthorizationSet& input_params,
                             AuthorizationSet* output_params) override;
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
 
     int evp_encrypt_mode() override { return 1; }
 
@@ -132,8 +136,9 @@ class AesEvpDecryptOperation : public AesEvpOperation {
     keymaster_error_t Update(const AuthorizationSet& additional_params, const Buffer& input,
                              AuthorizationSet* output_params, Buffer* output,
                              size_t* input_consumed) override;
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
 
     int evp_encrypt_mode() override { return 0; }
 
diff --git a/android_keymaster.cpp b/android_keymaster.cpp
index c2ff8e6..3a53394 100644
--- a/android_keymaster.cpp
+++ b/android_keymaster.cpp
@@ -310,7 +310,7 @@ void AndroidKeymaster::FinishOperation(const FinishOperationRequest& request,
         }
     }
 
-    response->error = operation->Finish(request.additional_params, request.signature,
+    response->error = operation->Finish(request.additional_params, request.input, request.signature,
                                         &response->output_params, &response->output);
     operation_table_->Delete(request.op_handle);
 }
diff --git a/authorization_set.cpp b/authorization_set.cpp
index b8f45e3..3a5b46e 100644
--- a/authorization_set.cpp
+++ b/authorization_set.cpp
@@ -196,21 +196,31 @@ int AuthorizationSet::find(keymaster_tag_t tag, int begin) const {
         return i;
 }
 
-keymaster_key_param_t empty;
+bool AuthorizationSet::erase(size_t index) {
+    if (index >= size())
+        return false;
+
+    --elems_size_;
+    for (size_t i = index; i < elems_size_; ++i)
+        elems_[i] = elems_[i + 1];
+    return true;
+}
+
+keymaster_key_param_t empty_set = {};
 keymaster_key_param_t& AuthorizationSet::operator[](int at) {
     if (is_valid() == OK && at < (int)elems_size_) {
         return elems_[at];
     }
-    memset(&empty, 0, sizeof(empty));
-    return empty;
+    empty_set = {};
+    return empty_set;
 }
 
 keymaster_key_param_t AuthorizationSet::operator[](int at) const {
     if (is_valid() == OK && at < (int)elems_size_) {
         return elems_[at];
     }
-    memset(&empty, 0, sizeof(empty));
-    return empty;
+    empty_set = {};
+    return empty_set;
 }
 
 bool AuthorizationSet::push_back(const keymaster_key_param_set_t& set) {
diff --git a/ecdsa_keymaster1_operation.h b/ecdsa_keymaster1_operation.h
index 7530fbb..6045686 100644
--- a/ecdsa_keymaster1_operation.h
+++ b/ecdsa_keymaster1_operation.h
@@ -69,12 +69,13 @@ template <typename BaseOperation> class EcdsaKeymaster1Operation : public BaseOp
         return super::Begin(input_params, output_params);
     }
 
-    keymaster_error_t Finish(const AuthorizationSet& input_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override {
+    keymaster_error_t Finish(const AuthorizationSet& input_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override {
         keymaster_error_t error = wrapped_operation_.PrepareFinish(super::ecdsa_key_, input_params);
         if (error != KM_ERROR_OK)
             return error;
-        error = super::Finish(input_params, signature, output_params, output);
+        error = super::Finish(input_params, input, signature, output_params, output);
         if (wrapped_operation_.GetError(super::ecdsa_key_) != KM_ERROR_OK)
             error = wrapped_operation_.GetError(super::ecdsa_key_);
         if (error == KM_ERROR_OK)
diff --git a/ecdsa_operation.cpp b/ecdsa_operation.cpp
index 0e56344..405dcb5 100644
--- a/ecdsa_operation.cpp
+++ b/ecdsa_operation.cpp
@@ -135,13 +135,17 @@ keymaster_error_t EcdsaSignOperation::Update(const AuthorizationSet& /* addition
     return KM_ERROR_OK;
 }
 
-keymaster_error_t EcdsaSignOperation::Finish(const AuthorizationSet& /* additional_params */,
-                                             const Buffer& /* signature */,
+keymaster_error_t EcdsaSignOperation::Finish(const AuthorizationSet& additional_params,
+                                             const Buffer& input, const Buffer& /* signature */,
                                              AuthorizationSet* /* output_params */,
                                              Buffer* output) {
     if (!output)
         return KM_ERROR_OUTPUT_PARAMETER_NULL;
 
+    keymaster_error_t error = UpdateForFinish(additional_params, input);
+    if (error != KM_ERROR_OK)
+        return error;
+
     size_t siglen;
     if (digest_ == KM_DIGEST_NONE) {
         UniquePtr<EC_KEY, EC_KEY_Delete> ecdsa(EVP_PKEY_get1_EC_KEY(ecdsa_key_));
@@ -196,10 +200,14 @@ keymaster_error_t EcdsaVerifyOperation::Update(const AuthorizationSet& /* additi
     return KM_ERROR_OK;
 }
 
-keymaster_error_t EcdsaVerifyOperation::Finish(const AuthorizationSet& /* additional_params */,
-                                               const Buffer& signature,
+keymaster_error_t EcdsaVerifyOperation::Finish(const AuthorizationSet& additional_params,
+                                               const Buffer& input, const Buffer& signature,
                                                AuthorizationSet* /* output_params */,
                                                Buffer* /* output */) {
+    keymaster_error_t error = UpdateForFinish(additional_params, input);
+    if (error != KM_ERROR_OK)
+        return error;
+
     if (digest_ == KM_DIGEST_NONE) {
         UniquePtr<EC_KEY, EC_KEY_Delete> ecdsa(EVP_PKEY_get1_EC_KEY(ecdsa_key_));
         if (!ecdsa.get())
diff --git a/ecdsa_operation.h b/ecdsa_operation.h
index fba743f..4b95dc9 100644
--- a/ecdsa_operation.h
+++ b/ecdsa_operation.h
@@ -56,8 +56,9 @@ class EcdsaSignOperation : public EcdsaOperation {
     keymaster_error_t Update(const AuthorizationSet& additional_params, const Buffer& input,
                              AuthorizationSet* output_params, Buffer* output,
                              size_t* input_consumed) override;
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
 };
 
 class EcdsaVerifyOperation : public EcdsaOperation {
@@ -69,8 +70,9 @@ class EcdsaVerifyOperation : public EcdsaOperation {
     keymaster_error_t Update(const AuthorizationSet& additional_params, const Buffer& input,
                              AuthorizationSet* output_params, Buffer* output,
                              size_t* input_consumed) override;
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
 };
 
 class EcdsaOperationFactory : public OperationFactory {
diff --git a/hmac_operation.cpp b/hmac_operation.cpp
index 6b3117e..7f21393 100644
--- a/hmac_operation.cpp
+++ b/hmac_operation.cpp
@@ -160,9 +160,13 @@ keymaster_error_t HmacOperation::Abort() {
     return KM_ERROR_OK;
 }
 
-keymaster_error_t HmacOperation::Finish(const AuthorizationSet& /* additional_params */,
-                                        const Buffer& signature,
+keymaster_error_t HmacOperation::Finish(const AuthorizationSet& additional_params,
+                                        const Buffer& input, const Buffer& signature,
                                         AuthorizationSet* /* output_params */, Buffer* output) {
+    keymaster_error_t error = UpdateForFinish(additional_params, input);
+    if (error != KM_ERROR_OK)
+        return error;
+
     uint8_t digest[EVP_MAX_MD_SIZE];
     unsigned int digest_len;
     if (!HMAC_Final(&ctx_, digest, &digest_len))
diff --git a/hmac_operation.h b/hmac_operation.h
index 9c2d59b..83f2e09 100644
--- a/hmac_operation.h
+++ b/hmac_operation.h
@@ -35,7 +35,7 @@ class HmacOperation : public Operation {
                                      AuthorizationSet* output_params, Buffer* output,
                                      size_t* input_consumed);
     virtual keymaster_error_t Abort();
-    virtual keymaster_error_t Finish(const AuthorizationSet& additional_params,
+    virtual keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
                                      const Buffer& signature, AuthorizationSet* output_params,
                                      Buffer* output);
 
diff --git a/include/keymaster/authorization_set.h b/include/keymaster/authorization_set.h
index 5f0f8c3..74daa8d 100644
--- a/include/keymaster/authorization_set.h
+++ b/include/keymaster/authorization_set.h
@@ -118,6 +118,11 @@ class AuthorizationSet : public Serializable, public keymaster_key_param_set_t {
     size_t size() const { return elems_size_; }
 
     /**
+     * Returns true if the set is empty.
+     */
+    bool empty() const { return size() == 0; }
+
+    /**
      * Returns the total size of all indirect data referenced by set elements.
      */
     size_t indirect_size() const { return indirect_data_size_; }
@@ -152,6 +157,12 @@ class AuthorizationSet : public Serializable, public keymaster_key_param_set_t {
     int find(keymaster_tag_t tag, int begin = -1) const;
 
     /**
+     * Removes the entry at the specified index. Returns true if successful, false if the index was
+     * out of bounds.
+     */
+    bool erase(size_t index);
+
+    /**
      * Returns iterator (pointer) to beginning of elems array, to enable STL-style iteration
      */
     const keymaster_key_param_t* begin() const { return elems_; }
diff --git a/operation.cpp b/operation.cpp
index c36531d..410c9aa 100644
--- a/operation.cpp
+++ b/operation.cpp
@@ -134,4 +134,22 @@ bool OperationFactory::GetAndValidateDigest(const AuthorizationSet& begin_params
     return true;
 }
 
+keymaster_error_t Operation::UpdateForFinish(const AuthorizationSet& input_params,
+                                             const Buffer& input) {
+    if (!input_params.empty() || input.available_read()) {
+        size_t input_consumed;
+        Buffer output;
+        AuthorizationSet output_params;
+        keymaster_error_t error =
+            Update(input_params, input, &output_params, &output, &input_consumed);
+        if (error != KM_ERROR_OK)
+            return error;
+        assert(input_consumed == input.available_read());
+        assert(output_params.empty());
+        assert(output.available_read() == 0);
+    }
+
+    return KM_ERROR_OK;
+}
+
 }  // namespace keymaster
diff --git a/operation.h b/operation.h
index 74948fa..1b87e23 100644
--- a/operation.h
+++ b/operation.h
@@ -105,10 +105,16 @@ class Operation {
     virtual keymaster_error_t Update(const AuthorizationSet& input_params, const Buffer& input,
                                      AuthorizationSet* output_params, Buffer* output,
                                      size_t* input_consumed) = 0;
-    virtual keymaster_error_t Finish(const AuthorizationSet& input_params, const Buffer& signature,
-                                     AuthorizationSet* output_params, Buffer* output) = 0;
+    virtual keymaster_error_t Finish(const AuthorizationSet& input_params, const Buffer& input,
+                                     const Buffer& signature, AuthorizationSet* output_params,
+                                     Buffer* output) = 0;
     virtual keymaster_error_t Abort() = 0;
 
+protected:
+    // Helper function for implementing Finish() methods that need to call Update() to process
+    // input, but don't expect any output.
+    keymaster_error_t UpdateForFinish(const AuthorizationSet& input_params, const Buffer& input);
+
   private:
     const keymaster_purpose_t purpose_;
     AuthorizationSet key_auths_;
diff --git a/rsa_keymaster1_operation.h b/rsa_keymaster1_operation.h
index bdf1a4e..30123f0 100644
--- a/rsa_keymaster1_operation.h
+++ b/rsa_keymaster1_operation.h
@@ -22,8 +22,8 @@
 #include <hardware/keymaster1.h>
 #include <keymaster/android_keymaster_utils.h>
 
-#include "rsa_operation.h"
 #include "keymaster1_engine.h"
+#include "rsa_operation.h"
 
 namespace keymaster {
 
@@ -69,12 +69,13 @@ template <typename BaseOperation> class RsaKeymaster1Operation : public BaseOper
         return super::Begin(input_params, output_params);
     }
 
-    keymaster_error_t Finish(const AuthorizationSet& input_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override {
+    keymaster_error_t Finish(const AuthorizationSet& input_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override {
         keymaster_error_t error = wrapped_operation_.PrepareFinish(super::rsa_key_, input_params);
         if (error != KM_ERROR_OK)
             return error;
-        error = super::Finish(input_params, signature, output_params, output);
+        error = super::Finish(input_params, input, signature, output_params, output);
         if (wrapped_operation_.GetError(super::rsa_key_) != KM_ERROR_OK)
             error = wrapped_operation_.GetError(super::rsa_key_);
         if (error == KM_ERROR_OK)
diff --git a/rsa_operation.cpp b/rsa_operation.cpp
index 846eaa3..ce8d2f3 100644
--- a/rsa_operation.cpp
+++ b/rsa_operation.cpp
@@ -289,11 +289,15 @@ keymaster_error_t RsaSignOperation::Update(const AuthorizationSet& additional_pa
     return KM_ERROR_OK;
 }
 
-keymaster_error_t RsaSignOperation::Finish(const AuthorizationSet& /* additional_params */,
-                                           const Buffer& /* signature */,
+keymaster_error_t RsaSignOperation::Finish(const AuthorizationSet& additional_params,
+                                           const Buffer& input, const Buffer& /* signature */,
                                            AuthorizationSet* /* output_params */, Buffer* output) {
     assert(output);
 
+    keymaster_error_t error = UpdateForFinish(additional_params, input);
+    if (error != KM_ERROR_OK)
+        return error;
+
     if (digest_ == KM_DIGEST_NONE)
         return SignUndigested(output);
     else
@@ -408,10 +412,14 @@ keymaster_error_t RsaVerifyOperation::Update(const AuthorizationSet& additional_
     return KM_ERROR_OK;
 }
 
-keymaster_error_t RsaVerifyOperation::Finish(const AuthorizationSet& /* additional_params */,
-                                             const Buffer& signature,
+keymaster_error_t RsaVerifyOperation::Finish(const AuthorizationSet& additional_params,
+                                             const Buffer& input, const Buffer& signature,
                                              AuthorizationSet* /* output_params */,
                                              Buffer* /* output */) {
+    keymaster_error_t error = UpdateForFinish(additional_params, input);
+    if (error != KM_ERROR_OK)
+        return error;
+
     if (digest_ == KM_DIGEST_NONE)
         return VerifyUndigested(signature);
     else
@@ -509,11 +517,16 @@ struct EVP_PKEY_CTX_Delete {
     void operator()(EVP_PKEY_CTX* p) { EVP_PKEY_CTX_free(p); }
 };
 
-keymaster_error_t RsaEncryptOperation::Finish(const AuthorizationSet& /* additional_params */,
-                                              const Buffer& /* signature */,
+keymaster_error_t RsaEncryptOperation::Finish(const AuthorizationSet& additional_params,
+                                              const Buffer& input, const Buffer& /* signature */,
                                               AuthorizationSet* /* output_params */,
                                               Buffer* output) {
-    assert(output);
+    if (!output)
+        return KM_ERROR_OUTPUT_PARAMETER_NULL;
+
+    keymaster_error_t error = UpdateForFinish(additional_params, input);
+    if (error != KM_ERROR_OK)
+        return error;
 
     UniquePtr<EVP_PKEY_CTX, EVP_PKEY_CTX_Delete> ctx(
         EVP_PKEY_CTX_new(rsa_key_, nullptr /* engine */));
@@ -523,7 +536,7 @@ keymaster_error_t RsaEncryptOperation::Finish(const AuthorizationSet& /* additio
     if (EVP_PKEY_encrypt_init(ctx.get()) <= 0)
         return TranslateLastOpenSslError();
 
-    keymaster_error_t error = SetRsaPaddingInEvpContext(ctx.get(), false /* signing */);
+    error = SetRsaPaddingInEvpContext(ctx.get(), false /* signing */);
     if (error != KM_ERROR_OK)
         return error;
     error = SetOaepDigestIfRequired(ctx.get());
@@ -557,11 +570,16 @@ keymaster_error_t RsaEncryptOperation::Finish(const AuthorizationSet& /* additio
     return KM_ERROR_OK;
 }
 
-keymaster_error_t RsaDecryptOperation::Finish(const AuthorizationSet& /* additional_params */,
-                                              const Buffer& /* signature */,
+keymaster_error_t RsaDecryptOperation::Finish(const AuthorizationSet& additional_params,
+                                              const Buffer& input, const Buffer& /* signature */,
                                               AuthorizationSet* /* output_params */,
                                               Buffer* output) {
-    assert(output);
+    if (!output)
+        return KM_ERROR_OUTPUT_PARAMETER_NULL;
+
+    keymaster_error_t error = UpdateForFinish(additional_params, input);
+    if (error != KM_ERROR_OK)
+        return error;
 
     UniquePtr<EVP_PKEY_CTX, EVP_PKEY_CTX_Delete> ctx(
         EVP_PKEY_CTX_new(rsa_key_, nullptr /* engine */));
@@ -571,7 +589,7 @@ keymaster_error_t RsaDecryptOperation::Finish(const AuthorizationSet& /* additio
     if (EVP_PKEY_decrypt_init(ctx.get()) <= 0)
         return TranslateLastOpenSslError();
 
-    keymaster_error_t error = SetRsaPaddingInEvpContext(ctx.get(), false /* signing */);
+    error = SetRsaPaddingInEvpContext(ctx.get(), false /* signing */);
     if (error != KM_ERROR_OK)
         return error;
     error = SetOaepDigestIfRequired(ctx.get());
diff --git a/rsa_operation.h b/rsa_operation.h
index cbf71a5..669b21a 100644
--- a/rsa_operation.h
+++ b/rsa_operation.h
@@ -95,8 +95,9 @@ class RsaSignOperation : public RsaDigestingOperation {
     keymaster_error_t Update(const AuthorizationSet& additional_params, const Buffer& input,
                              AuthorizationSet* output_params, Buffer* output,
                              size_t* input_consumed) override;
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
 
   private:
     keymaster_error_t SignUndigested(Buffer* output);
@@ -116,8 +117,9 @@ class RsaVerifyOperation : public RsaDigestingOperation {
     keymaster_error_t Update(const AuthorizationSet& additional_params, const Buffer& input,
                              AuthorizationSet* output_params, Buffer* output,
                              size_t* input_consumed) override;
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
 
   private:
     keymaster_error_t VerifyUndigested(const Buffer& signature);
@@ -148,8 +150,9 @@ class RsaEncryptOperation : public RsaCryptOperation {
   public:
     RsaEncryptOperation(keymaster_digest_t digest, keymaster_padding_t padding, EVP_PKEY* key)
         : RsaCryptOperation(KM_PURPOSE_ENCRYPT, digest, padding, key) {}
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
 };
 
 /**
@@ -159,8 +162,9 @@ class RsaDecryptOperation : public RsaCryptOperation {
   public:
     RsaDecryptOperation(keymaster_digest_t digest, keymaster_padding_t padding, EVP_PKEY* key)
         : RsaCryptOperation(KM_PURPOSE_DECRYPT, digest, padding, key) {}
-    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& signature,
-                             AuthorizationSet* output_params, Buffer* output) override;
+    keymaster_error_t Finish(const AuthorizationSet& additional_params, const Buffer& input,
+                             const Buffer& signature, AuthorizationSet* output_params,
+                             Buffer* output) override;
 };
 
 /**