Keystore 2.0 km_compat: Buffer incomplete updates. am: b7f303146f

Original change: https://googleplex-android-review.googlesource.com/c/platform/system/security/+/15942712

Change-Id: I551bff3569dfa4e62f9b9c718bfe6f8f8e9161c6
diff --git a/keystore2/src/km_compat/km_compat.cpp b/keystore2/src/km_compat/km_compat.cpp
index 64849c1..8d59a5a 100644
--- a/keystore2/src/km_compat/km_compat.cpp
+++ b/keystore2/src/km_compat/km_compat.cpp
@@ -762,7 +762,21 @@
     return convertErrorCode(errorCode);
 }
 
-ScopedAStatus KeyMintOperation::update(const std::vector<uint8_t>& input,
+void KeyMintOperation::setUpdateBuffer(std::vector<uint8_t> data) {
+    mUpdateBuffer = std::move(data);
+}
+
+const std::vector<uint8_t>&
+KeyMintOperation::getExtendedUpdateBuffer(const std::vector<uint8_t>& suffix) {
+    if (mUpdateBuffer.empty()) {
+        return suffix;
+    } else {
+        mUpdateBuffer.insert(mUpdateBuffer.end(), suffix.begin(), suffix.end());
+        return mUpdateBuffer;
+    }
+}
+
+ScopedAStatus KeyMintOperation::update(const std::vector<uint8_t>& input_raw,
                                        const std::optional<HardwareAuthToken>& optAuthToken,
                                        const std::optional<TimeStampToken>& optTimeStampToken,
                                        std::vector<uint8_t>* out_output) {
@@ -772,8 +786,10 @@
     size_t inputPos = 0;
     *out_output = {};
     KMV1::ErrorCode errorCode = KMV1::ErrorCode::OK;
+    auto input = getExtendedUpdateBuffer(input_raw);
 
     while (inputPos < input.size() && errorCode == KMV1::ErrorCode::OK) {
+        uint32_t consumed = 0;
         auto result =
             mDevice->update(mOperationHandle, {} /* inParams */,
                             {input.begin() + inputPos, input.end()}, authToken, verificationToken,
@@ -781,13 +797,22 @@
                                 const hidl_vec<uint8_t>& output) {
                                 errorCode = convert(error);
                                 out_output->insert(out_output->end(), output.begin(), output.end());
-                                inputPos += inputConsumed;
+                                consumed = inputConsumed;
                             });
 
         if (!result.isOk()) {
             LOG(ERROR) << __func__ << " transaction failed. " << result.description();
             errorCode = KMV1::ErrorCode::UNKNOWN_ERROR;
         }
+
+        if (errorCode == KMV1::ErrorCode::OK && consumed == 0) {
+            // Some very old KM implementations do not buffer sub blocks in certain block modes,
+            // instead, the simply return consumed == 0. So we buffer the input here in the
+            // hope that we complete the bock in a future call to update.
+            setUpdateBuffer({input.begin() + inputPos, input.end()});
+            return convertErrorCode(errorCode);
+        }
+        inputPos += consumed;
     }
 
     if (errorCode != KMV1::ErrorCode::OK) mOperationSlot.freeSlot();
@@ -802,7 +827,8 @@
                          const std::optional<TimeStampToken>& in_timeStampToken,
                          const std::optional<std::vector<uint8_t>>& in_confirmationToken,
                          std::vector<uint8_t>* out_output) {
-    auto input = in_input.value_or(std::vector<uint8_t>());
+    auto input_raw = in_input.value_or(std::vector<uint8_t>());
+    auto input = getExtendedUpdateBuffer(input_raw);
     auto signature = in_signature.value_or(std::vector<uint8_t>());
     V4_0_HardwareAuthToken authToken = convertAuthTokenToLegacy(in_authToken);
     V4_0_VerificationToken verificationToken = convertTimestampTokenToLegacy(in_timeStampToken);
diff --git a/keystore2/src/km_compat/km_compat.h b/keystore2/src/km_compat/km_compat.h
index 2d892da..70c7b86 100644
--- a/keystore2/src/km_compat/km_compat.h
+++ b/keystore2/src/km_compat/km_compat.h
@@ -140,11 +140,6 @@
 };
 
 class KeyMintOperation : public aidl::android::hardware::security::keymint::BnKeyMintOperation {
-  private:
-    ::android::sp<Keymaster> mDevice;
-    uint64_t mOperationHandle;
-    OperationSlot mOperationSlot;
-
   public:
     KeyMintOperation(::android::sp<Keymaster> device, uint64_t operationHandle,
                      OperationSlots* slots, bool isActive)
@@ -168,6 +163,25 @@
                          std::vector<uint8_t>* output) override;
 
     ScopedAStatus abort();
+
+  private:
+    /**
+     * Sets mUpdateBuffer to the given value.
+     * @param data
+     */
+    void setUpdateBuffer(std::vector<uint8_t> data);
+    /**
+     * If mUpdateBuffer is not empty, suffix is appended to mUpdateBuffer, and a reference to
+     * mUpdateBuffer is returned. Otherwise a reference to suffix is returned.
+     * @param suffix
+     * @return
+     */
+    const std::vector<uint8_t>& getExtendedUpdateBuffer(const std::vector<uint8_t>& suffix);
+
+    std::vector<uint8_t> mUpdateBuffer;
+    ::android::sp<Keymaster> mDevice;
+    uint64_t mOperationHandle;
+    OperationSlot mOperationSlot;
 };
 
 class SharedSecret : public aidl::android::hardware::security::sharedsecret::BnSharedSecret {