Merge 6e1e2646851b278c55e39e480d51fc1434567dfd on remote branch

Change-Id: I033b6eb046f4ac571311303ef8658e0f7124fb33

19 files changed, 355 insertions, 113 deletions

diff --git a/Android.bp b/Android.bp
index 1991c496..c321de8d 100644
--- a/Android.bp
+++ b/Android.bp
@@ -207,6 +207,7 @@ cc_defaults {
         "libcow_operation_convert",
         "lz4diff-protos",
         "liblz4patch",
+        "libzstd",
     ],
     shared_libs: [
         "libbase",
@@ -557,7 +558,10 @@ cc_binary {
         "libupdate_engine_android_exports",
     ],
 
-    static_libs: ["libupdate_engine_android"],
+    static_libs: [
+        "libupdate_engine_android",
+        "libgflags",
+    ],
     required: [
         "cacerts_google",
         "otacerts",
@@ -699,6 +703,7 @@ cc_defaults {
         "libselinux",
         "lz4diff-protos",
         "liblz4diff",
+        "libzstd",
     ],
     shared_libs: [
         "libbase",
@@ -856,6 +861,7 @@ cc_binary_host {
         "libavb_host_sysdeps",
         "libpayload_consumer",
         "libpayload_generator",
+        "libgflags",
     ],
 
     srcs: ["payload_generator/generate_delta_main.cc"],
@@ -873,6 +879,7 @@ cc_test {
     static_libs: [
         "libpayload_consumer",
         "libpayload_generator",
+        "libgflags",
     ],
 
     srcs: ["payload_generator/generate_delta_main.cc"],
@@ -1068,6 +1075,7 @@ cc_test {
         "libdm",
         "libgmock",
         "libz",
+        "libzstd",
     ],
     shared_libs: [
         "libssl",
diff --git a/aosp/apex_handler_android.cc b/aosp/apex_handler_android.cc
index f47889ba..acc6becf 100644
--- a/aosp/apex_handler_android.cc
+++ b/aosp/apex_handler_android.cc
@@ -22,7 +22,6 @@
 #include <ApexProperties.sysprop.h>
 
 #include "update_engine/aosp/apex_handler_android.h"
-#include "update_engine/common/utils.h"
 
 namespace chromeos_update_engine {
 
diff --git a/aosp/logging_android.cc b/aosp/logging_android.cc
index 5ccf7bc5..5940f785 100644
--- a/aosp/logging_android.cc
+++ b/aosp/logging_android.cc
@@ -22,6 +22,7 @@
 #include <algorithm>
 #include <functional>
 #include <iomanip>
+#include <sstream>
 #include <string>
 #include <string_view>
 #include <vector>
@@ -35,6 +36,7 @@
 #include <base/strings/stringprintf.h>
 #include <log/log.h>
 
+#include "android/log.h"
 #include "update_engine/common/utils.h"
 
 using std::string;
@@ -204,8 +206,23 @@ class CombinedLogger {
     }
   }
   void operator()(const struct __android_log_message* log_message) {
-    for (auto&& logger : loggers_) {
-      logger(log_message);
+    if (log_message->file != nullptr && log_message->line != 0) {
+      __android_log_message formatted = *log_message;
+      std::stringstream ss;
+      ss << "[" << LogPriorityToCString(formatted.priority) << ":"
+         << formatted.file << "(" << formatted.line << ")] "
+         << formatted.message;
+      formatted.file = nullptr;
+      formatted.line = 0;
+      const auto str = ss.str();
+      formatted.message = str.c_str();
+      for (auto&& logger : loggers_) {
+        logger(&formatted);
+      }
+    } else {
+      for (auto&& logger : loggers_) {
+        logger(log_message);
+      }
     }
   }
 
@@ -248,7 +265,17 @@ bool RedirectToLiblog(int severity,
   } else {
     // This will eventually be redirected to CombinedLogger.
     // Use nullptr as tag so that liblog infers log tag from getprogname().
-    __android_log_write(priority, nullptr /* tag */, str.c_str());
+    if (file == nullptr || file[0] == 0 || line == 0 || message_start != 0) {
+      __android_log_write(priority, nullptr /* tag */, str.c_str());
+    } else {
+      __android_log_print(priority,
+                          nullptr,
+                          "[%s:%s(%d)] %s",
+                          LogPriorityToCString(priority),
+                          file,
+                          line,
+                          str.c_str());
+    }
   }
   return true;
 }
diff --git a/aosp/update_attempter_android.cc b/aosp/update_attempter_android.cc
index 555ed333..56281092 100644
--- a/aosp/update_attempter_android.cc
+++ b/aosp/update_attempter_android.cc
@@ -55,6 +55,7 @@
 #include "update_engine/payload_consumer/payload_verifier.h"
 #include "update_engine/payload_consumer/postinstall_runner_action.h"
 #include "update_engine/update_boot_flags_action.h"
+#include "update_engine/update_status.h"
 #include "update_engine/update_status_utils.h"
 
 #ifndef _UE_SIDELOAD
@@ -281,6 +282,11 @@ bool UpdateAttempterAndroid::ApplyPayload(
   install_plan_.powerwash_required =
       GetHeaderAsBool(headers[kPayloadPropertyPowerwash], false);
 
+  if (!IsProductionBuild()) {
+    install_plan_.disable_vabc =
+        GetHeaderAsBool(headers[kPayloadDisableVABC], false);
+  }
+
   install_plan_.switch_slot_on_reboot =
       GetHeaderAsBool(headers[kPayloadPropertySwitchSlotOnReboot], true);
 
@@ -422,6 +428,13 @@ bool UpdateAttempterAndroid::ResetStatus(brillo::ErrorPtr* error) {
     return LogAndSetError(
         error, FROM_HERE, "Already processing an update, cancel it first.");
   }
+  if (status_ != UpdateStatus::IDLE &&
+      status_ != UpdateStatus::UPDATED_NEED_REBOOT) {
+    return LogAndSetError(error,
+                          FROM_HERE,
+                          "Status reset not allowed in this state, please "
+                          "cancel on going OTA first.");
+  }
 
   if (apex_handler_android_ != nullptr) {
     LOG(INFO) << "Cleaning up reserved space for compressed APEX (if any)";
@@ -437,30 +450,18 @@ bool UpdateAttempterAndroid::ResetStatus(brillo::ErrorPtr* error) {
                           "Failed to reset the status because "
                           "ClearUpdateCompletedMarker() failed");
   }
-
+  if (status_ == UpdateStatus::UPDATED_NEED_REBOOT) {
+    if (!resetShouldSwitchSlotOnReboot(error)) {
+      LOG(INFO) << "Failed to reset slot switch.";
+      return false;
+    }
+    LOG(INFO) << "Slot switch reset successful";
+  }
   if (!boot_control_->GetDynamicPartitionControl()->ResetUpdate(prefs_)) {
     LOG(WARNING) << "Failed to reset snapshots. UpdateStatus is IDLE but"
                  << "space might not be freed.";
   }
-  switch (status_) {
-    case UpdateStatus::IDLE: {
-      return true;
-    }
-
-    case UpdateStatus::UPDATED_NEED_REBOOT: {
-      const bool ret_value = resetShouldSwitchSlotOnReboot(error);
-      if (ret_value) {
-        LOG(INFO) << "Reset status successful";
-      }
-      return ret_value;
-    }
-
-    default:
-      return LogAndSetError(
-          error,
-          FROM_HERE,
-          "Reset not allowed in this state. Cancel the ongoing update first");
-  }
+  return true;
 }
 
 bool UpdateAttempterAndroid::VerifyPayloadParseManifest(
@@ -1318,4 +1319,13 @@ void UpdateAttempterAndroid::RemoveCleanupPreviousUpdateCallback(
       end_it, cleanup_previous_update_callbacks_.end());
 }
 
+bool UpdateAttempterAndroid::IsProductionBuild() {
+  if (android::base::GetProperty("ro.build.type", "") != "userdebug" ||
+      android::base::GetProperty("ro.build.tags", "") == "release-keys" ||
+      android::base::GetProperty("ro.boot.verifiedbootstate", "") == "green") {
+    return true;
+  }
+  return false;
+}
+
 }  // namespace chromeos_update_engine
diff --git a/aosp/update_attempter_android.h b/aosp/update_attempter_android.h
index 956bd25a..c2226b20 100644
--- a/aosp/update_attempter_android.h
+++ b/aosp/update_attempter_android.h
@@ -234,6 +234,8 @@ class UpdateAttempterAndroid
   void RemoveCleanupPreviousUpdateCallback(
       CleanupSuccessfulUpdateCallbackInterface* callback);
 
+  bool IsProductionBuild();
+
   DaemonStateInterface* daemon_state_;
 
   // DaemonStateAndroid pointers.
diff --git a/common/constants.h b/common/constants.h
index dc45e72b..004d9d95 100644
--- a/common/constants.h
+++ b/common/constants.h
@@ -179,6 +179,8 @@ static constexpr const auto& kPayloadPropertyNetworkProxy = "NETWORK_PROXY";
 // Set Virtual AB Compression's compression algorithm to "none", but still use
 // userspace snapshots and snapuserd for update installation.
 static constexpr const auto& kPayloadVABCNone = "VABC_NONE";
+// Enable/Disable VABC, falls back on plain VAB
+static constexpr const auto& kPayloadDisableVABC = "DISABLE_VABC";
 // Enable multi-threaded compression for VABC
 static constexpr const auto& kPayloadEnableThreading = "ENABLE_THREADING";
 // Enable batched writes for VABC
diff --git a/main.cc b/main.cc
index a23a08bf..103e1a10 100644
--- a/main.cc
+++ b/main.cc
@@ -21,25 +21,24 @@
 #include <base/at_exit.h>
 #include <base/command_line.h>
 #include <base/logging.h>
-#include <brillo/flag_helper.h>
+#include <gflags/gflags.h>
 
 #include "update_engine/common/daemon_base.h"
 #include "update_engine/common/logging.h"
 #include "update_engine/common/subprocess.h"
 #include "update_engine/common/terminator.h"
-#include "update_engine/common/utils.h"
 
 using std::string;
+DEFINE_bool(logtofile, false, "Write logs to a file in log_dir.");
+DEFINE_bool(logtostderr,
+            false,
+            "Write logs to stderr instead of to a file in log_dir.");
+DEFINE_bool(foreground, false, "Don't daemon()ize; run in foreground.");
 
 int main(int argc, char** argv) {
-  DEFINE_bool(logtofile, false, "Write logs to a file in log_dir.");
-  DEFINE_bool(logtostderr,
-              false,
-              "Write logs to stderr instead of to a file in log_dir.");
-  DEFINE_bool(foreground, false, "Don't daemon()ize; run in foreground.");
-
   chromeos_update_engine::Terminator::Init();
-  brillo::FlagHelper::Init(argc, argv, "A/B Update Engine");
+  gflags::SetUsageMessage("A/B Update Engine");
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
 
   // We have two logging flags "--logtostderr" and "--logtofile"; and the logic
   // to choose the logging destination is:
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index 26c3d591..3b9f2b66 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -41,6 +41,7 @@
 #include <google/protobuf/repeated_field.h>
 #include <puffin/puffpatch.h>
 
+#include "libsnapshot/cow_format.h"
 #include "update_engine/common/constants.h"
 #include "update_engine/common/download_action.h"
 #include "update_engine/common/error_code.h"
@@ -501,12 +502,28 @@ bool DeltaPerformer::Write(const void* bytes, size_t count, ErrorCode* error) {
                 operation.dst_extent().num_blocks() * manifest_.block_size();
           }
         }
-        // Adding extra 8MB headroom. OTA will sometimes write labels/metadata
-        // to COW image. If we overrun reserved COW size, entire OTA will fail
+        // Every block written to COW device will come with a header which
+        // stores src/dst block info along with other data.
+        const auto cow_metadata_size = partition.new_partition_info().size() /
+                                       manifest_.block_size() *
+                                       sizeof(android::snapshot::CowOperation);
+        // update_engine will emit a label op every op or every two seconds,
+        // whichever one is longer. In the worst case, we add 1 label per
+        // InstallOp. So take size of label ops into account.
+        const auto label_ops_size = partition.operations_size() *
+                                    sizeof(android::snapshot::CowOperation);
+        // Adding extra 2MB headroom just for any unexpected space usage.
+        // If we overrun reserved COW size, entire OTA will fail
         // and no way for user to retry OTA
-        partition.set_estimate_cow_size(new_cow_size + (1024 * 1024 * 8));
+        partition.set_estimate_cow_size(new_cow_size + (1024 * 1024 * 2) +
+                                        cow_metadata_size + label_ops_size);
+        LOG(INFO) << "New COW size for partition " << partition.partition_name()
+                  << " is " << partition.estimate_cow_size();
       }
     }
+    if (install_plan_->disable_vabc) {
+      manifest_.mutable_dynamic_partition_metadata()->set_vabc_enabled(false);
+    }
     if (install_plan_->enable_threading) {
       manifest_.mutable_dynamic_partition_metadata()
           ->mutable_vabc_feature_set()
diff --git a/payload_consumer/install_plan.h b/payload_consumer/install_plan.h
index 8fe104a3..93aebcee 100644
--- a/payload_consumer/install_plan.h
+++ b/payload_consumer/install_plan.h
@@ -74,6 +74,7 @@ struct InstallPlan {
 
   bool is_resume{false};
   bool vabc_none{false};
+  bool disable_vabc{false};
   std::string download_url;  // url to download from
   std::string version;       // version we are installing.
 
diff --git a/payload_consumer/postinstall_runner_action.cc b/payload_consumer/postinstall_runner_action.cc
index 6beca56a..a6fa4b30 100644
--- a/payload_consumer/postinstall_runner_action.cc
+++ b/payload_consumer/postinstall_runner_action.cc
@@ -430,12 +430,6 @@ void PostinstallRunnerAction::CompletePartitionPostinstall(
 }
 
 void PostinstallRunnerAction::CompletePostinstall(ErrorCode error_code) {
-  if (!install_plan_.partitions.empty()) {
-    auto dynamic_control = boot_control_->GetDynamicPartitionControl();
-    CHECK(dynamic_control);
-    dynamic_control->UnmapAllPartitions();
-    LOG(INFO) << "Unmapped all partitions.";
-  }
   // We only attempt to mark the new slot as active if all the postinstall
   // steps succeeded.
   if (error_code == ErrorCode::kSuccess) {
@@ -454,6 +448,12 @@ void PostinstallRunnerAction::CompletePostinstall(ErrorCode error_code) {
       error_code = ErrorCode::kUpdatedButNotActive;
     }
   }
+  if (!install_plan_.partitions.empty()) {
+    auto dynamic_control = boot_control_->GetDynamicPartitionControl();
+    CHECK(dynamic_control);
+    dynamic_control->UnmapAllPartitions();
+    LOG(INFO) << "Unmapped all partitions.";
+  }
 
   ScopedActionCompleter completer(processor_, this);
   completer.set_code(error_code);
diff --git a/payload_consumer/vabc_partition_writer.cc b/payload_consumer/vabc_partition_writer.cc
index df9c0a66..17b7d50d 100644
--- a/payload_consumer/vabc_partition_writer.cc
+++ b/payload_consumer/vabc_partition_writer.cc
@@ -95,6 +95,42 @@ VABCPartitionWriter::VABCPartitionWriter(
     }
     copy_blocks_.AddExtent(cow_op.dst_extent());
   }
+  LOG(INFO) << "Partition `" << partition_update.partition_name() << " has "
+            << copy_blocks_.blocks() << " copy blocks";
+}
+
+bool VABCPartitionWriter::DoesDeviceSupportsXor() {
+  return dynamic_control_->GetVirtualAbCompressionXorFeatureFlag().IsEnabled();
+}
+
+bool VABCPartitionWriter::WriteAllCopyOps() {
+  const bool userSnapshots = android::base::GetBoolProperty(
+      "ro.virtual_ab.userspace.snapshots.enabled", false);
+  for (const auto& cow_op : partition_update_.merge_operations()) {
+    if (cow_op.type() != CowMergeOperation::COW_COPY) {
+      continue;
+    }
+    if (cow_op.dst_extent() == cow_op.src_extent()) {
+      continue;
+    }
+    if (userSnapshots) {
+      TEST_AND_RETURN_FALSE(cow_op.src_extent().num_blocks() != 0);
+      TEST_AND_RETURN_FALSE(
+          cow_writer_->AddCopy(cow_op.dst_extent().start_block(),
+                               cow_op.src_extent().start_block(),
+                               cow_op.src_extent().num_blocks()));
+    } else {
+      // Add blocks in reverse order, because snapused specifically prefers
+      // this ordering. Since we already eliminated all self-overlapping
+      // SOURCE_COPY during delta generation, this should be safe to do.
+      for (size_t i = cow_op.src_extent().num_blocks(); i > 0; i--) {
+        TEST_AND_RETURN_FALSE(
+            cow_writer_->AddCopy(cow_op.dst_extent().start_block() + i - 1,
+                                 cow_op.src_extent().start_block() + i - 1));
+      }
+    }
+  }
+  return true;
 }
 
 bool VABCPartitionWriter::Init(const InstallPlan* install_plan,
@@ -144,35 +180,21 @@ bool VABCPartitionWriter::Init(const InstallPlan* install_plan,
     if (IsXorEnabled()) {
       LOG(INFO) << "VABC XOR enabled for partition "
                 << partition_update_.partition_name();
+    }
+    // When merge sequence is present in COW, snapuserd will merge blocks in
+    // order specified by the merge seuqnece op. Hence we have the freedom of
+    // writing COPY operations out of order. Delay processing of copy ops so
+    // that update_engine can be more responsive in progress updates.
+    if (DoesDeviceSupportsXor()) {
+      LOG(INFO) << "Snapuserd supports XOR and merge sequence, writing merge "
+                   "sequence and delay writing COPY operations";
       TEST_AND_RETURN_FALSE(WriteMergeSequence(
           partition_update_.merge_operations(), cow_writer_.get()));
-    }
-    const bool userSnapshots = android::base::GetBoolProperty(
-        "ro.virtual_ab.userspace.snapshots.enabled", false);
-
-    for (const auto& cow_op : partition_update_.merge_operations()) {
-      if (cow_op.type() != CowMergeOperation::COW_COPY) {
-        continue;
-      }
-      if (cow_op.dst_extent() == cow_op.src_extent()) {
-        continue;
-      }
-      if (userSnapshots) {
-        TEST_AND_RETURN_FALSE(cow_op.src_extent().num_blocks() != 0);
-        TEST_AND_RETURN_FALSE(
-            cow_writer_->AddCopy(cow_op.dst_extent().start_block(),
-                                 cow_op.src_extent().start_block(),
-                                 cow_op.src_extent().num_blocks()));
-      } else {
-        // Add blocks in reverse order, because snapused specifically prefers
-        // this ordering. Since we already eliminated all self-overlapping
-        // SOURCE_COPY during delta generation, this should be safe to do.
-        for (size_t i = cow_op.src_extent().num_blocks(); i > 0; i--) {
-          TEST_AND_RETURN_FALSE(
-              cow_writer_->AddCopy(cow_op.dst_extent().start_block() + i - 1,
-                                   cow_op.src_extent().start_block() + i - 1));
-        }
-      }
+    } else {
+      LOG(INFO) << "Snapuserd does not support merge sequence, writing all "
+                   "COPY operations up front, this may take few "
+                   "minutes.";
+      TEST_AND_RETURN_FALSE(WriteAllCopyOps());
     }
     cow_writer_->AddLabel(0);
   }
@@ -259,6 +281,11 @@ std::unique_ptr<ExtentWriter> VABCPartitionWriter::CreateBaseExtentWriter() {
   const auto& dst_extents = operation.dst_extents();
   BlockIterator it1{src_extents};
   BlockIterator it2{dst_extents};
+  const bool userSnapshots = android::base::GetBoolProperty(
+      "ro.virtual_ab.userspace.snapshots.enabled", false);
+  // For devices not supporting XOR, sequence op is not supported, so all COPY
+  // operations are written up front in strict merge order.
+  const auto sequence_op_supported = DoesDeviceSupportsXor();
   while (!it1.is_end() && !it2.is_end()) {
     const auto src_block = *it1;
     const auto dst_block = *it2;
@@ -267,13 +294,32 @@ std::unique_ptr<ExtentWriter> VABCPartitionWriter::CreateBaseExtentWriter() {
     if (src_block == dst_block) {
       continue;
     }
-    if (!copy_blocks_.ContainsBlock(dst_block)) {
+    if (copy_blocks_.ContainsBlock(dst_block)) {
+      if (sequence_op_supported) {
+        push_back(&converted, {CowOperation::CowCopy, src_block, dst_block, 1});
+      }
+    } else {
       push_back(&converted,
                 {CowOperation::CowReplace, src_block, dst_block, 1});
     }
   }
   std::vector<uint8_t> buffer;
   for (const auto& cow_op : converted) {
+    if (cow_op.op == CowOperation::CowCopy) {
+      if (userSnapshots) {
+        cow_writer_->AddCopy(
+            cow_op.dst_block, cow_op.src_block, cow_op.block_count);
+      } else {
+        // Add blocks in reverse order, because snapused specifically prefers
+        // this ordering. Since we already eliminated all self-overlapping
+        // SOURCE_COPY during delta generation, this should be safe to do.
+        for (size_t i = cow_op.block_count; i > 0; i--) {
+          TEST_AND_RETURN_FALSE(cow_writer_->AddCopy(cow_op.dst_block + i - 1,
+                                                     cow_op.src_block + i - 1));
+        }
+      }
+      continue;
+    }
     buffer.resize(block_size_ * cow_op.block_count);
     ssize_t bytes_read = 0;
     TEST_AND_RETURN_FALSE(utils::ReadAll(source_fd,
@@ -312,7 +358,11 @@ bool VABCPartitionWriter::PerformDiffOperation(
 
   std::unique_ptr<ExtentWriter> writer =
       IsXorEnabled() ? std::make_unique<XORExtentWriter>(
-                           operation, source_fd, cow_writer_.get(), xor_map_)
+                           operation,
+                           source_fd,
+                           cow_writer_.get(),
+                           xor_map_,
+                           partition_update_.old_partition_info().size())
                      : CreateBaseExtentWriter();
   return executor_.ExecuteDiffOperation(
       operation, std::move(writer), source_fd, data, count);
diff --git a/payload_consumer/vabc_partition_writer.h b/payload_consumer/vabc_partition_writer.h
index 8393fe5b..889f3763 100644
--- a/payload_consumer/vabc_partition_writer.h
+++ b/payload_consumer/vabc_partition_writer.h
@@ -70,7 +70,9 @@ class VABCPartitionWriter final : public PartitionWriterInterface {
       android::snapshot::ICowWriter* cow_writer);
 
  private:
+  [[nodiscard]] bool DoesDeviceSupportsXor();
   bool IsXorEnabled() const noexcept { return xor_map_.size() > 0; }
+  [[nodiscard]] bool WriteAllCopyOps();
   std::unique_ptr<android::snapshot::ISnapshotWriter> cow_writer_;
 
   [[nodiscard]] std::unique_ptr<ExtentWriter> CreateBaseExtentWriter();
diff --git a/payload_consumer/vabc_partition_writer_unittest.cc b/payload_consumer/vabc_partition_writer_unittest.cc
index 0bfa23b2..3cdf3bca 100644
--- a/payload_consumer/vabc_partition_writer_unittest.cc
+++ b/payload_consumer/vabc_partition_writer_unittest.cc
@@ -24,11 +24,13 @@
 #include <libsnapshot/cow_writer.h>
 #include <libsnapshot/mock_snapshot_writer.h>
 
+#include "update_engine/common/error_code.h"
 #include "update_engine/common/hash_calculator.h"
 #include "update_engine/common/mock_dynamic_partition_control.h"
 #include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/vabc_partition_writer.h"
 #include "update_engine/payload_generator/delta_diff_generator.h"
+#include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
@@ -55,6 +57,7 @@ class VABCPartitionWriterTest : public ::testing::Test {
   }
 
  protected:
+  void EmitBlockTest(bool xor_enabled);
   CowMergeOperation* AddMergeOp(PartitionUpdate* partition,
                                 std::array<size_t, 2> src_extent,
                                 std::array<size_t, 2> dst_extent,
@@ -136,7 +139,33 @@ TEST_F(VABCPartitionWriterTest, MergeSequenceXorSameBlock) {
   ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
 }
 
-TEST_F(VABCPartitionWriterTest, EmitBlockTest) {
+TEST_F(VABCPartitionWriterTest, EmitBlockTestXor) {
+  return EmitBlockTest(true);
+}
+
+TEST_F(VABCPartitionWriterTest, EmitBlockTestNoXor) {
+  return EmitBlockTest(false);
+}
+
+void VABCPartitionWriterTest::EmitBlockTest(bool xor_enabled) {
+  if (xor_enabled) {
+    ON_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
+        .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::LAUNCH)));
+  } else {
+    ON_CALL(dynamic_control_, GetVirtualAbCompressionXorFeatureFlag())
+        .WillByDefault(Return(FeatureFlag(FeatureFlag::Value::NONE)));
+  }
+  InstallOperation& install_op = *partition_update_.add_operations();
+  install_op.set_type(InstallOperation::SOURCE_COPY);
+  *install_op.add_src_extents() = ExtentForRange(5, 1);
+  *install_op.add_src_extents() = ExtentForRange(10, 1);
+  *install_op.add_src_extents() = ExtentForRange(15, 2);
+  *install_op.add_src_extents() = ExtentForRange(20, 2);
+
+  *install_op.add_dst_extents() = ExtentForRange(10, 1);
+  *install_op.add_dst_extents() = ExtentForRange(15, 1);
+  *install_op.add_dst_extents() = ExtentForRange(20, 2);
+  *install_op.add_dst_extents() = ExtentForRange(25, 2);
   AddMergeOp(&partition_update_, {5, 1}, {10, 1}, CowMergeOperation::COW_COPY);
   AddMergeOp(&partition_update_, {10, 1}, {15, 1}, CowMergeOperation::COW_COPY);
   AddMergeOp(&partition_update_, {15, 2}, {20, 2}, CowMergeOperation::COW_COPY);
@@ -144,25 +173,45 @@ TEST_F(VABCPartitionWriterTest, EmitBlockTest) {
   VABCPartitionWriter writer_{
       partition_update_, install_part_, &dynamic_control_, kBlockSize};
   EXPECT_CALL(dynamic_control_, OpenCowWriter(fake_part_name, _, false))
-      .WillOnce(Invoke(
-          [](const std::string&, const std::optional<std::string>&, bool) {
-            auto cow_writer =
-                std::make_unique<android::snapshot::MockSnapshotWriter>(
-                    android::snapshot::CowOptions{});
-            Sequence s;
-            ON_CALL(*cow_writer, EmitCopy(_, _, _)).WillByDefault(Return(true));
-            ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
-            ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
-            EXPECT_CALL(*cow_writer, Initialize()).InSequence(s);
-            EXPECT_CALL(*cow_writer, EmitCopy(10, 5, 1)).InSequence(s);
-            EXPECT_CALL(*cow_writer, EmitCopy(15, 10, 1)).InSequence(s);
-            // libsnapshot want blocks in reverser order, so 21 goes before 20
-            EXPECT_CALL(*cow_writer, EmitCopy(20, 15, 2)).InSequence(s);
-
-            EXPECT_CALL(*cow_writer, EmitCopy(25, 20, 1)).InSequence(s);
-            return cow_writer;
-          }));
+      .WillOnce(Invoke([xor_enabled](const std::string&,
+                                     const std::optional<std::string>&,
+                                     bool) {
+        auto cow_writer =
+            std::make_unique<android::snapshot::MockSnapshotWriter>(
+                android::snapshot::CowOptions{});
+        ON_CALL(*cow_writer, EmitCopy(_, _, _)).WillByDefault(Return(true));
+        ON_CALL(*cow_writer, EmitLabel(_)).WillByDefault(Return(true));
+        ON_CALL(*cow_writer, Initialize()).WillByDefault(Return(true));
+        EXPECT_CALL(*cow_writer, Initialize());
+        if (xor_enabled) {
+          EXPECT_CALL(*cow_writer, EmitSequenceData(_, _))
+              .WillOnce(Return(true));
+          EXPECT_CALL(*cow_writer, EmitCopy(10, 5, 1));
+          EXPECT_CALL(*cow_writer, EmitCopy(15, 10, 1));
+          // libsnapshot want blocks in reverser order, so 21 goes before 20
+          EXPECT_CALL(*cow_writer, EmitCopy(20, 15, 2));
+
+          EXPECT_CALL(*cow_writer, EmitCopy(25, 20, 1));
+          EXPECT_CALL(*cow_writer, EmitRawBlocks(26, _, 4096))
+              .WillOnce(Return(true));
+          EXPECT_CALL(*cow_writer, Finalize());
+        } else {
+          Sequence s;
+          EXPECT_CALL(*cow_writer, EmitCopy(10, 5, 1)).InSequence(s);
+          EXPECT_CALL(*cow_writer, EmitCopy(15, 10, 1)).InSequence(s);
+          // libsnapshot want blocks in reverser order, so 21 goes before 20
+          EXPECT_CALL(*cow_writer, EmitCopy(20, 15, 2)).InSequence(s);
+
+          EXPECT_CALL(*cow_writer, EmitCopy(25, 20, 1)).InSequence(s);
+          EXPECT_CALL(*cow_writer, EmitRawBlocks(26, _, 4096))
+              .InSequence(s)
+              .WillOnce(Return(true));
+        }
+        return cow_writer;
+      }));
   ASSERT_TRUE(writer_.Init(&install_plan_, true, 0));
+  ErrorCode error{};
+  ASSERT_TRUE(writer_.PerformSourceCopyOperation(install_op, &error));
 }
 
 std::string GetNoopBSDIFF(size_t data_size) {
diff --git a/payload_consumer/xor_extent_writer.cc b/payload_consumer/xor_extent_writer.cc
index 1a02a2d9..4534c05e 100644
--- a/payload_consumer/xor_extent_writer.cc
+++ b/payload_consumer/xor_extent_writer.cc
@@ -64,6 +64,23 @@ bool XORExtentWriter::WriteExtent(const void* bytes,
     const auto src_block = merge_op->src_extent().start_block() +
                            xor_ext.start_block() -
                            merge_op->dst_extent().start_block();
+    const auto i = xor_ext.start_block() - extent.start_block();
+    const auto dst_block_data =
+        static_cast<const unsigned char*>(bytes) + i * BlockSize();
+    const auto is_out_of_bound_read =
+        (src_block + xor_ext.num_blocks()) * BlockSize() + src_offset >
+            partition_size_ &&
+        partition_size_ != 0;
+    if (is_out_of_bound_read) {
+      LOG(INFO) << "Getting partial read for last block, converting "
+                   "XOR operation to a regular replace "
+                << xor_ext;
+      TEST_AND_RETURN_FALSE(
+          cow_writer_->AddRawBlocks(xor_ext.start_block(),
+                                    dst_block_data,
+                                    xor_ext.num_blocks() * BlockSize()));
+      continue;
+    }
     xor_block_data.resize(BlockSize() * xor_ext.num_blocks());
     ssize_t bytes_read = 0;
     TEST_AND_RETURN_FALSE_ERRNO(
@@ -73,14 +90,12 @@ bool XORExtentWriter::WriteExtent(const void* bytes,
                         src_offset + src_block * BlockSize(),
                         &bytes_read));
     if (bytes_read != static_cast<ssize_t>(xor_block_data.size())) {
-      LOG(ERROR) << "bytes_read: " << bytes_read;
+      LOG(ERROR) << "bytes_read: " << bytes_read << ", expected to read "
+                 << xor_block_data.size() << " at block " << src_block
+                 << " offset " << src_offset;
       return false;
     }
 
-    const auto i = xor_ext.start_block() - extent.start_block();
-
-    const auto dst_block_data =
-        static_cast<const unsigned char*>(bytes) + i * BlockSize();
     std::transform(xor_block_data.cbegin(),
                    xor_block_data.cbegin() + xor_block_data.size(),
                    dst_block_data,
diff --git a/payload_consumer/xor_extent_writer.h b/payload_consumer/xor_extent_writer.h
index 35565eab..57c99c22 100644
--- a/payload_consumer/xor_extent_writer.h
+++ b/payload_consumer/xor_extent_writer.h
@@ -19,12 +19,9 @@
 
 #include <vector>
 
+#include "common/utils.h"
 #include "update_engine/payload_consumer/block_extent_writer.h"
 #include "update_engine/payload_consumer/extent_map.h"
-#include "update_engine/payload_consumer/extent_reader.h"
-#include "update_engine/payload_consumer/extent_writer.h"
-#include "update_engine/payload_generator/extent_ranges.h"
-#include "update_engine/payload_generator/extent_utils.h"
 
 #include <update_engine/update_metadata.pb.h>
 #include <libsnapshot/cow_writer.h>
@@ -38,11 +35,13 @@ class XORExtentWriter : public BlockExtentWriter {
   XORExtentWriter(const InstallOperation& op,
                   FileDescriptorPtr source_fd,
                   android::snapshot::ICowWriter* cow_writer,
-                  const ExtentMap<const CowMergeOperation*>& xor_map)
+                  const ExtentMap<const CowMergeOperation*>& xor_map,
+                  size_t partition_size)
       : src_extents_(op.src_extents()),
         source_fd_(source_fd),
         xor_map_(xor_map),
-        cow_writer_(cow_writer) {
+        cow_writer_(cow_writer),
+        partition_size_(partition_size) {
     CHECK(source_fd->IsOpen());
   }
   ~XORExtentWriter() = default;
@@ -61,6 +60,7 @@ class XORExtentWriter : public BlockExtentWriter {
   const FileDescriptorPtr source_fd_;
   const ExtentMap<const CowMergeOperation*>& xor_map_;
   android::snapshot::ICowWriter* cow_writer_;
+  const size_t partition_size_;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/xor_extent_writer_unittest.cc b/payload_consumer/xor_extent_writer_unittest.cc
index 015c9ab1..55c3c6c0 100644
--- a/payload_consumer/xor_extent_writer_unittest.cc
+++ b/payload_consumer/xor_extent_writer_unittest.cc
@@ -94,7 +94,8 @@ TEST_F(XorExtentWriterTest, StreamTest) {
   ASSERT_TRUE(xor_map_.AddExtent(op3.dst_extent(), &op3));
   *op_.add_src_extents() = ExtentForRange(12, 4);
   *op_.add_dst_extents() = ExtentForRange(320, 4);
-  XORExtentWriter writer_{op_, source_fd_, &cow_writer_, xor_map_};
+  XORExtentWriter writer_{
+      op_, source_fd_, &cow_writer_, xor_map_, NUM_BLOCKS * kBlockSize};
 
   // OTA op:
   // [5-6] => [5-6], [45-47] => [455-457], [12-15] => [320-323]
@@ -146,7 +147,8 @@ TEST_F(XorExtentWriterTest, SubsetExtentTest) {
   *op_.add_dst_extents() = ExtentForRange(420, 3);
   *op_.add_src_extents() = ExtentForRange(15, 1);
   *op_.add_dst_extents() = ExtentForRange(323, 1);
-  XORExtentWriter writer_{op_, source_fd_, &cow_writer_, xor_map_};
+  XORExtentWriter writer_{
+      op_, source_fd_, &cow_writer_, xor_map_, NUM_BLOCKS * kBlockSize};
 
   // OTA op:
   // [12-14] => [320-322], [20-22] => [420-422], [15-16] => [323-324]
@@ -173,4 +175,54 @@ TEST_F(XorExtentWriterTest, SubsetExtentTest) {
   ASSERT_TRUE(writer_.Write(zeros->data(), zeros->size()));
 }
 
+TEST_F(XorExtentWriterTest, LastBlockTest) {
+  constexpr auto COW_XOR = CowMergeOperation::COW_XOR;
+  ON_CALL(cow_writer_, EmitXorBlocks(_, _, _, _, _))
+      .WillByDefault(Return(true));
+
+  const auto op3 = CreateCowMergeOperation(
+      ExtentForRange(NUM_BLOCKS - 1, 1), ExtentForRange(2, 1), COW_XOR, 777);
+  ASSERT_TRUE(xor_map_.AddExtent(op3.dst_extent(), &op3));
+
+  *op_.add_src_extents() = ExtentForRange(12, 3);
+  *op_.add_dst_extents() = ExtentForRange(320, 3);
+
+  *op_.add_src_extents() = ExtentForRange(20, 3);
+  *op_.add_dst_extents() = ExtentForRange(420, 3);
+
+  *op_.add_src_extents() = ExtentForRange(NUM_BLOCKS - 3, 3);
+  *op_.add_dst_extents() = ExtentForRange(2, 3);
+  XORExtentWriter writer_{
+      op_, source_fd_, &cow_writer_, xor_map_, NUM_BLOCKS * kBlockSize};
+
+  // OTA op:
+  // [12-14] => [320-322], [20-22] => [420-422], [NUM_BLOCKS-3] => [2-5]
+
+  // merge op:
+  // [NUM_BLOCKS-1] => [2-3]
+
+  // Expected result:
+  // [12-16] should be REPLACE blocks
+  // [420-422] should be REPLACE blocks
+  // [2-4] should be REPLACE blocks
+
+  auto zeros = utils::GetReadonlyZeroBlock(kBlockSize * 9);
+  EXPECT_CALL(cow_writer_, EmitRawBlocks(320, zeros->data(), kBlockSize * 3))
+      .WillOnce(Return(true));
+  EXPECT_CALL(
+      cow_writer_,
+      EmitRawBlocks(420, zeros->data() + 3 * kBlockSize, kBlockSize * 3))
+      .WillOnce(Return(true));
+
+  EXPECT_CALL(cow_writer_,
+              EmitRawBlocks(2, zeros->data() + 6 * kBlockSize, kBlockSize))
+      .WillOnce(Return(true));
+  EXPECT_CALL(cow_writer_,
+              EmitRawBlocks(3, zeros->data() + 7 * kBlockSize, kBlockSize * 2))
+      .WillOnce(Return(true));
+
+  ASSERT_TRUE(writer_.Init(op_.dst_extents(), kBlockSize));
+  ASSERT_TRUE(writer_.Write(zeros->data(), zeros->size()));
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/generate_delta_main.cc b/payload_generator/generate_delta_main.cc
index 64e72495..3a8d91a9 100644
--- a/payload_generator/generate_delta_main.cc
+++ b/payload_generator/generate_delta_main.cc
@@ -18,6 +18,7 @@
 #include <string>
 #include <vector>
 
+#include <android-base/strings.h>
 #include <base/bind.h>
 #include <base/files/file_path.h>
 #include <base/files/file_util.h>
@@ -25,10 +26,10 @@
 #include <base/strings/string_number_conversions.h>
 #include <base/strings/string_split.h>
 #include <base/strings/string_util.h>
-#include <brillo/flag_helper.h>
 #include <brillo/key_value_store.h>
 #include <brillo/message_loops/base_message_loop.h>
 #include <xz.h>
+#include <gflags/gflags.h>
 
 #include "update_engine/common/download_action.h"
 #include "update_engine/common/fake_boot_control.h"
@@ -165,7 +166,7 @@ class ApplyPayloadProcessorDelegate : public ActionProcessorDelegate {
   void ProcessingStopped(const ActionProcessor* processor) override {
     brillo::MessageLoop::current()->BreakLoop();
   }
-  ErrorCode code_;
+  ErrorCode code_{};
 };
 
 // TODO(deymo): Move this function to a new file and make the delta_performer
@@ -303,7 +304,6 @@ bool ParsePerPartitionTimestamps(const string& partition_timestamps,
   return true;
 }
 
-int Main(int argc, char** argv) {
   DEFINE_string(old_image, "", "Path to the old rootfs");
   DEFINE_string(new_image, "", "Path to the new rootfs");
   DEFINE_string(old_kernel, "", "Path to the old kernel partition image");
@@ -453,14 +453,17 @@ int Main(int argc, char** argv) {
                 "",
                 "Compression parameter passed to mkfs.erofs's -z option. "
                 "Example: lz4 lz4hc,9");
-
-  brillo::FlagHelper::Init(
-      argc,
-      argv,
+  int Main(int argc, char** argv) {
+  gflags::SetUsageMessage(
       "Generates a payload to provide to ChromeOS' update_engine.\n\n"
       "This tool can create full payloads and also delta payloads if the src\n"
       "image is provided. It also provides debugging options to apply, sign\n"
       "and verify payloads.");
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+  CHECK_EQ(argc, 1) << " Unused args: "
+                    << android::base::Join(
+                           std::vector<char*>(argv + 1, argv + argc), " ");
+
   Terminator::Init();
 
   logging::LoggingSettings log_settings;
@@ -770,7 +773,7 @@ int Main(int argc, char** argv) {
                            metadata_size_string.size()));
   }
   return 0;
-}
+  }
 
 }  // namespace
 
diff --git a/scripts/update_device.py b/scripts/update_device.py
index 2985c38e..f94774b1 100755
--- a/scripts/update_device.py
+++ b/scripts/update_device.py
@@ -483,6 +483,8 @@ def main():
                       help='Wipe userdata after installing OTA')
   parser.add_argument('--vabc-none', action='store_true',
                       help='Set Virtual AB Compression algorithm to none, but still use Android COW format')
+  parser.add_argument('--disable-vabc', action='store_true',
+                      help='Option to enable or disable vabc. If set to false, will fall back on A/B')
   parser.add_argument('--enable-threading', action='store_true',
                       help='Enable multi-threaded compression for VABC')
   parser.add_argument('--batched-writes', action='store_true',
@@ -549,6 +551,8 @@ def main():
     args.extra_headers += "\nPOWERWASH=1"
   if args.vabc_none:
     args.extra_headers += "\nVABC_NONE=1"
+  if args.disable_vabc:
+    args.extra_headers += "\nDISABLE_VABC=1"
   if args.enable_threading:
     args.extra_headers += "\nENABLE_THREADING=1"
   if args.batched_writes:
diff --git a/scripts/update_payload/payload.py b/scripts/update_payload/payload.py
index b513f72d..4abd63e6 100644
--- a/scripts/update_payload/payload.py
+++ b/scripts/update_payload/payload.py
@@ -138,6 +138,8 @@ class Payload(object):
     else:
       self.name = payload_file.name
       self.payload_file = payload_file
+    self.payload_file_size = self.payload_file.seek(0, io.SEEK_END)
+    self.payload_file.seek(0, io.SEEK_SET)
     self.payload_file_offset = payload_file_offset
     self.manifest_hasher = None
     self.is_init = False
@@ -266,7 +268,7 @@ class Payload(object):
     self.metadata_size = self.header.size + self.header.manifest_len
     self.data_offset = self.metadata_size + self.header.metadata_signature_len
 
-    if self.manifest.signatures_offset and self.manifest.signatures_size:
+    if self.manifest.signatures_offset and self.manifest.signatures_size and self.manifest.signatures_offset + self.manifest.signatures_size <= self.payload_file_size:
       payload_signature_blob = self.ReadDataBlob(
           self.manifest.signatures_offset, self.manifest.signatures_size)
       payload_signature = update_metadata_pb2.Signatures()