Fetch a/b ota partitions from the build prop, instead of walking through
/dev/block/by-name

Test: administer a gki partial update on cuttlefish
Bug: 162148770
Change-Id: I173bd3ee1c462428ed02a9421c87ebed8dde636d
diff --git a/common/fake_boot_control.h b/common/fake_boot_control.h
index adbacd6..5d8823a 100644
--- a/common/fake_boot_control.h
+++ b/common/fake_boot_control.h
@@ -57,6 +57,9 @@
     if (part_it == devices_[slot].end())
       return false;
     *device = part_it->second;
+    if (is_dynamic != nullptr) {
+      *is_dynamic = false;
+    }
     return true;
   }
 
diff --git a/payload_consumer/partition_update_generator_android.cc b/payload_consumer/partition_update_generator_android.cc
index 5768dd6..d5d5313 100644
--- a/payload_consumer/partition_update_generator_android.cc
+++ b/payload_consumer/partition_update_generator_android.cc
@@ -18,30 +18,23 @@
 
 #include <filesystem>
 #include <memory>
-#include <set>
-#include <string_view>
 #include <utility>
 
+#include <android-base/properties.h>
 #include <android-base/strings.h>
 #include <base/logging.h>
+#include <base/strings/string_split.h>
 
+#include "update_engine/common/boot_control_interface.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::
@@ -50,22 +43,57 @@
         BootControlInterface::Slot target_slot,
         const std::set<std::string>& partitions_in_payload,
         std::vector<PartitionUpdate>* update_list) {
-  auto ab_partitions = GetStaticAbPartitionsOnDevice();
-  if (!ab_partitions.has_value()) {
+  auto ab_partitions = GetAbPartitionsOnDevice();
+  if (ab_partitions.empty()) {
     LOG(ERROR) << "Failed to load static a/b partitions";
     return false;
   }
 
   std::vector<PartitionUpdate> partition_updates;
-  for (const auto& partition_name : ab_partitions.value()) {
+  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;
     }
+    bool is_source_dynamic = false;
+    std::string source_device;
 
-    auto partition_update =
-        CreatePartitionUpdate(partition_name, source_slot, target_slot);
+    TEST_AND_RETURN_FALSE(
+        boot_control_->GetPartitionDevice(partition_name,
+                                          source_slot,
+                                          true, /* not_in_payload */
+                                          &source_device,
+                                          &is_source_dynamic));
+    bool is_target_dynamic = false;
+    std::string target_device;
+    TEST_AND_RETURN_FALSE(boot_control_->GetPartitionDevice(
+        partition_name, target_slot, true, &target_device, &is_target_dynamic));
+
+    if (is_source_dynamic || is_target_dynamic) {
+      if (is_source_dynamic != is_target_dynamic) {
+        LOG(ERROR) << "Partition " << partition_name << " is expected to be a"
+                   << " static partition. source slot is "
+                   << (is_source_dynamic ? "" : "not")
+                   << " dynamic, and target slot " << target_slot << " is "
+                   << (is_target_dynamic ? "" : "not") << " dynamic.";
+        return false;
+      } else {
+        continue;
+      }
+    }
+
+    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 false;
+    }
+
+    auto partition_update = CreatePartitionUpdate(
+        partition_name, source_device, target_device, source_size);
     if (!partition_update.has_value()) {
       LOG(ERROR) << "Failed to create partition update for " << partition_name;
       return false;
@@ -76,98 +104,14 @@
   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) << "Partition " << partition_name << " is expected to be a"
-               << " static partition. source slot is "
-               << (is_source_dynamic ? "" : "not")
-               << " dynamic, and 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);
+std::vector<std::string>
+PartitionUpdateGeneratorAndroid::GetAbPartitionsOnDevice() const {
+  auto partition_list_str =
+      android::base::GetProperty("ro.product.ab_ota_partitions", "");
+  return base::SplitString(partition_list_str,
+                           ",",
+                           base::TRIM_WHITESPACE,
+                           base::SPLIT_WANT_NONEMPTY);
 }
 
 std::optional<PartitionUpdate>
@@ -183,6 +127,8 @@
 
   auto raw_hash = CalculateHashForPartition(source_device, partition_size);
   if (!raw_hash.has_value()) {
+    LOG(ERROR) << "Failed to calculate hash for partition " << source_device
+               << " size: " << partition_size;
     return {};
   }
   old_partition_info->set_hash(raw_hash->data(), raw_hash->size());
@@ -225,16 +171,9 @@
 std::unique_ptr<PartitionUpdateGeneratorInterface> Create(
     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));
