update_engine: Split fragmented REPLACE_BZ ops.

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

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

Change-Id: Ia5c055047c4bb863871507f32d028e4e28cb6af6
Reviewed-on: https://chromium-review.googlesource.com/269130
Reviewed-by: Allie Wood <alliewood@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 b6e7585..2d2aabc 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -1044,8 +1044,14 @@
   }
 
   // Fragment operations so we can sort them later.
-  TEST_AND_RETURN_FALSE(FragmentOperations(rootfs_ops));
-  TEST_AND_RETURN_FALSE(FragmentOperations(kernel_ops));
+  TEST_AND_RETURN_FALSE(FragmentOperations(rootfs_ops,
+                                           config.target.rootfs_part,
+                                           data_file_fd,
+                                           data_file_size));
+  TEST_AND_RETURN_FALSE(FragmentOperations(kernel_ops,
+                                           config.target.rootfs_part,
+                                           data_file_fd,
+                                           data_file_size));
 
   return true;
 }
@@ -1334,7 +1340,10 @@
   *extents = new_extents;
 }
 
-bool DeltaDiffGenerator::FragmentOperations(vector<AnnotatedOperation>* aops) {
+bool DeltaDiffGenerator::FragmentOperations(vector<AnnotatedOperation>* aops,
+                                            const string& target_rootfs_part,
+                                            int data_fd,
+                                            off_t* data_file_size) {
   vector<AnnotatedOperation> fragmented_aops;
   for (const AnnotatedOperation& aop : *aops) {
     if (aop.op.type() ==
@@ -1343,6 +1352,13 @@
     } else if (aop.op.type() ==
                DeltaArchiveManifest_InstallOperation_Type_REPLACE) {
       TEST_AND_RETURN_FALSE(SplitReplace(aop, &fragmented_aops));
+    } else if (aop.op.type() ==
+               DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ) {
+      TEST_AND_RETURN_FALSE(SplitReplaceBz(aop,
+                                           &fragmented_aops,
+                                           target_rootfs_part,
+                                           data_fd,
+                                           data_file_size));
     } else {
       fragmented_aops.push_back(aop);
     }
@@ -1408,16 +1424,16 @@
 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 ==
+  TEST_AND_RETURN_FALSE(original_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);
+    new_op.set_type(original_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);
@@ -1430,8 +1446,61 @@
     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.";
+    LOG(FATAL) << "Incorrectly split REPLACE operation. New data lengths do "
+               << "not sum to original data length.";
+  }
+  return true;
+}
+
+bool DeltaDiffGenerator::SplitReplaceBz(
+    const AnnotatedOperation& original_aop,
+    vector<AnnotatedOperation>* result_aops,
+    const string& target_rootfs_part,
+    int data_fd,
+    off_t* data_file_size) {
+  DeltaArchiveManifest_InstallOperation original_op = original_aop.op;
+  TEST_AND_RETURN_FALSE(original_op.type() ==
+                        DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ);
+
+  int target_rootfs_fd = open(target_rootfs_part.c_str(), O_RDONLY, 000);
+  TEST_AND_RETURN_FALSE_ERRNO(target_rootfs_fd >= 0);
+  ScopedFdCloser target_rootfs_fd_closer(&target_rootfs_fd);
+
+  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(original_op.type());
+    uint32_t uncompressed_data_size = dst_ext.num_blocks() * kBlockSize;
+    new_op.set_dst_length(uncompressed_data_size);
+
+    // Get the original uncompressed data for this extent.
+    ssize_t bytes_read;
+    chromeos::Blob uncompressed_data(uncompressed_data_size);
+    TEST_AND_RETURN_FALSE(utils::PReadAll(target_rootfs_fd,
+                                          uncompressed_data.data(),
+                                          uncompressed_data_size,
+                                          kBlockSize * dst_ext.start_block(),
+                                          &bytes_read));
+    TEST_AND_RETURN_FALSE(bytes_read ==
+                          static_cast<ssize_t>(uncompressed_data_size));
+
+    chromeos::Blob new_data_bz;
+    TEST_AND_RETURN_FALSE(BzipCompress(uncompressed_data, &new_data_bz));
+    CHECK(!new_data_bz.empty());
+    new_op.set_data_length(new_data_bz.size());
+    new_op.set_data_offset(*data_file_size);
+    TEST_AND_RETURN_FALSE(utils::PWriteAll(data_fd,
+                                           new_data_bz.data(),
+                                           new_data_bz.size(),
+                                           *data_file_size));
+    *data_file_size += new_data_bz.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);
   }
   return true;
 }
