Add InstallOperationExecutor which hanldes interpretation of install op

PartitionWriter and VABC Partition writer both need to "execute" install
ops by reconstructing target data(decompress payload data, apply bsdiff,
etc). Right now we achieve code sharing for this via inheritance and
protected members. But this is error prone. Instead, add a separate
class for handling executeion logic and achieve code sharing via
composition.

Test: th
Change-Id: If978790ce87b0d037839505bf8797c7024ce4c6d
diff --git a/payload_consumer/file_descriptor_utils.cc b/payload_consumer/file_descriptor_utils.cc
index 9a6a601..91b5673 100644
--- a/payload_consumer/file_descriptor_utils.cc
+++ b/payload_consumer/file_descriptor_utils.cc
@@ -35,9 +35,12 @@
 // Size of the buffer used to copy blocks.
 const uint64_t kMaxCopyBufferSize = 1024 * 1024;
 
+}  // namespace
+namespace fd_utils {
+
 bool CommonHashExtents(FileDescriptorPtr source,
                        const RepeatedPtrField<Extent>& src_extents,
-                       DirectExtentWriter* writer,
+                       ExtentWriter* writer,
                        uint64_t block_size,
                        brillo::Blob* hash_out) {
   auto total_blocks = utils::BlocksInExtents(src_extents);
@@ -72,10 +75,6 @@
   return true;
 }
 
-}  // namespace
-
-namespace fd_utils {
-
 bool CopyAndHashExtents(FileDescriptorPtr source,
                         const RepeatedPtrField<Extent>& src_extents,
                         FileDescriptorPtr target,
diff --git a/payload_consumer/file_descriptor_utils.h b/payload_consumer/file_descriptor_utils.h
index 68fb001..3a4027d 100644
--- a/payload_consumer/file_descriptor_utils.h
+++ b/payload_consumer/file_descriptor_utils.h
@@ -19,12 +19,20 @@
 
 #include <brillo/secure_blob.h>
 
+#include "update_engine/payload_consumer/extent_writer.h"
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
 namespace fd_utils {
 
+bool CommonHashExtents(
+    FileDescriptorPtr source,
+    const google::protobuf::RepeatedPtrField<Extent>& src_extents,
+    ExtentWriter* writer,
+    uint64_t block_size,
+    brillo::Blob* hash_out);
+
 // Copy blocks from the |source| file to the |target| file and hashes the
 // contents. The blocks to copy from the |source| to the |target| files are
 // specified by the |src_extents| and |tgt_extents| list of Extents, which
diff --git a/payload_consumer/partition_writer.cc b/payload_consumer/partition_writer.cc
index f2022a1..24ea7d8 100644
--- a/payload_consumer/partition_writer.cc
+++ b/payload_consumer/partition_writer.cc
@@ -17,6 +17,7 @@
 
 #include <fcntl.h>
 #include <linux/fs.h>
+#include <sys/mman.h>
 
 #include <algorithm>
 #include <initializer_list>
@@ -24,6 +25,7 @@
 #include <utility>
 #include <vector>
 
+#include <base/files/memory_mapped_file.h>
 #include <base/strings/string_number_conversions.h>
 #include <bsdiff/bspatch.h>
 #include <puffin/puffpatch.h>
@@ -250,7 +252,8 @@
       install_part_(install_part),
       dynamic_control_(dynamic_control),
       interactive_(is_interactive),
-      block_size_(block_size) {}
+      block_size_(block_size),
+      install_op_executor_(block_size) {}
 
 PartitionWriter::~PartitionWriter() {
   Close();
@@ -317,24 +320,57 @@
   return true;
 }
 
-bool PartitionWriter::PerformReplaceOperation(const InstallOperation& operation,
-                                              const void* data,
-                                              size_t count) {
+bool InstallOperationExecutor::ExecuteReplaceOperation(
+    const InstallOperation& operation,
+    std::unique_ptr<ExtentWriter> writer,
+    const void* data,
+    size_t count) {
+  TEST_AND_RETURN_FALSE(operation.type() == InstallOperation::REPLACE ||
+                        operation.type() == InstallOperation::REPLACE_BZ ||
+                        operation.type() == InstallOperation::REPLACE_XZ);
   // Setup the ExtentWriter stack based on the operation type.
-  std::unique_ptr<ExtentWriter> writer = CreateBaseExtentWriter();
-
   if (operation.type() == InstallOperation::REPLACE_BZ) {
     writer.reset(new BzipExtentWriter(std::move(writer)));
   } else if (operation.type() == InstallOperation::REPLACE_XZ) {
     writer.reset(new XzExtentWriter(std::move(writer)));
   }
-
   TEST_AND_RETURN_FALSE(writer->Init(operation.dst_extents(), block_size_));
   TEST_AND_RETURN_FALSE(writer->Write(data, operation.data_length()));
 
   return true;
 }
 
+bool PartitionWriter::PerformReplaceOperation(const InstallOperation& operation,
+                                              const void* data,
+                                              size_t count) {
+  // Setup the ExtentWriter stack based on the operation type.
+  std::unique_ptr<ExtentWriter> writer = CreateBaseExtentWriter();
+  writer->Init(operation.dst_extents(), block_size_);
+  return install_op_executor_.ExecuteReplaceOperation(
+      operation, std::move(writer), data, count);
+}
+
+bool InstallOperationExecutor::ExecuteZeroOrDiscardOperation(
+    const InstallOperation& operation, ExtentWriter* writer) {
+  TEST_AND_RETURN_FALSE(operation.type() == InstallOperation::ZERO ||
+                        operation.type() == InstallOperation::DISCARD);
+  for (const auto& extent : operation.dst_extents()) {
+    // Mmap a region of /dev/zero, as we don't need any actual memory to store
+    // these 0s, so mmap a region of "free memory".
+    base::File dev_zero(base::FilePath("/dev/zero"),
+                        base::File::FLAG_OPEN | base::File::FLAG_READ);
+    base::MemoryMappedFile buffer;
+    TEST_AND_RETURN_FALSE_ERRNO(buffer.Initialize(
+        std::move(dev_zero),
+        base::MemoryMappedFile::Region{
+            0, static_cast<size_t>(extent.num_blocks() * block_size_)},
+        base::MemoryMappedFile::Access::READ_ONLY));
+    TEST_AND_RETURN_FALSE(buffer.data() != nullptr);
+    writer->Write(buffer.data(), buffer.length());
+  }
+  return true;
+}
+
 bool PartitionWriter::PerformZeroOrDiscardOperation(
     const InstallOperation& operation) {
 #ifdef BLKZEROOUT
@@ -385,6 +421,17 @@
   return out;
 }
 
+bool InstallOperationExecutor::ExecuteSourceCopyOperation(
+    const InstallOperation& operation,
+    ExtentWriter* writer,
+    FileDescriptorPtr source_fd) {
+  TEST_AND_RETURN_FALSE(operation.type() == InstallOperation::SOURCE_COPY);
+  TEST_AND_RETURN_FALSE(source_fd != nullptr);
+
+  return fd_utils::CommonHashExtents(
+      source_fd, operation.src_extents(), writer, block_size_, nullptr);
+}
+
 bool PartitionWriter::PerformSourceCopyOperation(
     const InstallOperation& operation, ErrorCode* error) {
   TEST_AND_RETURN_FALSE(source_fd_ != nullptr);
@@ -410,20 +457,21 @@
     return false;
   }
 
-  return fd_utils::CopyAndHashExtents(source_fd,
-                                      optimized.src_extents(),
-                                      target_fd_,
-                                      optimized.dst_extents(),
-                                      block_size_,
-                                      nullptr);
+  auto writer = CreateBaseExtentWriter();
+  writer->Init(optimized.dst_extents(), block_size_);
+  return install_op_executor_.ExecuteSourceCopyOperation(
+      optimized, writer.get(), source_fd);
 }
 
-bool PartitionWriter::PerformSourceBsdiffOperation(
+bool InstallOperationExecutor::ExecuteSourceBsdiffOperation(
     const InstallOperation& operation,
-    ErrorCode* error,
+    std::unique_ptr<ExtentWriter> writer,
+    FileDescriptorPtr source_fd,
     const void* data,
     size_t count) {
-  FileDescriptorPtr source_fd = ChooseSourceFD(operation, error);
+  TEST_AND_RETURN_FALSE(operation.type() == InstallOperation::SOURCE_BSDIFF ||
+                        operation.type() == InstallOperation::BROTLI_BSDIFF ||
+                        operation.type() == InstallOperation::BSDIFF);
   TEST_AND_RETURN_FALSE(source_fd != nullptr);
 
   auto reader = std::make_unique<DirectExtentReader>();
@@ -433,7 +481,6 @@
       std::move(reader),
       utils::BlocksInExtents(operation.src_extents()) * block_size_);
 
-  auto writer = CreateBaseExtentWriter();
   TEST_AND_RETURN_FALSE(writer->Init(operation.dst_extents(), block_size_));
   auto dst_file = std::make_unique<BsdiffExtentFile>(
       std::move(writer),
@@ -446,7 +493,7 @@
   return true;
 }
 
-bool PartitionWriter::PerformPuffDiffOperation(
+bool PartitionWriter::PerformSourceBsdiffOperation(
     const InstallOperation& operation,
     ErrorCode* error,
     const void* data,
@@ -454,6 +501,21 @@
   FileDescriptorPtr source_fd = ChooseSourceFD(operation, error);
   TEST_AND_RETURN_FALSE(source_fd != nullptr);
 
+  auto writer = CreateBaseExtentWriter();
+  writer->Init(operation.dst_extents(), block_size_);
+  return install_op_executor_.ExecuteSourceBsdiffOperation(
+      operation, std::move(writer), source_fd, data, count);
+}
+
+bool InstallOperationExecutor::ExecutePuffDiffOperation(
+    const InstallOperation& operation,
+    std::unique_ptr<ExtentWriter> writer,
+    FileDescriptorPtr source_fd,
+    const void* data,
+    size_t count) {
+  TEST_AND_RETURN_FALSE(operation.type() == InstallOperation::PUFFDIFF);
+  TEST_AND_RETURN_FALSE(source_fd != nullptr);
+
   auto reader = std::make_unique<DirectExtentReader>();
   TEST_AND_RETURN_FALSE(
       reader->Init(source_fd, operation.src_extents(), block_size_));
@@ -461,7 +523,6 @@
       std::move(reader),
       utils::BlocksInExtents(operation.src_extents()) * block_size_));
 
-  auto writer = CreateBaseExtentWriter();
   TEST_AND_RETURN_FALSE(writer->Init(operation.dst_extents(), block_size_));
   puffin::UniqueStreamPtr dst_stream(new PuffinExtentStream(
       std::move(writer),
@@ -477,6 +538,20 @@
   return true;
 }
 
+bool PartitionWriter::PerformPuffDiffOperation(
+    const InstallOperation& operation,
+    ErrorCode* error,
+    const void* data,
+    size_t count) {
+  FileDescriptorPtr source_fd = ChooseSourceFD(operation, error);
+  TEST_AND_RETURN_FALSE(source_fd != nullptr);
+
+  auto writer = CreateBaseExtentWriter();
+  writer->Init(operation.dst_extents(), block_size_);
+  return install_op_executor_.ExecutePuffDiffOperation(
+      operation, std::move(writer), source_fd, data, count);
+}
+
 FileDescriptorPtr PartitionWriter::ChooseSourceFD(
     const InstallOperation& operation, ErrorCode* error) {
   if (source_fd_ == nullptr) {
@@ -600,4 +675,34 @@
   return std::make_unique<DirectExtentWriter>(target_fd_);
 }
 
+bool InstallOperationExecutor::ExecuteInstallOp(
+    const InstallOperation& op,
+    std::unique_ptr<ExtentWriter> writer,
+    FileDescriptorPtr source_fd,
+    const void* data,
+    size_t size) {
+  switch (op.type()) {
+    case InstallOperation::REPLACE:
+    case InstallOperation::REPLACE_BZ:
+    case InstallOperation::REPLACE_XZ:
+      return ExecuteReplaceOperation(op, std::move(writer), data, size);
+    case InstallOperation::ZERO:
+    case InstallOperation::DISCARD:
+      return ExecuteZeroOrDiscardOperation(op, writer.get());
+    case InstallOperation::SOURCE_COPY:
+      return ExecuteSourceCopyOperation(op, writer.get(), source_fd);
+    case InstallOperation::SOURCE_BSDIFF:
+    case InstallOperation::BROTLI_BSDIFF:
+      return ExecuteSourceBsdiffOperation(
+          op, std::move(writer), source_fd, data, size);
+    case InstallOperation::PUFFDIFF:
+      return ExecutePuffDiffOperation(
+          op, std::move(writer), source_fd, data, size);
+      break;
+    default:
+      return false;
+  }
+  return false;
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/partition_writer.h b/payload_consumer/partition_writer.h
index 82e557a..dee30e8 100644
--- a/payload_consumer/partition_writer.h
+++ b/payload_consumer/partition_writer.h
@@ -31,6 +31,46 @@
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
+
+// A reference class for interpretation of different OTA ops.
+// Different partition writers(VABCPartitionWriter and the regular
+// PartitionWriter) will use this class via composition, and optionally optimize
+// for some specific operations. For example, in VABC, copy operations are
+// handled with special care, but other operations are defaulted to this class.
+class InstallOperationExecutor {
+ public:
+  explicit InstallOperationExecutor(size_t block_size)
+      : block_size_(block_size) {}
+
+  bool ExecuteInstallOp(const InstallOperation& op,
+                        std::unique_ptr<ExtentWriter> writer,
+                        FileDescriptorPtr source_fd,
+                        const void* data,
+                        size_t size);
+  bool ExecuteReplaceOperation(const InstallOperation& operation,
+                               std::unique_ptr<ExtentWriter> writer,
+                               const void* data,
+                               size_t count);
+  bool ExecuteZeroOrDiscardOperation(const InstallOperation& operation,
+                                     ExtentWriter* writer);
+  bool ExecuteSourceCopyOperation(const InstallOperation& operation,
+                                  ExtentWriter* writer,
+                                  FileDescriptorPtr source_fd);
+  bool ExecuteSourceBsdiffOperation(const InstallOperation& operation,
+                                    std::unique_ptr<ExtentWriter> writer,
+                                    FileDescriptorPtr source_fd,
+                                    const void* data,
+                                    size_t count);
+  bool ExecutePuffDiffOperation(const InstallOperation& operation,
+                                std::unique_ptr<ExtentWriter> writer,
+                                FileDescriptorPtr source_fd,
+                                const void* data,
+                                size_t count);
+
+ private:
+  size_t block_size_;
+};
+
 class PartitionWriter {
  public:
   PartitionWriter(const PartitionUpdate& partition_update,
@@ -129,6 +169,11 @@
   // Used to avoid re-opening the same source partition if it is not actually
   // error corrected.
   bool source_ecc_open_failure_{false};
+
+  // This instance handles decompression/bsdfif/puffdiff. It's responsible for
+  // constructing data which should be written to target partition, actual
+  // "writing" is handled by |PartitionWriter|
+  InstallOperationExecutor install_op_executor_;
 };
 
 namespace partition_writer {