Refactor ECC related code to a separate class

Both PartitionWriter and VABC partition writer need to deal with
hashes/ecc. Refactor out common code to a free function.

Test: th
Change-Id: I40033a1671a2c3a63e7d2d8266c4a0087d067100
diff --git a/Android.bp b/Android.bp
index e07681b..c519280 100644
--- a/Android.bp
+++ b/Android.bp
@@ -244,6 +244,7 @@
         "payload_consumer/vabc_partition_writer.cc",
         "payload_consumer/snapshot_extent_writer.cc",
         "payload_consumer/postinstall_runner_action.cc",
+        "payload_consumer/verified_source_fd.cc",
         "payload_consumer/verity_writer_android.cc",
         "payload_consumer/xz_extent_writer.cc",
         "payload_consumer/fec_file_descriptor.cc",
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index a57169b..7c07f1d 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -32,8 +32,6 @@
 #include <base/format_macros.h>
 #include <base/metrics/histogram_macros.h>
 #include <base/strings/string_number_conversions.h>
-#include <base/strings/string_util.h>
-#include <base/strings/stringprintf.h>
 #include <base/time/time.h>
 #include <brillo/data_encoding.h>
 #include <bsdiff/bspatch.h>
@@ -873,44 +871,6 @@
   return partition_writer_->PerformZeroOrDiscardOperation(operation);
 }
 
-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) {
-    LOG(ERROR) << "The hash of the source data on disk for this operation "
-               << "doesn't match the expected value. This could mean that the "
-               << "delta update payload was targeted for another version, or "
-               << "that the source partition was modified after it was "
-               << "installed, for example, by mounting a filesystem.";
-    LOG(ERROR) << "Expected:   sha256|hex = "
-               << base::HexEncode(expected_source_hash.data(),
-                                  expected_source_hash.size());
-    LOG(ERROR) << "Calculated: sha256|hex = "
-               << base::HexEncode(calculated_hash.data(),
-                                  calculated_hash.size());
-
-    vector<string> source_extents;
-    for (const Extent& ext : operation.src_extents()) {
-      source_extents.push_back(
-          base::StringPrintf("%" PRIu64 ":%" PRIu64,
-                             static_cast<uint64_t>(ext.start_block()),
-                             static_cast<uint64_t>(ext.num_blocks())));
-    }
-    LOG(ERROR) << "Operation source (offset:size) in blocks: "
-               << base::JoinString(source_extents, ",");
-
-    // Log remount history if this device is an ext4 partition.
-    LogMountHistory(source_fd);
-
-    *error = ErrorCode::kDownloadStateInitializationError;
-    return false;
-  }
-  return true;
-}
-
 bool DeltaPerformer::PerformSourceCopyOperation(
     const InstallOperation& operation, ErrorCode* error) {
   if (operation.has_src_length())
diff --git a/payload_consumer/delta_performer.h b/payload_consumer/delta_performer.h
index c54316b..1d95e1e 100644
--- a/payload_consumer/delta_performer.h
+++ b/payload_consumer/delta_performer.h
@@ -176,14 +176,6 @@
   // Exposed for testing purposes.
   bool CheckpointUpdateProgress(bool force);
 
-  // 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.
diff --git a/payload_consumer/install_operation_executor.cc b/payload_consumer/install_operation_executor.cc
index 3e477e5..4bfcb5d 100644
--- a/payload_consumer/install_operation_executor.cc
+++ b/payload_consumer/install_operation_executor.cc
@@ -248,7 +248,7 @@
     ExtentWriter* writer,
     FileDescriptorPtr source_fd) {
   TEST_AND_RETURN_FALSE(operation.type() == InstallOperation::SOURCE_COPY);
-  writer->Init(operation.dst_extents(), block_size_);
+  TEST_AND_RETURN_FALSE(writer->Init(operation.dst_extents(), block_size_));
   return fd_utils::CommonHashExtents(
       source_fd, operation.src_extents(), writer, block_size_, nullptr);
 }
