Add partition writer class. am: 9bd519d4f5 am: 8a188256ba am: 88e2d6d48c am: 95169d81ac

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

Change-Id: Ifdbbbf920f3c0600c2975b01650bc9f228ac78eb
diff --git a/Android.bp b/Android.bp
index c5e66e3..acd3633 100644
--- a/Android.bp
+++ b/Android.bp
@@ -690,6 +690,7 @@
         "payload_consumer/certificate_parser_android_unittest.cc",
         "payload_consumer/delta_performer_integration_test.cc",
         "payload_consumer/delta_performer_unittest.cc",
+        "payload_consumer/partition_writer_unittest.cc",
         "payload_consumer/download_action_android_unittest.cc",
         "payload_consumer/extent_reader_unittest.cc",
         "payload_consumer/extent_writer_unittest.cc",
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index d9efc30..b49139e 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -282,6 +282,15 @@
 }
 
 int DeltaPerformer::CloseCurrentPartition() {
+  if (!partition_writer_) {
+    return 0;
+  }
+  int err = partition_writer_->Close();
+  partition_writer_ = nullptr;
+  return err;
+}
+
+int PartitionWriter::Close() {
   int err = 0;
   if (source_fd_ && !source_fd_->Close()) {
     err = errno;
@@ -290,14 +299,6 @@
       err = 1;
   }
   source_fd_.reset();
-  if (source_ecc_fd_ && !source_ecc_fd_->Close()) {
-    err = errno;
-    PLOG(ERROR) << "Error closing ECC source partition";
-    if (!err)
-      err = 1;
-  }
-  source_ecc_fd_.reset();
-  source_ecc_open_failure_ = false;
   source_path_.clear();
 
   if (target_fd_ && !target_fd_->Close()) {
@@ -308,6 +309,15 @@
   }
   target_fd_.reset();
   target_path_.clear();
+
+  if (source_ecc_fd_ && !source_ecc_fd_->Close()) {
+    err = errno;
+    PLOG(ERROR) << "Error closing ECC source partition";
+    if (!err)
+      err = 1;
+  }
+  source_ecc_fd_.reset();
+  source_ecc_open_failure_ = false;
   return -err;
 }
 
@@ -320,27 +330,43 @@
       install_plan_->partitions.size() - partitions_.size();
   const InstallPlan::Partition& install_part =
       install_plan_->partitions[num_previous_partitions + current_partition_];
+  partition_writer_ = std::make_unique<PartitionWriter>(
+      partition,
+      install_part,
+      boot_control_->GetDynamicPartitionControl(),
+      block_size_,
+      interactive_);
+
   // Open source fds if we have a delta payload, or for partitions in the
   // partial update.
   bool source_may_exist = manifest_.partial_update() ||
                           payload_->type == InstallPayloadType::kDelta;
