From 4d307b0975d6b4f70a5d2e41f25552aa54c853ea Mon Sep 17 00:00:00 2001
From: David Anderson <dvander@google.com>
Date: Mon, 17 Dec 2018 17:07:34 -0800
Subject: [PATCH] fastbootd: Only flash slots listed by the boot control HAL.

Bug: N/A
Test: flash when metadata slot count is >2
Change-Id: I67481be0de162cab5da8d32c2e318489427f1932
Merged-In: I67481be0de162cab5da8d32c2e318489427f1932
(cherry picked from commit 8568dcb057d63023feca09b031e456592c133f0e)
---
 fastboot/device/commands.cpp |  6 ++++--
 fastboot/device/flashing.cpp |  2 +-
 fastboot/device/utility.cpp  | 10 ++++++++--
 fastboot/device/utility.h    |  2 +-
 4 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/fastboot/device/commands.cpp b/fastboot/device/commands.cpp
index 71d2a1d61..e91598d01 100644
--- a/fastboot/device/commands.cpp
+++ b/fastboot/device/commands.cpp
@@ -329,12 +329,14 @@ class PartitionBuilder {
     MetadataBuilder* operator->() const { return builder_.get(); }
 
   private:
+    FastbootDevice* device_;
     std::string super_device_;
     uint32_t slot_number_;
     std::unique_ptr<MetadataBuilder> builder_;
 };
 
-PartitionBuilder::PartitionBuilder(FastbootDevice* device, const std::string& partition_name) {
+PartitionBuilder::PartitionBuilder(FastbootDevice* device, const std::string& partition_name)
+    : device_(device) {
     std::string slot_suffix = GetSuperSlotSuffix(device, partition_name);
     slot_number_ = SlotNumberForSlotSuffix(slot_suffix);
     auto super_device = FindPhysicalPartition(fs_mgr_get_super_partition_name(slot_number_));
@@ -350,7 +352,7 @@ bool PartitionBuilder::Write() {
     if (!metadata) {
         return false;
     }
-    return UpdateAllPartitionMetadata(super_device_, *metadata.get());
+    return UpdateAllPartitionMetadata(device_, super_device_, *metadata.get());
 }
 
 bool CreatePartitionHandler(FastbootDevice* device, const std::vector<std::string>& args) {
diff --git a/fastboot/device/flashing.cpp b/fastboot/device/flashing.cpp
index 7b99884a8..343661d4f 100644
--- a/fastboot/device/flashing.cpp
+++ b/fastboot/device/flashing.cpp
@@ -183,7 +183,7 @@ bool UpdateSuper(FastbootDevice* device, const std::string& super_name, bool wip
     }
 
     // Write the new table to every metadata slot.
-    if (!UpdateAllPartitionMetadata(super_name, *new_metadata.get())) {
+    if (!UpdateAllPartitionMetadata(device, super_name, *new_metadata.get())) {
         return device->WriteFail("Unable to write new partition table");
     }
     return device->WriteOkay("Successfully updated partition table");
diff --git a/fastboot/device/utility.cpp b/fastboot/device/utility.cpp
index 2ae9ac55b..2ebd57d12 100644
--- a/fastboot/device/utility.cpp
+++ b/fastboot/device/utility.cpp
@@ -200,10 +200,16 @@ bool GetDeviceLockStatus() {
     return cmdline.find("androidboot.verifiedbootstate=orange") == std::string::npos;
 }
 
-bool UpdateAllPartitionMetadata(const std::string& super_name,
+bool UpdateAllPartitionMetadata(FastbootDevice* device, const std::string& super_name,
                                 const android::fs_mgr::LpMetadata& metadata) {
+    size_t num_slots = 1;
+    auto boot_control_hal = device->boot_control_hal();
+    if (boot_control_hal) {
+        num_slots = boot_control_hal->getNumberSlots();
+    }
+
     bool ok = true;
-    for (size_t i = 0; i < metadata.geometry.metadata_slot_count; i++) {
+    for (size_t i = 0; i < num_slots; i++) {
         ok &= UpdatePartitionTable(super_name, metadata, i);
     }
     return ok;
diff --git a/fastboot/device/utility.h b/fastboot/device/utility.h
index 4c6aa07ad..bfeeb742f 100644
--- a/fastboot/device/utility.h
+++ b/fastboot/device/utility.h
@@ -68,5 +68,5 @@ std::vector<std::string> ListPartitions(FastbootDevice* device);
 bool GetDeviceLockStatus();
 
 // Update all copies of metadata.
-bool UpdateAllPartitionMetadata(const std::string& super_name,
+bool UpdateAllPartitionMetadata(FastbootDevice* device, const std::string& super_name,
                                 const android::fs_mgr::LpMetadata& metadata);