Support zucchini patch

Adding code to handle the zucchini patches.

Bug: 197361113
Test: TH, brillo_update_payload generate & verify
Change-Id: I939b10621f13fcae8bf981fefe35971635a6f965
diff --git a/payload_consumer/install_operation_executor.cc b/payload_consumer/install_operation_executor.cc
index eb4efe6..47c6b5c 100644
--- a/payload_consumer/install_operation_executor.cc
+++ b/payload_consumer/install_operation_executor.cc
@@ -15,18 +15,22 @@
 //
 
 #include "update_engine/payload_consumer/install_operation_executor.h"
-#include <memory>
-#include <utility>
-#include <vector>
 
 #include <fcntl.h>
 #include <glob.h>
 #include <linux/fs.h>
 
+#include <memory>
+#include <utility>
+#include <vector>
+
 #include <base/files/memory_mapped_file.h>
+#include <base/files/file_util.h>
 #include <bsdiff/bspatch.h>
+#include <puffin/brotli_util.h>
 #include <puffin/puffpatch.h>
-#include <sys/mman.h>
+#include <zucchini/patch_reader.h>
+#include <zucchini/zucchini.h>
 
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/bzip_extent_writer.h"
@@ -36,7 +40,6 @@
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/file_descriptor_utils.h"
 #include "update_engine/payload_consumer/xz_extent_writer.h"
-#include "update_engine/payload_generator/delta_diff_generator.h"
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
@@ -309,8 +312,46 @@
     FileDescriptorPtr source_fd,
     const void* data,
     size_t count) {
-  LOG(ERROR) << "zucchini operation isn't supported";
-  return false;
+  uint64_t src_size =
+      utils::BlocksInExtents(operation.src_extents()) * block_size_;
+  brillo::Blob source_bytes(src_size);
+
+  // TODO(197361113) either make zucchini stream the read, or use memory mapped
+  // files.
+  auto reader = std::make_unique<DirectExtentReader>();
+  TEST_AND_RETURN_FALSE(
+      reader->Init(source_fd, operation.src_extents(), block_size_));
+  TEST_AND_RETURN_FALSE(reader->Seek(0));
+  TEST_AND_RETURN_FALSE(reader->Read(source_bytes.data(), src_size));
+
+  brillo::Blob zucchini_patch;
+  TEST_AND_RETURN_FALSE(puffin::BrotliDecode(
+      static_cast<const uint8_t*>(data), count, &zucchini_patch));
+  auto patch_reader = zucchini::EnsemblePatchReader::Create(
+      {zucchini_patch.data(), zucchini_patch.size()});
+  if (!patch_reader.has_value()) {
+    LOG(ERROR) << "Failed to parse the zucchini patch.";
+    return false;
+  }
+
+  auto dst_size = patch_reader->header().new_size;
+  TEST_AND_RETURN_FALSE(dst_size ==
+                        utils::BlocksInExtents(operation.dst_extents()) *
+                            block_size_);
+
+  brillo::Blob patched_data(dst_size);
+  auto status =
+      zucchini::ApplyBuffer({source_bytes.data(), source_bytes.size()},
+                            *patch_reader,
+                            {patched_data.data(), patched_data.size()});
+  if (status != zucchini::status::kStatusSuccess) {
+    LOG(ERROR) << "Failed to apply the zucchini patch: " << status;
+    return false;
+  }
+
+  TEST_AND_RETURN_FALSE(
+      writer->Write(patched_data.data(), patched_data.size()));
+  return true;
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/install_operation_executor_unittest.cc b/payload_consumer/install_operation_executor_unittest.cc
index e4135fc..a7ab7e9 100644
--- a/payload_consumer/install_operation_executor_unittest.cc
+++ b/payload_consumer/install_operation_executor_unittest.cc
@@ -31,11 +31,17 @@
 #include <brillo/secure_blob.h>
 #include <gtest/gtest.h>
 #include <update_engine/update_metadata.pb.h>
+#include <zucchini/buffer_view.h>
+#include <zucchini/patch_writer.h>
+#include <zucchini/zucchini.h>
+#include <puffin/brotli_util.h>
 
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/extent_writer.h"
+#include "update_engine/payload_consumer/fake_extent_writer.h"
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/payload_constants.h"
+#include "update_engine/payload_generator/delta_diff_utils.h"
 #include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/payload_generator/extent_utils.h"
 
@@ -57,18 +63,27 @@
   static constexpr size_t BLOCK_SIZE = 4096;
   void SetUp() override {
     // Fill source partition with arbitrary data.
-    std::array<uint8_t, BLOCK_SIZE> buffer{};
+    source_data_.resize(NUM_BLOCKS * BLOCK_SIZE);
+    target_data_.resize(NUM_BLOCKS * BLOCK_SIZE);
     for (size_t i = 0; i < NUM_BLOCKS; i++) {
       // Fill block with arbitrary data. We don't care about what data is being
       // written to source partition, so as long as each block is slightly
       // different.
-      std::fill(buffer.begin(), buffer.end(), i);
-      ASSERT_TRUE(utils::WriteAll(source_.fd(), buffer.data(), buffer.size()))
-          << "Failed to write to source partition file: " << strerror(errno);
-      std::fill(buffer.begin(), buffer.end(), NUM_BLOCKS + i);
-      ASSERT_TRUE(utils::WriteAll(target_.fd(), buffer.data(), buffer.size()))
-          << "Failed to write to target partition file: " << strerror(errno);
+      uint32_t offset = i * BLOCK_SIZE;
+      std::fill(source_data_.begin() + offset,
+                source_data_.begin() + offset + BLOCK_SIZE,
+                i);
+      std::fill(target_data_.begin() + offset,
+                target_data_.begin() + offset + BLOCK_SIZE,
+                NUM_BLOCKS + i);
     }
+
+    ASSERT_TRUE(
+        utils::WriteAll(source_.fd(), source_data_.data(), source_data_.size()))
+        << "Failed to write to source partition file: " << strerror(errno);
+    ASSERT_TRUE(
+        utils::WriteAll(target_.fd(), target_data_.data(), target_data_.size()))
+        << "Failed to write to target partition file: " << strerror(errno);
     fsync(source_.fd());
     fsync(target_.fd());
 
@@ -111,6 +126,9 @@
   ScopedTempFile target_{"target_partition.XXXXXXXX", true};
   FileDescriptorPtr source_fd_ = std::make_shared<EintrSafeFileDescriptor>();
   FileDescriptorPtr target_fd_ = std::make_shared<EintrSafeFileDescriptor>();
+  std::vector<uint8_t> source_data_;
+  std::vector<uint8_t> target_data_;
+
   InstallOperationExecutor executor_{BLOCK_SIZE};
 };
 
@@ -206,6 +224,42 @@
   VerityUntouchedExtents(op);
 }
 