diff --git a/payload_consumer/partition_writer.cc b/payload_consumer/partition_writer.cc
index 4df0af6..9db7ae0 100644
--- a/payload_consumer/partition_writer.cc
+++ b/payload_consumer/partition_writer.cc
@@ -19,13 +19,18 @@
 #include <linux/fs.h>
 #include <sys/mman.h>
 
+#include <inttypes.h>
+
 #include <algorithm>
 #include <initializer_list>
 #include <memory>
+#include <string>
 #include <utility>
 #include <vector>
 
 #include <base/strings/string_number_conversions.h>
+#include <base/strings/string_util.h>
+#include <base/strings/stringprintf.h>
 
 #include "update_engine/common/terminator.h"
 #include "update_engine/common/utils.h"
@@ -33,7 +38,6 @@
 #include "update_engine/payload_consumer/cached_file_descriptor.h"
 #include "update_engine/payload_consumer/extent_reader.h"
 #include "update_engine/payload_consumer/extent_writer.h"
-#include "update_engine/payload_consumer/fec_file_descriptor.h"
 #include "update_engine/payload_consumer/file_descriptor_utils.h"
 #include "update_engine/payload_consumer/install_plan.h"
 #include "update_engine/payload_consumer/mount_history.h"
@@ -117,6 +121,7 @@
     : partition_update_(partition_update),
       install_part_(install_part),
       dynamic_control_(dynamic_control),
+      verified_source_fd_(block_size, install_part.source_path),
       interactive_(is_interactive),
       block_size_(block_size),
       install_op_executor_(block_size) {}
