Re-map source partitions with RW permission during OTA

Test: th
Bug: 351928254
Change-Id: Ia58b62f038b3bd0d73b429f761f4ff05ab93e072
diff --git a/aosp/dynamic_partition_control_android.cc b/aosp/dynamic_partition_control_android.cc
index 78cda20..363a433 100644
--- a/aosp/dynamic_partition_control_android.cc
+++ b/aosp/dynamic_partition_control_android.cc
@@ -19,6 +19,7 @@
 #include <algorithm>
 #include <chrono>  // NOLINT(build/c++11) - using libsnapshot / liblp API
 #include <cstdint>
+#include <iterator>
 #include <map>
 #include <memory>
 #include <set>
@@ -100,7 +101,11 @@
 constexpr std::chrono::milliseconds kMapSnapshotTimeout{10000};
 
 DynamicPartitionControlAndroid::~DynamicPartitionControlAndroid() {
-  UnmapAllPartitions();
+  std::set<std::string> mapped = mapped_devices_;
+  LOG(INFO) << "Destroying [" << Join(mapped, ", ") << "] from device mapper";
+  for (const auto& device_name : mapped) {
+    ignore_result(UnmapPartitionOnDeviceMapper(device_name));
+  }
   metadata_device_.reset();
 }
 
@@ -184,18 +189,34 @@
   return false;
 }
 
+constexpr auto&& kRWSourcePartitionSuffix = "_ota";
+std::string DynamicPartitionControlAndroid::GetDeviceName(
+    std::string partition_name, uint32_t slot) const {
+  if (partition_name.ends_with(kRWSourcePartitionSuffix)) {
+    return partition_name;
+  }
+  if (!partition_name.ends_with("_a") && !partition_name.ends_with("_b")) {
+    partition_name += slot ? "_b" : "_a";
+  }
+  if (slot == source_slot_) {
+    return partition_name + kRWSourcePartitionSuffix;
+  }
+  return partition_name;
+}
+
 bool DynamicPartitionControlAndroid::MapPartitionInternal(
     const std::string& super_device,
     const std::string& target_partition_name,
     uint32_t slot,
     bool force_writable,
     std::string* path) {
+  auto device_name = GetDeviceName(target_partition_name, slot);
   CreateLogicalPartitionParams params = {
       .block_device = super_device,
       .metadata_slot = slot,
       .partition_name = target_partition_name,
       .force_writable = force_writable,
-  };
+      .device_name = device_name};
   bool success = false;
   if (GetVirtualAbFeatureFlag().IsEnabled() && target_supports_snapshot_ &&
       slot != source_slot_ && force_writable && ExpectMetadataMounted()) {
@@ -219,7 +240,7 @@
   LOG(INFO) << "Succesfully mapped " << target_partition_name
             << " to device mapper (force_writable = " << force_writable
             << "); device path at " << *path;
-  mapped_devices_.insert(target_partition_name);
+  mapped_devices_.insert(params.device_name);
   return true;
 }
 