+TEST_F(InstallOperationExecutorTest, ZucchiniOpTest) {
+  InstallOperation op;
+  op.set_type(InstallOperation::ZUCCHINI);
+  *op.mutable_src_extents()->Add() = ExtentForRange(0, NUM_BLOCKS);
+  *op.mutable_dst_extents()->Add() = ExtentForRange(0, NUM_BLOCKS);
+
+  // Make a zucchini patch
+  std::vector<Extent> src_extents{ExtentForRange(0, NUM_BLOCKS)};
+  std::vector<Extent> dst_extents{ExtentForRange(0, NUM_BLOCKS)};
+  PayloadGenerationConfig config{
+      .version = PayloadVersion(kBrilloMajorPayloadVersion,
+                                kZucchiniMinorPayloadVersion)};
+  diff_utils::BestDiffGenerator best_diff_generator(
+      source_data_, target_data_, src_extents, dst_extents, {}, {}, config);
+  std::vector<uint8_t> patch_data = target_data_;  // Fake the full operation
+  AnnotatedOperation aop;
+  ASSERT_TRUE(best_diff_generator.GenerateBestDiffOperation(
+      {{InstallOperation::ZUCCHINI, 1024 * BLOCK_SIZE}}, &aop, &patch_data));
+  ASSERT_EQ(InstallOperation::ZUCCHINI, aop.op.type());
+
+  // Call the executor
+  ScopedTempFile patched{"patched.XXXXXXXX", true};
+  FileDescriptorPtr patched_fd = std::make_shared<EintrSafeFileDescriptor>();
+  patched_fd->Open(patched.path().c_str(), O_RDWR);
+  std::unique_ptr<ExtentWriter> writer(new DirectExtentWriter(patched_fd));
+  writer->Init(op.dst_extents(), BLOCK_SIZE);
+  ASSERT_TRUE(executor_.ExecuteDiffOperation(
+      op, std::move(writer), source_fd_, patch_data.data(), patch_data.size()));
+
+  // Compare the result
+  std::vector<uint8_t> patched_data;
+  ASSERT_TRUE(utils::ReadFile(patched.path(), &patched_data));
+  ASSERT_EQ(NUM_BLOCKS * BLOCK_SIZE, patched_data.size());
+  ASSERT_EQ(target_data_, patched_data);
+}
+
 TEST_F(InstallOperationExecutorTest, GetNthBlockTest) {
   std::vector<Extent> extents;
   extents.emplace_back(ExtentForRange(10, 3));