+  return partition_writer_->Init(install_plan_, source_may_exist);
+}
+
+bool PartitionWriter::Init(const InstallPlan* install_plan,
+                           bool source_may_exist) {
+  const PartitionUpdate& partition = partition_update_;
+  uint32_t source_slot = install_plan->source_slot;
+  uint32_t target_slot = install_plan->target_slot;
+
   // We shouldn't open the source partition in certain cases, e.g. some dynamic
   // partitions in delta payload, partitions included in the full payload for
   // partial updates. Use the source size as the indicator.
-  if (source_may_exist && install_part.source_size > 0) {
-    source_path_ = install_part.source_path;
+  if (source_may_exist && install_part_.source_size > 0) {
+    source_path_ = install_part_.source_path;
     int err;
     source_fd_ = OpenFile(source_path_.c_str(), O_RDONLY, false, &err);
     if (!source_fd_) {
       LOG(ERROR) << "Unable to open source partition "
                  << partition.partition_name() << " on slot "
-                 << BootControlInterface::SlotName(install_plan_->source_slot)
-                 << ", file " << source_path_;
+                 << BootControlInterface::SlotName(source_slot) << ", file "
+                 << source_path_;
       return false;
     }
   }
 
-  target_path_ = install_part.target_path;
+  target_path_ = install_part_.target_path;
   int err;
 
   int flags = O_RDWR;
@@ -354,8 +380,8 @@
   if (!target_fd_) {
     LOG(ERROR) << "Unable to open target partition "
                << partition.partition_name() << " on slot "
-               << BootControlInterface::SlotName(install_plan_->target_slot)
-               << ", file " << target_path_;
+               << BootControlInterface::SlotName(target_slot) << ", file "
+               << target_path_;
     return false;
   }
 
@@ -364,38 +390,28 @@
             << "\"";
 
   // Discard the end of the partition, but ignore failures.
-  DiscardPartitionTail(target_fd_, install_part.target_size);
+  DiscardPartitionTail(target_fd_, install_part_.target_size);
 
   return true;
 }
 
-bool DeltaPerformer::OpenCurrentECCPartition() {
+bool PartitionWriter::OpenCurrentECCPartition() {
+  // No support for ECC for full payloads.
+  // Full Paylods should not have any operation that requires ECCPartition.
   if (source_ecc_fd_)
     return true;
 
   if (source_ecc_open_failure_)
     return false;
 
-  if (current_partition_ >= partitions_.size())
-    return false;
-
-  // No support for ECC for full payloads.
-  if (payload_->type == InstallPayloadType::kFull)
-    return false;
-
 #if USE_FEC
-  const PartitionUpdate& partition = partitions_[current_partition_];
-  size_t num_previous_partitions =
-      install_plan_->partitions.size() - partitions_.size();
-  const InstallPlan::Partition& install_part =
-      install_plan_->partitions[num_previous_partitions + current_partition_];
-  string path = install_part.source_path;
+  const PartitionUpdate& partition = partition_update_;
+  const InstallPlan::Partition& install_part = install_part_;
+  std::string path = install_part.source_path;
   FileDescriptorPtr fd(new FecFileDescriptor());
   if (!fd->Open(path.c_str(), O_RDONLY, 0)) {
     PLOG(ERROR) << "Unable to open ECC source partition "
-                << partition.partition_name() << " on slot "
-                << BootControlInterface::SlotName(install_plan_->source_slot)
-                << ", file " << path;
+                << partition.partition_name() << ", file " << path;
     source_ecc_open_failure_ = true;
     return false;
   }
@@ -733,10 +749,6 @@
     if (!HandleOpResult(op_result, InstallOperationTypeName(op.type()), error))
       return false;
 
-    if (!target_fd_->Flush()) {
-      return false;
-    }
-
     next_operation_num_++;
     UpdateOverallProgress(false, "Completed ");
     CheckpointUpdateProgress(false);
@@ -1003,9 +1015,18 @@
 
   // Since we delete data off the beginning of the buffer as we use it,
   // the data we need should be exactly at the beginning of the buffer.
-  TEST_AND_RETURN_FALSE(buffer_offset_ == operation.data_offset());
   TEST_AND_RETURN_FALSE(buffer_.size() >= operation.data_length());
 
+  TEST_AND_RETURN_FALSE(partition_writer_->PerformReplaceOperation(
+      operation, buffer_.data(), buffer_.size()));
+  // Update buffer
+  DiscardBuffer(true, buffer_.size());
+  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 = std::make_unique<DirectExtentWriter>();
 
@@ -1017,11 +1038,9 @@
 
   TEST_AND_RETURN_FALSE(
       writer->Init(target_fd_, operation.dst_extents(), block_size_));
-  TEST_AND_RETURN_FALSE(writer->Write(buffer_.data(), operation.data_length()));
+  TEST_AND_RETURN_FALSE(writer->Write(data, operation.data_length()));
 
-  // Update buffer
-  DiscardBuffer(true, buffer_.size());
-  return true;
+  return target_fd_->Flush();
 }
 
 bool DeltaPerformer::PerformZeroOrDiscardOperation(
@@ -1032,7 +1051,11 @@
   // These operations have no blob.
   TEST_AND_RETURN_FALSE(!operation.has_data_offset());
   TEST_AND_RETURN_FALSE(!operation.has_data_length());
+  return partition_writer_->PerformZeroOrDiscardOperation(operation);
+}
 
+bool PartitionWriter::PerformZeroOrDiscardOperation(
+    const InstallOperation& operation) {
 #ifdef BLKZEROOUT
   bool attempt_ioctl = true;
   int request =
@@ -1061,13 +1084,13 @@
           target_fd_, zeros.data(), chunk_length, start + offset));
     }
   }
-  return true;
+  return target_fd_->Flush();
 }
 
