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

1 files changed, 54 insertions, 30 deletions

diff --git a/aosp/dynamic_partition_control_android.cc b/aosp/dynamic_partition_control_android.cc
index 1db8da8d..444fe42c 100644
--- a/aosp/dynamic_partition_control_android.cc
+++ b/aosp/dynamic_partition_control_android.cc
@@ -805,9 +805,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();
@@ -824,21 +822,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;
@@ -877,29 +922,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;
@@ -1167,7 +1191,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);