Check the super partiton size in VAB case

When the snapshot is used, we should check that the maximum size
of all dynamic partition groups doesn't exceed the super partition
size.

Bug: 182431975
Test: primary payload fails as expected Pixel21
Change-Id: I5df8976e6b7e011284b29fd554dda80e31305698

3 files changed, 68 insertions, 30 deletions

diff --git a/dynamic_partition_control_android.cc b/dynamic_partition_control_android.cc
index 3103a381..5f688a26 100644
--- a/dynamic_partition_control_android.cc
+++ b/dynamic_partition_control_android.cc
@@ -721,9 +721,7 @@ bool DynamicPartitionControlAndroid::PrepareDynamicPartitionsForUpdate(
   }
 
   std::string device_dir_str;
-  if (!GetDeviceDir(&device_dir_str)) {
-    return false;
-  }
+  TEST_AND_RETURN_FALSE(GetDeviceDir(&device_dir_str));
   base::FilePath device_dir(device_dir_str);
   auto source_device =
       device_dir.Append(GetSuperPartitionName(source_slot)).value();
@@ -740,21 +738,68 @@ bool DynamicPartitionControlAndroid::PrepareDynamicPartitionsForUpdate(
         DeleteSourcePartitions(builder.get(), source_slot, manifest));
   }
 
-  if (!UpdatePartitionMetadata(builder.get(), target_slot, manifest)) {
-    return false;
-  }
+  TEST_AND_RETURN_FALSE(
+      UpdatePartitionMetadata(builder.get(), target_slot, manifest));
 
   auto target_device =
       device_dir.Append(GetSuperPartitionName(target_slot)).value();
   return StoreMetadata(target_device, builder.get(), target_slot);
 }
 