-bool DeltaPerformer::ValidateSourceHash(const brillo::Blob& calculated_hash,
-                                        const InstallOperation& operation,
-                                        const FileDescriptorPtr source_fd,
-                                        ErrorCode* error) {
+bool PartitionWriter::ValidateSourceHash(const brillo::Blob& calculated_hash,
+                                         const InstallOperation& operation,
+                                         const FileDescriptorPtr source_fd,
+                                         ErrorCode* error) {
   brillo::Blob expected_source_hash(operation.src_sha256_hash().begin(),
                                     operation.src_sha256_hash().end());
   if (calculated_hash != expected_source_hash) {
@@ -1108,14 +1131,18 @@
     TEST_AND_RETURN_FALSE(operation.src_length() % block_size_ == 0);
   if (operation.has_dst_length())
     TEST_AND_RETURN_FALSE(operation.dst_length() % block_size_ == 0);
+  return partition_writer_->PerformSourceCopyOperation(operation, error);
+}
 
+bool PartitionWriter::PerformSourceCopyOperation(
+    const InstallOperation& operation, ErrorCode* error) {
   TEST_AND_RETURN_FALSE(source_fd_ != nullptr);
 
   // The device may optimize the SOURCE_COPY operation.
   // Being this a device-specific optimization let DynamicPartitionController
   // decide it the operation should be skipped.
-  const PartitionUpdate& partition = partitions_[current_partition_];
-  const auto& partition_control = boot_control_->GetDynamicPartitionControl();
+  const PartitionUpdate& partition = partition_update_;
+  const auto& partition_control = dynamic_control_;
 
   InstallOperation buf;
   bool should_optimize = partition_control->OptimizeOperation(
@@ -1189,7 +1216,7 @@
     }
     TEST_AND_RETURN_FALSE(
         ValidateSourceHash(source_hash, operation, source_ecc_fd_, error));
-    // At this point reading from the the error corrected device worked, but
+    // At this point reading from the error corrected device worked, but
     // reading from the raw device failed, so this is considered a recovered
     // failure.
     source_ecc_recovered_failures_++;
@@ -1215,10 +1242,10 @@
                                                        block_size_,
                                                        nullptr));
   }
-  return true;
+  return target_fd_->Flush();
 }
 
-FileDescriptorPtr DeltaPerformer::ChooseSourceFD(
+FileDescriptorPtr PartitionWriter::ChooseSourceFD(
     const InstallOperation& operation, ErrorCode* error) {
   if (source_fd_ == nullptr) {
     LOG(ERROR) << "ChooseSourceFD fail: source_fd_ == nullptr";
@@ -1264,7 +1291,7 @@
   if (fd_utils::ReadAndHashExtents(
           source_ecc_fd_, operation.src_extents(), block_size_, &source_hash) &&
       ValidateSourceHash(source_hash, operation, source_ecc_fd_, error)) {
-    // At this point reading from the the error corrected device worked, but
+    // At this point reading from the error corrected device worked, but
     // reading from the raw device failed, so this is considered a recovered
     // failure.
     source_ecc_recovered_failures_++;
@@ -1369,6 +1396,17 @@
   if (operation.has_dst_length())
     TEST_AND_RETURN_FALSE(operation.dst_length() % block_size_ == 0);
 
+  TEST_AND_RETURN_FALSE(partition_writer_->PerformSourceBsdiffOperation(
+      operation, error, buffer_.data(), buffer_.size()));
+  DiscardBuffer(true, buffer_.size());
+  return true;
+}
+
+bool PartitionWriter::PerformSourceBsdiffOperation(
+    const InstallOperation& operation,
+    ErrorCode* error,
+    const void* data,
+    size_t count) {
   FileDescriptorPtr source_fd = ChooseSourceFD(operation, error);
   TEST_AND_RETURN_FALSE(source_fd != nullptr);
 
@@ -1388,10 +1426,9 @@
 
   TEST_AND_RETURN_FALSE(bsdiff::bspatch(std::move(src_file),
                                         std::move(dst_file),
-                                        buffer_.data(),
-                                        buffer_.size()) == 0);
-  DiscardBuffer(true, buffer_.size());
-  return true;
+                                        static_cast<const uint8_t*>(data),
+                                        count) == 0);
+  return target_fd_->Flush();
 }
 
 namespace {
@@ -1475,7 +1512,17 @@
   // the data we need should be exactly at the beginning of the buffer.
   TEST_AND_RETURN_FALSE(buffer_offset_ == operation.data_offset());
   TEST_AND_RETURN_FALSE(buffer_.size() >= operation.data_length());
+  TEST_AND_RETURN_FALSE(partition_writer_->PerformPuffDiffOperation(
+      operation, error, buffer_.data(), buffer_.size()));
+  DiscardBuffer(true, buffer_.size());
+  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);
 
@@ -1496,11 +1543,10 @@
   const size_t kMaxCacheSize = 5 * 1024 * 1024;  // Total 5MB cache.
   TEST_AND_RETURN_FALSE(puffin::PuffPatch(std::move(src_stream),
                                           std::move(dst_stream),
-                                          buffer_.data(),
-                                          buffer_.size(),
+                                          static_cast<const uint8_t*>(data),
+                                          count,
                                           kMaxCacheSize));
-  DiscardBuffer(true, buffer_.size());
-  return true;
+  return target_fd_->Flush();
 }
 
 bool DeltaPerformer::ExtractSignatureMessage() {
@@ -2040,4 +2086,20 @@
   return true;
 }
 