@@ -133,9 +138,7 @@
   }
   if (install_part_.source_size > 0 && !install_part_.source_path.empty()) {
     source_path_ = install_part_.source_path;
-    int err;
-    source_fd_ = OpenFile(source_path_.c_str(), O_RDONLY, false, &err);
-    if (source_fd_ == nullptr) {
+    if (!verified_source_fd_.Open()) {
       LOG(ERROR) << "Unable to open source partition " << install_part_.name
                  << " on slot " << BootControlInterface::SlotName(source_slot)
                  << ", file " << source_path_;
@@ -244,8 +247,6 @@
 
 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.
@@ -303,98 +304,12 @@
 
 FileDescriptorPtr PartitionWriter::ChooseSourceFD(
     const InstallOperation& operation, ErrorCode* error) {
-  if (source_fd_ == nullptr) {
-    LOG(ERROR) << "ChooseSourceFD fail: source_fd_ == nullptr";
-    return nullptr;
-  }
-
-  if (!operation.has_src_sha256_hash()) {
-    // When the operation doesn't include a source hash, we attempt the error
-    // corrected device first since we can't verify the block in the raw
-    // device at this point, but we first need to make sure all extents are
-    // readable since the error corrected device can be shorter or not
-    // available.
-    if (OpenCurrentECCPartition() &&
-        fd_utils::ReadAndHashExtents(
-            source_ecc_fd_, operation.src_extents(), block_size_, nullptr)) {
-      return source_ecc_fd_;
-    }
-    return source_fd_;
-  }
-
-  brillo::Blob source_hash;
-  brillo::Blob expected_source_hash(operation.src_sha256_hash().begin(),
-                                    operation.src_sha256_hash().end());
-  if (fd_utils::ReadAndHashExtents(
-          source_fd_, operation.src_extents(), block_size_, &source_hash) &&
-      source_hash == expected_source_hash) {
-    return source_fd_;
-  }
-  // We fall back to use the error corrected device if the hash of the raw
-  // device doesn't match or there was an error reading the source partition.
-  if (!OpenCurrentECCPartition()) {
-    // The following function call will return false since the source hash
-    // mismatches, but we still want to call it so it prints the appropriate
-    // log message.
-    ValidateSourceHash(source_hash, operation, source_fd_, error);
-    return nullptr;
-  }
-  LOG(WARNING) << "Source hash from RAW device mismatched: found "
-               << base::HexEncode(source_hash.data(), source_hash.size())
-               << ", expected "
-               << base::HexEncode(expected_source_hash.data(),
-                                  expected_source_hash.size());
-
-  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 error corrected device worked, but
-    // reading from the raw device failed, so this is considered a recovered
-    // failure.
-    source_ecc_recovered_failures_++;
-    return source_ecc_fd_;
-  }
-  return nullptr;
-}
-
-bool PartitionWriter::OpenCurrentECCPartition() {
-  // No support for ECC for full payloads.
-  // Full payload should not have any opeartion that requires ECC partitions.
-  if (source_ecc_fd_)
-    return true;
-
-  if (source_ecc_open_failure_)
-    return false;
-
-#if USE_FEC
-  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() << ", file " << path;
-    source_ecc_open_failure_ = true;
-    return false;
-  }
-  source_ecc_fd_ = fd;
-#else
-  // No support for ECC compiled.
-  source_ecc_open_failure_ = true;
-#endif  // USE_FEC
-
-  return !source_ecc_open_failure_;
+  return verified_source_fd_.ChooseSourceFD(operation, error);
 }
 
 int PartitionWriter::Close() {
   int err = 0;
-  if (source_fd_ && !source_fd_->Close()) {
-    err = errno;
-    PLOG(ERROR) << "Error closing source partition";
-    if (!err)
-      err = 1;
-  }
-  source_fd_.reset();
+
   source_path_.clear();
 
   if (target_fd_ && !target_fd_->Close()) {
@@ -406,14 +321,6 @@
   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;
 }
 
@@ -425,4 +332,44 @@
   return std::make_unique<DirectExtentWriter>(target_fd_);
 }
 
+bool PartitionWriter::ValidateSourceHash(const brillo::Blob& calculated_hash,
+                                         const InstallOperation& operation,
+                                         const FileDescriptorPtr source_fd,
+                                         ErrorCode* error) {
+  using std::string;
+  using std::vector;
+  brillo::Blob expected_source_hash(operation.src_sha256_hash().begin(),
+                                    operation.src_sha256_hash().end());
+  if (calculated_hash != expected_source_hash) {
+    LOG(ERROR) << "The hash of the source data on disk for this operation "
+               << "doesn't match the expected value. This could mean that the "
+               << "delta update payload was targeted for another version, or "
+               << "that the source partition was modified after it was "
+               << "installed, for example, by mounting a filesystem.";
+    LOG(ERROR) << "Expected:   sha256|hex = "
+               << base::HexEncode(expected_source_hash.data(),
+                                  expected_source_hash.size());
+    LOG(ERROR) << "Calculated: sha256|hex = "
+               << base::HexEncode(calculated_hash.data(),
+                                  calculated_hash.size());
+
+    vector<string> source_extents;
+    for (const Extent& ext : operation.src_extents()) {
+      source_extents.push_back(
+          base::StringPrintf("%" PRIu64 ":%" PRIu64,
+                             static_cast<uint64_t>(ext.start_block()),
+                             static_cast<uint64_t>(ext.num_blocks())));
+    }
+    LOG(ERROR) << "Operation source (offset:size) in blocks: "
+               << base::JoinString(source_extents, ",");
+
+    // Log remount history if this device is an ext4 partition.
+    LogMountHistory(source_fd);
+
+    *error = ErrorCode::kDownloadStateInitializationError;
+    return false;
+  }
+  return true;
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/partition_writer.h b/payload_consumer/partition_writer.h
index 554e590..14bd18a 100644
--- a/payload_consumer/partition_writer.h
+++ b/payload_consumer/partition_writer.h
@@ -29,6 +29,7 @@
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_consumer/install_operation_executor.h"
 #include "update_engine/payload_consumer/install_plan.h"
+#include "update_engine/payload_consumer/verified_source_fd.h"
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
@@ -94,7 +95,8 @@
   friend class PartitionWriterTest;
   FRIEND_TEST(PartitionWriterTest, ChooseSourceFDTest);
 
-  bool OpenSourcePartition(uint32_t source_slot, bool source_may_exist);
+  [[nodiscard]] bool OpenSourcePartition(uint32_t source_slot,
+                                         bool source_may_exist);
 
   bool OpenCurrentECCPartition();
   // For a given operation, choose the source fd to be used (raw device or error
@@ -110,27 +112,12 @@
   DynamicPartitionControlInterface* dynamic_control_;
   // Path to source partition
   std::string source_path_;
+  VerifiedSourceFd verified_source_fd_;
   // Path to target partition
   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};
 
   // This instance handles decompression/bsdfif/puffdiff. It's responsible for
   // constructing data which should be written to target partition, actual
diff --git a/payload_consumer/partition_writer_unittest.cc b/payload_consumer/partition_writer_unittest.cc
index 564d8d4..331a061 100644
--- a/payload_consumer/partition_writer_unittest.cc
+++ b/payload_consumer/partition_writer_unittest.cc
@@ -46,18 +46,19 @@
   // 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.";
+    EXPECT_FALSE(writer_.verified_source_fd_.source_ecc_fd_)
+        << "source_ecc_fdb 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_;
+    writer_.verified_source_fd_.source_ecc_fd_ = fake_ecc_fd_;
     return ret;
   }
 
   uint64_t GetSourceEccRecoveredFailures() const {
-    return writer_.source_ecc_recovered_failures_;
+    return writer_.verified_source_fd_.source_ecc_recovered_failures_;
   }
 
   AnnotatedOperation GenerateSourceCopyOp(const brillo::Blob& copied_data,
@@ -81,22 +82,31 @@
 
   brillo::Blob PerformSourceCopyOp(const InstallOperation& op,
                                    const brillo::Blob blob_data) {
-    ScopedTempFile source_partition("Blob-XXXXXX");
+    LOG(INFO) << "Using source part " << source_partition.path();
     FileDescriptorPtr fd(new EintrSafeFileDescriptor());
     DirectExtentWriter extent_writer{fd};
     EXPECT_TRUE(fd->Open(source_partition.path().c_str(), O_RDWR));
+    if (HasFailure()) {
+      return {};
+    }
     EXPECT_TRUE(extent_writer.Init(op.src_extents(), kBlockSize));
+    if (HasFailure()) {
+      return {};
+    }
     EXPECT_TRUE(extent_writer.Write(blob_data.data(), blob_data.size()));
+    if (HasFailure()) {
+      return {};
+    }
+    fd->Flush();
 
-    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, 0));
+    if (HasFailure()) {
+      return {};
+    }
     EXPECT_TRUE(writer_.PerformSourceCopyOperation(op, &error));
     writer_.CheckpointUpdateProgress(1);
 