diff --git a/payload_generator/delta_diff_generator.h b/payload_generator/delta_diff_generator.h
index d046546..6b4dc63 100644
--- a/payload_generator/delta_diff_generator.h
+++ b/payload_generator/delta_diff_generator.h
@@ -236,9 +236,10 @@
   // such that there is only one dst extent per operation. Sets |aops| to a
   // vector of the new fragmented operations. This is only called when delta
   // minor version is 2.
-  // Currently, this only modifies SOURCE_COPY operations, but it will
-  // eventually fragment all operations.
-  static bool FragmentOperations(std::vector<AnnotatedOperation>* aops);
+  static bool FragmentOperations(std::vector<AnnotatedOperation>* aops,
+                                 const std::string& target_rootfs_part,
+                                 int data_fd,
+                                 off_t* data_file_size);
 
   // 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
@@ -258,6 +259,15 @@
   static bool SplitReplace(const AnnotatedOperation& original_aop,
                            std::vector<AnnotatedOperation>* result_aops);
 
+  // Takes a REPLACE_BZ 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 SplitReplaceBz(const AnnotatedOperation& original_aop,
+                             std::vector<AnnotatedOperation>* result_aops,
+                             const std::string& target_rootfs_part,
+                             int data_fd,
+                             off_t* data_file_size);
+
  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 4d8966d..fe15516 100644
--- a/payload_generator/delta_diff_generator_unittest.cc
+++ b/payload_generator/delta_diff_generator_unittest.cc
@@ -21,6 +21,7 @@
 #include <base/strings/string_util.h>
 #include <gtest/gtest.h>
 
+#include "update_engine/bzip.h"
 #include "update_engine/delta_performer.h"
 #include "update_engine/extent_ranges.h"
 #include "update_engine/payload_constants.h"
@@ -840,12 +841,12 @@
             first_op.type());
   EXPECT_EQ(kBlockSize * 2, first_op.src_length());
   EXPECT_EQ(1, first_op.src_extents().size());
-  EXPECT_EQ(2, first_op.src_extents().Get(0).start_block());
-  EXPECT_EQ(2, first_op.src_extents().Get(0).num_blocks());
+  EXPECT_EQ(2, first_op.src_extents(0).start_block());
+  EXPECT_EQ(2, first_op.src_extents(0).num_blocks());
   EXPECT_EQ(kBlockSize * 2, first_op.dst_length());
   EXPECT_EQ(1, first_op.dst_extents().size());
-  EXPECT_EQ(10, first_op.dst_extents().Get(0).start_block());
-  EXPECT_EQ(2, first_op.dst_extents().Get(0).num_blocks());
+  EXPECT_EQ(10, first_op.dst_extents(0).start_block());
+  EXPECT_EQ(2, first_op.dst_extents(0).num_blocks());
 
   EXPECT_EQ("SplitSourceCopyTestOp:1", result_ops[1].name);
   DeltaArchiveManifest_InstallOperation second_op = result_ops[1].op;
@@ -853,16 +854,16 @@
             second_op.type());
   EXPECT_EQ(kBlockSize * 3, second_op.src_length());
   EXPECT_EQ(3, second_op.src_extents().size());
