1 files changed, 55 insertions, 29 deletions

diff --git a/payload_consumer/file_descriptor_utils.cc b/payload_consumer/file_descriptor_utils.cc
index 73f86dff..4ffded2b 100644
--- a/payload_consumer/file_descriptor_utils.cc
+++ b/payload_consumer/file_descriptor_utils.cc
@@ -33,53 +33,79 @@ namespace chromeos_update_engine {
 namespace {
 
 // Size of the buffer used to copy blocks.
-const int kMaxCopyBufferSize = 1024 * 1024;
-
-}  // namespace
-
-namespace fd_utils {
-
-bool CopyAndHashExtents(FileDescriptorPtr source,
-                        const RepeatedPtrField<Extent>& src_extents,
-                        FileDescriptorPtr target,
-                        const RepeatedPtrField<Extent>& tgt_extents,
-                        uint32_t block_size,
-                        brillo::Blob* hash_out) {
-  uint64_t total_blocks = utils::BlocksInExtents(src_extents);
-  TEST_AND_RETURN_FALSE(total_blocks == utils::BlocksInExtents(tgt_extents));
-
-  DirectExtentReader reader;
-  TEST_AND_RETURN_FALSE(reader.Init(source, src_extents, block_size));
-  DirectExtentWriter writer;
-  TEST_AND_RETURN_FALSE(writer.Init(target, tgt_extents, block_size));
-
-  uint64_t buffer_blocks = kMaxCopyBufferSize / block_size;
+const uint64_t kMaxCopyBufferSize = 1024 * 1024;
+
+bool CommonHashExtents(FileDescriptorPtr source,
+                       const RepeatedPtrField<Extent>& src_extents,
+                       FileDescriptorPtr target,
+                       const RepeatedPtrField<Extent>& tgt_extents,
+                       uint64_t block_size,
+                       brillo::Blob* hash_out) {
+  auto total_blocks = utils::BlocksInExtents(src_extents);
+  auto buffer_blocks = kMaxCopyBufferSize / block_size;
   // Ensure we copy at least one block at a time.
   if (buffer_blocks < 1)
     buffer_blocks = 1;
   brillo::Blob buf(buffer_blocks * block_size);
 
+  DirectExtentReader reader;
+  TEST_AND_RETURN_FALSE(reader.Init(source, src_extents, block_size));
+  DirectExtentWriter writer;
+  if (target) {
+    TEST_AND_RETURN_FALSE(writer.Init(target, tgt_extents, block_size));
+    TEST_AND_RETURN_FALSE(total_blocks == utils::BlocksInExtents(tgt_extents));
+  }
+
   HashCalculator source_hasher;
-  uint64_t blocks_left = total_blocks;
-  while (blocks_left > 0) {
-    uint64_t read_blocks = std::min(blocks_left, buffer_blocks);
+  while (total_blocks > 0) {
+    auto read_blocks = std::min(total_blocks, buffer_blocks);
     TEST_AND_RETURN_FALSE(reader.Read(buf.data(), read_blocks * block_size));
-    if (hash_out) {
+    if (hash_out != nullptr) {
       TEST_AND_RETURN_FALSE(
           source_hasher.Update(buf.data(), read_blocks * block_size));
     }
-    TEST_AND_RETURN_FALSE(writer.Write(buf.data(), read_blocks * block_size));
-    blocks_left -= read_blocks;
+    if (target) {
+      TEST_AND_RETURN_FALSE(writer.Write(buf.data(), read_blocks * block_size));
+    }
+    total_blocks -= read_blocks;
+  }
+  if (target) {
+    TEST_AND_RETURN_FALSE(writer.End());
   }
-  TEST_AND_RETURN_FALSE(writer.End());
 
-  if (hash_out) {
+  if (hash_out != nullptr) {
     TEST_AND_RETURN_FALSE(source_hasher.Finalize());
     *hash_out = source_hasher.raw_hash();
   }
   return true;
 }
 
+}  // namespace
+
+namespace fd_utils {
+
+bool CopyAndHashExtents(FileDescriptorPtr source,
+                        const RepeatedPtrField<Extent>& src_extents,
+                        FileDescriptorPtr target,
+                        const RepeatedPtrField<Extent>& tgt_extents,
+                        uint64_t block_size,
+                        brillo::Blob* hash_out) {
+  TEST_AND_RETURN_FALSE(target);
+  TEST_AND_RETURN_FALSE(CommonHashExtents(
+      source, src_extents, target, tgt_extents, block_size, hash_out));
+  return true;
+}
+
+bool ReadAndHashExtents(FileDescriptorPtr source,
+                        const RepeatedPtrField<Extent>& extents,
+                        uint64_t block_size,
+                        brillo::Blob* hash_out) {
+  TEST_AND_RETURN_FALSE(hash_out != nullptr);
+  TEST_AND_RETURN_FALSE(
+      CommonHashExtents(source, extents, nullptr, {}, block_size, hash_out));
+  return true;
+}
+
 }  // namespace fd_utils
 
 }  // namespace chromeos_update_engine