@@ -111,8 +121,11 @@
   DynamicPartitionControlStub dynamic_control_{};
   FileDescriptorPtr fake_ecc_fd_{};
   DeltaArchiveManifest manifest_{};
+  ScopedTempFile source_partition{"source-part-XXXXXX"};
+  ScopedTempFile target_partition{"target-part-XXXXXX"};
+  InstallPlan::Partition install_part_{.source_path = source_partition.path(),
+                                       .target_path = target_partition.path()};
   PartitionUpdate partition_update_{};
-  InstallPlan::Partition install_part_{};
   PartitionWriter writer_{
       partition_update_, install_part_, &dynamic_control_, kBlockSize, false};
 };
@@ -124,7 +137,7 @@
   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));
+  ASSERT_TRUE(test_utils::WriteFileVector(source.path(), expected_data));
 
   // Setup the fec file descriptor as the fake stream, with smaller data than
   // the expected.
@@ -136,17 +149,18 @@
 
   // The payload operation doesn't include an operation hash.
   auto source_copy_op = GenerateSourceCopyOp(expected_data, false, &old_part);
-
+  ASSERT_NO_FATAL_FAILURE();
   auto output_data = PerformSourceCopyOp(source_copy_op.op, expected_data);
+  ASSERT_NO_FATAL_FAILURE();
   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());
+  ASSERT_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());
+  ASSERT_EQ(0U, GetSourceEccRecoveredFailures());
 }
 
 // Test that the error-corrected file descriptor is used to read the partition
@@ -163,7 +177,9 @@
   brillo::Blob expected_data = FakeFileDescriptorData(kCopyOperationSize);
 
   auto source_copy_op = GenerateSourceCopyOp(expected_data, true);
+  ASSERT_NO_FATAL_FAILURE();
   auto output_data = PerformSourceCopyOp(source_copy_op.op, invalid_data);
