liblp: Add support for updateable partition groups.

This patch introduces a new "groups" table in the super partition
metadata. Each entry denotes a named partition group with a maximum
size. All partitions now belong to a group, and the total size of
partitions in a group must not exceed its maximum size. This is enforced
by MetadataBuilder.

There is also a "default" group with no size restriction. This is used
for one-off partitions that aren't restricted by updates, for example,
the scratch partition for overlayfs, or partitions created through
fastbootd.

Bug: 116817738
Test: liblp gtest
Change-Id: I7049ffd35d326e41e25d01b1748cb53a584783a7

8 files changed, 214 insertions, 16 deletions

diff --git a/fs_mgr/liblp/builder.cpp b/fs_mgr/liblp/builder.cpp
index 352647b8e..d9a0e9c55 100644
--- a/fs_mgr/liblp/builder.cpp
+++ b/fs_mgr/liblp/builder.cpp
@@ -79,8 +79,9 @@ void ZeroExtent::AddTo(LpMetadata* out) const {
     out->extents.push_back(LpMetadataExtent{num_sectors_, LP_TARGET_TYPE_ZERO, 0});
 }
 
-Partition::Partition(const std::string& name, const std::string& guid, uint32_t attributes)
-    : name_(name), guid_(guid), attributes_(attributes), size_(0) {}
+Partition::Partition(const std::string& name, const std::string& group_name,
+                     const std::string& guid, uint32_t attributes)
+    : name_(name), group_name_(group_name), guid_(guid), attributes_(attributes), size_(0) {}
 
 void Partition::AddExtent(std::unique_ptr<Extent>&& extent) {
     size_ += extent->num_sectors() * LP_SECTOR_SIZE;
@@ -128,6 +129,17 @@ void Partition::ShrinkTo(uint64_t aligned_size) {
     DCHECK(size_ == aligned_size);
 }
 
+uint64_t Partition::BytesOnDisk() const {
+    uint64_t sectors = 0;
+    for (const auto& extent : extents_) {
+        if (!extent->AsLinearExtent()) {
+            continue;
+        }
+        sectors += extent->num_sectors();
+    }
+    return sectors * LP_SECTOR_SIZE;
+}
+
 std::unique_ptr<MetadataBuilder> MetadataBuilder::New(const std::string& block_device,
                                                       uint32_t slot_number) {
     std::unique_ptr<LpMetadata> metadata = ReadMetadata(block_device.c_str(), slot_number);
@@ -175,14 +187,23 @@ MetadataBuilder::MetadataBuilder() {
     header_.header_size = sizeof(header_);
     header_.partitions.entry_size = sizeof(LpMetadataPartition);
     header_.extents.entry_size = sizeof(LpMetadataExtent);
+    header_.groups.entry_size = sizeof(LpMetadataPartitionGroup);
 }
 
 bool MetadataBuilder::Init(const LpMetadata& metadata) {
     geometry_ = metadata.geometry;
 
+    for (const auto& group : metadata.groups) {
+        std::string group_name = GetPartitionGroupName(group);
+        if (!AddGroup(group_name, group.maximum_size)) {
+            return false;
+        }
+    }
+
     for (const auto& partition : metadata.partitions) {
-        Partition* builder = AddPartition(GetPartitionName(partition), GetPartitionGuid(partition),
-                                          partition.attributes);
+        std::string group_name = GetPartitionGroupName(metadata.groups[partition.group_index]);
+        Partition* builder = AddPartition(GetPartitionName(partition), group_name,
+                                          GetPartitionGuid(partition), partition.attributes);
         if (!builder) {
             return false;
         }
@@ -292,11 +313,29 @@ bool MetadataBuilder::Init(const BlockDeviceInfo& device_info, uint32_t metadata
     geometry_.alignment_offset = device_info_.alignment_offset;
     geometry_.block_device_size = device_info_.size;
     geometry_.logical_block_size = device_info.logical_block_size;
+
+    if (!AddGroup("default", 0)) {
+        return false;
+    }
+    return true;
+}
+
+bool MetadataBuilder::AddGroup(const std::string& group_name, uint64_t maximum_size) {
+    if (FindGroup(group_name)) {
+        LERROR << "Group already exists: " << group_name;
+        return false;
+    }
+    groups_.push_back(std::make_unique<PartitionGroup>(group_name, maximum_size));
     return true;
 }
 
 Partition* MetadataBuilder::AddPartition(const std::string& name, const std::string& guid,
                                          uint32_t attributes) {
+    return AddPartition(name, "default", guid, attributes);
+}
+
+Partition* MetadataBuilder::AddPartition(const std::string& name, const std::string& group_name,
+                                         const std::string& guid, uint32_t attributes) {
     if (name.empty()) {
         LERROR << "Partition must have a non-empty name.";
         return nullptr;
@@ -305,7 +344,11 @@ Partition* MetadataBuilder::AddPartition(const std::string& name, const std::str
         LERROR << "Attempting to create duplication partition with name: " << name;
         return nullptr;
     }
-    partitions_.push_back(std::make_unique<Partition>(name, guid, attributes));
+    if (!FindGroup(group_name)) {
+        LERROR << "Could not find partition group: " << group_name;
+        return nullptr;
+    }
+    partitions_.push_back(std::make_unique<Partition>(name, group_name, guid, attributes));
     return partitions_.back().get();
 }
 
@@ -318,6 +361,26 @@ Partition* MetadataBuilder::FindPartition(const std::string& name) {
     return nullptr;
 }
 
+PartitionGroup* MetadataBuilder::FindGroup(const std::string& group_name) const {
+    for (const auto& group : groups_) {
+        if (group->name() == group_name) {
+            return group.get();
+        }
+    }
+    return nullptr;
+}
+
+uint64_t MetadataBuilder::TotalSizeOfGroup(PartitionGroup* group) const {
+    uint64_t total = 0;
+    for (const auto& partition : partitions_) {
+        if (partition->group_name() != group->name()) {
+            continue;
+        }
+        total += partition->BytesOnDisk();
+    }
+    return total;
+}
+
 void MetadataBuilder::RemovePartition(const std::string& name) {
     for (auto iter = partitions_.begin(); iter != partitions_.end(); iter++) {
         if ((*iter)->name() == name) {
@@ -328,8 +391,23 @@ void MetadataBuilder::RemovePartition(const std::string& name) {
 }
 
 bool MetadataBuilder::GrowPartition(Partition* partition, uint64_t aligned_size) {
-    // Figure out how much we need to allocate.
+    PartitionGroup* group = FindGroup(partition->group_name());
+    CHECK(group);
+
+    // Figure out how much we need to allocate, and whether our group has
+    // enough space remaining.
     uint64_t space_needed = aligned_size - partition->size();
+    if (group->maximum_size() > 0) {
+        uint64_t group_size = TotalSizeOfGroup(group);
+        if (group_size >= group->maximum_size() ||
+            group->maximum_size() - group_size < space_needed) {
+            LERROR << "Partition " << partition->name() << " is part of group " << group->name()
+                   << " which does not have enough space free (" << space_needed << "requested, "
+                   << group_size << " used out of " << group->maximum_size();
+            return false;
+        }
+    }
+
     uint64_t sectors_needed = space_needed / LP_SECTOR_SIZE;
     DCHECK(sectors_needed * LP_SECTOR_SIZE == space_needed);
 
@@ -441,6 +519,20 @@ std::unique_ptr<LpMetadata> MetadataBuilder::Export() {
     metadata->header = header_;
     metadata->geometry = geometry_;
 
+    std::map<std::string, size_t> group_indices;
+    for (const auto& group : groups_) {
+        LpMetadataPartitionGroup out = {};
+
+        if (group->name().size() > sizeof(out.name)) {
+            LERROR << "Partition group name is too long: " << group->name();
+            return nullptr;
+        }
+        strncpy(out.name, group->name().c_str(), sizeof(out.name));
+        out.maximum_size = group->maximum_size();
+
+        metadata->groups.push_back(out);
+    }
+
     // Flatten the partition and extent structures into an LpMetadata, which
     // makes it very easy to validate, serialize, or pass on to device-mapper.
     for (const auto& partition : partitions_) {
@@ -475,6 +567,7 @@ std::unique_ptr<LpMetadata> MetadataBuilder::Export() {
 
     metadata->header.partitions.num_entries = static_cast<uint32_t>(metadata->partitions.size());
     metadata->header.extents.num_entries = static_cast<uint32_t>(metadata->extents.size());
+    metadata->header.groups.num_entries = static_cast<uint32_t>(metadata->groups.size());
     return metadata;
 }
 
@@ -530,8 +623,10 @@ bool MetadataBuilder::ResizePartition(Partition* partition, uint64_t requested_s
         ShrinkPartition(partition, aligned_size);
     }
 
-    LINFO << "Partition " << partition->name() << " will resize from " << old_size << " bytes to "
-          << aligned_size << " bytes";
+    if (partition->size() != old_size) {
+        LINFO << "Partition " << partition->name() << " will resize from " << old_size
+              << " bytes to " << aligned_size << " bytes";
+    }
     return true;
 }
 
diff --git a/fs_mgr/liblp/builder_test.cpp b/fs_mgr/liblp/builder_test.cpp
index 0c7e43d24..711cc64ba 100644
--- a/fs_mgr/liblp/builder_test.cpp
+++ b/fs_mgr/liblp/builder_test.cpp
@@ -495,3 +495,28 @@ TEST(liblp, AlignedFreeSpace) {
     unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 512, 1);
     ASSERT_EQ(builder, nullptr);
 }
+
+TEST(liblp, HasDefaultGroup) {
+    BlockDeviceInfo device_info(1024 * 1024, 0, 0, 4096);
+    unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
+    ASSERT_NE(builder, nullptr);
+
+    EXPECT_FALSE(builder->AddGroup("default", 0));
+}
+
+TEST(liblp, GroupSizeLimits) {
+    BlockDeviceInfo device_info(1024 * 1024, 0, 0, 4096);
+    unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(device_info, 1024, 1);
+    ASSERT_NE(builder, nullptr);
+
+    ASSERT_TRUE(builder->AddGroup("google", 16384));
+
+    Partition* partition = builder->AddPartition("system", "google", TEST_GUID, 0);
+    ASSERT_NE(partition, nullptr);
+    EXPECT_TRUE(builder->ResizePartition(partition, 8192));
+    EXPECT_EQ(partition->size(), 8192);
+    EXPECT_TRUE(builder->ResizePartition(partition, 16384));
+    EXPECT_EQ(partition->size(), 16384);
+    EXPECT_FALSE(builder->ResizePartition(partition, 32768));
+    EXPECT_EQ(partition->size(), 16384);
+}
diff --git a/fs_mgr/liblp/include/liblp/builder.h b/fs_mgr/liblp/include/liblp/builder.h
index 27808256f..f1edee7ea 100644
--- a/fs_mgr/liblp/include/liblp/builder.h
+++ b/fs_mgr/liblp/include/liblp/builder.h
@@ -95,11 +95,25 @@ class ZeroExtent final : public Extent {
     void AddTo(LpMetadata* out) const override;
 };
 
+class PartitionGroup final {
+  public:
+    explicit PartitionGroup(const std::string& name, uint64_t maximum_size)
+        : name_(name), maximum_size_(maximum_size) {}
+
+    const std::string& name() const { return name_; }
+    uint64_t maximum_size() const { return maximum_size_; }
+
+  private:
+    std::string name_;
+    uint64_t maximum_size_;
+};
+
 class Partition final {
     friend class MetadataBuilder;
 
   public:
-    Partition(const std::string& name, const std::string& guid, uint32_t attributes);
+    Partition(const std::string& name, const std::string& group_name, const std::string& guid,
+              uint32_t attributes);
 
     // Add a raw extent.
     void AddExtent(std::unique_ptr<Extent>&& extent);
@@ -107,7 +121,12 @@ class Partition final {
     // Remove all extents from this partition.
     void RemoveExtents();
 
+    // Compute the size used by linear extents. This is the same as size(),
+    // but does not factor in extents which do not take up space.
+    uint64_t BytesOnDisk() const;
+
     const std::string& name() const { return name_; }
+    const std::string& group_name() const { return group_name_; }
     uint32_t attributes() const { return attributes_; }
     const std::string& guid() const { return guid_; }
     const std::vector<std::unique_ptr<Extent>>& extents() const { return extents_; }
@@ -117,6 +136,7 @@ class Partition final {
     void ShrinkTo(uint64_t aligned_size);
 
     std::string name_;
+    std::string group_name_;
     std::string guid_;
     std::vector<std::unique_ptr<Extent>> extents_;
     uint32_t attributes_;
@@ -156,12 +176,24 @@ class MetadataBuilder {
         return New(device_info, metadata_max_size, metadata_slot_count);
     }
 
+    // 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.
+    //
+    // This can fail and return false if the group already exists.
+    bool AddGroup(const std::string& group_name, uint64_t maximum_size);
+
     // Export metadata so it can be serialized to an image, to disk, or mounted
     // via device-mapper.
     std::unique_ptr<LpMetadata> Export();
 
     // Add a partition, returning a handle so it can be sized as needed. If a
     // partition with the given name already exists, nullptr is returned.
+    Partition* AddPartition(const std::string& name, const std::string& group_name,
+                            const std::string& guid, uint32_t attributes);
+
+    // Same as AddPartition above, but uses the default partition group which
+    // has no size restrictions.
     Partition* AddPartition(const std::string& name, const std::string& guid, uint32_t attributes);
 
     // Delete a partition by name if it exists.
@@ -202,10 +234,13 @@ class MetadataBuilder {
     bool GrowPartition(Partition* partition, uint64_t aligned_size);
     void ShrinkPartition(Partition* partition, uint64_t aligned_size);
     uint64_t AlignSector(uint64_t sector);
+    PartitionGroup* FindGroup(const std::string& group_name) const;
+    uint64_t TotalSizeOfGroup(PartitionGroup* group) const;
 
     LpMetadataGeometry geometry_;
     LpMetadataHeader header_;
     std::vector<std::unique_ptr<Partition>> partitions_;
+    std::vector<std::unique_ptr<PartitionGroup>> groups_;
     BlockDeviceInfo device_info_;
 };
 
diff --git a/fs_mgr/liblp/include/liblp/liblp.h b/fs_mgr/liblp/include/liblp/liblp.h
index 6da24f6ac..51d262b37 100644
--- a/fs_mgr/liblp/include/liblp/liblp.h
+++ b/fs_mgr/liblp/include/liblp/liblp.h
@@ -36,6 +36,7 @@ struct LpMetadata {
     LpMetadataHeader header;
     std::vector<LpMetadataPartition> partitions;
     std::vector<LpMetadataExtent> extents;
+    std::vector<LpMetadataPartitionGroup> groups;
 };
 
 // Place an initial partition table on the device. This will overwrite the
@@ -68,6 +69,7 @@ std::unique_ptr<LpMetadata> ReadFromImageBlob(const void* data, size_t bytes);
 // Helper to extract safe C++ strings from partition info.
 std::string GetPartitionName(const LpMetadataPartition& partition);
 std::string GetPartitionGuid(const LpMetadataPartition& partition);
+std::string GetPartitionGroupName(const LpMetadataPartitionGroup& group);
 
 // Helper to return a slot number for a slot suffix.
 uint32_t SlotNumberForSlotSuffix(const std::string& suffix);
diff --git a/fs_mgr/liblp/include/liblp/metadata_format.h b/fs_mgr/liblp/include/liblp/metadata_format.h
index 52c80f7f4..79ef2ab58 100644
--- a/fs_mgr/liblp/include/liblp/metadata_format.h
+++ b/fs_mgr/liblp/include/liblp/metadata_format.h
@@ -38,7 +38,7 @@ extern "C" {
 #define LP_METADATA_HEADER_MAGIC 0x414C5030
 
 /* Current metadata version. */
-#define LP_METADATA_MAJOR_VERSION 1
+#define LP_METADATA_MAJOR_VERSION 2
 #define LP_METADATA_MINOR_VERSION 0
 
 /* Attributes for the LpMetadataPartition::attributes field.
@@ -216,6 +216,8 @@ typedef struct LpMetadataHeader {
     LpMetadataTableDescriptor partitions;
     /* 92: Extent table descriptor. */
     LpMetadataTableDescriptor extents;
+    /* 104: Updateable group descriptor. */
+    LpMetadataTableDescriptor groups;
 } __attribute__((packed)) LpMetadataHeader;
 
 /* This struct defines a logical partition entry, similar to what would be
@@ -245,6 +247,9 @@ typedef struct LpMetadataPartition {
      * least one extent.
      */
     uint32_t num_extents;
+
+    /* 64: Group this partition belongs to. */
+    uint32_t group_index;
 } __attribute__((packed)) LpMetadataPartition;
 
 /* This extent is a dm-linear target, and the index is an index into the
@@ -271,6 +276,19 @@ typedef struct LpMetadataExtent {
     uint64_t target_data;
 } __attribute__((packed)) LpMetadataExtent;
 
+/* This struct defines an entry in the groups table. Each group has a maximum
+ * size, and partitions in a group must not exceed that size. There is always
+ * a "default" group of unlimited size, which is used when not using update
+ * groups or when using overlayfs or fastbootd.
+ */
+typedef struct LpMetadataPartitionGroup {
+    /*  0: Name of this group. Any unused characters must be 0. */
+    char name[36];
+
+    /* 36: Maximum size in bytes. If 0, the group has no maximum size. */
+    uint64_t maximum_size;
+} LpMetadataPartitionGroup;
+
 #ifdef __cplusplus
 } /* extern "C" */
 #endif
diff --git a/fs_mgr/liblp/io_test.cpp b/fs_mgr/liblp/io_test.cpp
index eda68fd66..f93852b1d 100644
--- a/fs_mgr/liblp/io_test.cpp
+++ b/fs_mgr/liblp/io_test.cpp
@@ -343,9 +343,6 @@ TEST(liblp, TooManyPartitions) {
         ASSERT_NE(partition, nullptr);
     }
     ASSERT_NE(partition, nullptr);
-    // Add one extent to any partition to fill up more space - we're at 508
-    // bytes after this, out of 512.
-    ASSERT_TRUE(builder->ResizePartition(partition, 1024));
 
     unique_ptr<LpMetadata> exported = builder->Export();
     ASSERT_NE(exported, nullptr);
diff --git a/fs_mgr/liblp/reader.cpp b/fs_mgr/liblp/reader.cpp
index 190c65098..005d493a5 100644
--- a/fs_mgr/liblp/reader.cpp
+++ b/fs_mgr/liblp/reader.cpp
@@ -188,7 +188,8 @@ static bool ValidateMetadataHeader(const LpMetadataHeader& header) {
         return false;
     }
     if (!ValidateTableBounds(header, header.partitions) ||
-        !ValidateTableBounds(header, header.extents)) {
+        !ValidateTableBounds(header, header.extents) ||
+        !ValidateTableBounds(header, header.groups)) {
         LERROR << "Logical partition metadata has invalid table bounds.";
         return false;
     }
@@ -202,6 +203,10 @@ static bool ValidateMetadataHeader(const LpMetadataHeader& header) {
         LERROR << "Logical partition metadata has invalid extent table entry size.";
         return false;
     }
+    if (header.groups.entry_size != sizeof(LpMetadataPartitionGroup)) {
+        LERROR << "Logical partition metadata has invalid group table entry size.";
+        return false;
+    }
     return true;
 }
 
@@ -257,6 +262,10 @@ static std::unique_ptr<LpMetadata> ParseMetadata(const LpMetadataGeometry& geome
             LERROR << "Logical partition has invalid extent list.";
             return nullptr;
         }
+        if (partition.group_index >= header.groups.num_entries) {
+            LERROR << "Logical partition has invalid group index.";
+            return nullptr;
+        }
 
         metadata->partitions.push_back(partition);
     }
@@ -269,6 +278,16 @@ static std::unique_ptr<LpMetadata> ParseMetadata(const LpMetadataGeometry& geome
 
         metadata->extents.push_back(extent);
     }
+
+    cursor = buffer.get() + header.groups.offset;
+    for (size_t i = 0; i < header.groups.num_entries; i++) {
+        LpMetadataPartitionGroup group = {};
+        memcpy(&group, cursor, sizeof(group));
+        cursor += header.groups.entry_size;
+
+        metadata->groups.push_back(group);
+    }
+
     return metadata;
 }
 
@@ -345,5 +364,9 @@ std::string GetPartitionName(const LpMetadataPartition& partition) {
     return NameFromFixedArray(partition.name, sizeof(partition.name));
 }
 
+std::string GetPartitionGroupName(const LpMetadataPartitionGroup& group) {
+    return NameFromFixedArray(group.name, sizeof(group.name));
+}
+
 }  // namespace fs_mgr
 }  // namespace android
diff --git a/fs_mgr/liblp/writer.cpp b/fs_mgr/liblp/writer.cpp
index 9dd274511..241562999 100644
--- a/fs_mgr/liblp/writer.cpp
+++ b/fs_mgr/liblp/writer.cpp
@@ -56,14 +56,17 @@ std::string SerializeMetadata(const LpMetadata& input) {
                            metadata.partitions.size() * sizeof(LpMetadataPartition));
     std::string extents(reinterpret_cast<const char*>(metadata.extents.data()),
                         metadata.extents.size() * sizeof(LpMetadataExtent));
+    std::string groups(reinterpret_cast<const char*>(metadata.groups.data()),
+                       metadata.groups.size() * sizeof(LpMetadataPartitionGroup));
 
     // Compute positions of tables.
     header.partitions.offset = 0;
     header.extents.offset = header.partitions.offset + partitions.size();
-    header.tables_size = header.extents.offset + extents.size();
+    header.groups.offset = header.extents.offset + extents.size();
+    header.tables_size = header.groups.offset + groups.size();
 
     // Compute payload checksum.
-    std::string tables = partitions + extents;
+    std::string tables = partitions + extents + groups;
     SHA256(tables.data(), tables.size(), header.tables_checksum);
 
     // Compute header checksum.