Implement PartitionUpdateGenerator for partial updates am: 99d570d67b am: d7ef9ad1f8 am: 485f36ac6c am: d9947187f0 am: a5250974e2

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

Change-Id: I95648552cce8322983bec1ea04bd66118879c8ca

11 files changed, 486 insertions, 12 deletions

diff --git a/Android.bp b/Android.bp
index 59d698d8..a5223c76 100644
--- a/Android.bp
+++ b/Android.bp
@@ -123,15 +123,16 @@ cc_defaults {
         "libbz",
         "libbspatch",
         "libbrotli",
+        "libc++fs",
         "libfec_rs",
         "libpuffpatch",
         "libverity_tree",
     ],
     shared_libs: [
-        "libziparchive",
         "libbase",
         "libcrypto",
         "libfec",
+        "libziparchive",
     ],
 }
 
@@ -731,6 +732,7 @@ cc_test {
         "payload_consumer/file_descriptor_utils_unittest.cc",
         "payload_consumer/file_writer_unittest.cc",
         "payload_consumer/filesystem_verifier_action_unittest.cc",
+        "payload_consumer/partition_update_generator_android_unittest.cc",
         "payload_consumer/postinstall_runner_action_unittest.cc",
         "payload_consumer/verity_writer_android_unittest.cc",
         "payload_consumer/xz_extent_writer_unittest.cc",
diff --git a/common/dynamic_partition_control_interface.h b/common/dynamic_partition_control_interface.h
index 58ebfe46..7289deee 100644
--- a/common/dynamic_partition_control_interface.h
+++ b/common/dynamic_partition_control_interface.h
@@ -21,6 +21,7 @@
 
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "update_engine/common/action.h"
 #include "update_engine/common/cleanup_previous_update_action_delegate.h"
@@ -118,6 +119,17 @@ class DynamicPartitionControlInterface {
   // progress, while ResetUpdate() forcefully free previously
   // allocated space for snapshot updates.
   virtual bool ResetUpdate(PrefsInterface* prefs) = 0;
+
+  // Reads the dynamic partitions metadata from the current slot, and puts the
+  // name of the dynamic partitions with the current suffix to |partitions|.
+  // Returns true on success.
+  virtual bool ListDynamicPartitionsForSlot(
+      uint32_t current_slot, std::vector<std::string>* partitions) = 0;
+
+  // Finds a possible location that list all block devices by name; and puts
+  // the result in |path|. Returns true on success.
+  // Sample result: /dev/block/by-name/
+  virtual bool GetDeviceDir(std::string* path) = 0;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/common/dynamic_partition_control_stub.cc b/common/dynamic_partition_control_stub.cc
index 903b7ee0..cde36afc 100644
--- a/common/dynamic_partition_control_stub.cc
+++ b/common/dynamic_partition_control_stub.cc
@@ -67,4 +67,13 @@ bool DynamicPartitionControlStub::ResetUpdate(PrefsInterface* prefs) {
   return false;
 }
 
+bool DynamicPartitionControlStub::ListDynamicPartitionsForSlot(
+    uint32_t current_slot, std::vector<std::string>* partitions) {
+  return true;
+}
+
+bool DynamicPartitionControlStub::GetDeviceDir(std::string* path) {
+  return true;
+}
+
 }  // namespace chromeos_update_engine
diff --git a/common/dynamic_partition_control_stub.h b/common/dynamic_partition_control_stub.h
index d8e254ea..28e3e6a5 100644
--- a/common/dynamic_partition_control_stub.h
+++ b/common/dynamic_partition_control_stub.h
@@ -21,6 +21,7 @@
 
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "update_engine/common/dynamic_partition_control_interface.h"
 
@@ -46,6 +47,10 @@ class DynamicPartitionControlStub : public DynamicPartitionControlInterface {
       PrefsInterface* prefs,
       CleanupPreviousUpdateActionDelegateInterface* delegate) override;
   bool ResetUpdate(PrefsInterface* prefs) override;
+
+  bool ListDynamicPartitionsForSlot(
+      uint32_t current_slot, std::vector<std::string>* partitions) override;
+  bool GetDeviceDir(std::string* path) override;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/dynamic_partition_control_android.cc b/dynamic_partition_control_android.cc
index 79c269c4..6817c21e 100644
--- a/dynamic_partition_control_android.cc
+++ b/dynamic_partition_control_android.cc
@@ -21,6 +21,8 @@
 #include <memory>
 #include <set>
 #include <string>
+#include <string_view>
+#include <utility>
 #include <vector>
 
 #include <android-base/properties.h>
@@ -1081,6 +1083,36 @@ bool DynamicPartitionControlAndroid::ResetUpdate(PrefsInterface* prefs) {
   return true;
 }
 
+bool DynamicPartitionControlAndroid::ListDynamicPartitionsForSlot(
+    uint32_t current_slot, std::vector<std::string>* partitions) {
+  if (!GetDynamicPartitionsFeatureFlag().IsEnabled()) {
+    LOG(ERROR) << "Dynamic partition is not enabled";
+    return false;
+  }
+
+  std::string device_dir_str;
+  TEST_AND_RETURN_FALSE(GetDeviceDir(&device_dir_str));
+  base::FilePath device_dir(device_dir_str);
+  auto super_device =
+      device_dir.Append(GetSuperPartitionName(current_slot)).value();
+  auto builder = LoadMetadataBuilder(super_device, current_slot);
+  TEST_AND_RETURN_FALSE(builder != nullptr);
+
+  std::vector<std::string> result;
+  auto suffix = SlotSuffixForSlotNumber(current_slot);
+  for (const auto& group : builder->ListGroups()) {
+    for (const auto& partition : builder->ListPartitionsInGroup(group)) {
+      std::string_view partition_name = partition->name();
+      if (!android::base::ConsumeSuffix(&partition_name, suffix)) {
+        continue;
+      }
+      result.emplace_back(partition_name);
+    }
+  }
+  *partitions = std::move(result);
+  return true;
+}
+
 bool DynamicPartitionControlAndroid::ExpectMetadataMounted() {
   // No need to mount metadata for non-Virtual A/B devices.
   if (!GetVirtualAbFeatureFlag().IsEnabled()) {
diff --git a/dynamic_partition_control_android.h b/dynamic_partition_control_android.h
index e3bedbcb..69026a40 100644
--- a/dynamic_partition_control_android.h
+++ b/dynamic_partition_control_android.h
@@ -20,6 +20,7 @@
 #include <memory>
 #include <set>
 #include <string>
+#include <vector>
 
 #include <base/files/file_util.h>
 #include <libsnapshot/auto_device.h>
@@ -53,6 +54,11 @@ class DynamicPartitionControlAndroid : public DynamicPartitionControlInterface {
 
   bool ResetUpdate(PrefsInterface* prefs) override;
 
+  bool ListDynamicPartitionsForSlot(
+      uint32_t current_slot, std::vector<std::string>* partitions) override;
+
+  bool GetDeviceDir(std::string* path) override;
+
   // Return the device for partition |partition_name| at slot |slot|.
   // |current_slot| should be set to the current active slot.
   // Note: this function is only used by BootControl*::GetPartitionDevice.
@@ -131,9 +137,6 @@ class DynamicPartitionControlAndroid : public DynamicPartitionControlInterface {
   virtual std::unique_ptr<android::fs_mgr::MetadataBuilder> LoadMetadataBuilder(
       const std::string& super_device, uint32_t source_slot);
 
-  // Return a possible location for devices listed by name.
-  virtual bool GetDeviceDir(std::string* path);
-
   // Return the name of the super partition (which stores super partition
   // metadata) for a given slot.
   virtual std::string GetSuperPartitionName(uint32_t slot);
diff --git a/payload_consumer/delta_performer.cc b/payload_consumer/delta_performer.cc
index d1de9f4d..d9b739de 100644
--- a/payload_consumer/delta_performer.cc
+++ b/payload_consumer/delta_performer.cc
@@ -879,7 +879,8 @@ bool DeltaPerformer::ParseManifestPartitions(ErrorCode* error) {
       touched_partitions.insert(partition_update.partition_name());
     }
 
-    auto generator = partition_update_generator::Create(boot_control_);
+    auto generator = partition_update_generator::Create(boot_control_,
+                                                        manifest_.block_size());
     std::vector<PartitionUpdate> other_partitions;
     TEST_AND_RETURN_FALSE(
         generator->GenerateOperationsForPartitionsNotInPayload(
diff --git a/payload_consumer/partition_update_generator_android.cc b/payload_consumer/partition_update_generator_android.cc
index fcacc86d..aa3f2e57 100644
--- a/payload_consumer/partition_update_generator_android.cc
+++ b/payload_consumer/partition_update_generator_android.cc
@@ -16,25 +16,241 @@
 
 #include "update_engine/payload_consumer/partition_update_generator_android.h"
 
+#include <filesystem>
 #include <memory>
+#include <set>
+#include <string_view>
+#include <utility>
+
+#include <android-base/strings.h>
+#include <base/logging.h>
+
+#include "update_engine/common/hash_calculator.h"
+#include "update_engine/common/utils.h"
+
+namespace {
+// TODO(xunchang) use definition in fs_mgr, e.g. fs_mgr_get_slot_suffix
+const char* SUFFIX_A = "_a";
+const char* SUFFIX_B = "_b";
+}  // namespace
 
 namespace chromeos_update_engine {
 
+PartitionUpdateGeneratorAndroid::PartitionUpdateGeneratorAndroid(
+    BootControlInterface* boot_control,
+    std::string device_dir,
+    size_t block_size)
+    : boot_control_(boot_control),
+      block_device_dir_(std::move(device_dir)),
+      block_size_(block_size) {}
+
 bool PartitionUpdateGeneratorAndroid::
     GenerateOperationsForPartitionsNotInPayload(
         BootControlInterface::Slot source_slot,
         BootControlInterface::Slot target_slot,
         const std::set<std::string>& partitions_in_payload,
         std::vector<PartitionUpdate>* update_list) {
-  // TODO(xunchang) implement the function
-  CHECK(boot_control_);
+  auto ret = GetStaticAbPartitionsOnDevice();
+  if (!ret.has_value()) {
+    LOG(ERROR) << "Failed to load static a/b partitions";
+    return false;
+  }
+  auto ab_partitions = ret.value();
+
+  // Add the dynamic partitions.
+  auto dynamic_control = boot_control_->GetDynamicPartitionControl();
+  std::vector<std::string> dynamic_partitions;
+  if (!dynamic_control->ListDynamicPartitionsForSlot(source_slot,
+                                                     &dynamic_partitions)) {
+    LOG(ERROR) << "Failed to load dynamic partitions from slot " << source_slot;
+    return false;
+  }
+  ab_partitions.insert(dynamic_partitions.begin(), dynamic_partitions.end());
+
+  std::vector<PartitionUpdate> partition_updates;
+  for (const auto& partition_name : ab_partitions) {
+    if (partitions_in_payload.find(partition_name) !=
+        partitions_in_payload.end()) {
+      LOG(INFO) << partition_name << " has included in payload";
+      continue;
+    }
+
+    auto partition_update =
+        CreatePartitionUpdate(partition_name, source_slot, target_slot);
+    if (!partition_update.has_value()) {
+      LOG(ERROR) << "Failed to create partition update for " << partition_name;
+      return false;
+    }
+    partition_updates.push_back(partition_update.value());
+  }
+  *update_list = std::move(partition_updates);
   return true;
 }
 
+std::optional<std::set<std::string>>
+PartitionUpdateGeneratorAndroid::GetStaticAbPartitionsOnDevice() {
+  if (std::error_code error_code;
+      !std::filesystem::exists(block_device_dir_, error_code) || error_code) {
+    LOG(ERROR) << "Failed to find " << block_device_dir_ << " "
+               << error_code.message();
+    return std::nullopt;
+  }
+
+  std::error_code error_code;
+  auto it = std::filesystem::directory_iterator(block_device_dir_, error_code);
+  if (error_code) {
+    LOG(ERROR) << "Failed to iterate " << block_device_dir_ << " "
+               << error_code.message();
+    return std::nullopt;
+  }
+
+  std::set<std::string> partitions_with_suffix;
+  for (const auto& entry : it) {
+    auto partition_name = entry.path().filename().string();
+    if (android::base::EndsWith(partition_name, SUFFIX_A) ||
+        android::base::EndsWith(partition_name, SUFFIX_B)) {
+      partitions_with_suffix.insert(partition_name);
+    }
+  }
+
+  // Second iteration to add the partition name without suffixes.
+  std::set<std::string> ab_partitions;
+  for (std::string_view name : partitions_with_suffix) {
+    if (!android::base::ConsumeSuffix(&name, SUFFIX_A)) {
+      continue;
+    }
+
+    // Add to the output list if the partition exist for both slot a and b.
+    auto base_name = std::string(name);
+    if (partitions_with_suffix.find(base_name + SUFFIX_B) !=
+        partitions_with_suffix.end()) {
+      ab_partitions.insert(base_name);
+    } else {
+      LOG(WARNING) << "Failed to find the b partition for " << base_name;
+    }
+  }
+
+  return ab_partitions;
+}
+
+std::optional<PartitionUpdate>
+PartitionUpdateGeneratorAndroid::CreatePartitionUpdate(
+    const std::string& partition_name,
+    BootControlInterface::Slot source_slot,
+    BootControlInterface::Slot target_slot) {
+  bool is_source_dynamic = false;
+  std::string source_device;
+  if (!boot_control_->GetPartitionDevice(partition_name,
+                                         source_slot,
+                                         true, /* not_in_payload */
+                                         &source_device,
+                                         &is_source_dynamic)) {
+    LOG(ERROR) << "Failed to load source " << partition_name;
+    return std::nullopt;
+  }
+  bool is_target_dynamic = false;
+  std::string target_device;
+  if (!boot_control_->GetPartitionDevice(partition_name,
+                                         target_slot,
+                                         true,
+                                         &target_device,
+                                         &is_target_dynamic)) {
+    LOG(ERROR) << "Failed to load target " << partition_name;
+    return std::nullopt;
+  }
+
+  if (is_source_dynamic != is_target_dynamic) {
+    LOG(ERROR) << "Source slot " << source_slot << " for partition "
+               << partition_name << " is " << (is_source_dynamic ? "" : "not")
+               << " dynamic, but target slot " << target_slot << " is "
+               << (is_target_dynamic ? "" : "not") << " dynamic.";
+    return std::nullopt;
+  }
+  auto source_size = utils::FileSize(source_device);
+  auto target_size = utils::FileSize(target_device);
+  if (source_size == -1 || target_size == -1 || source_size != target_size ||
+      source_size % block_size_ != 0) {
+    LOG(ERROR) << "Invalid partition size. source size " << source_size
+               << ", target size " << target_size;
+    return std::nullopt;
+  }
+
+  return CreatePartitionUpdate(partition_name,
+                               source_device,
+                               target_device,
+                               source_size,
+                               is_source_dynamic);
+}
+
+std::optional<PartitionUpdate>
+PartitionUpdateGeneratorAndroid::CreatePartitionUpdate(
+    const std::string& partition_name,
+    const std::string& source_device,
+    const std::string& target_device,
+    int64_t partition_size,
+    bool is_dynamic) {
+  PartitionUpdate partition_update;
+  partition_update.set_partition_name(partition_name);
+  auto old_partition_info = partition_update.mutable_old_partition_info();
+  old_partition_info->set_size(partition_size);
+
+  auto raw_hash = CalculateHashForPartition(source_device, partition_size);
+  if (!raw_hash.has_value()) {
+    return {};
+  }
+  old_partition_info->set_hash(raw_hash->data(), raw_hash->size());
+  auto new_partition_info = partition_update.mutable_new_partition_info();
+  new_partition_info->set_size(partition_size);
+  new_partition_info->set_hash(raw_hash->data(), raw_hash->size());
+  // TODO(xunchang) TBD, should we skip hashing and verification of the
+  // dynamic partitions not in payload?
+  if (!is_dynamic) {
+    auto copy_operation = partition_update.add_operations();
+    copy_operation->set_type(InstallOperation::SOURCE_COPY);
+    Extent copy_extent;
+    copy_extent.set_start_block(0);
+    copy_extent.set_num_blocks(partition_size / block_size_);
+
+    *copy_operation->add_src_extents() = copy_extent;
+    *copy_operation->add_dst_extents() = copy_extent;
+  }
+
+  return partition_update;
+}
+
+std::optional<brillo::Blob>
+PartitionUpdateGeneratorAndroid::CalculateHashForPartition(
+    const std::string& block_device, int64_t partition_size) {
+  // TODO(xunchang) compute the hash with ecc partitions first, the hashing
+  // behavior should match the one in SOURCE_COPY. Also, we don't have the
+  // correct hash for source partition.
+  // An alternative way is to verify the written bytes match the read bytes
+  // during filesystem verification. This could probably save us a read of
+  // partitions here.
+  brillo::Blob raw_hash;
+  if (HashCalculator::RawHashOfFile(block_device, partition_size, &raw_hash) !=
+      partition_size) {
+    LOG(ERROR) << "Failed to calculate hash for " << block_device;
+    return std::nullopt;
+  }
+
+  return raw_hash;
+}
+
 namespace partition_update_generator {
 std::unique_ptr<PartitionUpdateGeneratorInterface> Create(
-    BootControlInterface* boot_control) {
-  return std::make_unique<PartitionUpdateGeneratorAndroid>(boot_control);
+    BootControlInterface* boot_control, size_t block_size) {
+  CHECK(boot_control);
+  auto dynamic_control = boot_control->GetDynamicPartitionControl();
+  CHECK(dynamic_control);
+  std::string dir_path;
+  if (!dynamic_control->GetDeviceDir(&dir_path)) {
+    return nullptr;
+  }
+
+  return std::unique_ptr<PartitionUpdateGeneratorInterface>(
+      new PartitionUpdateGeneratorAndroid(
+          boot_control, std::move(dir_path), block_size));
 }
 }  // namespace partition_update_generator
 
diff --git a/payload_consumer/partition_update_generator_android.h b/payload_consumer/partition_update_generator_android.h
index bb50133a..8f33077f 100644
--- a/payload_consumer/partition_update_generator_android.h
+++ b/payload_consumer/partition_update_generator_android.h
@@ -17,10 +17,14 @@
 #ifndef UPDATE_ENGINE_PAYLOAD_CONSUMER_PARTITION_UPDATE_GENERATOR_ANDROID_H_
 #define UPDATE_ENGINE_PAYLOAD_CONSUMER_PARTITION_UPDATE_GENERATOR_ANDROID_H_
 
+#include <optional>
 #include <set>
 #include <string>
 #include <vector>
 
+#include <brillo/secure_blob.h>
+#include <gtest/gtest_prod.h>  // for FRIEND_TEST
+
 #include "update_engine/common/boot_control_interface.h"
 #include "update_engine/payload_consumer/partition_update_generator_interface.h"
 
@@ -28,8 +32,9 @@ namespace chromeos_update_engine {
 class PartitionUpdateGeneratorAndroid
     : public PartitionUpdateGeneratorInterface {
  public:
-  explicit PartitionUpdateGeneratorAndroid(BootControlInterface* boot_control)
-      : boot_control_(boot_control) {}
+  PartitionUpdateGeneratorAndroid(BootControlInterface* boot_control,
+                                  std::string device_dir,
+                                  size_t block_size);
 
   bool GenerateOperationsForPartitionsNotInPayload(
       BootControlInterface::Slot source_slot,
@@ -38,7 +43,34 @@ class PartitionUpdateGeneratorAndroid
       std::vector<PartitionUpdate>* update_list) override;
 
  private:
+  friend class PartitionUpdateGeneratorAndroidTest;
+  FRIEND_TEST(PartitionUpdateGeneratorAndroidTest, GetStaticPartitions);
+  FRIEND_TEST(PartitionUpdateGeneratorAndroidTest, CreatePartitionUpdate);
+
+  // Gets the name of the static a/b partitions on the device.
+  std::optional<std::set<std::string>> GetStaticAbPartitionsOnDevice();
+
+  // Creates a PartitionUpdate object for a given partition to update from
+  // source to target. Returns std::nullopt on failure.
+  std::optional<PartitionUpdate> CreatePartitionUpdate(
+      const std::string& partition_name,
+      const std::string& source_device,
+      const std::string& target_device,
+      int64_t partition_size,
+      bool is_dynamic);
+
+  std::optional<PartitionUpdate> CreatePartitionUpdate(
+      const std::string& partition_name,
+      BootControlInterface::Slot source_slot,
+      BootControlInterface::Slot target_slot);
+
+  std::optional<brillo::Blob> CalculateHashForPartition(
+      const std::string& block_device, int64_t partition_size);
+
   BootControlInterface* boot_control_;
+  // Path to look for a/b partitions
+  std::string block_device_dir_;
+  size_t block_size_;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/partition_update_generator_android_unittest.cc b/payload_consumer/partition_update_generator_android_unittest.cc
new file mode 100644
index 00000000..c3be9dbc
--- /dev/null
+++ b/payload_consumer/partition_update_generator_android_unittest.cc
@@ -0,0 +1,162 @@
+//
+// 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 "update_engine/payload_consumer/partition_update_generator_android.h"
+
+#include <map>
+#include <memory>
+#include <set>
+#include <vector>
+
+#include <android-base/strings.h>
+#include <brillo/secure_blob.h>
+#include <gtest/gtest.h>
+
+#include "update_engine/common/fake_boot_control.h"
+#include "update_engine/common/hash_calculator.h"
+#include "update_engine/common/test_utils.h"
+#include "update_engine/common/utils.h"
+
+namespace chromeos_update_engine {
+
+class PartitionUpdateGeneratorAndroidTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    ASSERT_TRUE(device_dir_.CreateUniqueTempDir());
+    boot_control_ = std::make_unique<FakeBootControl>();
+    boot_control_->SetNumSlots(2);
+    auto generator =
+        partition_update_generator::Create(boot_control_.get(), 4096);
+    generator_.reset(
+        static_cast<PartitionUpdateGeneratorAndroid*>(generator.release()));
+    ASSERT_TRUE(boot_control_);
+    ASSERT_TRUE(generator_);
+    generator_->block_device_dir_ = device_dir_.GetPath().value();
+  }
+
+  std::unique_ptr<PartitionUpdateGeneratorAndroid> generator_;
+  std::unique_ptr<FakeBootControl> boot_control_;
+
+  base::ScopedTempDir device_dir_;
+
+  void SetUpBlockDevice(const std::map<std::string, std::string>& contents) {
+    for (const auto& [name, content] : contents) {
+      auto path = generator_->block_device_dir_ + "/" + name;
+      ASSERT_TRUE(
+          utils::WriteFile(path.c_str(), content.data(), content.size()));
+
+      if (android::base::EndsWith(name, "_a")) {
+        boot_control_->SetPartitionDevice(
+            name.substr(0, name.size() - 2), 0, path);
+      } else if (android::base::EndsWith(name, "_b")) {
+        boot_control_->SetPartitionDevice(
+            name.substr(0, name.size() - 2), 1, path);
+      }
+    }
+  }
+
+  void CheckPartitionUpdate(const std::string& name,
+                            const std::string& content,
+                            const PartitionUpdate& partition_update) {
+    ASSERT_EQ(name, partition_update.partition_name());
+
+    brillo::Blob out_hash;
+    ASSERT_TRUE(HashCalculator::RawHashOfBytes(
+        content.data(), content.size(), &out_hash));
+    ASSERT_EQ(std::string(out_hash.begin(), out_hash.end()),
+              partition_update.old_partition_info().hash());
+    ASSERT_EQ(std::string(out_hash.begin(), out_hash.end()),
+              partition_update.new_partition_info().hash());
+
+    ASSERT_EQ(1, partition_update.operations_size());
+    const auto& operation = partition_update.operations(0);
+    ASSERT_EQ(InstallOperation::SOURCE_COPY, operation.type());
+
+    ASSERT_EQ(1, operation.src_extents_size());
+    ASSERT_EQ(0u, operation.src_extents(0).start_block());
+    ASSERT_EQ(content.size() / 4096, operation.src_extents(0).num_blocks());
+
+    ASSERT_EQ(1, operation.dst_extents_size());
+    ASSERT_EQ(0u, operation.dst_extents(0).start_block());
+    ASSERT_EQ(content.size() / 4096, operation.dst_extents(0).num_blocks());
+  }
+};
+
+TEST_F(PartitionUpdateGeneratorAndroidTest, GetStaticPartitions) {
+  std::map<std::string, std::string> contents = {
+      {"system_a", ""},
+      {"system_b", ""},
+      {"vendor_a", ""},
+      {"vendor_b", ""},
+      {"persist", ""},
+      {"vbmeta_a", ""},
+      {"vbmeta_b", ""},
+      {"boot_a", ""},
+      {"boot_b", ""},
+  };
+
+  SetUpBlockDevice(contents);
+  auto partitions = generator_->GetStaticAbPartitionsOnDevice();
+  ASSERT_EQ(std::set<std::string>({"system", "vendor", "vbmeta", "boot"}),
+            partitions);
+}
+
+TEST_F(PartitionUpdateGeneratorAndroidTest, CreatePartitionUpdate) {
+  auto system_contents = std::string(4096 * 2, '1');
+  auto boot_contents = std::string(4096 * 5, 'b');
+  std::map<std::string, std::string> contents = {
+      {"system_a", system_contents},
+      {"system_b", std::string(4096 * 2, 0)},
+      {"boot_a", boot_contents},
+      {"boot_b", std::string(4096 * 5, 0)},
+  };
+  SetUpBlockDevice(contents);
+
+  auto system_partition_update =
+      generator_->CreatePartitionUpdate("system", 0, 1);
+  ASSERT_TRUE(system_partition_update.has_value());
+  CheckPartitionUpdate(
+      "system", system_contents, system_partition_update.value());
+
+  auto boot_partition_update = generator_->CreatePartitionUpdate("boot", 0, 1);
+  ASSERT_TRUE(boot_partition_update.has_value());
+  CheckPartitionUpdate("boot", boot_contents, boot_partition_update.value());
+}
+
+TEST_F(PartitionUpdateGeneratorAndroidTest, GenerateOperations) {
+  auto system_contents = std::string(4096 * 10, '2');
+  auto boot_contents = std::string(4096 * 5, 'b');
+  std::map<std::string, std::string> contents = {
+      {"system_a", system_contents},
+      {"system_b", std::string(4096 * 10, 0)},
+      {"boot_a", boot_contents},
+      {"boot_b", std::string(4096 * 5, 0)},
+      {"vendor_a", ""},
+      {"vendor_b", ""},
+      {"persist", ""},
+  };
+  SetUpBlockDevice(contents);
+
+  std::vector<PartitionUpdate> update_list;
+  ASSERT_TRUE(generator_->GenerateOperationsForPartitionsNotInPayload(
+      0, 1, std::set<std::string>{"vendor"}, &update_list));
+
+  ASSERT_EQ(2u, update_list.size());
+  CheckPartitionUpdate("boot", boot_contents, update_list[0]);
+  CheckPartitionUpdate("system", system_contents, update_list[1]);
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_consumer/partition_update_generator_interface.h b/payload_consumer/partition_update_generator_interface.h
index 0341d40b..3fa3dfbc 100644
--- a/payload_consumer/partition_update_generator_interface.h
+++ b/payload_consumer/partition_update_generator_interface.h
@@ -47,7 +47,7 @@ class PartitionUpdateGeneratorInterface {
 
 namespace partition_update_generator {
 std::unique_ptr<PartitionUpdateGeneratorInterface> Create(
-    BootControlInterface* boot_control);
+    BootControlInterface* boot_control, size_t block_size);
 }
 
 }  // namespace chromeos_update_engine