-  EXPECT_EQ(4, second_op.src_extents().Get(0).start_block());
-  EXPECT_EQ(1, second_op.src_extents().Get(0).num_blocks());
-  EXPECT_EQ(6, second_op.src_extents().Get(1).start_block());
-  EXPECT_EQ(1, second_op.src_extents().Get(1).num_blocks());
-  EXPECT_EQ(8, second_op.src_extents().Get(2).start_block());
-  EXPECT_EQ(1, second_op.src_extents().Get(2).num_blocks());
+  EXPECT_EQ(4, second_op.src_extents(0).start_block());
+  EXPECT_EQ(1, second_op.src_extents(0).num_blocks());
+  EXPECT_EQ(6, second_op.src_extents(1).start_block());
+  EXPECT_EQ(1, second_op.src_extents(1).num_blocks());
+  EXPECT_EQ(8, second_op.src_extents(2).start_block());
+  EXPECT_EQ(1, second_op.src_extents(2).num_blocks());
   EXPECT_EQ(kBlockSize * 3, second_op.dst_length());
   EXPECT_EQ(1, second_op.dst_extents().size());
-  EXPECT_EQ(14, second_op.dst_extents().Get(0).start_block());
-  EXPECT_EQ(3, second_op.dst_extents().Get(0).num_blocks());
+  EXPECT_EQ(14, second_op.dst_extents(0).start_block());
+  EXPECT_EQ(3, second_op.dst_extents(0).num_blocks());
 
   EXPECT_EQ("SplitSourceCopyTestOp:2", result_ops[2].name);
   DeltaArchiveManifest_InstallOperation third_op = result_ops[2].op;
@@ -870,12 +871,12 @@
             third_op.type());
   EXPECT_EQ(kBlockSize * 3, third_op.src_length());
   EXPECT_EQ(1, third_op.src_extents().size());
-  EXPECT_EQ(9, third_op.src_extents().Get(0).start_block());
-  EXPECT_EQ(3, third_op.src_extents().Get(0).num_blocks());
+  EXPECT_EQ(9, third_op.src_extents(0).start_block());
+  EXPECT_EQ(3, third_op.src_extents(0).num_blocks());
   EXPECT_EQ(kBlockSize * 3, third_op.dst_length());
   EXPECT_EQ(1, third_op.dst_extents().size());
-  EXPECT_EQ(18, third_op.dst_extents().Get(0).start_block());
-  EXPECT_EQ(3, third_op.dst_extents().Get(0).num_blocks());
+  EXPECT_EQ(18, third_op.dst_extents(0).start_block());
+  EXPECT_EQ(3, third_op.dst_extents(0).num_blocks());
 }
 
 TEST_F(DeltaDiffGeneratorTest, SplitReplaceTest) {
@@ -903,8 +904,8 @@
   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(2, first_op.dst_extents(0).start_block());
+  EXPECT_EQ(2, first_op.dst_extents(0).num_blocks());
 
   EXPECT_EQ("SplitReplaceTestOp:1", result_ops[1].name);
   DeltaArchiveManifest_InstallOperation second_op = result_ops[1].op;
@@ -914,8 +915,91 @@
   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());
+  EXPECT_EQ(6, second_op.dst_extents(0).start_block());
+  EXPECT_EQ(1, second_op.dst_extents(0).num_blocks());
 }
 
