Write XOR ops for bsdiff and puffdiff

Test: th
Bug: 177104308

Change-Id: I10f79f1a2de910e5ba57a72a5dba5901c9348d77
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index 3cdd4a4..a5d03d9 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -16,20 +16,26 @@
 
 #include "update_engine/payload_consumer/vabc_partition_writer.h"
 
+#include <algorithm>
+#include <map>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include <brillo/secure_blob.h>
 #include <libsnapshot/cow_writer.h>
 
 #include "update_engine/common/cow_operation_convert.h"
 #include "update_engine/common/utils.h"
-#include "update_engine/payload_consumer/extent_writer.h"
+#include "update_engine/payload_consumer/block_extent_writer.h"
+#include "update_engine/payload_consumer/extent_map.h"
+#include "update_engine/payload_consumer/extent_reader.h"
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/partition_writer.h"
 #include "update_engine/payload_consumer/snapshot_extent_writer.h"
+#include "update_engine/payload_consumer/xor_extent_writer.h"
 #include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/payload_generator/extent_utils.h"
 #include "update_engine/update_metadata.pb.h"
@@ -62,6 +68,18 @@
 using android::snapshot::ICowWriter;
 using ::google::protobuf::RepeatedPtrField;
 
+// Compute XOR map, a map from dst extent to corresponding merge operation
+static ExtentMap<const CowMergeOperation*, ExtentLess> ComputeXorMap(
+    const RepeatedPtrField<CowMergeOperation>& merge_ops) {
+  ExtentMap<const CowMergeOperation*, ExtentLess> xor_map;
+  for (const auto& merge_op : merge_ops) {
+    if (merge_op.type() == CowMergeOperation::COW_XOR) {
+      xor_map.AddExtent(merge_op.dst_extent(), &merge_op);
+    }
+  }
+  return xor_map;
+}
+
 VABCPartitionWriter::VABCPartitionWriter(
     const PartitionUpdate& partition_update,
     const InstallPlan::Partition& install_part,
@@ -77,6 +95,7 @@
 bool VABCPartitionWriter::Init(const InstallPlan* install_plan,
                                bool source_may_exist,
                                size_t next_op_index) {
+  xor_map_ = ComputeXorMap(partition_update_.merge_operations());
   TEST_AND_RETURN_FALSE(install_plan != nullptr);
   if (source_may_exist) {
     TEST_AND_RETURN_FALSE(verified_source_fd_.Open());
@@ -123,7 +142,8 @@
     auto source_fd = std::make_shared<EintrSafeFileDescriptor>();
     TEST_AND_RETURN_FALSE_ERRNO(
         source_fd->Open(install_part_.source_path.c_str(), O_RDONLY));
-    WriteAllCowOps(block_size_, converted, cow_writer_.get(), source_fd);
+    TEST_AND_RETURN_FALSE(WriteSourceCopyCowOps(
+        block_size_, converted, cow_writer_.get(), source_fd));
     cow_writer_->AddLabel(0);
   }
   return true;
@@ -141,12 +161,6 @@
         merge_op.src_extent() == merge_op.dst_extent()) {
       continue;
     }
-    // libsnapshot doesn't like us include REPLACE blocks in merge sequence,
-    // since we don't support XOR ops at this CL, skip them! Remove once we
-    // write XOR ops.
-    if (merge_op.type() == CowMergeOperation::COW_XOR) {
-      continue;
-    }
     // libsnapshot prefers blocks in reverse order
     for (int i = dst_extent.num_blocks() - 1; i >= 0; i--) {
       blocks_merge_order.push_back(dst_extent.start_block() + i);
@@ -156,7 +170,7 @@
                                      blocks_merge_order.data());
 }
 