+  ASSERT_NO_FATAL_FAILURE();
   ASSERT_EQ(output_data, expected_data);
 
   // Verify that the fake_fec was actually used.
@@ -177,10 +193,11 @@
   // 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));
+  ASSERT_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);
+  writer_.verified_source_fd_.source_fd_ =
+      std::make_shared<EintrSafeFileDescriptor>();
+  writer_.verified_source_fd_.source_fd_->Open(source.path().c_str(), O_RDONLY);
 
   // Setup the fec file descriptor as the fake stream, which matches
   // |expected_data|.
@@ -190,15 +207,16 @@
   InstallOperation op;
   *(op.add_src_extents()) = ExtentForRange(0, kSourceSize / 4096);
   brillo::Blob src_hash;
-  EXPECT_TRUE(HashCalculator::RawHashOfData(expected_data, &src_hash));
+  ASSERT_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);
+  ASSERT_EQ(writer_.verified_source_fd_.source_ecc_fd_,
+            writer_.ChooseSourceFD(op, &error));
+  ASSERT_EQ(ErrorCode::kSuccess, error);
   // Verify that the fake_fec was actually used.
-  EXPECT_EQ(1U, fake_fec->GetReadOps().size());
-  EXPECT_EQ(1U, GetSourceEccRecoveredFailures());
+  ASSERT_EQ(1U, fake_fec->GetReadOps().size());
+  ASSERT_EQ(1U, GetSourceEccRecoveredFailures());
 }
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index 0843fff..a27847c 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -90,7 +90,15 @@
   auto converted = ConvertToCowOperations(partition_update_.operations(),
                                           partition_update_.merge_operations());
 
-  WriteAllCowOps(block_size_, converted, cow_writer_.get(), source_fd_);
+  if (!converted.empty()) {
+    // Use source fd directly. Ideally we want to verify all extents used in
+    // source copy, but then what do we do if some extents contain correct
+    // hashes and some don't?
+    auto source_fd = std::make_shared<EintrSafeFileDescriptor>();
+    TEST_AND_RETURN_FALSE_ERRNO(
+        source_fd->Open(install_part_.source_path.c_str(), O_RDONLY));
+    WriteAllCowOps(block_size_, converted, cow_writer_.get(), source_fd);
+  }
   return true;
 }
 
