Fix xor writer failure am: cb9932ff76

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

Change-Id: Ia4f388155f17bed48f1f750540246b07af4a5c9f
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/aosp/boot_control_android.h b/aosp/boot_control_android.h
index 018c00c..51923e2 100644
--- a/aosp/boot_control_android.h
+++ b/aosp/boot_control_android.h
@@ -26,7 +26,7 @@
 #include <BootControlClient.h>
 
 #include "update_engine/aosp/dynamic_partition_control_android.h"
-#include "update_engine/common/boot_control.h"
+#include "update_engine/common/boot_control_interface.h"
 #include "update_engine/common/dynamic_partition_control_interface.h"
 
 namespace chromeos_update_engine {
diff --git a/common/prefs.cc b/common/prefs.cc
index 41a7b24..a070302 100644
--- a/common/prefs.cc
+++ b/common/prefs.cc
@@ -25,7 +25,6 @@
 #include <base/strings/string_split.h>
 #include <base/strings/string_util.h>
 
-#include "brillo/file_utils.h"
 #include "update_engine/common/utils.h"
 
 using std::string;
diff --git a/payload_consumer/extent_map.h b/payload_consumer/extent_map.h
index a985cb5..a83bf0f 100644
--- a/payload_consumer/extent_map.h
+++ b/payload_consumer/extent_map.h
@@ -56,6 +56,15 @@
     const auto it = map_.find(extent);
     if (it == map_.end()) {
       for (const auto& ext : set_.GetCandidateRange(extent)) {
+        // Sometimes there are operations like
+        // map.AddExtent({0, 5}, 42);
+        // map.Get({2, 1})
+        // If the querying extent is completely covered within the key, we still
+        // consdier this to be a valid query.
+
+        if (ExtentContains(ext, extent)) {
+          return map_.at(ext);
+        }
         if (ExtentRanges::ExtentsOverlap(ext, extent)) {
           LOG(WARNING) << "Looking up a partially intersecting extent isn't "
                           "supported by "
diff --git a/payload_consumer/extent_map_unittest.cc b/payload_consumer/extent_map_unittest.cc
index d8137a0..8e79b33 100644
--- a/payload_consumer/extent_map_unittest.cc
+++ b/payload_consumer/extent_map_unittest.cc
@@ -20,7 +20,6 @@
 
 #include "update_engine/payload_consumer/extent_map.h"
 #include "update_engine/payload_generator/extent_ranges.h"
-#include "update_engine/payload_generator/extent_utils.h"
 
 namespace chromeos_update_engine {
 
@@ -41,6 +40,17 @@
   ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
   ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 1));
   auto ret = map_.Get(ExtentForRange(1, 2));
+  ASSERT_EQ(ret, 7);
+}
+
+TEST_F(ExtentMapTest, QueryNoMerge) {
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
+  ASSERT_TRUE(map_.AddExtent(ExtentForRange(5, 5), 1));
+  auto ret = map_.Get(ExtentForRange(1, 2));
+  ASSERT_EQ(ret, 7);
+  ret = map_.Get(ExtentForRange(0, 10));
+  ASSERT_EQ(ret, std::nullopt);
+  ret = map_.Get(ExtentForRange(4, 3));
   ASSERT_EQ(ret, std::nullopt);
 }
 
@@ -48,9 +58,9 @@
   ASSERT_TRUE(map_.AddExtent(ExtentForRange(0, 5), 7));
   ASSERT_TRUE(map_.AddExtent(ExtentForRange(10, 5), 1));
   auto ret = map_.Get(ExtentForRange(3, 2));
-  ASSERT_EQ(ret, std::nullopt);
+  ASSERT_EQ(ret, 7);
   ret = map_.Get(ExtentForRange(4, 1));
-  ASSERT_EQ(ret, std::nullopt);
+  ASSERT_EQ(ret, 7);
   ret = map_.Get(ExtentForRange(5, 5));
   ASSERT_EQ(ret, std::nullopt);
   ret = map_.Get(ExtentForRange(5, 6));
diff --git a/payload_consumer/payload_verifier.cc b/payload_consumer/payload_verifier.cc
index 8a3ea65..83ef8c9 100644
--- a/payload_consumer/payload_verifier.cc
+++ b/payload_consumer/payload_verifier.cc
@@ -23,7 +23,6 @@
 #include <openssl/pem.h>
 
 #include "update_engine/common/constants.h"
-#include "update_engine/common/hash_calculator.h"
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/certificate_parser_interface.h"
 #include "update_engine/update_metadata.pb.h"
diff --git a/payload_consumer/xor_extent_writer.cc b/payload_consumer/xor_extent_writer.cc
index 31567f2..1a02a2d 100644
--- a/payload_consumer/xor_extent_writer.cc
+++ b/payload_consumer/xor_extent_writer.cc
@@ -48,21 +48,22 @@
     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()) {
+    if (!ExtentContains(extent, xor_ext)) {
       LOG(ERROR) << "CowXor merge op extent should be completely inside "
                     "InstallOp's extent. merge op extent: "
                  << xor_ext << " InstallOp extent: " << extent;
       return false;
     }
+    if (!ExtentContains(merge_op->dst_extent(), xor_ext)) {
+      LOG(ERROR) << "CowXor op extent should be completely inside "
+                    "xor_map's extent. merge op extent: "
+                 << xor_ext << " xor_map extent: " << merge_op->dst_extent();
+      return false;
+    }
     const auto src_offset = merge_op->src_offset();
-    const auto src_block = merge_op->src_extent().start_block();
+    const auto src_block = merge_op->src_extent().start_block() +
+                           xor_ext.start_block() -
+                           merge_op->dst_extent().start_block();
     xor_block_data.resize(BlockSize() * xor_ext.num_blocks());
     ssize_t bytes_read = 0;
     TEST_AND_RETURN_FALSE_ERRNO(
diff --git a/payload_consumer/xor_extent_writer_unittest.cc b/payload_consumer/xor_extent_writer_unittest.cc
index 7f35bc2..015c9ab 100644
--- a/payload_consumer/xor_extent_writer_unittest.cc
+++ b/payload_consumer/xor_extent_writer_unittest.cc
@@ -131,4 +131,46 @@
   ASSERT_TRUE(writer_.Write(zeros->data(), 9 * kBlockSize));
 }
 
+TEST_F(XorExtentWriterTest, SubsetExtentTest) {
+  constexpr auto COW_XOR = CowMergeOperation::COW_XOR;
+  ON_CALL(cow_writer_, EmitXorBlocks(_, _, _, _, _))
+      .WillByDefault(Return(true));
+
+  const auto op3 = CreateCowMergeOperation(
+      ExtentForRange(12, 4), ExtentForRange(320, 4), COW_XOR, 777);
+  ASSERT_TRUE(xor_map_.AddExtent(op3.dst_extent(), &op3));
+
+  *op_.add_src_extents() = ExtentForRange(12, 3);
+  *op_.add_dst_extents() = ExtentForRange(320, 3);
+  *op_.add_src_extents() = ExtentForRange(20, 3);
+  *op_.add_dst_extents() = ExtentForRange(420, 3);
+  *op_.add_src_extents() = ExtentForRange(15, 1);
+  *op_.add_dst_extents() = ExtentForRange(323, 1);
+  XORExtentWriter writer_{op_, source_fd_, &cow_writer_, xor_map_};
+
+  // OTA op:
+  // [12-14] => [320-322], [20-22] => [420-422], [15-16] => [323-324]
+
+  // merge op:
+  // [12-16] => [321-322]
+
+  // Expected result:
+  // [12-16] should be XOR blocks
+  // [420-422] should be regular replace blocks
+
+  auto zeros = utils::GetReadonlyZeroBlock(kBlockSize * 7);
+  EXPECT_CALL(
+      cow_writer_,
+      EmitRawBlocks(420, zeros->data() + 3 * kBlockSize, kBlockSize * 3))
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(cow_writer_, EmitXorBlocks(320, _, kBlockSize * 3, 12, 777))
+      .WillOnce(Return(true));
+  EXPECT_CALL(cow_writer_, EmitXorBlocks(323, _, kBlockSize, 15, 777))
+      .WillOnce(Return(true));
+
+  ASSERT_TRUE(writer_.Init(op_.dst_extents(), kBlockSize));
+  ASSERT_TRUE(writer_.Write(zeros->data(), zeros->size()));
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/extent_utils.h b/payload_generator/extent_utils.h
index 17995b1..f8d36e7 100644
--- a/payload_generator/extent_utils.h
+++ b/payload_generator/extent_utils.h
@@ -151,6 +151,13 @@
          block < extent.start_block() + extent.num_blocks();
 }
 
+// return true iff |big| extent contains |small| extent
+constexpr bool ExtentContains(const Extent& big, const Extent& small) {
+  return big.start_block() <= small.start_block() &&
+         small.start_block() + small.num_blocks() <=
+             big.start_block() + big.num_blocks();
+}
+
 }  // namespace chromeos_update_engine
 
 #endif  // UPDATE_ENGINE_PAYLOAD_GENERATOR_EXTENT_UTILS_H_