Add XorExtentWriter am: 4bb4920a57

Original change: https://android-review.googlesource.com/c/platform/system/update_engine/+/1761165

Change-Id: If2140625a35079de56b0743ce422861ca5a7f8ba
diff --git a/Android.bp b/Android.bp
index b3fbccf..83be871 100644
--- a/Android.bp
+++ b/Android.bp
@@ -242,6 +242,7 @@
         "payload_consumer/partition_writer.cc",
         "payload_consumer/partition_writer_factory_android.cc",
         "payload_consumer/vabc_partition_writer.cc",
+        "payload_consumer/xor_extent_writer.cc",
         "payload_consumer/block_extent_writer.cc",
         "payload_consumer/snapshot_extent_writer.cc",
         "payload_consumer/postinstall_runner_action.cc",
@@ -806,6 +807,7 @@
         "payload_consumer/extent_writer_unittest.cc",
         "payload_consumer/extent_map_unittest.cc",
         "payload_consumer/snapshot_extent_writer_unittest.cc",
+        "payload_consumer/xor_extent_writer_unittest.cc",
         "payload_consumer/fake_file_descriptor.cc",
         "payload_consumer/file_descriptor_utils_unittest.cc",
         "payload_consumer/file_writer_unittest.cc",
diff --git a/common/utils.cc b/common/utils.cc
index 5dbb445..25df540 100644
--- a/common/utils.cc
+++ b/common/utils.cc
@@ -1040,6 +1040,11 @@
   return ErrorCode::kSuccess;
 }
 
+std::unique_ptr<android::base::MappedFile> GetReadonlyZeroBlock(size_t size) {
+  android::base::unique_fd fd{HANDLE_EINTR(open("/dev/zero", O_RDONLY))};
+  return android::base::MappedFile::FromFd(fd, 0, size, PROT_READ);
+}
+
 }  // namespace utils
 
 }  // namespace chromeos_update_engine
diff --git a/common/utils.h b/common/utils.h
index 4e78097..9ed0c7b 100644
--- a/common/utils.h
+++ b/common/utils.h
@@ -36,6 +36,7 @@
 #include <brillo/key_value_store.h>
 #include <brillo/secure_blob.h>
 
+#include "android-base/mapped_file.h"
 #include "update_engine/common/action.h"
 #include "update_engine/common/action_processor.h"
 #include "update_engine/common/constants.h"
@@ -357,6 +358,8 @@
 ErrorCode IsTimestampNewer(const std::string& old_version,
                            const std::string& new_version);
 
+std::unique_ptr<android::base::MappedFile> GetReadonlyZeroBlock(size_t size);
+
 }  // namespace utils
 
 // Utility class to close a file descriptor
