libsnapshot: Implement MapAllSnapshots and UnmapAllSnapshots.

Bug: 168554689
Test: vts_libsnapshot_test
Change-Id: I6809e226741dabcf337c3a5cfaba56afdb9edd64
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index 9dbbbd7..8bed1b9 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -559,6 +559,9 @@
     std::string GetSnapshotDeviceName(const std::string& snapshot_name,
                                       const SnapshotStatus& status);
 
+    bool MapAllPartitions(LockedFile* lock, const std::string& super_device, uint32_t slot,
+                          const std::chrono::milliseconds& timeout_ms);
+
     // Reason for calling MapPartitionWithSnapshot.
     enum class SnapshotContext {
         // For writing or verification (during update_engine).
@@ -632,9 +635,12 @@
             const LpMetadata* exported_target_metadata, const std::string& target_suffix,
             const std::map<std::string, SnapshotStatus>& all_snapshot_status);
 
+    // Implementation of UnmapAllSnapshots(), with the lock provided.
+    bool UnmapAllSnapshots(LockedFile* lock);
+
     // Unmap all partitions that were mapped by CreateLogicalAndSnapshotPartitions.
     // This should only be called in recovery.
-    bool UnmapAllPartitions();
+    bool UnmapAllPartitionsInRecovery();
 
     // Check no snapshot overflows. Note that this returns false negatives if the snapshot
     // overflows, then is remapped and not written afterwards.
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index b6c0a29..f9bb0dd 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -1709,8 +1709,13 @@
     auto lock = LockExclusive();
     if (!lock) return false;
 
-    const auto& opener = device_->GetPartitionOpener();
     uint32_t slot = SlotNumberForSlotSuffix(device_->GetSlotSuffix());
+    return MapAllPartitions(lock.get(), super_device, slot, timeout_ms);
+}
+
+bool SnapshotManager::MapAllPartitions(LockedFile* lock, const std::string& super_device,
+                                       uint32_t slot, const std::chrono::milliseconds& timeout_ms) {
+    const auto& opener = device_->GetPartitionOpener();
     auto metadata = android::fs_mgr::ReadMetadata(opener, super_device, slot);
     if (!metadata) {
         LOG(ERROR) << "Could not read dynamic partition metadata for device: " << super_device;
@@ -1731,8 +1736,7 @@
                 .partition_opener = &opener,
                 .timeout_ms = timeout_ms,
         };
-        if (!MapPartitionWithSnapshot(lock.get(), std::move(params), SnapshotContext::Mount,
-                                      nullptr)) {
+        if (!MapPartitionWithSnapshot(lock, std::move(params), SnapshotContext::Mount, nullptr)) {
             return false;
         }
     }
@@ -2078,14 +2082,49 @@
     return true;
 }
 
-bool SnapshotManager::MapAllSnapshots(const std::chrono::milliseconds&) {
-    LOG(ERROR) << "Not yet implemented.";
-    return false;
+bool SnapshotManager::MapAllSnapshots(const std::chrono::milliseconds& timeout_ms) {
+    auto lock = LockExclusive();
+    if (!lock) return false;
+
+    auto state = ReadUpdateState(lock.get());
+    if (state == UpdateState::Unverified) {
+        if (GetCurrentSlot() == Slot::Target) {
+            LOG(ERROR) << "Cannot call MapAllSnapshots when booting from the target slot.";
+            return false;
+        }
+    } else if (state != UpdateState::Initiated) {
+        LOG(ERROR) << "Cannot call MapAllSnapshots from update state: " << state;
+        return false;
+    }
+
+    if (!UnmapAllSnapshots(lock.get())) {
+        return false;
+    }
+
+    uint32_t slot = SlotNumberForSlotSuffix(device_->GetOtherSlotSuffix());
+    return MapAllPartitions(lock.get(), device_->GetSuperDevice(slot), slot, timeout_ms);
 }
 
 bool SnapshotManager::UnmapAllSnapshots() {
-    LOG(ERROR) << "Not yet implemented.";
-    return false;
+    auto lock = LockExclusive();
+    if (!lock) return false;
+
+    return UnmapAllSnapshots(lock.get());
+}
+
+bool SnapshotManager::UnmapAllSnapshots(LockedFile* lock) {
+    std::vector<std::string> snapshots;
+    if (!ListSnapshots(lock, &snapshots)) {
+        return false;
+    }
+
+    for (const auto& snapshot : snapshots) {
+        if (!UnmapPartitionWithSnapshot(lock, snapshot)) {
+            LOG(ERROR) << "Failed to unmap snapshot: " << snapshot;
+            return false;
+        }
+    }
+    return true;
 }
 
 auto SnapshotManager::OpenFile(const std::string& file, int lock_flags)
@@ -2868,7 +2907,7 @@
     return UnmapPartitionWithSnapshot(lock.get(), target_partition_name);
 }
 
-bool SnapshotManager::UnmapAllPartitions() {
+bool SnapshotManager::UnmapAllPartitionsInRecovery() {
     auto lock = LockExclusive();
     if (!lock) return false;
 
@@ -3012,7 +3051,7 @@
     }
 
     // Nothing should be depending on partitions now, so unmap them all.
-    if (!UnmapAllPartitions()) {
+    if (!UnmapAllPartitionsInRecovery()) {
         LOG(ERROR) << "Unable to unmap all partitions; fastboot may fail to flash.";
     }
     return true;
@@ -3043,7 +3082,7 @@
     }
 
     // Nothing should be depending on partitions now, so unmap them all.
-    if (!UnmapAllPartitions()) {
+    if (!UnmapAllPartitionsInRecovery()) {
         LOG(ERROR) << "Unable to unmap all partitions; fastboot may fail to flash.";
     }
     return true;
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 4454281..7fc64e5 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -1793,6 +1793,21 @@
     ASSERT_EQ(access("/dev/dm-user/sys_b-user-cow", F_OK), 0);
 }
 
+TEST_F(SnapshotUpdateTest, MapAllSnapshots) {
+    AddOperationForPartitions();
+    // Execute the update.
+    ASSERT_TRUE(sm->BeginUpdate());
+    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
+    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
+        ASSERT_TRUE(WriteSnapshotAndHash(name));
+    }
+    ASSERT_TRUE(sm->FinishedSnapshotWrites(false));
+    ASSERT_TRUE(sm->MapAllSnapshots(10s));
+
+    // Read bytes back and verify they match the cache.
+    ASSERT_TRUE(IsPartitionUnchanged("sys_b"));
+}
+
 class FlashAfterUpdateTest : public SnapshotUpdateTest,
                              public WithParamInterface<std::tuple<uint32_t, bool>> {
   public: