liblp: Add helpers for modifying groups.

These are needed for non-A/B OTAs.

Bug: 122473283
Test: liblp_test gtest
Change-Id: Ib30614f1691dbea0a56c5a98aadc84fc26d1e639
diff --git a/fs_mgr/liblp/builder.cpp b/fs_mgr/liblp/builder.cpp
index 07e3c8a..b99ff8f 100644
--- a/fs_mgr/liblp/builder.cpp
+++ b/fs_mgr/liblp/builder.cpp
@@ -33,6 +33,8 @@
 bool MetadataBuilder::sABOverrideSet;
 bool MetadataBuilder::sABOverrideValue;
 
+static const std::string kDefaultGroup = "default";
+
 bool LinearExtent::AddTo(LpMetadata* out) const {
     if (device_index_ >= out->block_devices.size()) {
         LERROR << "Extent references unknown block device.";
@@ -403,7 +405,7 @@
     geometry_.metadata_slot_count = metadata_slot_count;
     geometry_.logical_block_size = logical_block_size;
 
-    if (!AddGroup("default", 0)) {
+    if (!AddGroup(kDefaultGroup, 0)) {
         return false;
     }
     return true;
@@ -419,7 +421,7 @@
 }
 
 Partition* MetadataBuilder::AddPartition(const std::string& name, uint32_t attributes) {
-    return AddPartition(name, "default", attributes);
+    return AddPartition(name, kDefaultGroup, attributes);
 }
 
 Partition* MetadataBuilder::AddPartition(const std::string& name, const std::string& group_name,
@@ -675,6 +677,10 @@
 }
 
 std::unique_ptr<LpMetadata> MetadataBuilder::Export() {
+    if (!ValidatePartitionGroups()) {
+        return nullptr;
+    }
+
     std::unique_ptr<LpMetadata> metadata = std::make_unique<LpMetadata>();
     metadata->header = header_;
     metadata->geometry = geometry_;
@@ -695,7 +701,7 @@
             LERROR << "Partition group name is too long: " << group->name();
             return nullptr;
         }
-        if (auto_slot_suffixing_ && group->name() != "default") {
+        if (auto_slot_suffixing_ && group->name() != kDefaultGroup) {
             out.flags |= LP_GROUP_SLOT_SUFFIXED;
         }
         strncpy(out.name, group->name().c_str(), sizeof(out.name));
@@ -877,7 +883,7 @@
 }
 
 void MetadataBuilder::RemoveGroupAndPartitions(const std::string& group_name) {
-    if (group_name == "default") {
+    if (group_name == kDefaultGroup) {
         // Cannot remove the default group.
         return;
     }
@@ -1000,5 +1006,53 @@
     return true;
 }
 
+std::vector<Partition*> MetadataBuilder::ListPartitionsInGroup(const std::string& group_name) {
+    std::vector<Partition*> partitions;
+    for (const auto& partition : partitions_) {
+        if (partition->group_name() == group_name) {
+            partitions.emplace_back(partition.get());
+        }
+    }
+    return partitions;
+}
+
+bool MetadataBuilder::ChangePartitionGroup(Partition* partition, const std::string& group_name) {
+    if (!FindGroup(group_name)) {
+        LERROR << "Partition cannot change to unknown group: " << group_name;
+        return false;
+    }
+    partition->set_group_name(group_name);
+    return true;
+}
+
+bool MetadataBuilder::ValidatePartitionGroups() const {
+    for (const auto& group : groups_) {
+        if (!group->maximum_size()) {
+            continue;
+        }
+        uint64_t used = TotalSizeOfGroup(group.get());
+        if (used > group->maximum_size()) {
+            LERROR << "Partition group " << group->name() << " exceeds maximum size (" << used
+                   << " bytes used, maximum " << group->maximum_size() << ")";
+            return false;
+        }
+    }
+    return true;
+}
+
+bool MetadataBuilder::ChangeGroupSize(const std::string& group_name, uint64_t maximum_size) {
+    if (group_name == kDefaultGroup) {
+        LERROR << "Cannot change the size of the default group";
+        return false;
+    }
+    PartitionGroup* group = FindGroup(group_name);
+    if (!group) {
+        LERROR << "Cannot change size of unknown partition group: " << group_name;
+        return false;
+    }
+    group->set_maximum_size(maximum_size);
+    return true;
+}
+
 }  // namespace fs_mgr
 }  // namespace android
diff --git a/fs_mgr/liblp/builder_test.cpp b/fs_mgr/liblp/builder_test.cpp
index 3793964..7d615a3 100644
--- a/fs_mgr/liblp/builder_test.cpp
+++ b/fs_mgr/liblp/builder_test.cpp
@@ -549,6 +549,58 @@
     EXPECT_EQ(partition->size(), 16384);
 }
 
