Fix partial read failure on last block of src partition

When an XOR op is at the very end of partition with >0 src_offest,
update_engine will attempt to read beyond partition size and fail.
Convert such blocks to REPLACE instead.

Bug: 276366819

Change-Id: I6e50037c192599f4d0c78da35ef2aa26cf2e92de
diff --git a/aosp/apex_handler_android.cc b/aosp/apex_handler_android.cc
index f47889b..acc6bec 100644
--- a/aosp/apex_handler_android.cc
+++ b/aosp/apex_handler_android.cc
@@ -22,7 +22,6 @@
 #include <ApexProperties.sysprop.h>
 
 #include "update_engine/aosp/apex_handler_android.h"
-#include "update_engine/common/utils.h"
 
 namespace chromeos_update_engine {
 
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index 92f27a9..17b7d50 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -358,7 +358,11 @@
 
   std::unique_ptr<ExtentWriter> writer =
       IsXorEnabled() ? std::make_unique<XORExtentWriter>(
-                           operation, source_fd, cow_writer_.get(), xor_map_)
+                           operation,
+                           source_fd,
+                           cow_writer_.get(),
+                           xor_map_,
+                           partition_update_.old_partition_info().size())
                      : CreateBaseExtentWriter();
   return executor_.ExecuteDiffOperation(
       operation, std::move(writer), source_fd, data, count);
diff --git a/payload_consumer/xor_extent_writer.cc b/payload_consumer/xor_extent_writer.cc
index 1a02a2d..4534c05 100644
--- a/payload_consumer/xor_extent_writer.cc
+++ b/payload_consumer/xor_extent_writer.cc
@@ -64,6 +64,23 @@
     const auto src_block = merge_op->src_extent().start_block() +
                            xor_ext.start_block() -
                            merge_op->dst_extent().start_block();
+    const auto i = xor_ext.start_block() - extent.start_block();
+    const auto dst_block_data =
+        static_cast<const unsigned char*>(bytes) + i * BlockSize();
+    const auto is_out_of_bound_read =
+        (src_block + xor_ext.num_blocks()) * BlockSize() + src_offset >
+            partition_size_ &&
+        partition_size_ != 0;
+    if (is_out_of_bound_read) {
+      LOG(INFO) << "Getting partial read for last block, converting "
+                   "XOR operation to a regular replace "
+                << xor_ext;
+      TEST_AND_RETURN_FALSE(
+          cow_writer_->AddRawBlocks(xor_ext.start_block(),
+                                    dst_block_data,
+                                    xor_ext.num_blocks() * BlockSize()));
+      continue;
+    }
     xor_block_data.resize(BlockSize() * xor_ext.num_blocks());
     ssize_t bytes_read = 0;
     TEST_AND_RETURN_FALSE_ERRNO(
@@ -73,14 +90,12 @@
                         src_offset + src_block * BlockSize(),
                         &bytes_read));
     if (bytes_read != static_cast<ssize_t>(xor_block_data.size())) {
-      LOG(ERROR) << "bytes_read: " << bytes_read;
+      LOG(ERROR) << "bytes_read: " << bytes_read << ", expected to read "
+                 << xor_block_data.size() << " at block " << src_block
+                 << " offset " << src_offset;
       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,
diff --git a/payload_consumer/xor_extent_writer.h b/payload_consumer/xor_extent_writer.h
index 35565ea..57c99c2 100644
--- a/payload_consumer/xor_extent_writer.h
+++ b/payload_consumer/xor_extent_writer.h
@@ -19,12 +19,9 @@
 
 #include <vector>
 
+#include "common/utils.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>
@@ -38,11 +35,13 @@
   XORExtentWriter(const InstallOperation& op,
                   FileDescriptorPtr source_fd,
                   android::snapshot::ICowWriter* cow_writer,
-                  const ExtentMap<const CowMergeOperation*>& xor_map)
+                  const ExtentMap<const CowMergeOperation*>& xor_map,
+                  size_t partition_size)
       : src_extents_(op.src_extents()),
         source_fd_(source_fd),
         xor_map_(xor_map),