+PartitionWriter::PartitionWriter(
+    const PartitionUpdate& partition_update,
+    const InstallPlan::Partition& install_part,
+    DynamicPartitionControlInterface* dynamic_control,
+    size_t block_size,
+    bool is_interactive)
+    : partition_update_(partition_update),
+      install_part_(install_part),
+      dynamic_control_(dynamic_control),
+      interactive_(is_interactive),
+      block_size_(block_size) {}
+
+PartitionWriter::~PartitionWriter() {
+  Close();
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/delta_performer.h b/payload_consumer/delta_performer.h
index 88076af..bee7fde 100644
--- a/payload_consumer/delta_performer.h
+++ b/payload_consumer/delta_performer.h
@@ -46,6 +46,9 @@
 class HardwareInterface;
 class PrefsInterface;
 
+// At the bottom of this file.
+class PartitionWriter;
+
 // This class performs the actions in a delta update synchronously. The delta
 // update itself should be passed in in chunks as it is received.
 class DeltaPerformer : public FileWriter {
@@ -101,10 +104,6 @@
   // work. Returns whether the required file descriptors were successfully open.
   bool OpenCurrentPartition();
 
-  // Attempt to open the error-corrected device for the current partition.
-  // Returns whether the operation succeeded.
-  bool OpenCurrentECCPartition();
-
   // Closes the current partition file descriptors if open. Returns 0 on success
   // or -errno on error.
   int CloseCurrentPartition();
@@ -177,14 +176,6 @@
   // it returns that value, otherwise it returns the default value.
   uint32_t GetMinorVersion() const;
 
-  // Compare |calculated_hash| with source hash in |operation|, return false and
-  // dump hash and set |error| if don't match.
-  // |source_fd| is the file descriptor of the source partition.
-  static bool ValidateSourceHash(const brillo::Blob& calculated_hash,
-                                 const InstallOperation& operation,
-                                 const FileDescriptorPtr source_fd,
-                                 ErrorCode* error);
-
   // Initialize partitions and allocate required space for an update with the
   // given |manifest|. |update_check_response_hash| is used to check if the
   // previous call to this function corresponds to the same payload.
@@ -208,7 +199,6 @@
   friend class DeltaPerformerIntegrationTest;
   FRIEND_TEST(DeltaPerformerTest, BrilloMetadataSignatureSizeTest);
   FRIEND_TEST(DeltaPerformerTest, BrilloParsePayloadMetadataTest);
-  FRIEND_TEST(DeltaPerformerTest, ChooseSourceFDTest);
   FRIEND_TEST(DeltaPerformerTest, UsePublicKeyFromResponse);
 
   // Parse and move the update instructions of all partitions into our local
@@ -262,13 +252,6 @@
   bool PerformPuffDiffOperation(const InstallOperation& operation,
                                 ErrorCode* error);
 
-  // For a given operation, choose the source fd to be used (raw device or error
-  // correction device) based on the source operation hash.
-  // Returns nullptr if the source hash mismatch cannot be corrected, and set
-  // the |error| accordingly.
-  FileDescriptorPtr ChooseSourceFD(const InstallOperation& operation,
-                                   ErrorCode* error);
-
   // Extracts the payload signature message from the current |buffer_| if the
   // offset matches the one specified by the manifest. Returns whether the
   // signature was extracted.
@@ -335,34 +318,6 @@
   // Pointer to the current payload in install_plan_.payloads.
   InstallPlan::Payload* payload_{nullptr};
 
-  // File descriptor of the source partition. Only set while updating a
-  // partition when using a delta payload.
-  FileDescriptorPtr source_fd_{nullptr};
-
-  // File descriptor of the error corrected source partition. Only set while
-  // updating partition using a delta payload for a partition where error
-  // correction is available. The size of the error corrected device is smaller
-  // than the underlying raw device, since it doesn't include the error
-  // correction blocks.
-  FileDescriptorPtr source_ecc_fd_{nullptr};
-
-  // The total number of operations that failed source hash verification but
-  // passed after falling back to the error-corrected |source_ecc_fd_| device.
-  uint64_t source_ecc_recovered_failures_{0};
-
-  // Whether opening the current partition as an error-corrected device failed.
-  // Used to avoid re-opening the same source partition if it is not actually
-  // error corrected.
-  bool source_ecc_open_failure_{false};
-
-  // File descriptor of the target partition. Only set while performing the
-  // operations of a given partition.
-  FileDescriptorPtr target_fd_{nullptr};
-
-  // Paths the |source_fd_| and |target_fd_| refer to.
-  std::string source_path_;
-  std::string target_path_;
-
   PayloadMetadata payload_metadata_;
 
   // Parsed manifest. Set after enough bytes to parse the manifest were
@@ -452,9 +407,93 @@
       base::TimeDelta::FromSeconds(kCheckpointFrequencySeconds)};
   base::TimeTicks update_checkpoint_time_;
 
+  std::unique_ptr<PartitionWriter> partition_writer_;
+
   DISALLOW_COPY_AND_ASSIGN(DeltaPerformer);
 };
 
