Support input to "finish()" in AndroidKeymaster operations.

This CL does not yet take advantage of the simplifications that allowing
input to finish() provides. That will require updating the Java layer
first, to remove some assumptions and code that assume update() must
eventually consume all input.

Change-Id: Ie85896027a1d55ddec06750d19addbb1f5e462c8

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 7b6514c..2046a64 100644
--- a/rsa_operation.cpp
+++ b/rsa_operation.cpp
@@ -282,11 +282,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
@@ -401,10 +405,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
@@ -502,11 +510,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 */));
@@ -516,7 +529,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());
+    error = SetRsaPaddingInEvpContext(ctx.get());
     if (error != KM_ERROR_OK)
         return error;
     error = SetOaepDigestIfRequired(ctx.get());
@@ -550,11 +563,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 */));
@@ -564,7 +582,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());
+    error = SetRsaPaddingInEvpContext(ctx.get());
     if (error != KM_ERROR_OK)
         return error;
     error = SetOaepDigestIfRequired(ctx.get());
diff --git a/rsa_operation.h b/rsa_operation.h
index 3ad4f0d..8283f2e 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;
 };
 
 /**