-        cow_writer_(cow_writer) {
+        cow_writer_(cow_writer),
+        partition_size_(partition_size) {
     CHECK(source_fd->IsOpen());
   }
   ~XORExtentWriter() = default;
@@ -61,6 +60,7 @@
   const FileDescriptorPtr source_fd_;
   const ExtentMap<const CowMergeOperation*>& xor_map_;
   android::snapshot::ICowWriter* cow_writer_;
+  const size_t partition_size_;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/xor_extent_writer_unittest.cc b/payload_consumer/xor_extent_writer_unittest.cc
index 015c9ab..55c3c6c 100644
--- a/payload_consumer/xor_extent_writer_unittest.cc
+++ b/payload_consumer/xor_extent_writer_unittest.cc
@@ -94,7 +94,8 @@
   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_};
+  XORExtentWriter writer_{
+      op_, source_fd_, &cow_writer_, xor_map_, NUM_BLOCKS * kBlockSize};
 
   // OTA op:
   // [5-6] => [5-6], [45-47] => [455-457], [12-15] => [320-323]
@@ -146,7 +147,8 @@
   *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_};
+  XORExtentWriter writer_{
+      op_, source_fd_, &cow_writer_, xor_map_, NUM_BLOCKS * kBlockSize};
 
   // OTA op:
   // [12-14] => [320-322], [20-22] => [420-422], [15-16] => [323-324]
@@ -173,4 +175,54 @@
   ASSERT_TRUE(writer_.Write(zeros->data(), zeros->size()));
 }
 
+TEST_F(XorExtentWriterTest, LastBlockTest) {
+  constexpr auto COW_XOR = CowMergeOperation::COW_XOR;
+  ON_CALL(cow_writer_, EmitXorBlocks(_, _, _, _, _))
+      .WillByDefault(Return(true));
+
+  const auto op3 = CreateCowMergeOperation(
+      ExtentForRange(NUM_BLOCKS - 1, 1), ExtentForRange(2, 1), 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(NUM_BLOCKS - 3, 3);
+  *op_.add_dst_extents() = ExtentForRange(2, 3);
+  XORExtentWriter writer_{
+      op_, source_fd_, &cow_writer_, xor_map_, NUM_BLOCKS * kBlockSize};
+
+  // OTA op:
+  // [12-14] => [320-322], [20-22] => [420-422], [NUM_BLOCKS-3] => [2-5]
+
+  // merge op:
+  // [NUM_BLOCKS-1] => [2-3]
+
+  // Expected result:
+  // [12-16] should be REPLACE blocks
+  // [420-422] should be REPLACE blocks
+  // [2-4] should be REPLACE blocks
+
+  auto zeros = utils::GetReadonlyZeroBlock(kBlockSize * 9);
+  EXPECT_CALL(cow_writer_, EmitRawBlocks(320, zeros->data(), kBlockSize * 3))
+      .WillOnce(Return(true));
+  EXPECT_CALL(
+      cow_writer_,
+      EmitRawBlocks(420, zeros->data() + 3 * kBlockSize, kBlockSize * 3))
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(cow_writer_,
+              EmitRawBlocks(2, zeros->data() + 6 * kBlockSize, kBlockSize))
+      .WillOnce(Return(true));
+  EXPECT_CALL(cow_writer_,
+              EmitRawBlocks(3, zeros->data() + 7 * kBlockSize, kBlockSize * 2))
+      .WillOnce(Return(true));
+
+  ASSERT_TRUE(writer_.Init(op_.dst_extents(), kBlockSize));
+  ASSERT_TRUE(writer_.Write(zeros->data(), zeros->size()));
+}
+
 }  // namespace chromeos_update_engine