+TEST_F(BuilderTest, ListPartitionsInGroup) {
+    BlockDeviceInfo device_info("super", 1024 * 1024, 0, 0, 4096);
+    unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
+    ASSERT_NE(builder, nullptr);
+
+    ASSERT_TRUE(builder->AddGroup("groupA", 16384));
+    ASSERT_TRUE(builder->AddGroup("groupB", 16384));
+
+    Partition* system = builder->AddPartition("system", "groupA", 0);
+    Partition* vendor = builder->AddPartition("vendor", "groupA", 0);
+    Partition* product = builder->AddPartition("product", "groupB", 0);
+    ASSERT_NE(system, nullptr);
+    ASSERT_NE(vendor, nullptr);
+    ASSERT_NE(product, nullptr);
+
+    auto groupA = builder->ListPartitionsInGroup("groupA");
+    auto groupB = builder->ListPartitionsInGroup("groupB");
+    auto groupC = builder->ListPartitionsInGroup("groupC");
+    ASSERT_THAT(groupA, ElementsAre(system, vendor));
+    ASSERT_THAT(groupB, ElementsAre(product));
+    ASSERT_TRUE(groupC.empty());
+}
+
+TEST_F(BuilderTest, ChangeGroups) {
+    BlockDeviceInfo device_info("super", 1024 * 1024, 0, 0, 4096);
+    unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
+    ASSERT_NE(builder, nullptr);
+
+    ASSERT_TRUE(builder->AddGroup("groupA", 16384));
+    ASSERT_TRUE(builder->AddGroup("groupB", 32768));
+
+    Partition* system = builder->AddPartition("system", "groupA", 0);
+    Partition* vendor = builder->AddPartition("vendor", "groupB", 0);
+    ASSERT_NE(system, nullptr);
+    ASSERT_NE(vendor, nullptr);
+    ASSERT_NE(builder->Export(), nullptr);
+
+    ASSERT_FALSE(builder->ChangePartitionGroup(system, "groupXYZ"));
+    ASSERT_TRUE(builder->ChangePartitionGroup(system, "groupB"));
+    ASSERT_NE(builder->Export(), nullptr);
+
+    // Violate group constraint by reassigning groups.
+    ASSERT_TRUE(builder->ResizePartition(system, 16384 + 4096));
+    ASSERT_TRUE(builder->ChangePartitionGroup(system, "groupA"));
+    ASSERT_EQ(builder->Export(), nullptr);
+
+    ASSERT_FALSE(builder->ChangeGroupSize("default", 2));
+    ASSERT_FALSE(builder->ChangeGroupSize("unknown", 2));
+    ASSERT_TRUE(builder->ChangeGroupSize("groupA", 32768));
+    ASSERT_NE(builder->Export(), nullptr);
+}
+
 constexpr unsigned long long operator"" _GiB(unsigned long long x) {  // NOLINT
     return x << 30;
 }
diff --git a/fs_mgr/liblp/include/liblp/builder.h b/fs_mgr/liblp/include/liblp/builder.h
index 57cce21..53f480f 100644
--- a/fs_mgr/liblp/include/liblp/builder.h
+++ b/fs_mgr/liblp/include/liblp/builder.h
@@ -80,6 +80,8 @@
 };
 
 class PartitionGroup final {
+    friend class MetadataBuilder;
+
   public:
     explicit PartitionGroup(const std::string& name, uint64_t maximum_size)
         : name_(name), maximum_size_(maximum_size) {}
@@ -88,6 +90,8 @@
     uint64_t maximum_size() const { return maximum_size_; }
 
   private:
+    void set_maximum_size(uint64_t maximum_size) { maximum_size_ = maximum_size; }
+
     std::string name_;
     uint64_t maximum_size_;
 };
@@ -116,6 +120,7 @@
 
   private:
     void ShrinkTo(uint64_t aligned_size);
+    void set_group_name(const std::string& group_name) { group_name_ = group_name; }
 
     std::string name_;
     std::string group_name_;
@@ -235,6 +240,21 @@
     // underlying filesystem or contents of the partition on disk.
     bool ResizePartition(Partition* partition, uint64_t requested_size);
 
+    // Return the list of partitions belonging to a group.
+    std::vector<Partition*> ListPartitionsInGroup(const std::string& group_name);
+
+    // Changes a partition's group. Size constraints will not be checked until
+    // the metadata is exported, to avoid errors during potential group and
+    // size shuffling operations. This will return false if the new group does
+    // not exist.
+    bool ChangePartitionGroup(Partition* partition, const std::string& group_name);
+
+    // Changes the size of a partition group. Size constraints will not be
+    // checked until metadata is exported, to avoid errors during group
+    // reshuffling. This will return false if the group does not exist, or if
+    // the group name is "default".
+    bool ChangeGroupSize(const std::string& group_name, uint64_t maximum_size);
+
     // Amount of space that can be allocated to logical partitions.
     uint64_t AllocatableSpace() const;
     uint64_t UsedSpace() const;
@@ -283,6 +303,7 @@
     bool ImportPartition(const LpMetadata& metadata, const LpMetadataPartition& source);
     bool IsABDevice() const;
     bool IsRetrofitDevice() const;
+    bool ValidatePartitionGroups() const;
 
     struct Interval {
         uint32_t device_index;