update_engine: Split fragmented REPLACE operations.

Splits REPLACE operations with multiple destination extents into REPLACE
operations with only one destination extent each.

BUG=chromium:461643
TEST=`cros flash --src-image-to-delta` and unit tests.
CQ-DEPEND=CL:267697

Change-Id: Id14ec9f23a462f94821454289857e9f25e756d24
Reviewed-on: https://chromium-review.googlesource.com/268520
Reviewed-by: Alex Deymo <deymo@chromium.org>
Commit-Queue: Allie Wood <alliewood@chromium.org>
Trybot-Ready: Allie Wood <alliewood@chromium.org>
Tested-by: Allie Wood <alliewood@chromium.org>
diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index 19e7aae..b6e7585 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -975,10 +975,8 @@
     DeltaArchiveManifest_InstallOperation* op,
     const chromeos::Blob& buf) {
   OmahaHashCalculator hasher;
-
   TEST_AND_RETURN_FALSE(hasher.Update(buf.data(), buf.size()));
   TEST_AND_RETURN_FALSE(hasher.Finalize());
-
   const chromeos::Blob& hash = hasher.raw_hash();
   op->set_data_sha256_hash(hash.data(), hash.size());
   return true;
@@ -1046,8 +1044,8 @@
   }
 
   // Fragment operations so we can sort them later.
-  FragmentOperations(rootfs_ops);
-  FragmentOperations(kernel_ops);
+  TEST_AND_RETURN_FALSE(FragmentOperations(rootfs_ops));
+  TEST_AND_RETURN_FALSE(FragmentOperations(kernel_ops));
 
   return true;
 }
@@ -1336,22 +1334,29 @@
   *extents = new_extents;
 }
 
-void DeltaDiffGenerator::FragmentOperations(vector<AnnotatedOperation>* aops) {
+bool DeltaDiffGenerator::FragmentOperations(vector<AnnotatedOperation>* aops) {
   vector<AnnotatedOperation> fragmented_aops;
   for (const AnnotatedOperation& aop : *aops) {
     if (aop.op.type() ==
-            DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY) {
-      SplitSourceCopy(aop.op, &fragmented_aops);
+        DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY) {
+      TEST_AND_RETURN_FALSE(SplitSourceCopy(aop, &fragmented_aops));
+    } else if (aop.op.type() ==
+               DeltaArchiveManifest_InstallOperation_Type_REPLACE) {
+      TEST_AND_RETURN_FALSE(SplitReplace(aop, &fragmented_aops));
     } else {
       fragmented_aops.push_back(aop);
     }
   }
   *aops = fragmented_aops;
+  return true;
 }
 
-void DeltaDiffGenerator::SplitSourceCopy(
-    const DeltaArchiveManifest_InstallOperation& original_op,
+bool DeltaDiffGenerator::SplitSourceCopy(
+    const AnnotatedOperation& original_aop,
     vector<AnnotatedOperation>* result_aops) {
+  DeltaArchiveManifest_InstallOperation original_op = original_aop.op;
+  TEST_AND_RETURN_FALSE(original_op.type() ==
+                        DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY);
   // Keeps track of the index of curr_src_ext.
   int curr_src_ext_index = 0;
   Extent curr_src_ext = original_op.src_extents(curr_src_ext_index);
@@ -1390,12 +1395,45 @@
 
     AnnotatedOperation new_aop;
     new_aop.op = new_op;
+    new_aop.name = base::StringPrintf("%s:%d", original_aop.name.c_str(), i);
     result_aops->push_back(new_aop);
   }
   if (curr_src_ext_index != original_op.src_extents().size() - 1) {
     LOG(FATAL) << "Incorrectly split SOURCE_COPY operation. Did not use all "
                << "source extents.";
   }
+  return true;
+}
+
+bool DeltaDiffGenerator::SplitReplace(const AnnotatedOperation& original_aop,
+                                      vector<AnnotatedOperation>* result_aops) {
+  DeltaArchiveManifest_InstallOperation original_op = original_aop.op;
+  DeltaArchiveManifest_InstallOperation_Type op_type = original_op.type();
+  TEST_AND_RETURN_FALSE(op_type ==
+                        DeltaArchiveManifest_InstallOperation_Type_REPLACE);
+  uint32_t data_offset = original_op.data_offset();
+  for (int i = 0; i < original_op.dst_extents_size(); i++) {
+    Extent dst_ext = original_op.dst_extents(i);
+    // Make a new operation with only one dst extent.
+    DeltaArchiveManifest_InstallOperation new_op;
+    *(new_op.add_dst_extents()) = dst_ext;
+    new_op.set_type(op_type);
+    uint32_t data_size = dst_ext.num_blocks() * kBlockSize;
+    new_op.set_dst_length(data_size);
+    new_op.set_data_length(data_size);
+    new_op.set_data_offset(data_offset);
+    data_offset += data_size;
+
+    AnnotatedOperation new_aop;
+    new_aop.op = new_op;
+    new_aop.name = base::StringPrintf("%s:%d", original_aop.name.c_str(), i);
+    result_aops->push_back(new_aop);
+  }
+  if (data_offset != original_op.data_offset() + original_op.data_length()) {
+    LOG(FATAL) << "Incorrectly split REPLACE/REPLACE_BZ operation. New data "
+               << "lengths do not sum to original data length.";
+  }
+  return true;
 }
 
 };  // namespace chromeos_update_engine
