Add a library function for performing a cow dry run am: 4ab946d067

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

MUST ONLY BE SUBMITTED BY AUTOMERGER

Change-Id: I68ff2f1506a33c555a65ea0672becedb4489d6fe
diff --git a/payload_generator/cow_size_estimator.cc b/payload_generator/cow_size_estimator.cc
index f3f5264..fc66e67 100644
--- a/payload_generator/cow_size_estimator.cc
+++ b/payload_generator/cow_size_estimator.cc
@@ -24,13 +24,14 @@
 #include <libsnapshot/cow_writer.h>
 
 #include "update_engine/common/cow_operation_convert.h"
-#include "update_engine/payload_consumer/vabc_partition_writer.h"
+#include "update_engine/common/utils.h"
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
 using android::snapshot::CowWriter;
 
-void PerformReplaceOp(const InstallOperation& op,
+namespace {
+bool PerformReplaceOp(const InstallOperation& op,
                       CowWriter* writer,
                       FileDescriptorPtr target_fd,
                       size_t block_size) {
@@ -45,22 +46,62 @@
                                   buffer.size(),
                                   extent.start_block() * block_size,
                                   &bytes_read);
-    CHECK(success);
+    TEST_AND_RETURN_FALSE(success);
     CHECK_EQ(static_cast<size_t>(bytes_read), buffer.size());
-    writer->AddRawBlocks(extent.start_block(), buffer.data(), buffer.size());
+    TEST_AND_RETURN_FALSE(writer->AddRawBlocks(
+        extent.start_block(), buffer.data(), buffer.size()));
   }
+  return true;
 }
 
-void PerformZeroOp(const InstallOperation& op,
+bool PerformZeroOp(const InstallOperation& op,
                    CowWriter* writer,
                    size_t block_size) {
   for (const auto& extent : op.dst_extents()) {
-    writer->AddZeroBlocks(extent.start_block(), extent.num_blocks());
+    TEST_AND_RETURN_FALSE(
+        writer->AddZeroBlocks(extent.start_block(), extent.num_blocks()));
   }
+  return true;
 }
 
+bool WriteAllCowOps(size_t block_size,
+                    const std::vector<CowOperation>& converted,
+                    android::snapshot::ICowWriter* cow_writer,
+                    FileDescriptorPtr target_fd) {
+  std::vector<uint8_t> buffer(block_size);
+
+  for (const auto& cow_op : converted) {
+    switch (cow_op.op) {
+      case CowOperation::CowCopy:
+        if (cow_op.src_block == cow_op.dst_block) {
+          continue;
+        }
+        TEST_AND_RETURN_FALSE(
+            cow_writer->AddCopy(cow_op.dst_block, cow_op.src_block));
+        break;
+      case CowOperation::CowReplace:
+        ssize_t bytes_read = 0;
+        TEST_AND_RETURN_FALSE(chromeos_update_engine::utils::ReadAll(
+            target_fd,
+            buffer.data(),
+            block_size,
+            cow_op.dst_block * block_size,
+            &bytes_read));
+        if (bytes_read <= 0 || static_cast<size_t>(bytes_read) != block_size) {
+          LOG(ERROR) << "source_fd->Read failed: " << bytes_read;
+          return false;
+        }
+        TEST_AND_RETURN_FALSE(cow_writer->AddRawBlocks(
+            cow_op.dst_block, buffer.data(), block_size));
+        break;
+    }
+  }
+
+  return true;
+}
+}  // namespace
+
 size_t EstimateCowSize(
-    FileDescriptorPtr source_fd,
     FileDescriptorPtr target_fd,
     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
@@ -73,21 +114,32 @@
   // CowWriter treats -1 as special value, will discard all the data but still
   // reports Cow size. Good for estimation purposes
   cow_writer.Initialize(android::base::borrowed_fd{-1});
+  CHECK(CowDryRun(
+      target_fd, operations, merge_operations, block_size, &cow_writer));
+  return cow_writer.GetCowSize();
+}
 