+      new PartitionUpdateGeneratorAndroid(boot_control, 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 97b7d83..0330c99 100644
--- a/payload_consumer/partition_update_generator_android.h
+++ b/payload_consumer/partition_update_generator_android.h
@@ -29,11 +29,11 @@
 #include "update_engine/payload_consumer/partition_update_generator_interface.h"
 
 namespace chromeos_update_engine {
+
 class PartitionUpdateGeneratorAndroid
     : public PartitionUpdateGeneratorInterface {
  public:
   PartitionUpdateGeneratorAndroid(BootControlInterface* boot_control,
-                                  std::string device_dir,
                                   size_t block_size);
 
   bool GenerateOperationsForPartitionsNotInPayload(
@@ -41,15 +41,13 @@
       BootControlInterface::Slot target_slot,
       const std::set<std::string>& partitions_in_payload,
       std::vector<PartitionUpdate>* update_list) override;
+  virtual std::vector<std::string> GetAbPartitionsOnDevice() const;
 
  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(
@@ -58,17 +56,10 @@
       const std::string& target_device,
       int64_t partition_size);
 
-  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_;
 };
 
diff --git a/payload_consumer/partition_update_generator_android_unittest.cc b/payload_consumer/partition_update_generator_android_unittest.cc
index c3be9db..86d025e 100644
--- a/payload_consumer/partition_update_generator_android_unittest.cc
+++ b/payload_consumer/partition_update_generator_android_unittest.cc
@@ -19,12 +19,14 @@
 #include <map>
 #include <memory>
 #include <set>
+#include <utility>
 #include <vector>
 
 #include <android-base/strings.h>
 #include <brillo/secure_blob.h>
 #include <gtest/gtest.h>
 
+#include "update_engine/common/boot_control_interface.h"
 #include "update_engine/common/fake_boot_control.h"
 #include "update_engine/common/hash_calculator.h"
 #include "update_engine/common/test_utils.h"
@@ -32,40 +34,53 @@
 
 namespace chromeos_update_engine {
 
+class FakePartitionUpdateGenerator : public PartitionUpdateGeneratorAndroid {
+ public:
+  std::vector<std::string> GetAbPartitionsOnDevice() const {
+    return ab_partitions_;
+  }
+  using PartitionUpdateGeneratorAndroid::PartitionUpdateGeneratorAndroid;
+  std::vector<std::string> ab_partitions_;
+};
+
 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_);
+    boot_control_->SetNumSlots(2);
+    generator_ = std::make_unique<FakePartitionUpdateGenerator>(
+        boot_control_.get(), 4096);
     ASSERT_TRUE(generator_);
-    generator_->block_device_dir_ = device_dir_.GetPath().value();
   }
 
-  std::unique_ptr<PartitionUpdateGeneratorAndroid> generator_;
+  std::unique_ptr<FakePartitionUpdateGenerator> generator_;
   std::unique_ptr<FakeBootControl> boot_control_;
 
   base::ScopedTempDir device_dir_;
+  std::map<std::string, std::string> device_map_;
 
   void SetUpBlockDevice(const std::map<std::string, std::string>& contents) {
+    std::set<std::string> partition_base_names;
     for (const auto& [name, content] : contents) {
-      auto path = generator_->block_device_dir_ + "/" + name;
+      auto path = device_dir_.GetPath().value() + "/" + 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);
+        auto prefix = name.substr(0, name.size() - 2);
+        boot_control_->SetPartitionDevice(prefix, 0, path);
+        partition_base_names.emplace(prefix);
       } else if (android::base::EndsWith(name, "_b")) {
-        boot_control_->SetPartitionDevice(
-            name.substr(0, name.size() - 2), 1, path);
+        auto prefix = name.substr(0, name.size() - 2);
+        boot_control_->SetPartitionDevice(prefix, 1, path);
+        partition_base_names.emplace(prefix);
       }
+      device_map_[name] = std::move(path);
     }
+    generator_->ab_partitions_ = {partition_base_names.begin(),
+                                  partition_base_names.end()};
   }
 
   void CheckPartitionUpdate(const std::string& name,
@@ -95,25 +110,6 @@
   }
 };
 
-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');
@@ -125,13 +121,14 @@
   };
   SetUpBlockDevice(contents);
 
-  auto system_partition_update =
-      generator_->CreatePartitionUpdate("system", 0, 1);
+  auto system_partition_update = generator_->CreatePartitionUpdate(
+      "system", device_map_["system_a"], device_map_["system_b"], 4096 * 2);
   ASSERT_TRUE(system_partition_update.has_value());
   CheckPartitionUpdate(
       "system", system_contents, system_partition_update.value());
 
-  auto boot_partition_update = generator_->CreatePartitionUpdate("boot", 0, 1);
+  auto boot_partition_update = generator_->CreatePartitionUpdate(
+      "boot", device_map_["boot_a"], device_map_["boot_b"], 4096 * 5);
   ASSERT_TRUE(boot_partition_update.has_value());
   CheckPartitionUpdate("boot", boot_contents, boot_partition_update.value());
 }