diff --git a/payload_generator/delta_diff_generator.h b/payload_generator/delta_diff_generator.h
index 77bca6a..d046546 100644
--- a/payload_generator/delta_diff_generator.h
+++ b/payload_generator/delta_diff_generator.h
@@ -238,20 +238,25 @@
   // minor version is 2.
   // Currently, this only modifies SOURCE_COPY operations, but it will
   // eventually fragment all operations.
-  static void FragmentOperations(std::vector<AnnotatedOperation>* aops);
+  static bool FragmentOperations(std::vector<AnnotatedOperation>* aops);
 
-  // Takes an SOURCE_COPY install operation, |original_op|, and adds one
-  // operation for each dst extent in |original_op| to |ops|. The new
-  // operations added to |ops| will have only one dst extent. The src extents
-  // are split so the number of blocks in the src and dst extents are equal.
+  // Takes an SOURCE_COPY install operation, |aop|, and adds one operation for
+  // each dst extent in |aop| to |ops|. The new operations added to |ops| will
+  // have only one dst extent. The src extents are split so the number of blocks
+  // in the src and dst extents are equal.
   // E.g. we have a SOURCE_COPY operation:
   //   src extents: [(1, 3), (5, 1), (7, 1)], dst extents: [(2, 2), (6, 3)]
   // Then we will get 2 new operations:
   //   1. src extents: [(1, 2)], dst extents: [(2, 2)]
   //   2. src extents: [(3, 1),(5, 1),(7, 1)], dst extents: [(6, 3)]
-  static void SplitSourceCopy(
-    const DeltaArchiveManifest_InstallOperation& original_op,
-    std::vector<AnnotatedOperation>* ops);
+  static bool SplitSourceCopy(const AnnotatedOperation& original_aop,
+                              std::vector<AnnotatedOperation>* result_aops);
+
+  // Takes a REPLACE operation, |aop|, and adds one operation for each dst
+  // extent in |aop| to |ops|. The new operations added to |ops| will have only
+  // one dst extent each.
+  static bool SplitReplace(const AnnotatedOperation& original_aop,
+                           std::vector<AnnotatedOperation>* result_aops);
 
  private:
   DISALLOW_COPY_AND_ASSIGN(DeltaDiffGenerator);
diff --git a/payload_generator/delta_diff_generator_unittest.cc b/payload_generator/delta_diff_generator_unittest.cc
index 99b5d08..4d8966d 100644
--- a/payload_generator/delta_diff_generator_unittest.cc
+++ b/payload_generator/delta_diff_generator_unittest.cc
@@ -827,10 +827,14 @@
   *(op.add_dst_extents()) = ExtentForRange(14, 3);
   *(op.add_dst_extents()) = ExtentForRange(18, 3);
 
+  AnnotatedOperation aop;
+  aop.op = op;
+  aop.name = "SplitSourceCopyTestOp";
   vector<AnnotatedOperation> result_ops;
