Merge "fastbootd: Support two super partitions for retrofit devices."
diff --git a/fastboot/device/commands.cpp b/fastboot/device/commands.cpp
index 08744ec..71d2a1d 100644
--- a/fastboot/device/commands.cpp
+++ b/fastboot/device/commands.cpp
@@ -322,7 +322,7 @@
 // partition table to the same place it was read.
 class PartitionBuilder {
   public:
-    explicit PartitionBuilder(FastbootDevice* device);
+    explicit PartitionBuilder(FastbootDevice* device, const std::string& partition_name);
 
     bool Write();
     bool Valid() const { return !!builder_; }
@@ -330,19 +330,19 @@
 
   private:
     std::string super_device_;
+    uint32_t slot_number_;
     std::unique_ptr<MetadataBuilder> builder_;
 };
 
-PartitionBuilder::PartitionBuilder(FastbootDevice* device) {
-    auto super_device = FindPhysicalPartition(fs_mgr_get_super_partition_name());
+PartitionBuilder::PartitionBuilder(FastbootDevice* device, const std::string& partition_name) {
+    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_));
     if (!super_device) {
         return;
     }
     super_device_ = *super_device;
-
-    std::string slot = device->GetCurrentSlot();
-    uint32_t slot_number = SlotNumberForSlotSuffix(slot);
-    builder_ = MetadataBuilder::New(super_device_, slot_number);
+    builder_ = MetadataBuilder::New(super_device_, slot_number_);
 }
 
 bool PartitionBuilder::Write() {
@@ -350,11 +350,7 @@
     if (!metadata) {
         return false;
     }
-    bool ok = true;
-    for (uint32_t i = 0; i < metadata->geometry.metadata_slot_count; i++) {
-        ok &= UpdatePartitionTable(super_device_, *metadata.get(), i);
-    }
-    return ok;
+    return UpdateAllPartitionMetadata(super_device_, *metadata.get());
 }
 
 bool CreatePartitionHandler(FastbootDevice* device, const std::vector<std::string>& args) {
@@ -372,7 +368,7 @@
         return device->WriteFail("Invalid partition size");
     }
 
-    PartitionBuilder builder(device);
+    PartitionBuilder builder(device, partition_name);
     if (!builder.Valid()) {
         return device->WriteFail("Could not open super partition");
     }
@@ -404,11 +400,13 @@
         return device->WriteStatus(FastbootResult::FAIL, "Command not available on locked devices");
     }
 
-    PartitionBuilder builder(device);
+    std::string partition_name = args[1];
+
+    PartitionBuilder builder(device, partition_name);
     if (!builder.Valid()) {
         return device->WriteFail("Could not open super partition");
     }
-    builder->RemovePartition(args[1]);
+    builder->RemovePartition(partition_name);
     if (!builder.Write()) {
         return device->WriteFail("Failed to write partition table");
     }
@@ -430,7 +428,7 @@
         return device->WriteFail("Invalid partition size");
     }
 
-    PartitionBuilder builder(device);
+    PartitionBuilder builder(device, partition_name);
     if (!builder.Valid()) {
         return device->WriteFail("Could not open super partition");
     }
diff --git a/fastboot/device/flashing.cpp b/fastboot/device/flashing.cpp
index 7c9e1d0..7b99884 100644
--- a/fastboot/device/flashing.cpp
+++ b/fastboot/device/flashing.cpp
@@ -183,12 +183,7 @@
     }
 
     // Write the new table to every metadata slot.