+bool CowDryRun(
+    FileDescriptorPtr target_fd,
+    const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
+    const google::protobuf::RepeatedPtrField<CowMergeOperation>&
+        merge_operations,
+    size_t block_size,
+    android::snapshot::CowWriter* cow_writer) {
   const auto converted = ConvertToCowOperations(operations, merge_operations);
-  VABCPartitionWriter::WriteAllCowOps(
-      block_size, converted, &cow_writer, source_fd);
-  cow_writer.AddLabel(0);
+  WriteAllCowOps(block_size, converted, cow_writer, target_fd);
+  cow_writer->AddLabel(0);
   for (const auto& op : operations) {
     switch (op.type()) {
       case InstallOperation::REPLACE:
       case InstallOperation::REPLACE_BZ:
       case InstallOperation::REPLACE_XZ:
-        PerformReplaceOp(op, &cow_writer, target_fd, block_size);
+        TEST_AND_RETURN_FALSE(
+            PerformReplaceOp(op, cow_writer, target_fd, block_size));
         break;
       case InstallOperation::ZERO:
       case InstallOperation::DISCARD:
-        PerformZeroOp(op, &cow_writer, block_size);
+        TEST_AND_RETURN_FALSE(PerformZeroOp(op, cow_writer, block_size));
         break;
       case InstallOperation::SOURCE_COPY:
       case InstallOperation::MOVE:
@@ -99,15 +151,16 @@
       case InstallOperation::BSDIFF:
         // We might do something special by adding CowBsdiff to CowWriter.
         // For now proceed the same way as normal REPLACE operation.
-        PerformReplaceOp(op, &cow_writer, target_fd, block_size);
+        TEST_AND_RETURN_FALSE(
+            PerformReplaceOp(op, cow_writer, target_fd, block_size));
         break;
     }
     // Arbitrary label number, we won't be resuming use these labels here.
     // They are emitted just to keep size estimates accurate. As update_engine
     // emits 1 label for every op.
-    cow_writer.AddLabel(2);
+    cow_writer->AddLabel(2);
   }
   // TODO(zhangkelvin) Take FEC extents into account once VABC stabilizes
-  return cow_writer.GetCowSize();
+  return true;
 }
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/cow_size_estimator.h b/payload_generator/cow_size_estimator.h
index cad4bad..a94df88 100644
--- a/payload_generator/cow_size_estimator.h
+++ b/payload_generator/cow_size_estimator.h
@@ -16,19 +16,19 @@
 #include <cstddef>
 #include <string>
 
+#include <libsnapshot/cow_writer.h>
 #include <update_engine/update_metadata.pb.h>
 
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/payload_generator/delta_diff_generator.h"
 
 namespace chromeos_update_engine {
-// Given file descriptor to the source image, target image, and list of
+// Given file descriptor to the target image, and list of
 // operations, estimate the size of COW image if the operations are applied on
 // Virtual AB Compression enabled device. This is intended to be used by update
 // generators to put an estimate cow size in OTA payload. When installing an OTA
 // update, libsnapshot will take this estimate as a hint to allocate spaces.
 size_t EstimateCowSize(
-    FileDescriptorPtr source_fd,
     FileDescriptorPtr target_fd,
     const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
     const google::protobuf::RepeatedPtrField<CowMergeOperation>&
@@ -36,4 +36,13 @@
     size_t block_size,
     std::string compression);
 
+// Convert InstallOps to CowOps and apply the converted cow op to |cow_writer|
+bool CowDryRun(
+    FileDescriptorPtr target_fd,
+    const google::protobuf::RepeatedPtrField<InstallOperation>& operations,
+    const google::protobuf::RepeatedPtrField<CowMergeOperation>&
+        merge_operations,
+    size_t block_size,
+    android::snapshot::CowWriter* cow_writer);
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index a43c782..a87dabf 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -132,7 +132,6 @@
       *operations.Add() = aop.op;
     }
     *cow_size_ = EstimateCowSize(
-        std::move(source_fd),
         std::move(target_fd),
         std::move(operations),
         {cow_merge_sequence_->begin(), cow_merge_sequence_->end()},