Add InstallOperationExecutor which hanldes interpretation of install op am: 1d402cbc5e

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

Change-Id: Ief166b174055bfdeabd355120bc575cf5a63138f
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 {