-    bool ok = true;
-    for (size_t i = 0; i < new_metadata->geometry.metadata_slot_count; i++) {
-        ok &= UpdatePartitionTable(super_name, *new_metadata.get(), i);
-    }
-
-    if (!ok) {
+    if (!UpdateAllPartitionMetadata(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 b844b9f..2ae9ac5 100644
--- a/fastboot/device/utility.cpp
+++ b/fastboot/device/utility.cpp
@@ -23,9 +23,11 @@
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/properties.h>
 #include <android-base/strings.h>
 #include <fs_mgr.h>
 #include <fs_mgr_dm_linear.h>
+#include <liblp/builder.h>
 #include <liblp/liblp.h>
 
 #include "fastboot_device.h"
@@ -35,7 +37,9 @@
 using android::base::unique_fd;
 using android::hardware::boot::V1_0::Slot;
 
-static bool OpenPhysicalPartition(const std::string& name, PartitionHandle* handle) {
+namespace {
+
+bool OpenPhysicalPartition(const std::string& name, PartitionHandle* handle) {
     std::optional<std::string> path = FindPhysicalPartition(name);
     if (!path) {
         return false;
@@ -44,28 +48,31 @@
     return true;
 }
 
-static bool OpenLogicalPartition(const std::string& name, const std::string& slot,
-                                 PartitionHandle* handle) {
-    std::optional<std::string> path = FindPhysicalPartition(fs_mgr_get_super_partition_name());
+bool OpenLogicalPartition(FastbootDevice* device, const std::string& partition_name,
+                          PartitionHandle* handle) {
+    std::string slot_suffix = GetSuperSlotSuffix(device, partition_name);
+    uint32_t slot_number = SlotNumberForSlotSuffix(slot_suffix);
+    auto path = FindPhysicalPartition(fs_mgr_get_super_partition_name(slot_number));
     if (!path) {
         return false;
     }
-    uint32_t slot_number = SlotNumberForSlotSuffix(slot);
     std::string dm_path;
-    if (!CreateLogicalPartition(path->c_str(), slot_number, name, true, 5s, &dm_path)) {
-        LOG(ERROR) << "Could not map partition: " << name;
+    if (!CreateLogicalPartition(path->c_str(), slot_number, partition_name, true, 5s, &dm_path)) {
+        LOG(ERROR) << "Could not map partition: " << partition_name;
         return false;
     }
-    auto closer = [name]() -> void { DestroyLogicalPartition(name, 5s); };
+    auto closer = [partition_name]() -> void { DestroyLogicalPartition(partition_name, 5s); };
     *handle = PartitionHandle(dm_path, std::move(closer));
     return true;
 }
 
+}  // namespace
+
 bool OpenPartition(FastbootDevice* device, const std::string& name, PartitionHandle* handle) {
     // We prioritize logical partitions over physical ones, and do this
     // consistently for other partition operations (like getvar:partition-size).
-    if (LogicalPartitionExists(name, device->GetCurrentSlot())) {
-        if (!OpenLogicalPartition(name, device->GetCurrentSlot(), handle)) {
+    if (LogicalPartitionExists(device, name)) {
+        if (!OpenLogicalPartition(device, name, handle)) {
             return false;
         }
     } else if (!OpenPhysicalPartition(name, handle)) {
@@ -104,14 +111,14 @@
     return nullptr;
 }
 
-bool LogicalPartitionExists(const std::string& name, const std::string& slot_suffix,
-                            bool* is_zero_length) {
-    auto path = FindPhysicalPartition(fs_mgr_get_super_partition_name());
+bool LogicalPartitionExists(FastbootDevice* device, const std::string& name, bool* is_zero_length) {
+    std::string slot_suffix = GetSuperSlotSuffix(device, name);
+    uint32_t slot_number = SlotNumberForSlotSuffix(slot_suffix);
+    auto path = FindPhysicalPartition(fs_mgr_get_super_partition_name(slot_number));
     if (!path) {
         return false;
     }
 
-    uint32_t slot_number = SlotNumberForSlotSuffix(slot_suffix);
     std::unique_ptr<LpMetadata> metadata = ReadMetadata(path->c_str(), slot_number);
     if (!metadata) {
         return false;
@@ -154,12 +161,29 @@
         }
     }
 
-    // Next get logical partitions.
-    if (auto path = FindPhysicalPartition(fs_mgr_get_super_partition_name())) {
-        uint32_t slot_number = SlotNumberForSlotSuffix(device->GetCurrentSlot());
-        if (auto metadata = ReadMetadata(path->c_str(), slot_number)) {
-            for (const auto& partition : metadata->partitions) {
-                std::string partition_name = GetPartitionName(partition);
+    // Find metadata in each super partition (on retrofit devices, there will
+    // be two).
+    std::vector<std::unique_ptr<LpMetadata>> metadata_list;
+
+    uint32_t current_slot = SlotNumberForSlotSuffix(device->GetCurrentSlot());
+    std::string super_name = fs_mgr_get_super_partition_name(current_slot);
+    if (auto metadata = ReadMetadata(super_name, current_slot)) {
+        metadata_list.emplace_back(std::move(metadata));
+    }
+
+    uint32_t other_slot = (current_slot == 0) ? 1 : 0;
+    std::string other_super = fs_mgr_get_super_partition_name(other_slot);
+    if (super_name != other_super) {
+        if (auto metadata = ReadMetadata(other_super, other_slot)) {
+            metadata_list.emplace_back(std::move(metadata));
+        }
+    }
+
+    for (const auto& metadata : metadata_list) {
+        for (const auto& partition : metadata->partitions) {
+            std::string partition_name = GetPartitionName(partition);
+            if (std::find(partitions.begin(), partitions.end(), partition_name) ==
+                partitions.end()) {
                 partitions.emplace_back(partition_name);
             }
         }
@@ -175,3 +199,30 @@
     }
     return cmdline.find("androidboot.verifiedbootstate=orange") == std::string::npos;
 }
+
+bool UpdateAllPartitionMetadata(const std::string& super_name,
+                                const android::fs_mgr::LpMetadata& metadata) {
+    bool ok = true;
+    for (size_t i = 0; i < metadata.geometry.metadata_slot_count; i++) {
+        ok &= UpdatePartitionTable(super_name, metadata, i);
+    }
+    return ok;
+}
+
+std::string GetSuperSlotSuffix(FastbootDevice* device, const std::string& partition_name) {
+    // If the super partition does not have a slot suffix, this is not a
+    // retrofit device, and we should take the current slot.
+    std::string current_slot_suffix = device->GetCurrentSlot();
+    uint32_t current_slot_number = SlotNumberForSlotSuffix(current_slot_suffix);
+    std::string super_partition = fs_mgr_get_super_partition_name(current_slot_number);
+    if (GetPartitionSlotSuffix(super_partition).empty()) {
+        return current_slot_suffix;
+    }
+
+    // Otherwise, infer the slot from the partition name.
+    std::string slot_suffix = GetPartitionSlotSuffix(partition_name);
+    if (!slot_suffix.empty()) {
+        return slot_suffix;
+    }
+    return current_slot_suffix;
+}
diff --git a/fastboot/device/utility.h b/fastboot/device/utility.h
index bb08f72..4c6aa07 100644
--- a/fastboot/device/utility.h
+++ b/fastboot/device/utility.h
@@ -20,6 +20,7 @@
 
 #include <android-base/unique_fd.h>
 #include <android/hardware/boot/1.0/IBootControl.h>
+#include <liblp/liblp.h>
 
 // Logical partitions are only mapped to a block device as needed, and
 // immediately unmapped when no longer needed. In order to enforce this we
@@ -52,10 +53,20 @@
 
 class FastbootDevice;
 
+// On normal devices, the super partition is always named "super". On retrofit
+// devices, the name must be derived from the partition name or current slot.
+// This helper assists in choosing the correct super for a given partition
+// name.
+std::string GetSuperSlotSuffix(FastbootDevice* device, const std::string& partition_name);
+
 std::optional<std::string> FindPhysicalPartition(const std::string& name);
-bool LogicalPartitionExists(const std::string& name, const std::string& slot_suffix,
+bool LogicalPartitionExists(FastbootDevice* device, const std::string& name,
                             bool* is_zero_length = nullptr);
 bool OpenPartition(FastbootDevice* device, const std::string& name, PartitionHandle* handle);
 bool GetSlotNumber(const std::string& slot, android::hardware::boot::V1_0::Slot* number);
 std::vector<std::string> ListPartitions(FastbootDevice* device);
 bool GetDeviceLockStatus();
+
+// Update all copies of metadata.
+bool UpdateAllPartitionMetadata(const std::string& super_name,
+                                const android::fs_mgr::LpMetadata& metadata);
diff --git a/fastboot/device/variables.cpp b/fastboot/device/variables.cpp
index 601af34..130a3cf 100644
--- a/fastboot/device/variables.cpp
+++ b/fastboot/device/variables.cpp
@@ -271,8 +271,7 @@
         return true;
     }
     std::string partition_name = args[0] + slot_suffix;
-    if (FindPhysicalPartition(partition_name) ||
-        LogicalPartitionExists(partition_name, slot_suffix)) {
+    if (FindPhysicalPartition(partition_name) || LogicalPartitionExists(device, partition_name)) {
         *message = "yes";
     } else {
         *message = "no";
@@ -289,8 +288,7 @@
     // Zero-length partitions cannot be created through device-mapper, so we
     // special case them here.
     bool is_zero_length;
-    if (LogicalPartitionExists(args[0], device->GetCurrentSlot(), &is_zero_length) &&
-        is_zero_length) {
+    if (LogicalPartitionExists(device, args[0], &is_zero_length) && is_zero_length) {
         *message = "0x0";
         return true;
     }
@@ -313,8 +311,7 @@
     }
 
     std::string partition_name = args[0];
-    if (!FindPhysicalPartition(partition_name) &&
-        !LogicalPartitionExists(partition_name, device->GetCurrentSlot())) {
+    if (!FindPhysicalPartition(partition_name) && !LogicalPartitionExists(device, partition_name)) {
         *message = "Invalid partition";
         return false;
     }
@@ -363,7 +360,7 @@
     // return "true", to be consistent with prefering to flash logical partitions
     // over physical ones.
     std::string partition_name = args[0];
-    if (LogicalPartitionExists(partition_name, device->GetCurrentSlot())) {
+    if (LogicalPartitionExists(device, partition_name)) {
         *message = "yes";
         return true;
     }
diff --git a/fs_mgr/liblp/builder.cpp b/fs_mgr/liblp/builder.cpp
index 7035f42..da86d75 100644
--- a/fs_mgr/liblp/builder.cpp
+++ b/fs_mgr/liblp/builder.cpp
@@ -20,6 +20,7 @@
 
 #include <algorithm>
 
+#include <android-base/properties.h>
 #include <android-base/unique_fd.h>
 
 #include "liblp/liblp.h"
@@ -29,6 +30,9 @@
 namespace android {
 namespace fs_mgr {
 
+bool MetadataBuilder::sABOverrideSet;
+bool MetadataBuilder::sABOverrideValue;
+
 bool LinearExtent::AddTo(LpMetadata* out) const {
     if (device_index_ >= out->block_devices.size()) {
         LERROR << "Extent references unknown block device.";
@@ -203,6 +207,11 @@
     return New(*metadata.get(), &opener);
 }
 
+void MetadataBuilder::OverrideABForTesting(bool ab_device) {
+    sABOverrideSet = true;
+    sABOverrideValue = ab_device;
+}
+
 MetadataBuilder::MetadataBuilder() : auto_slot_suffixing_(false) {
     memset(&geometry_, 0, sizeof(geometry_));
     geometry_.magic = LP_METADATA_GEOMETRY_MAGIC;
@@ -427,6 +436,11 @@
         LERROR << "Could not find partition group: " << group_name;
         return nullptr;
     }
+    if (IsABDevice() && !auto_slot_suffixing_ && name != "scratch" &&
+        GetPartitionSlotSuffix(name).empty()) {
+        LERROR << "Unsuffixed partition not allowed on A/B device: " << name;
+        return nullptr;
+    }
     partitions_.push_back(std::make_unique<Partition>(name, group_name, attributes));
     return partitions_.back().get();
 }
@@ -909,5 +923,12 @@
     auto_slot_suffixing_ = true;
 }
 
+bool MetadataBuilder::IsABDevice() const {
+    if (sABOverrideSet) {
+        return sABOverrideValue;
+    }
+    return android::base::GetBoolProperty("ro.build.ab_update", false);
+}
+
 }  // namespace fs_mgr
 }  // namespace android
diff --git a/fs_mgr/liblp/builder_test.cpp b/fs_mgr/liblp/builder_test.cpp
index 35cab38..926fe12 100644
--- a/fs_mgr/liblp/builder_test.cpp
+++ b/fs_mgr/liblp/builder_test.cpp
@@ -25,7 +25,25 @@
 using namespace android::fs_mgr;
 using ::testing::ElementsAre;
 
-TEST(liblp, BuildBasic) {
+class Environment : public ::testing::Environment {
+  public:
+    void SetUp() override { MetadataBuilder::OverrideABForTesting(false); }
+};
+
+int main(int argc, char** argv) {
+    std::unique_ptr<Environment> env(new Environment);
+    ::testing::AddGlobalTestEnvironment(env.get());
+    ::testing::InitGoogleTest(&argc, argv);
+    return RUN_ALL_TESTS();
+}
+
+class BuilderTest : public ::testing::Test {
+  public:
+    void SetUp() override { MetadataBuilder::OverrideABForTesting(false); }
+    void TearDown() override { MetadataBuilder::OverrideABForTesting(false); }
+};
+
+TEST_F(BuilderTest, BuildBasic) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
     ASSERT_NE(builder, nullptr);
 
@@ -40,7 +58,7 @@
     EXPECT_EQ(builder->FindPartition("system"), nullptr);
 }
 
-TEST(liblp, ResizePartition) {
+TEST_F(BuilderTest, ResizePartition) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
     ASSERT_NE(builder, nullptr);
 
@@ -94,7 +112,7 @@
     EXPECT_EQ(system->extents().size(), 0);
 }
 
-TEST(liblp, PartitionAlignment) {
+TEST_F(BuilderTest, PartitionAlignment) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
     ASSERT_NE(builder, nullptr);
 
@@ -110,7 +128,7 @@
     EXPECT_EQ(system->extents().size(), 1);
 }
 
-TEST(liblp, DiskAlignment) {
+TEST_F(BuilderTest, DiskAlignment) {
     static const uint64_t kDiskSize = 1000000;
     static const uint32_t kMetadataSize = 1024;
     static const uint32_t kMetadataSlots = 2;
@@ -120,7 +138,7 @@
     ASSERT_EQ(builder, nullptr);
 }
 
-TEST(liblp, MetadataAlignment) {
+TEST_F(BuilderTest, MetadataAlignment) {
     // Make sure metadata sizes get aligned up.
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1000, 2);
     ASSERT_NE(builder, nullptr);
@@ -129,7 +147,7 @@
     EXPECT_EQ(exported->geometry.metadata_max_size, 1024);
 }
 
-TEST(liblp, InternalAlignment) {
+TEST_F(BuilderTest, InternalAlignment) {
     // Test the metadata fitting within alignment.
     BlockDeviceInfo device_info("super", 1024 * 1024, 768 * 1024, 0, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 2);
@@ -177,7 +195,7 @@
     EXPECT_EQ(super_device->first_logical_sector, 160);
 }
 
-TEST(liblp, InternalPartitionAlignment) {
+TEST_F(BuilderTest, InternalPartitionAlignment) {
     BlockDeviceInfo device_info("super", 512 * 1024 * 1024, 768 * 1024, 753664, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 32 * 1024, 2);
 
@@ -211,7 +229,7 @@
     EXPECT_EQ(exported->extents.back().target_data, 30656);
 }
 
-TEST(liblp, UseAllDiskSpace) {
+TEST_F(BuilderTest, UseAllDiskSpace) {
     static constexpr uint64_t total = 1024 * 1024;
     static constexpr uint64_t metadata = 1024;
     static constexpr uint64_t slots = 2;
@@ -237,7 +255,7 @@
     EXPECT_EQ(builder->AllocatableSpace(), allocatable);
 }
 
-TEST(liblp, BuildComplex) {
+TEST_F(BuilderTest, BuildComplex) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
 
     Partition* system = builder->AddPartition("system", LP_PARTITION_ATTR_READONLY);
@@ -271,7 +289,7 @@
     EXPECT_EQ(vendor1->physical_sector() + vendor1->num_sectors(), system2->physical_sector());
 }
 
-TEST(liblp, AddInvalidPartition) {
+TEST_F(BuilderTest, AddInvalidPartition) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
 
     Partition* partition = builder->AddPartition("system", LP_PARTITION_ATTR_READONLY);
@@ -286,7 +304,7 @@
     EXPECT_EQ(partition, nullptr);
 }
 
-TEST(liblp, BuilderExport) {
+TEST_F(BuilderTest, BuilderExport) {
     static const uint64_t kDiskSize = 1024 * 1024;
     static const uint32_t kMetadataSize = 1024;
     static const uint32_t kMetadataSlots = 2;
@@ -344,7 +362,7 @@
     }
 }
 
-TEST(liblp, BuilderImport) {
+TEST_F(BuilderTest, BuilderImport) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
 
     Partition* system = builder->AddPartition("system", LP_PARTITION_ATTR_READONLY);
@@ -382,7 +400,7 @@
     EXPECT_EQ(vendor1->num_sectors(), 32768 / LP_SECTOR_SIZE);
 }
 
-TEST(liblp, ExportNameTooLong) {
+TEST_F(BuilderTest, ExportNameTooLong) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
 
     std::string name = "abcdefghijklmnopqrstuvwxyz0123456789";
@@ -393,7 +411,7 @@
     EXPECT_EQ(exported, nullptr);
 }
 
-TEST(liblp, MetadataTooLarge) {
+TEST_F(BuilderTest, MetadataTooLarge) {
     static const size_t kDiskSize = 128 * 1024;
     static const size_t kMetadataSize = 64 * 1024;
 
@@ -423,7 +441,7 @@
     EXPECT_EQ(builder, nullptr);
 }
 
-TEST(liblp, block_device_info) {
+TEST_F(BuilderTest, block_device_info) {
     std::unique_ptr<fstab, decltype(&fs_mgr_free_fstab)> fstab(fs_mgr_read_fstab_default(),
                                                                fs_mgr_free_fstab);
     ASSERT_NE(fstab, nullptr);
@@ -444,7 +462,7 @@
     ASSERT_LT(device_info.alignment_offset, device_info.alignment);
 }
 
-TEST(liblp, UpdateBlockDeviceInfo) {
+TEST_F(BuilderTest, UpdateBlockDeviceInfo) {
     BlockDeviceInfo device_info("super", 1024 * 1024, 4096, 1024, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
     ASSERT_NE(builder, nullptr);
@@ -482,13 +500,13 @@
     EXPECT_EQ(new_info.logical_block_size, 4096);
 }
 
-TEST(liblp, InvalidBlockSize) {
+TEST_F(BuilderTest, InvalidBlockSize) {
     BlockDeviceInfo device_info("super", 1024 * 1024, 0, 0, 513);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
     EXPECT_EQ(builder, nullptr);
 }
 
-TEST(liblp, AlignedExtentSize) {
+TEST_F(BuilderTest, AlignedExtentSize) {
     BlockDeviceInfo device_info("super", 1024 * 1024, 0, 0, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
     ASSERT_NE(builder, nullptr);
@@ -499,14 +517,14 @@
     EXPECT_EQ(partition->size(), 4096);
 }
 
-TEST(liblp, AlignedFreeSpace) {
+TEST_F(BuilderTest, AlignedFreeSpace) {
     // Only one sector free - at least one block is required.
     BlockDeviceInfo device_info("super", 10240, 0, 0, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 512, 1);
     ASSERT_EQ(builder, nullptr);
 }
 
-TEST(liblp, HasDefaultGroup) {
+TEST_F(BuilderTest, HasDefaultGroup) {
     BlockDeviceInfo device_info("super", 1024 * 1024, 0, 0, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
     ASSERT_NE(builder, nullptr);
@@ -514,7 +532,7 @@
     EXPECT_FALSE(builder->AddGroup("default", 0));
 }
 
-TEST(liblp, GroupSizeLimits) {
+TEST_F(BuilderTest, GroupSizeLimits) {
     BlockDeviceInfo device_info("super", 1024 * 1024, 0, 0, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
     ASSERT_NE(builder, nullptr);
@@ -538,7 +556,7 @@
     return x << 20;
 }
 
-TEST(liblp, RemoveAndAddFirstPartition) {
+TEST_F(BuilderTest, RemoveAndAddFirstPartition) {
     auto builder = MetadataBuilder::New(10_GiB, 65536, 2);
     ASSERT_NE(nullptr, builder);
     ASSERT_TRUE(builder->AddGroup("foo_a", 5_GiB));
@@ -561,7 +579,7 @@
     ASSERT_TRUE(p && builder->ResizePartition(p, 1_GiB));
 }
 
-TEST(liblp, ListGroups) {
+TEST_F(BuilderTest, ListGroups) {
     BlockDeviceInfo device_info("super", 1024 * 1024, 0, 0, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
     ASSERT_NE(builder, nullptr);
@@ -571,7 +589,7 @@
     ASSERT_THAT(groups, ElementsAre("default", "example"));
 }
 
-TEST(liblp, RemoveGroupAndPartitions) {
+TEST_F(BuilderTest, RemoveGroupAndPartitions) {
     BlockDeviceInfo device_info("super", 1024 * 1024, 0, 0, 4096);
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
     ASSERT_NE(builder, nullptr);
@@ -588,7 +606,7 @@
     ASSERT_NE(builder->FindPartition("system"), nullptr);
 }
 
-TEST(liblp, MultipleBlockDevices) {
+TEST_F(BuilderTest, MultipleBlockDevices) {
     std::vector<BlockDeviceInfo> partitions = {
             BlockDeviceInfo("system_a", 256_MiB, 786432, 229376, 4096),
             BlockDeviceInfo("vendor_a", 128_MiB, 786432, 753664, 4096),
@@ -633,7 +651,7 @@
     EXPECT_EQ(metadata->extents[2].target_source, 2);
 }
 
-TEST(liblp, ImportPartitionsOk) {
+TEST_F(BuilderTest, ImportPartitionsOk) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
     ASSERT_NE(builder, nullptr);
 
@@ -673,7 +691,7 @@
     EXPECT_EQ(extent_a.target_source, extent_b.target_source);
 }
 
-TEST(liblp, ImportPartitionsFail) {
+TEST_F(BuilderTest, ImportPartitionsFail) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
     ASSERT_NE(builder, nullptr);
 
@@ -693,3 +711,12 @@
     ASSERT_NE(builder, nullptr);
     EXPECT_FALSE(builder->ImportPartitions(*exported.get(), {"system"}));
 }
+
+TEST_F(BuilderTest, UnsuffixedPartitions) {
+    MetadataBuilder::OverrideABForTesting(true);
+    unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
+    ASSERT_NE(builder, nullptr);
+
+    ASSERT_EQ(builder->AddPartition("system", 0), nullptr);
+    ASSERT_NE(builder->AddPartition("system_a", 0), nullptr);
+}
diff --git a/fs_mgr/liblp/include/liblp/builder.h b/fs_mgr/liblp/include/liblp/builder.h
index 297611b..4bb38d6 100644
--- a/fs_mgr/liblp/include/liblp/builder.h
+++ b/fs_mgr/liblp/include/liblp/builder.h
@@ -22,6 +22,7 @@
 
 #include <map>
 #include <memory>
+#include <optional>
 #include <set>
 
 #include "liblp.h"
@@ -186,6 +187,9 @@
         return New(device_info, metadata_max_size, metadata_slot_count);
     }
 
+    // Used by the test harness to override whether the device is "A/B".
+    static void OverrideABForTesting(bool ab_device);
+
     // Define a new partition group. By default there is one group called
     // "default", with an unrestricted size. A non-zero size will restrict the
     // total space used by all partitions in the group.
@@ -270,6 +274,7 @@
     void ImportExtents(Partition* dest, const LpMetadata& metadata,
                        const LpMetadataPartition& source);
     bool ImportPartition(const LpMetadata& metadata, const LpMetadataPartition& source);
+    bool IsABDevice() const;
 
     struct Interval {
         uint32_t device_index;
@@ -290,6 +295,9 @@
     void ExtentsToFreeList(const std::vector<Interval>& extents,
                            std::vector<Interval>* free_regions) const;
 
+    static bool sABOverrideValue;
+    static bool sABOverrideSet;
+
     LpMetadataGeometry geometry_;
     LpMetadataHeader header_;
     std::vector<std::unique_ptr<Partition>> partitions_;