init: Make sure /dev/block/by-name/userdata is created in snapshot mode.

When userdata is used to store COW devices, we need to ensure we've
generated uevents for it.

This patch also refactors FirstStageMount to pass required devices
through arguments rather than a member variable.

Bug: 140761481
Test: manual test
Change-Id: Ie5e1c9699f084da5467a758eea41c7907fecc5ca
diff --git a/init/first_stage_mount.cpp b/init/first_stage_mount.cpp
index 6b4216f..4a1bc5a 100644
--- a/init/first_stage_mount.cpp
+++ b/init/first_stage_mount.cpp
@@ -77,8 +77,9 @@
     bool InitDevices();
 
   protected:
-    ListenerAction HandleBlockDevice(const std::string& name, const Uevent&);
-    bool InitRequiredDevices();
+    ListenerAction HandleBlockDevice(const std::string& name, const Uevent&,
+                                     std::set<std::string>* required_devices);
+    bool InitRequiredDevices(std::set<std::string> devices);
     bool InitMappedDevice(const std::string& verity_device);
     bool InitDeviceMapper();
     bool CreateLogicalPartitions();
@@ -89,14 +90,14 @@
     bool TrySwitchSystemAsRoot();
     bool TrySkipMountingPartitions();
     bool IsDmLinearEnabled();
-    bool GetDmLinearMetadataDevice();
+    void GetDmLinearMetadataDevice(std::set<std::string>* devices);
     bool InitDmLinearBackingDevices(const android::fs_mgr::LpMetadata& metadata);
     void UseGsiIfPresent();
 
-    ListenerAction UeventCallback(const Uevent& uevent);
+    ListenerAction UeventCallback(const Uevent& uevent, std::set<std::string>* required_devices);
 
     // Pure virtual functions.
-    virtual bool GetDmVerityDevices() = 0;
+    virtual bool GetDmVerityDevices(std::set<std::string>* devices) = 0;
     virtual bool SetUpDmVerity(FstabEntry* fstab_entry) = 0;
 
     bool need_dm_verity_;
@@ -104,7 +105,6 @@
 
     Fstab fstab_;
     std::string lp_metadata_partition_;
-    std::set<std::string> required_devices_partition_names_;
     std::string super_partition_name_;
     std::unique_ptr<DeviceHandler> device_handler_;
     UeventListener uevent_listener_;
@@ -116,7 +116,7 @@
     ~FirstStageMountVBootV1() override = default;
 
   protected:
-    bool GetDmVerityDevices() override;
+    bool GetDmVerityDevices(std::set<std::string>* devices) override;
     bool SetUpDmVerity(FstabEntry* fstab_entry) override;
 };
 
@@ -128,7 +128,7 @@
     ~FirstStageMountVBootV2() override = default;
 
   protected:
-    bool GetDmVerityDevices() override;
+    bool GetDmVerityDevices(std::set<std::string>* devices) override;
     bool SetUpDmVerity(FstabEntry* fstab_entry) override;
     bool InitAvbHandle();
 
@@ -252,7 +252,13 @@
 }
 
 bool FirstStageMount::InitDevices() {
-    return GetDmLinearMetadataDevice() && GetDmVerityDevices() && InitRequiredDevices();
+    std::set<std::string> devices;
+    GetDmLinearMetadataDevice(&devices);
+
+    if (!GetDmVerityDevices(&devices)) {
+        return false;
+    }
+    return InitRequiredDevices(std::move(devices));
 }
 
 bool FirstStageMount::IsDmLinearEnabled() {
@@ -262,45 +268,46 @@
     return false;
 }
 