@@ -229,9 +250,10 @@
     uint32_t slot,
     bool force_writable,
     std::string* path) {
-  DmDeviceState state = GetState(target_partition_name);
+  auto device_name = GetDeviceName(target_partition_name, slot);
+  DmDeviceState state = GetState(device_name);
   if (state == DmDeviceState::ACTIVE) {
-    if (mapped_devices_.find(target_partition_name) != mapped_devices_.end()) {
+    if (mapped_devices_.find(device_name) != mapped_devices_.end()) {
       if (GetDmDevicePathByName(target_partition_name, path)) {
         LOG(INFO) << target_partition_name
                   << " is mapped on device mapper: " << *path;
@@ -245,12 +267,13 @@
     // Note that for source partitions, if GetState() == ACTIVE, callers (e.g.
     // BootControlAndroid) should not call MapPartitionOnDeviceMapper, but
     // should directly call GetDmDevicePathByName.
-    if (!UnmapPartitionOnDeviceMapper(target_partition_name)) {
+    LOG(INFO) << "Destroying `" << device_name << "` from device mapper";
+    if (!UnmapPartitionOnDeviceMapper(device_name)) {
       LOG(ERROR) << target_partition_name
                  << " is mapped before the update, and it cannot be unmapped.";
       return false;
     }
-    state = GetState(target_partition_name);
+    state = GetState(device_name);
     if (state != DmDeviceState::INVALID) {
       LOG(ERROR) << target_partition_name << " is unmapped but state is "
                  << static_cast<std::underlying_type_t<DmDeviceState>>(state);
@@ -270,32 +293,38 @@
 
 bool DynamicPartitionControlAndroid::UnmapPartitionOnDeviceMapper(
     const std::string& target_partition_name) {
-  if (DeviceMapper::Instance().GetState(target_partition_name) !=
+  auto device_name = target_partition_name;
+  if (target_partition_name.ends_with("_a") ||
+      target_partition_name.ends_with("_b")) {
+    auto slot = target_partition_name.ends_with("_a") ? 0 : 1;
+    device_name = GetDeviceName(target_partition_name, slot);
+    return false;
+  }
+  if (DeviceMapper::Instance().GetState(device_name) !=
       DmDeviceState::INVALID) {
     // Partitions at target slot on non-Virtual A/B devices are mapped as
     // dm-linear. Also, on Virtual A/B devices, system_other may be mapped for
     // preopt apps as dm-linear.
     // Call DestroyLogicalPartition to handle these cases.
-    bool success = DestroyLogicalPartition(target_partition_name);
+    bool success = DestroyLogicalPartition(device_name);
 
     // On a Virtual A/B device, |target_partition_name| may be a leftover from
     // a paused update. Clean up any underlying devices.
-    if (ExpectMetadataMounted()) {
-      success &= snapshot_->UnmapUpdateSnapshot(target_partition_name);
+    if (ExpectMetadataMounted() &&
+        !device_name.ends_with(kRWSourcePartitionSuffix)) {
+      success &= snapshot_->UnmapUpdateSnapshot(device_name);
     } else {
-      LOG(INFO) << "Skip UnmapUpdateSnapshot(" << target_partition_name
-                << ") because metadata is not mounted";
+      LOG(INFO) << "Skip UnmapUpdateSnapshot(" << device_name << ")";
     }
 
     if (!success) {
-      LOG(ERROR) << "Cannot unmap " << target_partition_name
-                 << " from device mapper.";
+      LOG(ERROR) << "Cannot unmap " << device_name << " from device mapper.";
       return false;
     }
-    LOG(INFO) << "Successfully unmapped " << target_partition_name
+    LOG(INFO) << "Successfully unmapped " << device_name
               << " from device mapper.";
   }
-  mapped_devices_.erase(target_partition_name);
+  mapped_devices_.erase(device_name);
   return true;
 }
 
@@ -306,16 +335,27 @@
   }
   // UnmapPartitionOnDeviceMapper removes objects from mapped_devices_, hence
   // a copy is needed for the loop.
-  std::set<std::string> mapped = mapped_devices_;
+  std::set<std::string> mapped;
+  std::copy_if(mapped_devices_.begin(),
+               mapped_devices_.end(),
+               std::inserter(mapped, mapped.end()),
+               [](auto&& device_name) {
+                 return !std::string_view(device_name)
+                             .ends_with(kRWSourcePartitionSuffix);
+               });
   LOG(INFO) << "Destroying [" << Join(mapped, ", ") << "] from device mapper";
-  for (const auto& partition_name : mapped) {
-    ignore_result(UnmapPartitionOnDeviceMapper(partition_name));
+  for (const auto& device_name : mapped) {
+    ignore_result(UnmapPartitionOnDeviceMapper(device_name));
   }
   return true;
 }
 
 void DynamicPartitionControlAndroid::Cleanup() {
-  UnmapAllPartitions();
+  std::set<std::string> mapped = mapped_devices_;
+  LOG(INFO) << "Destroying [" << Join(mapped, ", ") << "] from device mapper";
+  for (const auto& device_name : mapped) {
+    ignore_result(UnmapPartitionOnDeviceMapper(device_name));
+  }
   LOG(INFO) << "UnmapAllPartitions done";
   metadata_device_.reset();
   if (GetVirtualAbFeatureFlag().IsEnabled()) {
@@ -772,6 +812,8 @@
   // In recovery, metadata might not be mounted, and
   // UnmapPartitionOnDeviceMapper might fail. However,
   // it is unusual that system_other has already been mapped. Hence, just skip.
+  LOG(INFO) << "Destroying `" << partition_name_suffix
+            << "` from device mapper";
   TEST_AND_RETURN_FALSE(UnmapPartitionOnDeviceMapper(partition_name_suffix));
   // Use CreateLogicalPartition directly to avoid mapping with existing
   // snapshots.
@@ -813,6 +855,8 @@
   // should be called. If DestroyLogicalPartition does fail, it is still okay
   // to skip the error here and let Prepare*() fail later.
   if (should_unmap) {
+    LOG(INFO) << "Destroying `" << partition_name_suffix
+              << "` from device mapper";
     TEST_AND_RETURN_FALSE(UnmapPartitionOnDeviceMapper(partition_name_suffix));
   }
 
@@ -1161,12 +1205,13 @@
   std::string device;
   if (GetDynamicPartitionsFeatureFlag().IsEnabled() &&
       (slot == current_slot || is_target_dynamic_)) {
-    switch (GetDynamicPartitionDevice(device_dir,
-                                      partition_name_suffix,
-                                      slot,
-                                      current_slot,
-                                      not_in_payload,
-                                      &device)) {
+    auto status = GetDynamicPartitionDevice(device_dir,
+                                            partition_name_suffix,
+                                            slot,
+                                            current_slot,
+                                            not_in_payload,
+                                            &device);
+    switch (status) {
       case DynamicPartitionDeviceStatus::SUCCESS:
         return {{.rw_device_path = device,
                  .readonly_device_path = device,
@@ -1176,6 +1221,7 @@
         break;
       case DynamicPartitionDeviceStatus::ERROR:  // fallthrough
       default:
+        LOG(ERROR) << "Unhandled dynamic partition status " << (int)status;
         return {};
     }
   }
@@ -1211,6 +1257,7 @@
     std::string* device) {
   std::string super_device =
       device_dir.Append(GetSuperPartitionName(slot)).value();
+  auto device_name = GetDeviceName(partition_name_suffix, slot);
 
   auto builder = LoadMetadataBuilder(super_device, slot);
   if (builder == nullptr) {
@@ -1233,13 +1280,12 @@
   }
 
   if (slot == current_slot) {
-    if (GetState(partition_name_suffix) != DmDeviceState::ACTIVE) {
-      LOG(WARNING) << partition_name_suffix << " is at current slot but it is "
+    if (GetState(device_name) != DmDeviceState::ACTIVE) {
+      LOG(WARNING) << device_name << " is at current slot but it is "
                    << "not mapped. Now try to map it.";
     } else {
-      if (GetDmDevicePathByName(partition_name_suffix, device)) {
-        LOG(INFO) << partition_name_suffix
-                  << " is mapped on device mapper: " << *device;
+      if (GetDmDevicePathByName(device_name, device)) {
+        LOG(INFO) << device_name << " is mapped on device mapper: " << *device;
         return DynamicPartitionDeviceStatus::SUCCESS;
       }
       LOG(ERROR) << partition_name_suffix << "is mapped but path is unknown.";