diff --git a/payload_consumer/verified_source_fd.cc b/payload_consumer/verified_source_fd.cc
new file mode 100644
index 0000000..002bd07
--- /dev/null
+++ b/payload_consumer/verified_source_fd.cc
@@ -0,0 +1,124 @@
+//
+// Copyright (C) 2021 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
+// limi
+
+#include "update_engine/payload_consumer/verified_source_fd.h"
+
+#include <fcntl.h>
+#include <sys/stat.h>
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include <base/strings/string_number_conversions.h>
+#include <base/strings/string_util.h>
+#include <base/strings/stringprintf.h>
+
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_consumer/fec_file_descriptor.h"
+#include "update_engine/payload_consumer/file_descriptor_utils.h"
+#include "update_engine/payload_consumer/mount_history.h"
+#include "update_engine/payload_consumer/partition_writer.h"
+
+namespace chromeos_update_engine {
+using std::string;
+
+bool VerifiedSourceFd::OpenCurrentECCPartition() {
+  // No support for ECC for full payloads.
+  // Full payload should not have any opeartion that requires ECC partitions.
+  if (source_ecc_fd_)
+    return true;
+
+  if (source_ecc_open_failure_)
+    return false;
+
+#if USE_FEC
+  FileDescriptorPtr fd(new FecFileDescriptor());
+  if (!fd->Open(source_path_.c_str(), O_RDONLY, 0)) {
+    PLOG(ERROR) << "Unable to open ECC source partition " << source_path_;
+    source_ecc_open_failure_ = true;
+    return false;
+  }
+  source_ecc_fd_ = fd;
+#else
+  // No support for ECC compiled.
+  source_ecc_open_failure_ = true;
+#endif  // USE_FEC
+
+  return !source_ecc_open_failure_;
+}
+
+FileDescriptorPtr VerifiedSourceFd::ChooseSourceFD(
+    const InstallOperation& operation, ErrorCode* error) {
+  if (source_fd_ == nullptr) {
+    LOG(ERROR) << "ChooseSourceFD fail: source_fd_ == nullptr";
+    return nullptr;
+  }
+  if (!operation.has_src_sha256_hash()) {
+    // When the operation doesn't include a source hash, we attempt the error
+    // corrected device first since we can't verify the block in the raw device
+    // at this point, but we first need to make sure all extents are readable
+    // since the error corrected device can be shorter or not available.
+    if (OpenCurrentECCPartition() &&
+        fd_utils::ReadAndHashExtents(
+            source_ecc_fd_, operation.src_extents(), block_size_, nullptr)) {
+      return source_ecc_fd_;
+    }
+    return source_fd_;
+  }
+
+  brillo::Blob source_hash;
+  brillo::Blob expected_source_hash(operation.src_sha256_hash().begin(),
+                                    operation.src_sha256_hash().end());
+  if (fd_utils::ReadAndHashExtents(
+          source_fd_, operation.src_extents(), block_size_, &source_hash) &&
+      source_hash == expected_source_hash) {
+    return source_fd_;
+  }
+  // We fall back to use the error corrected device if the hash of the raw
+  // device doesn't match or there was an error reading the source partition.
+  if (!OpenCurrentECCPartition()) {
+    // The following function call will return false since the source hash
+    // mismatches, but we still want to call it so it prints the appropriate
+    // log message.
+    PartitionWriter::ValidateSourceHash(
+        source_hash, operation, source_fd_, error);
+    return nullptr;
+  }
+  LOG(WARNING) << "Source hash from RAW device mismatched: found "
+               << base::HexEncode(source_hash.data(), source_hash.size())
+               << ", expected "
+               << base::HexEncode(expected_source_hash.data(),
+                                  expected_source_hash.size());
+
+  if (fd_utils::ReadAndHashExtents(
+          source_ecc_fd_, operation.src_extents(), block_size_, &source_hash) &&
+      PartitionWriter::ValidateSourceHash(
+          source_hash, operation, source_ecc_fd_, error)) {
+    source_ecc_recovered_failures_++;
+    return source_ecc_fd_;
+  }
+  return nullptr;
+}
+
+bool VerifiedSourceFd::Open() {
+  source_fd_ = std::make_shared<EintrSafeFileDescriptor>();
+  if (source_fd_ == nullptr)
+    return false;
+  TEST_AND_RETURN_FALSE_ERRNO(source_fd_->Open(source_path_.c_str(), O_RDONLY));
+  return true;
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_consumer/verified_source_fd.h b/payload_consumer/verified_source_fd.h
new file mode 100644
index 0000000..f7d0620
--- /dev/null
+++ b/payload_consumer/verified_source_fd.h
@@ -0,0 +1,61 @@
+//
+// Copyright (C) 2021 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
+// limi
+
+#ifndef UPDATE_ENGINE_VERIFIED_SOURCE_FD_H__
+#define UPDATE_ENGINE_VERIFIED_SOURCE_FD_H__
+
+#include <cstddef>
+
+#include <string>
+#include <utility>
+
+#include <gtest/gtest_prod.h>
+#include <update_engine/update_metadata.pb.h>
+
+#include "update_engine/common/error_code.h"
+#include "update_engine/payload_consumer/file_descriptor.h"
+
+namespace chromeos_update_engine {
+
+class VerifiedSourceFd {
+ public:
+  explicit VerifiedSourceFd(size_t block_size, std::string source_path)
+      : block_size_(block_size), source_path_(std::move(source_path)) {}
+  FileDescriptorPtr ChooseSourceFD(const InstallOperation& operation,
+                                   ErrorCode* error);
+
+  [[nodiscard]] bool Open();
+
+ private:
+  bool OpenCurrentECCPartition();
+  const size_t block_size_;
+  const std::string source_path_;
+  FileDescriptorPtr source_ecc_fd_;
+  FileDescriptorPtr source_fd_;
+
+  friend class PartitionWriterTest;
+  FRIEND_TEST(PartitionWriterTest, ChooseSourceFDTest);
+  // 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