+bool DynamicPartitionControlAndroid::CheckSuperPartitionAllocatableSpace(
+    android::fs_mgr::MetadataBuilder* builder,
+    const DeltaArchiveManifest& manifest,
+    bool use_snapshot) {
+  uint64_t total_size = 0;
+  for (const auto& group : manifest.dynamic_partition_metadata().groups()) {
+    total_size += group.size();
+  }
+
+  std::string expr;
+  uint64_t allocatable_space = builder->AllocatableSpace();
+  // On device retrofitting dynamic partitions, allocatable_space = super.
+  // On device launching dynamic partitions w/o VAB,
+  //   allocatable_space = super / 2.
+  // On device launching dynamic partitions with VAB, allocatable_space = super.
+  // For recovery sideload, allocatable_space = super.
+  if (!GetDynamicPartitionsFeatureFlag().IsRetrofit() && !use_snapshot &&
+      !IsRecovery()) {
+    allocatable_space /= 2;
+    expr = "half of ";
+  }
+  if (total_size > allocatable_space) {
+    LOG(ERROR) << "The maximum size of all groups for the target slot"
+               << " (" << total_size << ") has exceeded " << expr
+               << "allocatable space for dynamic partitions "
+               << allocatable_space << ".";
+    return false;
+  }
+
+  return true;
+}
+
 bool DynamicPartitionControlAndroid::PrepareSnapshotPartitionsForUpdate(
     uint32_t source_slot,
     uint32_t target_slot,
     const DeltaArchiveManifest& manifest,
     uint64_t* required_size) {
   TEST_AND_RETURN_FALSE(ExpectMetadataMounted());
+
+  std::string device_dir_str;
+  TEST_AND_RETURN_FALSE(GetDeviceDir(&device_dir_str));
+  base::FilePath device_dir(device_dir_str);
+  auto super_device =
+      device_dir.Append(GetSuperPartitionName(source_slot)).value();
+  auto builder = LoadMetadataBuilder(super_device, source_slot);
+  if (builder == nullptr) {
+    LOG(ERROR) << "No metadata at "
+               << BootControlInterface::SlotName(source_slot);
+    return false;
+  }
+
+  TEST_AND_RETURN_FALSE(
+      CheckSuperPartitionAllocatableSpace(builder.get(), manifest, true));
+
   if (!snapshot_->BeginUpdate()) {
     LOG(ERROR) << "Cannot begin new update.";
     return false;
@@ -793,29 +838,8 @@ bool DynamicPartitionControlAndroid::UpdatePartitionMetadata(
   const std::string target_suffix = SlotSuffixForSlotNumber(target_slot);
   DeleteGroupsWithSuffix(builder, target_suffix);
 
-  uint64_t total_size = 0;
-  for (const auto& group : manifest.dynamic_partition_metadata().groups()) {
-    total_size += group.size();
-  }
-
-  std::string expr;
-  uint64_t allocatable_space = builder->AllocatableSpace();
-  // On device retrofitting dynamic partitions, allocatable_space = super.
-  // On device launching dynamic partitions w/o VAB,
-  //   allocatable_space = super / 2.
-  // On device launching dynamic partitions with VAB, allocatable_space = super.
-  if (!GetDynamicPartitionsFeatureFlag().IsRetrofit() &&
-      !GetVirtualAbFeatureFlag().IsEnabled()) {
-    allocatable_space /= 2;
-    expr = "half of ";
-  }
-  if (total_size > allocatable_space) {
-    LOG(ERROR) << "The maximum size of all groups with suffix " << target_suffix
-               << " (" << total_size << ") has exceeded " << expr
-               << "allocatable space for dynamic partitions "
-               << allocatable_space << ".";
-    return false;
-  }
+  TEST_AND_RETURN_FALSE(
+      CheckSuperPartitionAllocatableSpace(builder, manifest, false));
 
   // name of partition(e.g. "system") -> size in bytes
   std::map<std::string, uint64_t> partition_sizes;
@@ -1015,7 +1039,7 @@ bool DynamicPartitionControlAndroid::DeleteSourcePartitions(
 
   LOG(INFO) << "Will overwrite existing partitions. Slot "
             << BootControlInterface::SlotName(source_slot)
-            << "may be unbootable until update finishes!";
+            << " may be unbootable until update finishes!";
   const std::string source_suffix = SlotSuffixForSlotNumber(source_slot);
   DeleteGroupsWithSuffix(builder, source_suffix);
 
diff --git a/dynamic_partition_control_android.h b/dynamic_partition_control_android.h
index c2bec1dc..28699643 100644
--- a/dynamic_partition_control_android.h
+++ b/dynamic_partition_control_android.h
@@ -208,6 +208,13 @@ class DynamicPartitionControlAndroid : public DynamicPartitionControlInterface {
                                           const DeltaArchiveManifest& manifest,
                                           uint64_t* required_size);
 
+  // Returns true if the allocatable space in super partition is larger than
+  // the size of dynamic partition groups in the manifest.
+  bool CheckSuperPartitionAllocatableSpace(
+      android::fs_mgr::MetadataBuilder* builder,
+      const DeltaArchiveManifest& manifest,
+      bool use_snapshot);
+
   enum class DynamicPartitionDeviceStatus {
     SUCCESS,
     ERROR,
diff --git a/dynamic_partition_control_android_unittest.cc b/dynamic_partition_control_android_unittest.cc
index ccc4a36b..de673639 100644
--- a/dynamic_partition_control_android_unittest.cc
+++ b/dynamic_partition_control_android_unittest.cc
@@ -938,6 +938,7 @@ class SnapshotPartitionTestP : public DynamicPartitionControlAndroidTestP {
 // Test happy path of PreparePartitionsForUpdate on a Virtual A/B device.
 TEST_P(SnapshotPartitionTestP, PreparePartitions) {
   ExpectCreateUpdateSnapshots(android::snapshot::Return::Ok());
+  SetMetadata(source(), {});
   uint64_t required_size = 0;
   EXPECT_TRUE(PreparePartitionsForUpdate(&required_size));
   EXPECT_EQ(0u, required_size);
@@ -948,6 +949,8 @@ TEST_P(SnapshotPartitionTestP, PreparePartitions) {
 TEST_P(SnapshotPartitionTestP, PreparePartitionsNoSpace) {
   ExpectCreateUpdateSnapshots(android::snapshot::Return::NoSpace(1_GiB));
   uint64_t required_size = 0;
+
+  SetMetadata(source(), {});
   EXPECT_FALSE(PreparePartitionsForUpdate(&required_size));
   EXPECT_EQ(1_GiB, required_size);
 }
@@ -957,6 +960,10 @@ TEST_P(SnapshotPartitionTestP, PreparePartitionsNoSpace) {
 TEST_P(SnapshotPartitionTestP, RecoveryUseSuperEmpty) {
   ExpectCreateUpdateSnapshots(android::snapshot::Return::Ok());
   EXPECT_CALL(dynamicControl(), IsRecovery()).WillRepeatedly(Return(true));
+
+  // Metadata is needed to perform super partition size check.
+  SetMetadata(source(), {});
+
   // Must not call PrepareDynamicPartitionsForUpdate if
   // PrepareSnapshotPartitionsForUpdate succeeds.
   EXPECT_CALL(dynamicControl(), PrepareDynamicPartitionsForUpdate(_, _, _, _))