+class PartitionWriter {
+ public:
+  PartitionWriter(const PartitionUpdate& partition_update,
+                  const InstallPlan::Partition& install_part,
+                  DynamicPartitionControlInterface* dynamic_control,
+                  size_t block_size,
+                  bool is_interactive);
+  ~PartitionWriter();
+  // Compare |calculated_hash| with source hash in |operation|, return false and
+  // dump hash and set |error| if don't match.
+  // |source_fd| is the file descriptor of the source partition.
+  static bool ValidateSourceHash(const brillo::Blob& calculated_hash,
+                                 const InstallOperation& operation,
+                                 const FileDescriptorPtr source_fd,
+                                 ErrorCode* error);
+
+  // Perform necessary initialization work before InstallOperation can be
+  // applied to this partition
+  [[nodiscard]] bool Init(const InstallPlan* install_plan,
+                          bool source_may_exist);
+
+  int Close();
+
+  // These perform a specific type of operation and return true on success.
+  // |error| will be set if source hash mismatch, otherwise |error| might not be
+  // set even if it fails.
+  [[nodiscard]] bool PerformReplaceOperation(const InstallOperation& operation,
+                                             const void* data,
+                                             size_t count);
+  [[nodiscard]] bool PerformZeroOrDiscardOperation(
+      const InstallOperation& operation);
+
+  [[nodiscard]] bool PerformSourceCopyOperation(
+      const InstallOperation& operation, ErrorCode* error);
+  [[nodiscard]] bool PerformSourceBsdiffOperation(
+      const InstallOperation& operation,
+      ErrorCode* error,
+      const void* data,
+      size_t count);
+  [[nodiscard]] bool PerformPuffDiffOperation(const InstallOperation& operation,
+                                              ErrorCode* error,
+                                              const void* data,
+                                              size_t count);
+
+ private:
+  friend class PartitionWriterTest;
+  FRIEND_TEST(PartitionWriterTest, ChooseSourceFDTest);
+
+  bool OpenCurrentECCPartition();
+  // For a given operation, choose the source fd to be used (raw device or error
+  // correction device) based on the source operation hash.
+  // Returns nullptr if the source hash mismatch cannot be corrected, and set
+  // the |error| accordingly.
+  FileDescriptorPtr ChooseSourceFD(const InstallOperation& operation,
+                                   ErrorCode* error);
+
+  const PartitionUpdate& partition_update_;
+  const InstallPlan::Partition& install_part_;
+  DynamicPartitionControlInterface* dynamic_control_;
+  std::string source_path_;
+  std::string target_path_;
+  FileDescriptorPtr source_fd_;
+  FileDescriptorPtr target_fd_;
+  const bool interactive_;
+  const size_t block_size_;
+  // File descriptor of the error corrected source partition. Only set while
+  // updating partition using a delta payload for a partition where error
+  // correction is available. The size of the error corrected device is smaller
+  // than the underlying raw device, since it doesn't include the error
+  // correction blocks.
+  FileDescriptorPtr source_ecc_fd_{nullptr};
+
+  // The total number of operations that failed source hash verification but
+  // passed after falling back to the error-corrected |source_ecc_fd_| device.
+  uint64_t source_ecc_recovered_failures_{0};
+
+  // Whether opening the current partition as an error-corrected device failed.
+  // Used to avoid re-opening the same source partition if it is not actually
+  // error corrected.
+  bool source_ecc_open_failure_{false};
+};
+
 }  // namespace chromeos_update_engine
 
 #endif  // UPDATE_ENGINE_PAYLOAD_CONSUMER_DELTA_PERFORMER_H_
diff --git a/payload_consumer/delta_performer_unittest.cc b/payload_consumer/delta_performer_unittest.cc
index 449201c..a5eb538 100644
--- a/payload_consumer/delta_performer_unittest.cc
+++ b/payload_consumer/delta_performer_unittest.cc
@@ -418,22 +418,7 @@
     EXPECT_EQ(payload_.metadata_size, performer_.metadata_size_);
   }
 
-  // Helper function to pretend that the ECC file descriptor was already opened.
-  // Returns a pointer to the created file descriptor.
-  FakeFileDescriptor* SetFakeECCFile(size_t size) {
-    EXPECT_FALSE(performer_.source_ecc_fd_) << "source_ecc_fd_ already open.";
-    FakeFileDescriptor* ret = new FakeFileDescriptor();
-    fake_ecc_fd_.reset(ret);
-    // Call open to simulate it was already opened.
-    ret->Open("", 0);
-    ret->SetFileSize(size);
-    performer_.source_ecc_fd_ = fake_ecc_fd_;
-    return ret;
-  }
 
-  uint64_t GetSourceEccRecoveredFailures() const {
-    return performer_.source_ecc_recovered_failures_;
-  }
 
   FakePrefs prefs_;
   InstallPlan install_plan_;
@@ -660,94 +645,8 @@
   EXPECT_EQ(actual_data, ApplyPayload(payload_data, source.path(), false));
 }
 