-bool FirstStageMount::GetDmLinearMetadataDevice() {
+void FirstStageMount::GetDmLinearMetadataDevice(std::set<std::string>* devices) {
     // Add any additional devices required for dm-linear mappings.
     if (!IsDmLinearEnabled()) {
-        return true;
+        return;
     }
 
-    required_devices_partition_names_.emplace(super_partition_name_);
-    return true;
+    devices->emplace(super_partition_name_);
 }
 
-// Creates devices with uevent->partition_name matching one in the member variable
-// required_devices_partition_names_. Found partitions will then be removed from it
-// for the subsequent member function to check which devices are NOT created.
-bool FirstStageMount::InitRequiredDevices() {
+// Creates devices with uevent->partition_name matching ones in the given set.
+// Found partitions will then be removed from it for the subsequent member
+// function to check which devices are NOT created.
+bool FirstStageMount::InitRequiredDevices(std::set<std::string> devices) {
     if (!InitDeviceMapper()) {
         return false;
     }
 
-    if (required_devices_partition_names_.empty()) {
+    if (devices.empty()) {
         return true;
     }
 
-    auto uevent_callback = [this](const Uevent& uevent) { return UeventCallback(uevent); };
+    auto uevent_callback = [&, this](const Uevent& uevent) {
+        return UeventCallback(uevent, &devices);
+    };
     uevent_listener_.RegenerateUevents(uevent_callback);
 
-    // UeventCallback() will remove found partitions from required_devices_partition_names_.
-    // So if it isn't empty here, it means some partitions are not found.
-    if (!required_devices_partition_names_.empty()) {
+    // UeventCallback() will remove found partitions from |devices|. So if it
+    // isn't empty here, it means some partitions are not found.
+    if (!devices.empty()) {
         LOG(INFO) << __PRETTY_FUNCTION__
                   << ": partition(s) not found in /sys, waiting for their uevent(s): "
-                  << android::base::Join(required_devices_partition_names_, ", ");
+                  << android::base::Join(devices, ", ");
         Timer t;
         uevent_listener_.Poll(uevent_callback, 10s);
         LOG(INFO) << "Wait for partitions returned after " << t;
     }
 
-    if (!required_devices_partition_names_.empty()) {
+    if (!devices.empty()) {
         LOG(ERROR) << __PRETTY_FUNCTION__ << ": partition(s) not found after polling timeout: "
-                   << android::base::Join(required_devices_partition_names_, ", ");
+                   << android::base::Join(devices, ", ");
         return false;
     }
 
@@ -333,27 +340,20 @@
 }
 
 bool FirstStageMount::InitDmLinearBackingDevices(const android::fs_mgr::LpMetadata& metadata) {
+    std::set<std::string> devices;
+
     auto partition_names = android::fs_mgr::GetBlockDevicePartitionNames(metadata);
     for (const auto& partition_name : partition_names) {
         // The super partition was found in the earlier pass.
         if (partition_name == super_partition_name_) {
             continue;
         }
-        required_devices_partition_names_.emplace(partition_name);
+        devices.emplace(partition_name);
     }
-    if (required_devices_partition_names_.empty()) {
+    if (devices.empty()) {
         return true;
     }
-
-    auto uevent_callback = [this](const Uevent& uevent) { return UeventCallback(uevent); };
-    uevent_listener_.RegenerateUevents(uevent_callback);
-
-    if (!required_devices_partition_names_.empty()) {
-        LOG(ERROR) << __PRETTY_FUNCTION__ << ": partition(s) not found after polling timeout: "
-                   << android::base::Join(required_devices_partition_names_, ", ");
-        return false;
-    }
-    return true;
+    return InitRequiredDevices(std::move(devices));
 }
 
 bool FirstStageMount::CreateLogicalPartitions() {
@@ -372,6 +372,11 @@
             return false;
         }
         if (sm->NeedSnapshotsInFirstStageMount()) {
+            // When COW images are present for snapshots, they are stored on
+            // the data partition.
+            if (!InitRequiredDevices({"userdata"})) {
+                return false;
+            }
             return sm->CreateLogicalAndSnapshotPartitions(lp_metadata_partition_);
         }
     }
@@ -387,20 +392,21 @@
     return android::fs_mgr::CreateLogicalPartitions(*metadata.get(), lp_metadata_partition_);
 }
 
-ListenerAction FirstStageMount::HandleBlockDevice(const std::string& name, const Uevent& uevent) {
+ListenerAction FirstStageMount::HandleBlockDevice(const std::string& name, const Uevent& uevent,
+                                                  std::set<std::string>* required_devices) {
     // Matches partition name to create device nodes.
     // Both required_devices_partition_names_ and uevent->partition_name have A/B
     // suffix when A/B is used.
-    auto iter = required_devices_partition_names_.find(name);
-    if (iter != required_devices_partition_names_.end()) {
+    auto iter = required_devices->find(name);
+    if (iter != required_devices->end()) {
         LOG(VERBOSE) << __PRETTY_FUNCTION__ << ": found partition: " << *iter;
         if (IsDmLinearEnabled() && name == super_partition_name_) {
             std::vector<std::string> links = device_handler_->GetBlockDeviceSymlinks(uevent);
             lp_metadata_partition_ = links[0];
         }
-        required_devices_partition_names_.erase(iter);
+        required_devices->erase(iter);
         device_handler_->HandleUevent(uevent);
-        if (required_devices_partition_names_.empty()) {
+        if (required_devices->empty()) {
             return ListenerAction::kStop;
         } else {
             return ListenerAction::kContinue;
@@ -409,18 +415,19 @@
     return ListenerAction::kContinue;
 }
 
-ListenerAction FirstStageMount::UeventCallback(const Uevent& uevent) {
+ListenerAction FirstStageMount::UeventCallback(const Uevent& uevent,
+                                               std::set<std::string>* required_devices) {
     // Ignores everything that is not a block device.
     if (uevent.subsystem != "block") {
         return ListenerAction::kContinue;
     }
 
     if (!uevent.partition_name.empty()) {
-        return HandleBlockDevice(uevent.partition_name, uevent);
+        return HandleBlockDevice(uevent.partition_name, uevent, required_devices);
     } else {
         size_t base_idx = uevent.path.rfind('/');
         if (base_idx != std::string::npos) {
-            return HandleBlockDevice(uevent.path.substr(base_idx + 1), uevent);
+            return HandleBlockDevice(uevent.path.substr(base_idx + 1), uevent, required_devices);
         }
     }
     // Not found a partition or find an unneeded partition, continue to find others.
@@ -577,17 +584,7 @@
     const auto devices = fs_mgr_overlayfs_required_devices(&fstab_);
     for (auto const& device : devices) {
         if (android::base::StartsWith(device, "/dev/block/by-name/")) {
-            required_devices_partition_names_.emplace(basename(device.c_str()));
-            auto uevent_callback = [this](const Uevent& uevent) { return UeventCallback(uevent); };
-            uevent_listener_.RegenerateUevents(uevent_callback);
-            if (!required_devices_partition_names_.empty()) {
-                uevent_listener_.Poll(uevent_callback, 10s);
-                if (!required_devices_partition_names_.empty()) {
-                    LOG(ERROR) << __PRETTY_FUNCTION__
-                               << ": partition(s) not found after polling timeout: "
-                               << android::base::Join(required_devices_partition_names_, ", ");
-                }
-            }
+            InitRequiredDevices({basename(device.c_str())});
         } else {
             InitMappedDevice(device);
         }
@@ -641,7 +638,7 @@
     gsi_not_on_userdata_ = (super_name != "userdata");
 }
 
-bool FirstStageMountVBootV1::GetDmVerityDevices() {
+bool FirstStageMountVBootV1::GetDmVerityDevices(std::set<std::string>* devices) {
     need_dm_verity_ = false;
 
     for (const auto& fstab_entry : fstab_) {
@@ -660,7 +657,7 @@
     // Notes that fstab_rec->blk_device has A/B suffix updated by fs_mgr when A/B is used.
     for (const auto& fstab_entry : fstab_) {
         if (!fstab_entry.fs_mgr_flags.logical) {
-            required_devices_partition_names_.emplace(basename(fstab_entry.blk_device.c_str()));
+            devices->emplace(basename(fstab_entry.blk_device.c_str()));
         }
     }
 
@@ -711,7 +708,7 @@
     }
 }
 
-bool FirstStageMountVBootV2::GetDmVerityDevices() {
+bool FirstStageMountVBootV2::GetDmVerityDevices(std::set<std::string>* devices) {
     need_dm_verity_ = false;
 
     std::set<std::string> logical_partitions;
@@ -725,7 +722,7 @@
             // Don't try to find logical partitions via uevent regeneration.
             logical_partitions.emplace(basename(fstab_entry.blk_device.c_str()));
         } else {
-            required_devices_partition_names_.emplace(basename(fstab_entry.blk_device.c_str()));
+            devices->emplace(basename(fstab_entry.blk_device.c_str()));
         }
     }
 
@@ -742,11 +739,11 @@
             if (logical_partitions.count(partition_name)) {
                 continue;
             }
-            // required_devices_partition_names_ is of type std::set so it's not an issue
-            // to emplace a partition twice. e.g., /vendor might be in both places:
+            // devices is of type std::set so it's not an issue to emplace a
+            // partition twice. e.g., /vendor might be in both places:
             //   - device_tree_vbmeta_parts_ = "vbmeta,boot,system,vendor"
             //   - mount_fstab_recs_: /vendor_a
-            required_devices_partition_names_.emplace(partition_name);
+            devices->emplace(partition_name);
         }
     }
     return true;