+TEST_F(DeltaDiffGeneratorTest, SplitReplaceBzTest) {
+  string data_path;
+  EXPECT_TRUE(utils::MakeTempFile(
+      "ReplaceBzTest_data.XXXXXX", &data_path, nullptr));
+  int data_fd = open(data_path.c_str(), O_RDWR, 000);
+  EXPECT_GE(data_fd, 0);
+  ScopedFdCloser data_fd_closer(&data_fd);
+  off_t data_file_size = 0;
+
+  string rootfs_path;
+  EXPECT_TRUE(utils::MakeTempFile(
+      "ReplaceBzTest_rootfs.XXXXXX", &rootfs_path, nullptr));
+  chromeos::Blob data(7 * kBlockSize);
+  test_utils::FillWithData(&data);
+  EXPECT_TRUE(utils::WriteFile(rootfs_path.c_str(), data.data(), data.size()));
+
+
+  DeltaArchiveManifest_InstallOperation op;
+  op.set_type(DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ);
+  *(op.add_dst_extents()) = ExtentForRange(2, 2);
+  *(op.add_dst_extents()) = ExtentForRange(6, 1);
+
+  AnnotatedOperation aop;
+  aop.op = op;
+  aop.name = "SplitReplaceBzTestOp";
+  vector<AnnotatedOperation> result_ops;
+  EXPECT_TRUE(DeltaDiffGenerator::SplitReplaceBz(aop, &result_ops, rootfs_path,
+                                                 data_fd, &data_file_size));
+  EXPECT_EQ(result_ops.size(), 2);
+
+  EXPECT_EQ("SplitReplaceBzTestOp:0", result_ops[0].name);
+  DeltaArchiveManifest_InstallOperation first_op = result_ops[0].op;
+  EXPECT_EQ(DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ,
+            first_op.type());
+  EXPECT_EQ(2 * kBlockSize, first_op.dst_length());
+  EXPECT_EQ(1, first_op.dst_extents().size());
+  EXPECT_EQ(2, first_op.dst_extents(0).start_block());
+  EXPECT_EQ(2, first_op.dst_extents(0).num_blocks());
+  // Get the blob corresponding to this extent and compress it.
+  chromeos::Blob first_blob(data.begin() + (kBlockSize * 2),
+                            data.begin() + (kBlockSize * 4));
+  chromeos::Blob first_blob_bz;
+  EXPECT_TRUE(BzipCompress(first_blob, &first_blob_bz));
+  EXPECT_EQ(first_blob_bz.size(), first_op.data_length());
+  EXPECT_EQ(0, first_op.data_offset());
+  // Check that the compressed blob matches what's in data_fd.
+  chromeos::Blob first_data_blob(first_op.data_length());
+  ssize_t bytes_read;
+  EXPECT_TRUE(utils::PReadAll(data_fd,
+                              first_data_blob.data(),
+                              first_op.data_length(),
+                              first_op.data_offset(),
+                              &bytes_read));
+  EXPECT_EQ(bytes_read, first_op.data_length());
+  EXPECT_EQ(first_data_blob, first_blob_bz);
+
+  EXPECT_EQ("SplitReplaceBzTestOp:1", result_ops[1].name);
+  DeltaArchiveManifest_InstallOperation second_op = result_ops[1].op;
+  EXPECT_EQ(DeltaArchiveManifest_InstallOperation_Type_REPLACE_BZ,
+            second_op.type());
+  EXPECT_EQ(kBlockSize, second_op.dst_length());
+  EXPECT_EQ(1, second_op.dst_extents().size());
+  EXPECT_EQ(6, second_op.dst_extents(0).start_block());
+  EXPECT_EQ(1, second_op.dst_extents(0).num_blocks());
+  chromeos::Blob second_blob(data.begin() + (kBlockSize * 6),
+                             data.begin() + (kBlockSize * 7));
+  chromeos::Blob second_blob_bz;
+  EXPECT_TRUE(BzipCompress(second_blob, &second_blob_bz));
+  EXPECT_EQ(second_blob_bz.size(), second_op.data_length());
+  EXPECT_EQ(first_op.data_length(), second_op.data_offset());
+  chromeos::Blob second_data_blob(second_op.data_length());
+  EXPECT_TRUE(utils::PReadAll(data_fd,
+                              second_data_blob.data(),
+                              second_op.data_length(),
+                              second_op.data_offset(),
+                              &bytes_read));
+  EXPECT_EQ(bytes_read, second_op.data_length());
+  EXPECT_EQ(second_data_blob, second_blob_bz);
+
+  EXPECT_EQ(data_file_size, second_op.data_offset() + second_op.data_length());
+}
+
+
 }  // namespace chromeos_update_engine