-  DeltaDiffGenerator::SplitSourceCopy(op, &result_ops);
+  EXPECT_TRUE(DeltaDiffGenerator::SplitSourceCopy(aop, &result_ops));
   EXPECT_EQ(result_ops.size(), 3);
 
+  EXPECT_EQ("SplitSourceCopyTestOp:0", result_ops[0].name);
   DeltaArchiveManifest_InstallOperation first_op = result_ops[0].op;
   EXPECT_EQ(DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY,
             first_op.type());
@@ -843,6 +847,7 @@
   EXPECT_EQ(10, first_op.dst_extents().Get(0).start_block());
   EXPECT_EQ(2, first_op.dst_extents().Get(0).num_blocks());
 
+  EXPECT_EQ("SplitSourceCopyTestOp:1", result_ops[1].name);
   DeltaArchiveManifest_InstallOperation second_op = result_ops[1].op;
   EXPECT_EQ(DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY,
             second_op.type());
@@ -859,6 +864,7 @@
   EXPECT_EQ(14, second_op.dst_extents().Get(0).start_block());
   EXPECT_EQ(3, second_op.dst_extents().Get(0).num_blocks());
 
+  EXPECT_EQ("SplitSourceCopyTestOp:2", result_ops[2].name);
   DeltaArchiveManifest_InstallOperation third_op = result_ops[2].op;
   EXPECT_EQ(DeltaArchiveManifest_InstallOperation_Type_SOURCE_COPY,
             third_op.type());
@@ -872,4 +878,44 @@
   EXPECT_EQ(3, third_op.dst_extents().Get(0).num_blocks());
 }
 
+TEST_F(DeltaDiffGeneratorTest, SplitReplaceTest) {
+  uint32_t data_offset = 5555;
+
+  DeltaArchiveManifest_InstallOperation op;
+  op.set_type(DeltaArchiveManifest_InstallOperation_Type_REPLACE);
+  *(op.add_dst_extents()) = ExtentForRange(2, 2);
+  *(op.add_dst_extents()) = ExtentForRange(6, 1);
+  op.set_data_length(3 * kBlockSize);
+  op.set_data_offset(data_offset);
+
+  AnnotatedOperation aop;
+  aop.op = op;
+  aop.name = "SplitReplaceTestOp";
+  vector<AnnotatedOperation> result_ops;
+  EXPECT_TRUE(DeltaDiffGenerator::SplitReplace(aop, &result_ops));
+  EXPECT_EQ(result_ops.size(), 2);
+
+  EXPECT_EQ("SplitReplaceTestOp:0", result_ops[0].name);
+  DeltaArchiveManifest_InstallOperation first_op = result_ops[0].op;
+  EXPECT_EQ(DeltaArchiveManifest_InstallOperation_Type_REPLACE,
+            first_op.type());
+  EXPECT_EQ(2 * kBlockSize, first_op.dst_length());
+  EXPECT_EQ(2 * kBlockSize, first_op.data_length());
+  EXPECT_EQ(data_offset, first_op.data_offset());
+  EXPECT_EQ(1, first_op.dst_extents().size());
+  EXPECT_EQ(2, first_op.dst_extents().Get(0).start_block());
+  EXPECT_EQ(2, first_op.dst_extents().Get(0).num_blocks());
+
+  EXPECT_EQ("SplitReplaceTestOp:1", result_ops[1].name);
+  DeltaArchiveManifest_InstallOperation second_op = result_ops[1].op;
+  EXPECT_EQ(DeltaArchiveManifest_InstallOperation_Type_REPLACE,
+            second_op.type());
+  EXPECT_EQ(kBlockSize, second_op.dst_length());
+  EXPECT_EQ(kBlockSize, second_op.data_length());
+  EXPECT_EQ(data_offset + (2 * kBlockSize), second_op.data_offset());
+  EXPECT_EQ(1, second_op.dst_extents().size());
+  EXPECT_EQ(6, second_op.dst_extents().Get(0).start_block());
+  EXPECT_EQ(1, second_op.dst_extents().Get(0).num_blocks());
+}
+
 }  // namespace chromeos_update_engine