-bool VABCPartitionWriter::WriteAllCowOps(
+bool VABCPartitionWriter::WriteSourceCopyCowOps(
     size_t block_size,
     const std::vector<CowOperation>& converted,
     ICowWriter* cow_writer,
@@ -229,8 +243,12 @@
   FileDescriptorPtr source_fd =
       verified_source_fd_.ChooseSourceFD(operation, error);
   TEST_AND_RETURN_FALSE(source_fd != nullptr);
+  TEST_AND_RETURN_FALSE(source_fd->IsOpen());
 
-  auto writer = CreateBaseExtentWriter();
+  std::unique_ptr<ExtentWriter> writer =
+      IsXorEnabled() ? std::make_unique<XORExtentWriter>(
+                           operation, source_fd, cow_writer_.get(), xor_map_)
+                     : CreateBaseExtentWriter();
   return executor_.ExecuteDiffOperation(
       operation, std::move(writer), source_fd, data, count);
 }
diff --git a/payload_consumer/vabc_partition_writer.h b/payload_consumer/vabc_partition_writer.h
index fb63f8c..4df5151 100644
--- a/payload_consumer/vabc_partition_writer.h
+++ b/payload_consumer/vabc_partition_writer.h
@@ -17,6 +17,7 @@
 #ifndef UPDATE_ENGINE_VABC_PARTITION_WRITER_H_
 #define UPDATE_ENGINE_VABC_PARTITION_WRITER_H_
 
+#include <map>
 #include <memory>
 #include <string>
 #include <vector>
@@ -24,9 +25,11 @@
 #include <libsnapshot/snapshot_writer.h>
 
 #include "update_engine/common/cow_operation_convert.h"
+#include "update_engine/payload_consumer/extent_map.h"
 #include "update_engine/payload_consumer/install_operation_executor.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/partition_writer.h"
+#include "update_engine/payload_generator/extent_ranges.h"
 
 namespace chromeos_update_engine {
 class VABCPartitionWriter final : public PartitionWriterInterface {
@@ -59,10 +62,11 @@
 
   void CheckpointUpdateProgress(size_t next_op_index) override;
 
-  static bool WriteAllCowOps(size_t block_size,
-                             const std::vector<CowOperation>& converted,
-                             android::snapshot::ICowWriter* cow_writer,
-                             FileDescriptorPtr source_fd);
+  [[nodiscard]] static bool WriteSourceCopyCowOps(
+      size_t block_size,
+      const std::vector<CowOperation>& converted,
+      android::snapshot::ICowWriter* cow_writer,
+      FileDescriptorPtr source_fd);
 
   [[nodiscard]] bool FinishedInstallOps() override;
   int Close() override;
@@ -72,19 +76,21 @@
       android::snapshot::ICowWriter* cow_writer);
 
  private:
+  bool IsXorEnabled() const noexcept { return xor_map_.size() > 0; }
   std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer_;
 
   [[nodiscard]] std::unique_ptr<ExtentWriter> CreateBaseExtentWriter();
 
   const PartitionUpdate& partition_update_;
   const InstallPlan::Partition& install_part_;
-  DynamicPartitionControlInterface* dynamic_control_;
+  DynamicPartitionControlInterface* const dynamic_control_;
   // Path to source partition
-  std::string source_path_;
+  const std::string source_path_;
 
   const size_t block_size_;
   InstallOperationExecutor executor_;
   VerifiedSourceFd verified_source_fd_;
+  ExtentMap<const CowMergeOperation*, ExtentLess> xor_map_;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/vabc_partition_writer_unittest.cc b/payload_consumer/vabc_partition_writer_unittest.cc
new file mode 100644
index 0000000..d77006e
--- /dev/null
+++ b/payload_consumer/vabc_partition_writer_unittest.cc
@@ -0,0 +1,201 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include <unistd.h>
+
+#include <android-base/file.h>
+#include <android-base/mapped_file.h>
+#include <bsdiff/bsdiff.h>
+#include <gtest/gtest.h>
+#include <libsnapshot/cow_writer.h>
+#include <libsnapshot/mock_snapshot_writer.h>
+
+#include "update_engine/common/hash_calculator.h"
+#include "update_engine/common/mock_dynamic_partition_control.h"
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_consumer/vabc_partition_writer.h"
+#include "update_engine/payload_generator/delta_diff_generator.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+
+using android::snapshot::CowOptions;
+using testing::_;
+using testing::Args;
+using testing::ElementsAreArray;
+using testing::Invoke;
+using testing::Return;
+using testing::Sequence;
+using utils::GetReadonlyZeroBlock;
+
+static constexpr auto& fake_part_name = "fake_part";
+class VABCPartitionWriterTest : public ::testing::Test {
+ public:
+  void SetUp() override {}
+
+ protected:
+  void AddMergeOp(PartitionUpdate* partition,
+                  std::array<size_t, 2> src_extent,
+                  std::array<size_t, 2> dst_extent,
+                  CowMergeOperation_Type type) {
+    auto merge_op = partition->add_merge_operations();
+    auto src = merge_op->mutable_src_extent();
+    src->set_start_block(src_extent[0]);
+    src->set_num_blocks(src_extent[1]);
+    auto dst = merge_op->mutable_dst_extent();
+    dst->set_start_block(dst_extent[0]);
+    dst->set_num_blocks(dst_extent[1]);
+    merge_op->set_type(type);
+  }
+
+  android::snapshot::CowOptions options_ = {
+      .block_size = static_cast<uint32_t>(kBlockSize)};
+  android::snapshot::MockSnapshotWriter cow_writer_{options_};
+  MockDynamicPartitionControl dynamic_control_;
+  PartitionUpdate partition_update_;
+  InstallPlan install_plan_;
+  TemporaryFile source_part_;
+  InstallPlan::Partition install_part_{
+      .name = fake_part_name,
+      .source_path = source_part_.path,
+  };
+};
+
+TEST_F(VABCPartitionWriterTest, MergeSequenceWriteTest) {
+  AddMergeOp(&partition_update_, {5, 1}, {10, 1}, CowMergeOperation::COW_COPY);
+  AddMergeOp(&partition_update_, {10, 1}, {15, 1}, CowMergeOperation::COW_XOR);
+  AddMergeOp(&partition_update_, {15, 1}, {20, 1}, CowMergeOperation::COW_COPY);
+  AddMergeOp(&partition_update_, {20, 1}, {25, 1}, CowMergeOperation::COW_COPY);
+  VABCPartitionWriter writer_{
+      partition_update_, install_part_, &dynamic_control_, kBlockSize};
+  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
+      .WillOnce(Invoke(
+          [](const std::string&, const std::optional<std::string>&, bool) {
+            auto cow_writer =
+                std::make_unique<android::snapshot::MockSnapshotWriter>(
+                    android::snapshot::CowOptions{});
+            auto expected_merge_sequence = {10, 15, 20, 25};
+            EXPECT_CALL(*cow_writer, Initialize()).WillOnce(Return(true));
+            EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+                .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
+                .WillOnce(Return(true));
+            ON_CALL(*cow_writer, EmitCopy(_, _)).WillByDefault(Return(true));
+            ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
+            return cow_writer;
+          }));
+  ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
+}
+
+TEST_F(VABCPartitionWriterTest, EmitBlockTest) {
+  AddMergeOp(&partition_update_, {5, 1}, {10, 1}, CowMergeOperation::COW_COPY);
+  AddMergeOp(&partition_update_, {10, 1}, {15, 1}, CowMergeOperation::COW_COPY);
+  AddMergeOp(&partition_update_, {15, 2}, {20, 2}, CowMergeOperation::COW_COPY);
+  AddMergeOp(&partition_update_, {20, 1}, {25, 1}, CowMergeOperation::COW_COPY);
+  VABCPartitionWriter writer_{
+      partition_update_, install_part_, &dynamic_control_, kBlockSize};
+  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
+      .WillOnce(Invoke(
+          [](const std::string&, const std::optional<std::string>&, bool) {
+            auto cow_writer =
+                std::make_unique<android::snapshot::MockSnapshotWriter>(
+                    android::snapshot::CowOptions{});
+            Sequence s;
+            ON_CALL(*cow_writer, EmitCopy(_, _)).WillByDefault(Return(true));
+            ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
+            ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
+            EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+                .Times(1)
+                .WillOnce(Return(true));
+            EXPECT_CALL(*cow_writer, Initialize()).InSequence(s);
+            EXPECT_CALL(*cow_writer, EmitCopy(10, 5)).InSequence(s);
+            EXPECT_CALL(*cow_writer, EmitCopy(15, 10)).InSequence(s);
+            // libsnapshot want blocks in reverser order, so 21 goes before 20
+            EXPECT_CALL(*cow_writer, EmitCopy(21, 16)).InSequence(s);
+            EXPECT_CALL(*cow_writer, EmitCopy(20, 15)).InSequence(s);
+
+            EXPECT_CALL(*cow_writer, EmitCopy(25, 20)).InSequence(s);
+            return cow_writer;
+          }));
+  ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
+}
+
+std::string GetNoopBSDIFF(size_t data_size) {
+  auto zeros = GetReadonlyZeroBlock(data_size);
+  TemporaryFile patch_file;
+  int error = bsdiff::bsdiff(reinterpret_cast<const uint8_t*>(zeros->data()),
+                             zeros->size(),
+                             reinterpret_cast<const uint8_t*>(zeros->data()),
+                             zeros->size(),
+                             patch_file.path,
+                             nullptr);
+  if (error) {
+    LOG(ERROR) << "Failed to generate BSDIFF patch " << error;
+    return {};
+  }
+  std::string patch_data;
+  if (!utils::ReadFile(patch_file.path, &patch_data)) {
+    return {};
+  }
+  return patch_data;
+}
+
+TEST_F(VABCPartitionWriterTest, StreamXORBlockTest) {
+  AddMergeOp(&partition_update_, {5, 2}, {10, 2}, CowMergeOperation::COW_XOR);
+  AddMergeOp(&partition_update_, {8, 2}, {13, 2}, CowMergeOperation::COW_XOR);
+  auto install_op = partition_update_.add_operations();
+  *install_op->add_src_extents() = ExtentForRange(5, 5);
+  *install_op->add_dst_extents() = ExtentForRange(10, 5);
+  install_op->set_type(InstallOperation::SOURCE_BSDIFF);
+  auto data_hash = install_op->mutable_src_sha256_hash();
+  auto zeros = GetReadonlyZeroBlock(kBlockSize * 5);
+  brillo::Blob expected_hash;
+  truncate64(source_part_.path, kBlockSize * 20);
+  HashCalculator::RawHashOfBytes(zeros->data(), zeros->size(), &expected_hash);
+  data_hash->assign(reinterpret_cast<const char*>(expected_hash.data()),
+                    expected_hash.size());
+
+  EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
+      .WillOnce(Invoke([](const std::string&,
+                          const std::optional<std::string>&,
+                          bool) {
+        auto cow_writer =
+            std::make_unique<android::snapshot::MockSnapshotWriter>(
+                android::snapshot::CowOptions{});
+        ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
+        auto expected_merge_sequence = {11, 10, 14, 13};
+        ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
+        EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+            .With(Args<1, 0>(ElementsAreArray(expected_merge_sequence)))
+            .WillOnce(Return(true));
+        EXPECT_CALL(*cow_writer, Initialize()).Times(1);
+        EXPECT_CALL(*cow_writer, EmitCopy(_, _)).Times(0);
+        EXPECT_CALL(*cow_writer, EmitRawBlocks(_, _, _)).WillOnce(Return(true));
+        EXPECT_CALL(*cow_writer, EmitXorBlocks(10, _, kBlockSize * 2, 5, 0))
+            .WillOnce(Return(true));
+        EXPECT_CALL(*cow_writer, EmitXorBlocks(13, _, kBlockSize * 2, 8, 0))
+            .WillOnce(Return(true));
+        return cow_writer;
+      }));
+  VABCPartitionWriter writer_{
+      partition_update_, install_part_, &dynamic_control_, kBlockSize};
+  ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
+  auto patch_data = GetNoopBSDIFF(kBlockSize * 5);
+  ASSERT_GT(patch_data.size(), 0UL);
+  ASSERT_TRUE(writer_.PerformDiffOperation(
+      *install_op, nullptr, patch_data.data(), patch_data.size()));
+}
+
+}  // namespace chromeos_update_engine