-// Test that the error-corrected file descriptor is used to read the partition
-// since the source partition doesn't match the operation hash.
-TEST_F(DeltaPerformerTest, ErrorCorrectionSourceCopyFallbackTest) {
-  constexpr size_t kCopyOperationSize = 4 * 4096;
-  test_utils::ScopedTempFile source("Source-XXXXXX");
-  // Write invalid data to the source image, which doesn't match the expected
-  // hash.
-  brillo::Blob invalid_data(kCopyOperationSize, 0x55);
-  EXPECT_TRUE(test_utils::WriteFileVector(source.path(), invalid_data));
 
-  // Setup the fec file descriptor as the fake stream, which matches
-  // |expected_data|.
-  FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize);
-  brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
 
-  PartitionConfig old_part(kPartitionNameRoot);
-  old_part.path = source.path();
-  old_part.size = invalid_data.size();
-
-  brillo::Blob payload_data =
-      GenerateSourceCopyPayload(expected_data, true, &old_part);
-  EXPECT_EQ(expected_data, ApplyPayload(payload_data, source.path(), true));
-  // Verify that the fake_fec was actually used.
-  EXPECT_EQ(1U, fake_fec->GetReadOps().size());
-  EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
-}
-
-// Test that the error-corrected file descriptor is used to read a partition
-// when no hash is available for SOURCE_COPY but it falls back to the normal
-// file descriptor when the size of the error corrected one is too small.
-TEST_F(DeltaPerformerTest, ErrorCorrectionSourceCopyWhenNoHashFallbackTest) {
-  constexpr size_t kCopyOperationSize = 4 * 4096;
-  test_utils::ScopedTempFile source("Source-XXXXXX");
-  // Setup the source path with the right expected data.
-  brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
-  EXPECT_TRUE(test_utils::WriteFileVector(source.path(), expected_data));
-
-  // Setup the fec file descriptor as the fake stream, with smaller data than
-  // the expected.
-  FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize / 2);
-
-  PartitionConfig old_part(kPartitionNameRoot);
-  old_part.path = source.path();
-  old_part.size = expected_data.size();
-
-  // The payload operation doesn't include an operation hash.
-  brillo::Blob payload_data =
-      GenerateSourceCopyPayload(expected_data, false, &old_part);
-  EXPECT_EQ(expected_data, ApplyPayload(payload_data, source.path(), true));
-  // Verify that the fake_fec was attempted to be used. Since the file
-  // descriptor is shorter it can actually do more than one read to realize it
-  // reached the EOF.
-  EXPECT_LE(1U, fake_fec->GetReadOps().size());
-  // This fallback doesn't count as an error-corrected operation since the
-  // operation hash was not available.
-  EXPECT_EQ(0U, GetSourceEccRecoveredFailures());
-}
-
-TEST_F(DeltaPerformerTest, ChooseSourceFDTest) {
-  constexpr size_t kSourceSize = 4 * 4096;
-  test_utils::ScopedTempFile source("Source-XXXXXX");
-  // Write invalid data to the source image, which doesn't match the expected
-  // hash.
-  brillo::Blob invalid_data(kSourceSize, 0x55);
-  EXPECT_TRUE(test_utils::WriteFileVector(source.path(), invalid_data));
-
-  performer_.source_fd_ = std::make_shared<EintrSafeFileDescriptor>();
-  performer_.source_fd_->Open(source.path().c_str(), O_RDONLY);
-  performer_.block_size_ = 4096;
-
-  // Setup the fec file descriptor as the fake stream, which matches
-  // |expected_data|.
-  FakeFileDescriptor* fake_fec = SetFakeECCFile(kSourceSize);
-  brillo::Blob expected_data = FakeFileDescriptorData(kSourceSize);
-
-  InstallOperation op;
-  *(op.add_src_extents()) = ExtentForRange(0, kSourceSize / 4096);
-  brillo::Blob src_hash;
-  EXPECT_TRUE(HashCalculator::RawHashOfData(expected_data, &src_hash));
-  op.set_src_sha256_hash(src_hash.data(), src_hash.size());
-
-  ErrorCode error = ErrorCode::kSuccess;
-  EXPECT_EQ(performer_.source_ecc_fd_, performer_.ChooseSourceFD(op, &error));
-  EXPECT_EQ(ErrorCode::kSuccess, error);
-  // Verify that the fake_fec was actually used.
-  EXPECT_EQ(1U, fake_fec->GetReadOps().size());
-  EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
-}
 
 TEST_F(DeltaPerformerTest, ExtentsToByteStringTest) {
   uint64_t test[] = {1, 1, 4, 2, 0, 1};
diff --git a/payload_consumer/partition_writer_unittest.cc b/payload_consumer/partition_writer_unittest.cc
new file mode 100644
index 0000000..c1ff4f4
--- /dev/null
+++ b/payload_consumer/partition_writer_unittest.cc
@@ -0,0 +1,203 @@
+//
+// Copyright (C) 2020 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include <memory>
+#include <vector>
+
+#include <brillo/secure_blob.h>
+#include <gtest/gtest.h>
+
+#include "update_engine/common/dynamic_partition_control_stub.h"
+#include "update_engine/common/error_code.h"
+#include "update_engine/common/fake_prefs.h"
+#include "update_engine/common/hash_calculator.h"
+#include "update_engine/common/test_utils.h"
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_consumer/delta_performer.h"
+#include "update_engine/payload_consumer/extent_reader.h"
+#include "update_engine/payload_consumer/extent_writer.h"
+#include "update_engine/payload_consumer/fake_file_descriptor.h"
+#include "update_engine/payload_consumer/file_descriptor.h"
+#include "update_engine/payload_consumer/install_plan.h"
+#include "update_engine/payload_generator/annotated_operation.h"
+#include "update_engine/payload_generator/delta_diff_generator.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+#include "update_engine/payload_generator/payload_file.h"
+#include "update_engine/payload_generator/payload_generation_config.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+
+class PartitionWriterTest : public testing::Test {
+ public:
+  // Helper function to pretend that the ECC file descriptor was already opened.
+  // Returns a pointer to the created file descriptor.
+  FakeFileDescriptor* SetFakeECCFile(size_t size) {
+    EXPECT_FALSE(writer_.source_ecc_fd_) << "source_ecc_fd_ already open.";
+    FakeFileDescriptor* ret = new FakeFileDescriptor();
+    fake_ecc_fd_.reset(ret);
+    // Call open to simulate it was already opened.
+    ret->Open("", 0);
+    ret->SetFileSize(size);
+    writer_.source_ecc_fd_ = fake_ecc_fd_;
+    return ret;
+  }
+
+  uint64_t GetSourceEccRecoveredFailures() const {
+    return writer_.source_ecc_recovered_failures_;
+  }
+
+  AnnotatedOperation GenerateSourceCopyOp(const brillo::Blob& copied_data,
+                                          bool add_hash,
+                                          PartitionConfig* old_part = nullptr) {
+    PayloadGenerationConfig config;
+    const uint64_t kDefaultBlockSize = config.block_size;
+    EXPECT_EQ(0U, copied_data.size() % kDefaultBlockSize);
+    uint64_t num_blocks = copied_data.size() / kDefaultBlockSize;
+    AnnotatedOperation aop;
+    *(aop.op.add_src_extents()) = ExtentForRange(0, num_blocks);
+    *(aop.op.add_dst_extents()) = ExtentForRange(0, num_blocks);
+    aop.op.set_type(InstallOperation::SOURCE_COPY);
+    brillo::Blob src_hash;
+    EXPECT_TRUE(HashCalculator::RawHashOfData(copied_data, &src_hash));
+    if (add_hash)
+      aop.op.set_src_sha256_hash(src_hash.data(), src_hash.size());
+
+    return aop;
+  }
+
+  brillo::Blob PerformSourceCopyOp(const InstallOperation& op,
+                                   const brillo::Blob blob_data) {
+    test_utils::ScopedTempFile source_partition("Blob-XXXXXX");
+    DirectExtentWriter extent_writer;
+    FileDescriptorPtr fd(new EintrSafeFileDescriptor());
+    EXPECT_TRUE(fd->Open(source_partition.path().c_str(), O_RDWR));
+    EXPECT_TRUE(extent_writer.Init(fd, op.src_extents(), kBlockSize));
+    EXPECT_TRUE(extent_writer.Write(blob_data.data(), blob_data.size()));
+
+    test_utils::ScopedTempFile target_partition("Blob-XXXXXX");
+
+    install_part_.source_path = source_partition.path();
+    install_part_.target_path = target_partition.path();
+    install_part_.source_size = blob_data.size();
+    install_part_.target_size = blob_data.size();
+
+    ErrorCode error;
+    EXPECT_TRUE(writer_.Init(&install_plan_, true));
+    EXPECT_TRUE(writer_.PerformSourceCopyOperation(op, &error));
+
+    brillo::Blob output_data;
+    EXPECT_TRUE(utils::ReadFile(target_partition.path(), &output_data));
+    return output_data;
+  }
+
+  FakePrefs prefs_{};
+  InstallPlan install_plan_{};
+  InstallPlan::Payload payload_{};
+  DynamicPartitionControlStub dynamic_control_{};
+  FileDescriptorPtr fake_ecc_fd_{};
+  DeltaArchiveManifest manifest_{};
+  PartitionUpdate partition_update_{};
+  InstallPlan::Partition install_part_{};
+  PartitionWriter writer_{
+      partition_update_, install_part_, &dynamic_control_, kBlockSize, false};
+};
+// Test that the error-corrected file descriptor is used to read a partition
+// when no hash is available for SOURCE_COPY but it falls back to the normal
+// file descriptor when the size of the error corrected one is too small.
+TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyWhenNoHashFallbackTest) {
+  constexpr size_t kCopyOperationSize = 4 * 4096;
+  test_utils::ScopedTempFile source("Source-XXXXXX");
+  // Setup the source path with the right expected data.
+  brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
+  EXPECT_TRUE(test_utils::WriteFileVector(source.path(), expected_data));
+
+  // Setup the fec file descriptor as the fake stream, with smaller data than
+  // the expected.
+  FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize / 2);
+
+  PartitionConfig old_part(kPartitionNameRoot);
+  old_part.path = source.path();
+  old_part.size = expected_data.size();
+
+  // The payload operation doesn't include an operation hash.
+  auto source_copy_op = GenerateSourceCopyOp(expected_data, false, &old_part);
+
+  auto output_data = PerformSourceCopyOp(source_copy_op.op, expected_data);
+  ASSERT_EQ(output_data, expected_data);
+
+  // Verify that the fake_fec was attempted to be used. Since the file
+  // descriptor is shorter it can actually do more than one read to realize it
+  // reached the EOF.
+  EXPECT_LE(1U, fake_fec->GetReadOps().size());
+  // This fallback doesn't count as an error-corrected operation since the
+  // operation hash was not available.
+  EXPECT_EQ(0U, GetSourceEccRecoveredFailures());
+}
+
+// Test that the error-corrected file descriptor is used to read the partition
+// since the source partition doesn't match the operation hash.
+TEST_F(PartitionWriterTest, ErrorCorrectionSourceCopyFallbackTest) {
+  constexpr size_t kCopyOperationSize = 4 * 4096;
+  // Write invalid data to the source image, which doesn't match the expected
+  // hash.
+  brillo::Blob invalid_data(kCopyOperationSize, 0x55);
+
+  // Setup the fec file descriptor as the fake stream, which matches
+  // |expected_data|.
+  FakeFileDescriptor* fake_fec = SetFakeECCFile(kCopyOperationSize);
+  brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
+
+  auto source_copy_op = GenerateSourceCopyOp(expected_data, true);
+  auto output_data = PerformSourceCopyOp(source_copy_op.op, invalid_data);
+  ASSERT_EQ(output_data, expected_data);
+
+  // Verify that the fake_fec was actually used.
+  EXPECT_EQ(1U, fake_fec->GetReadOps().size());
+  EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
+}
+
+TEST_F(PartitionWriterTest, ChooseSourceFDTest) {
+  constexpr size_t kSourceSize = 4 * 4096;
+  test_utils::ScopedTempFile source("Source-XXXXXX");
+  // Write invalid data to the source image, which doesn't match the expected
+  // hash.
+  brillo::Blob invalid_data(kSourceSize, 0x55);
+  EXPECT_TRUE(test_utils::WriteFileVector(source.path(), invalid_data));
+
+  writer_.source_fd_ = std::make_shared<EintrSafeFileDescriptor>();
+  writer_.source_fd_->Open(source.path().c_str(), O_RDONLY);
+
+  // Setup the fec file descriptor as the fake stream, which matches
+  // |expected_data|.
+  FakeFileDescriptor* fake_fec = SetFakeECCFile(kSourceSize);
+  brillo::Blob expected_data = FakeFileDescriptorData(kSourceSize);
+
+  InstallOperation op;
+  *(op.add_src_extents()) = ExtentForRange(0, kSourceSize / 4096);
+  brillo::Blob src_hash;
+  EXPECT_TRUE(HashCalculator::RawHashOfData(expected_data, &src_hash));
+  op.set_src_sha256_hash(src_hash.data(), src_hash.size());
+
+  ErrorCode error = ErrorCode::kSuccess;
+  EXPECT_EQ(writer_.source_ecc_fd_, writer_.ChooseSourceFD(op, &error));
+  EXPECT_EQ(ErrorCode::kSuccess, error);
+  // Verify that the fake_fec was actually used.
+  EXPECT_EQ(1U, fake_fec->GetReadOps().size());
+  EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
+}
+
+}  // namespace chromeos_update_engine
diff --git a/update_attempter_android.cc b/update_attempter_android.cc
index 7fc13e1..3578d95 100644
--- a/update_attempter_android.cc
+++ b/update_attempter_android.cc
@@ -507,7 +507,7 @@
         return LogAndSetError(
             error, FROM_HERE, "Failed to hash " + partition_path);
       }
-      if (!DeltaPerformer::ValidateSourceHash(
+      if (!PartitionWriter::ValidateSourceHash(
               source_hash, operation, fd, &errorcode)) {
         return false;
       }