diff --git a/payload_consumer/block_extent_writer.cc b/payload_consumer/block_extent_writer.cc
index a5c990c..e50526c 100644
--- a/payload_consumer/block_extent_writer.cc
+++ b/payload_consumer/block_extent_writer.cc
@@ -45,7 +45,7 @@
 
   if (buffer_.empty() && count >= cur_extent_size) {
     if (!WriteExtent(data, cur_extent, block_size_)) {
-      LOG(ERROR) << "AddRawBlocks(" << cur_extent.start_block() << ", " << data
+      LOG(ERROR) << "WriteExtent(" << cur_extent.start_block() << ", " << data
                  << ", " << cur_extent_size << ") failed.";
       // return value is expected to be greater than 0. Return 0 to signal error
       // condition
diff --git a/payload_consumer/block_extent_writer.h b/payload_consumer/block_extent_writer.h
index 7c90459..f9c7b15 100644
--- a/payload_consumer/block_extent_writer.h
+++ b/payload_consumer/block_extent_writer.h
@@ -41,6 +41,7 @@
   virtual bool WriteExtent(const void* bytes,
                            const Extent& extent,
                            size_t block_size) = 0;
+  size_t BlockSize() const { return block_size_; }
 
  private:
   bool NextExtent();
diff --git a/payload_consumer/xor_extent_writer.cc b/payload_consumer/xor_extent_writer.cc
new file mode 100644
index 0000000..5c89c93
--- /dev/null
+++ b/payload_consumer/xor_extent_writer.cc
@@ -0,0 +1,115 @@
+//
+// 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 <algorithm>
+#include <optional>
+
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_consumer/xor_extent_writer.h"
+#include "update_engine/payload_generator/extent_utils.h"
+
+namespace chromeos_update_engine {
+
+// Returns true on success.
+bool XORExtentWriter::WriteExtent(const void* bytes,
+                                  const Extent& extent,
+                                  const size_t size) {
+  brillo::Blob xor_block_data;
+  const auto xor_extents = xor_map_.GetIntersectingExtents(extent);
+  for (const auto& xor_ext : xor_extents) {
+    const auto merge_op_opt = xor_map_.Get(xor_ext);
+    if (!merge_op_opt.has_value()) {
+      LOG(ERROR)
+          << xor_ext
+          << " isn't in XOR map but it's returned by GetIntersectingExtents(), "
+             "this is a bug inside GetIntersectingExtents";
+      return false;
+    }
+
+    const auto merge_op = merge_op_opt.value();
+    TEST_AND_RETURN_FALSE(merge_op->has_src_extent());
+    TEST_AND_RETURN_FALSE(merge_op->has_dst_extent());
+    if (merge_op->dst_extent() != xor_ext) {
+      LOG(ERROR) << "Each xor extent is expected to correspond to a complete "
+                    "MergeOp, extent in value: "
+                 << merge_op->dst_extent() << " extent in key: " << xor_ext;
+      return false;
+    }
+    if (xor_ext.start_block() + xor_ext.num_blocks() >
+        extent.start_block() + extent.num_blocks()) {
+      LOG(ERROR) << "CowXor merge op extent should be completely inside "
+                    "InstallOp's extent. merge op extent: "
+                 << xor_ext << " InstallOp extent: " << extent;
+      return false;
+    }
+    const auto src_offset = merge_op->src_offset();
+    const auto src_block = merge_op->src_extent().start_block();
+    xor_block_data.resize(BlockSize() * xor_ext.num_blocks());
+    ssize_t bytes_read = 0;
+    TEST_AND_RETURN_FALSE_ERRNO(
+        utils::PReadAll(source_fd_,
+                        xor_block_data.data(),
+                        xor_block_data.size(),
+                        src_offset + src_block * BlockSize(),
+                        &bytes_read));
+    if (bytes_read != static_cast<ssize_t>(xor_block_data.size())) {
+      LOG(ERROR) << "bytes_read: " << bytes_read;
+      return false;
+    }
+
+    const auto i = xor_ext.start_block() - extent.start_block();
+
+    const auto dst_block_data =
+        static_cast<const unsigned char*>(bytes) + i * BlockSize();
+    std::transform(xor_block_data.cbegin(),
+                   xor_block_data.cbegin() + xor_block_data.size(),
+                   dst_block_data,
+                   xor_block_data.begin(),
+                   std::bit_xor<unsigned char>{});
+    TEST_AND_RETURN_FALSE(cow_writer_->AddXorBlocks(xor_ext.start_block(),
+                                                    xor_block_data.data(),
+                                                    xor_block_data.size(),
+                                                    src_block,
+                                                    src_offset));
+  }
+  const auto replace_extents = xor_map_.GetNonIntersectingExtents(extent);
+  return WriteReplaceExtents(replace_extents, extent, bytes, size);
+}
+
+bool XORExtentWriter::WriteReplaceExtents(
+    const std::vector<Extent>& replace_extents,
+    const Extent& extent,
+    const void* bytes,
+    size_t size) {
+  const uint64_t new_block_start = extent.start_block();
+  for (const auto& ext : replace_extents) {
+    if (ext.start_block() + ext.num_blocks() >
+        extent.start_block() + extent.num_blocks()) {
+      LOG(ERROR) << "CowReplace merge op extent should be completely inside "
+                    "InstallOp's extent. merge op extent: "
+                 << ext << " InstallOp extent: " << extent;
+      return false;
+    }
+    const auto i = ext.start_block() - new_block_start;
+    const auto dst_block_data =
+        static_cast<const unsigned char*>(bytes) + i * BlockSize();
+    TEST_AND_RETURN_FALSE(cow_writer_->AddRawBlocks(
+        ext.start_block(), dst_block_data, ext.num_blocks() * BlockSize()));
+  }
+  return true;
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_consumer/xor_extent_writer.h b/payload_consumer/xor_extent_writer.h
new file mode 100644
index 0000000..94b5a9e
--- /dev/null
+++ b/payload_consumer/xor_extent_writer.h
@@ -0,0 +1,66 @@
+//
+// 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.
+//
+
+#ifndef UPDATE_ENGINE_PAYLOAD_CONSUMER_XOR_EXTENT_WRITER_H_
+#define UPDATE_ENGINE_PAYLOAD_CONSUMER_XOR_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/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>
+#include <libsnapshot/cow_writer.h>
+
+namespace chromeos_update_engine {
+
+// An extent writer that will selectively convert some of the blocks into an XOR
+// block. All blocks that appear in |xor_map| will be converted,
+class XORExtentWriter : public BlockExtentWriter {
+ public:
+  XORExtentWriter(const InstallOperation& op,
+                  FileDescriptorPtr source_fd,
+                  android::snapshot::ICowWriter* cow_writer,
+                  const ExtentMap<const CowMergeOperation*>& xor_map)
+      : src_extents_(op.src_extents()),
+        source_fd_(source_fd),
+        xor_map_(xor_map),
+        cow_writer_(cow_writer) {
+    CHECK(source_fd->IsOpen());
+  }
+  ~XORExtentWriter() = default;
+
+  // Returns true on success.
+  bool WriteExtent(const void* bytes,
+                   const Extent& extent,
+                   size_t size) override;
+
+ private:
+  bool WriteReplaceExtents(const std::vector<Extent>& replace_extents,
+                           const Extent& extent,
+                           const void* bytes,
+                           size_t size);
+  const google::protobuf::RepeatedPtrField<Extent>& src_extents_;
+  const FileDescriptorPtr source_fd_;
+  const ExtentMap<const CowMergeOperation*>& xor_map_;
+  android::snapshot::ICowWriter* cow_writer_;
+};
+
+}  // namespace chromeos_update_engine
+
+#endif  // UPDATE_ENGINE_PAYLOAD_CONSUMER_XOR_EXTENT_WRITER_H_
diff --git a/payload_consumer/xor_extent_writer_unittest.cc b/payload_consumer/xor_extent_writer_unittest.cc
new file mode 100644
index 0000000..7f35bc2
--- /dev/null
+++ b/payload_consumer/xor_extent_writer_unittest.cc
@@ -0,0 +1,134 @@
+//
+// 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 <memory>
+
+#include <unistd.h>
+
+#include <android-base/file.h>
+#include <gtest/gtest.h>
+#include <libsnapshot/mock_snapshot_writer.h>
+
+#include "common/utils.h"
+#include "update_engine/payload_consumer/extent_map.h"
+#include "update_engine/payload_consumer/file_descriptor.h"
+#include "update_engine/payload_consumer/xor_extent_writer.h"
+#include "update_engine/payload_generator/delta_diff_generator.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/payload_generator/merge_sequence_generator.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+
+using testing::_;
+using testing::Args;
+using testing::Return;
+
+class XorExtentWriterTest : public ::testing::Test {
+ public:
+  static constexpr size_t NUM_BLOCKS = 50;
+  void SetUp() override {
+    ASSERT_EQ(ftruncate64(source_part_.fd, kBlockSize * NUM_BLOCKS), 0);
+    ASSERT_EQ(ftruncate64(target_part_.fd, kBlockSize * NUM_BLOCKS), 0);
+
+    // Fill source part with 1s, as we are computing XOR between source and
+    // target data later.
+    ASSERT_EQ(lseek(source_part_.fd, 0, SEEK_SET), 0);
+    brillo::Blob buffer(kBlockSize);
+    std::fill(buffer.begin(), buffer.end(), 1);
+    for (size_t i = 0; i < NUM_BLOCKS; i++) {
+      ASSERT_EQ(write(source_part_.fd, buffer.data(), buffer.size()),
+                static_cast<ssize_t>(buffer.size()));
+    }
+    ASSERT_EQ(fsync(source_part_.fd), 0);
+    ASSERT_EQ(fsync(target_part_.fd), 0);
+    ASSERT_TRUE(source_fd_->Open(source_part_.path, O_RDONLY | O_CREAT, 0644));
+  }
+  InstallOperation op_;
+  FileDescriptorPtr source_fd_ = std::make_shared<EintrSafeFileDescriptor>();
+  ExtentMap<const CowMergeOperation*> xor_map_;
+  android::snapshot::MockSnapshotWriter cow_writer_;
+  TemporaryFile source_part_;
+  TemporaryFile target_part_;
+};
+
+MATCHER_P2(BytesEqual,
+           bytes,
+           size,
+           "Check if args match expected value byte for byte") {
+  return std::get<1>(arg) == size && std::get<0>(arg) != nullptr &&
+         memcmp(std::get<0>(arg), bytes, size) == 0;
+}
+
+TEST_F(XorExtentWriterTest, StreamTest) {
+  constexpr auto COW_XOR = CowMergeOperation::COW_XOR;
+  ON_CALL(cow_writer_, EmitXorBlocks(_, _, _, _, _))
+      .WillByDefault(Return(true));
+  const auto op1 = CreateCowMergeOperation(
+      ExtentForRange(5, 2), ExtentForRange(5, 2), COW_XOR);
+  ASSERT_TRUE(xor_map_.AddExtent(op1.dst_extent(), &op1));
+  *op_.add_src_extents() = op1.src_extent();
+  *op_.add_dst_extents() = op1.dst_extent();
+
+  const auto op2 = CreateCowMergeOperation(
+      ExtentForRange(45, 2), ExtentForRange(456, 2), COW_XOR);
+  ASSERT_TRUE(xor_map_.AddExtent(op2.dst_extent(), &op2));
+  *op_.add_src_extents() = ExtentForRange(45, 3);
+  *op_.add_dst_extents() = ExtentForRange(455, 3);
+
+  const auto op3 = CreateCowMergeOperation(
+      ExtentForRange(12, 2), ExtentForRange(321, 2), COW_XOR, 777);
+  ASSERT_TRUE(xor_map_.AddExtent(op3.dst_extent(), &op3));
+  *op_.add_src_extents() = ExtentForRange(12, 4);
+  *op_.add_dst_extents() = ExtentForRange(320, 4);
+  XORExtentWriter writer_{op_, source_fd_, &cow_writer_, xor_map_};
+
+  // OTA op:
+  // [5-6] => [5-6], [45-47] => [455-457], [12-15] => [320-323]
+
+  // merge op:
+  // [5-6] => [5-6], [45-46] => [456-457], [12-13] => [321-322]
+
+  // Expected result:
+  // [5-7], [45-47], [12-14] should be XOR blocks
+  // [320], [323], [455] should be regular replace blocks
+
+  auto zeros = utils::GetReadonlyZeroBlock(kBlockSize * 10);
+  EXPECT_CALL(cow_writer_,
+              EmitRawBlocks(455, zeros->data() + 2 * kBlockSize, kBlockSize))
+      .With(Args<1, 2>(BytesEqual(zeros->data(), kBlockSize)))
+      .WillOnce(Return(true));
+  EXPECT_CALL(cow_writer_,
+              EmitRawBlocks(320, zeros->data() + 5 * kBlockSize, kBlockSize))
+      .With(Args<1, 2>(BytesEqual(zeros->data(), kBlockSize)))
+      .WillOnce(Return(true));
+  EXPECT_CALL(cow_writer_,
+              EmitRawBlocks(323, zeros->data() + 8 * kBlockSize, kBlockSize))
+      .With(Args<1, 2>(BytesEqual(zeros->data(), kBlockSize)))
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(cow_writer_, EmitXorBlocks(5, _, kBlockSize * 2, 5, 0))
+      .WillOnce(Return(true));
+  EXPECT_CALL(cow_writer_, EmitXorBlocks(456, _, kBlockSize * 2, 45, 0))
+      .WillOnce(Return(true));
+  EXPECT_CALL(cow_writer_, EmitXorBlocks(321, _, kBlockSize * 2, 12, 777))
+      .WillOnce(Return(true));
+
+  ASSERT_TRUE(writer_.Init(op_.dst_extents(), kBlockSize));
+  ASSERT_TRUE(writer_.Write(zeros->data(), 9 * kBlockSize));
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_utils.cc b/payload_generator/extent_utils.cc
index f4a9ff0..612fd67 100644
--- a/payload_generator/extent_utils.cc
+++ b/payload_generator/extent_utils.cc
@@ -26,6 +26,7 @@
 #include <base/macros.h>
 #include <base/strings/stringprintf.h>
 
+#include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/payload_constants.h"
 #include "update_engine/payload_generator/annotated_operation.h"
 #include "update_engine/payload_generator/extent_ranges.h"
@@ -161,14 +162,44 @@
   return result;
 }
 
-bool operator==(const Extent& a, const Extent& b) {
+bool operator==(const Extent& a, const Extent& b) noexcept {
   return a.start_block() == b.start_block() && a.num_blocks() == b.num_blocks();
 }
 
+bool operator!=(const Extent& a, const Extent& b) noexcept {
+  return !(a == b);
+}
+
 std::ostream& operator<<(std::ostream& out, const Extent& extent) {
   out << "[" << extent.start_block() << " - "
       << extent.start_block() + extent.num_blocks() - 1 << "]";
   return out;
 }
 
+template <typename T>
+std::ostream& PrintExtents(std::ostream& out, const T& extents) {
+  if (extents.begin() == extents.end()) {
+    out << "{}";
+    return out;
+  }
+  out << "{";
+  auto begin = extents.begin();
+  out << *begin;
+  for (const auto& ext : Range{++begin, extents.end()}) {
+    out << ", " << ext;
+  }
+  out << "}";
+  return out;
+}
+
+std::ostream& operator<<(std::ostream& out,
+                         const std::vector<Extent>& extents) {
+  return PrintExtents(out, extents);
+}
+std::ostream& operator<<(
+    std::ostream& out,
+    const google::protobuf::RepeatedPtrField<Extent>& extents) {
+  return PrintExtents(out, extents);
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_utils.h b/payload_generator/extent_utils.h
index 2bd6626..f5bb39c 100644
--- a/payload_generator/extent_utils.h
+++ b/payload_generator/extent_utils.h
@@ -86,7 +86,9 @@
                                    uint64_t block_offset,
                                    uint64_t block_count);
 
-bool operator==(const Extent& a, const Extent& b);
+bool operator==(const Extent& a, const Extent& b) noexcept;
+
+bool operator!=(const Extent& a, const Extent& b) noexcept;
 
 // TODO(zhangkelvin) This is ugly. Rewrite using C++20's coroutine once
 // that's available. Unfortunately with C++17 this is the best I could do.
@@ -126,6 +128,10 @@
 };
 
 std::ostream& operator<<(std::ostream& out, const Extent& extent);
+std::ostream& operator<<(std::ostream& out, const std::vector<Extent>& extent);
+std::ostream& operator<<(
+    std::ostream& out,
+    const google::protobuf::RepeatedPtrField<Extent>& extent);
 
 template <typename Container>
 size_t GetNthBlock(const Container& extents, const size_t n) {