Keystore 2.0 km_compat: Buffer incomplete updates.

Older KM implementations do not consume data if in certain block modes
when too little data is presented. However, km_compat update assumes
that the backend always consumes some data. If this assumption does not
hold it can get stuck in an infinite loop.

This patch adds some buffering, allowing the km_compat to buffer
unconsumed data and make it appear to the caller that the data was
indeed consumed.

Ignore-AOSP-First: b/200041882 ASA review.

Bug: 200041882
Test: CtsKeystoreTestCases for regression testing
      keystore2_km_compat_test
Merged-In: Icae44c6bc97507f192ec44c944c3bc0a9dd60ba7
Change-Id: Icae44c6bc97507f192ec44c944c3bc0a9dd60ba7
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 {