From 7f994f4cfbb2d81f224cd9580a96e8ddbb2c7b8a Mon Sep 17 00:00:00 2001
From: David Anderson <dvander@google.com>
Date: Wed, 1 Sep 2021 20:51:35 -0700
Subject: [PATCH] libsnapshot: Fix new partitions not transitioning in
 second-stage init.

Bug: 196922070
Test: vts_libsnapshot_test
Change-Id: If8a7afde218fd719e4426dc1dda41f53a4e6544b
---
 fs_mgr/libsnapshot/snapshot.cpp      | 11 +++++++++--
 fs_mgr/libsnapshot/snapshot_test.cpp | 12 +++++++++++-
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 52324ba63..3e8d9c1d3 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -1452,7 +1452,7 @@ bool SnapshotManager::PerformInitTransition(InitTransition transition,
                                             std::vector<std::string>* snapuserd_argv) {
     LOG(INFO) << "Performing transition for snapuserd.";
 
-    // Don't use EnsuerSnapuserdConnected() because this is called from init,
+    // Don't use EnsureSnapuserdConnected() because this is called from init,
     // and attempting to do so will deadlock.
     if (!snapuserd_client_ && transition != InitTransition::SELINUX_DETACH) {
         snapuserd_client_ = SnapuserdClient::Connect(kSnapuserdSocket, 10s);
@@ -1509,8 +1509,15 @@ bool SnapshotManager::PerformInitTransition(InitTransition transition,
             continue;
         }
 
+        std::string source_device_name;
+        if (snapshot_status.old_partition_size() > 0) {
+            source_device_name = GetSourceDeviceName(snapshot);
+        } else {
+            source_device_name = GetBaseDeviceName(snapshot);
+        }
+
         std::string source_device;
-        if (!dm.GetDmDevicePathByName(GetSourceDeviceName(snapshot), &source_device)) {
+        if (!dm.GetDmDevicePathByName(source_device_name, &source_device)) {
             LOG(ERROR) << "Could not get device path for " << GetSourceDeviceName(snapshot);
             continue;
         }
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 057e5b19a..43c7fe2b3 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -2118,14 +2118,24 @@ TEST_F(SnapshotUpdateTest, AddPartition) {
     // After reboot, init does first stage mount.
     auto init = NewManagerForFirstStageMount("_b");
     ASSERT_NE(init, nullptr);
+
+    ASSERT_TRUE(init->EnsureSnapuserdConnected());
+    init->set_use_first_stage_snapuserd(true);
+
     ASSERT_TRUE(init->NeedSnapshotsInFirstStageMount());
     ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
 
     // Check that the target partitions have the same content.
-    for (const auto& name : {"sys_b", "vnd_b", "prd_b", "dlkm_b"}) {
+    std::vector<std::string> partitions = {"sys_b", "vnd_b", "prd_b", "dlkm_b"};
+    for (const auto& name : partitions) {
         ASSERT_TRUE(IsPartitionUnchanged(name));
     }
 
+    ASSERT_TRUE(init->PerformInitTransition(SnapshotManager::InitTransition::SECOND_STAGE));
+    for (const auto& name : partitions) {
+        ASSERT_TRUE(init->snapuserd_client()->WaitForDeviceDelete(name + "-user-cow-init"));
+    }
+
     // Initiate the merge and wait for it to be completed.
     ASSERT_TRUE(init->InitiateMerge());
     ASSERT_EQ(UpdateState::MergeCompleted, init->ProcessUpdateState());