Small refactoring in first_stage_mount

The existing approach in first_stage_init/first_stage_mount makes it
harder to add conditional logic that should only be applied for
Microdroid. Additionally, it forces the FirstStageMount object to be
created twice.

This change refactors the control flow to make first_stage_init take the
ownership of the FirstStageMount object. It will help with the follow up
change (which will add logic to conditionally mount /vendor partition
while booting Microdroid). As a nice side effect, this refactoring also
fixes the problem of the FirstStageMount being created twice.

This change also merges the FirstStageMount and FirstStageMountVBootV2
in a single class, since nobody actually uses FirstStageMount.

Bug: 285855433
Test: device boots
Test: atest MicrodroidTestApp
Change-Id: I38a72c0f20e7c1ac70031498aeeca22b091fa827
diff --git a/init/first_stage_init.cpp b/init/first_stage_init.cpp
index bff80c5..e1d383a 100644
--- a/init/first_stage_init.cpp
+++ b/init/first_stage_init.cpp
@@ -257,6 +257,16 @@
     return BootMode::NORMAL_MODE;
 }
 
+static std::unique_ptr<FirstStageMount> CreateFirstStageMount() {
+    auto ret = FirstStageMount::Create();
+    if (ret.ok()) {
+        return std::move(*ret);
+    } else {
+        LOG(ERROR) << "Failed to create FirstStageMount : " << ret.error();
+        return nullptr;
+    }
+}
+
 int FirstStageMain(int argc, char** argv) {
     if (REBOOT_BOOTLOADER_ON_PANIC) {
         InstallRebootSignalHandlers();
@@ -381,12 +391,17 @@
                   << module_elapse_time.count() << " ms";
     }
 
+    std::unique_ptr<FirstStageMount> fsm;
+
     bool created_devices = false;
     if (want_console == FirstStageConsoleParam::CONSOLE_ON_FAILURE) {
         if (!IsRecoveryMode()) {
-            created_devices = DoCreateDevices();
-            if (!created_devices) {
-                LOG(ERROR) << "Failed to create device nodes early";
+            fsm = CreateFirstStageMount();
+            if (fsm) {
+                created_devices = fsm->DoCreateDevices();
+                if (!created_devices) {
+                    LOG(ERROR) << "Failed to create device nodes early";
+                }
             }
         }
         StartConsole(cmdline);
@@ -437,8 +452,23 @@
         SwitchRoot("/first_stage_ramdisk");
     }
 
-    if (!DoFirstStageMount(!created_devices)) {
-        LOG(FATAL) << "Failed to mount required partitions early ...";
+    if (IsRecoveryMode()) {
+        LOG(INFO) << "First stage mount skipped (recovery mode)";
+    } else {
+        if (!fsm) {
+            fsm = CreateFirstStageMount();
+        }
+        if (!fsm) {
+            LOG(FATAL) << "FirstStageMount not available";
+        }
+
+        if (!created_devices && !fsm->DoCreateDevices()) {
+            LOG(FATAL) << "Failed to create devices required for first stage mount";
+        }
+
+        if (!fsm->DoFirstStageMount()) {
+            LOG(FATAL) << "Failed to mount required partitions early ...";
+        }
     }
 
     struct stat new_root_info;
diff --git a/init/first_stage_mount.cpp b/init/first_stage_mount.cpp
index 07ce458..4961f6a 100644
--- a/init/first_stage_mount.cpp
+++ b/init/first_stage_mount.cpp
@@ -76,21 +76,21 @@
 
 // Class Declarations
 // ------------------
-class FirstStageMount {
+class FirstStageMountVBootV2 : public FirstStageMount {
   public:
-    FirstStageMount(Fstab fstab);
-    virtual ~FirstStageMount() = default;
+    friend void SetInitAvbVersionInRecovery();
 
-    // The factory method to create a FirstStageMountVBootV2 instance.
-    static Result<std::unique_ptr<FirstStageMount>> Create();
-    bool DoCreateDevices();    // Creates devices and logical partitions from storage devices
-    bool DoFirstStageMount();  // Mounts fstab entries read from device tree.
+    FirstStageMountVBootV2(Fstab fstab);
+    virtual ~FirstStageMountVBootV2() = default;
+
+    bool DoCreateDevices() override;
+    bool DoFirstStageMount() override;
+
+  private:
     bool InitDevices();
-
-  protected:
     bool InitRequiredDevices(std::set<std::string> devices);
     bool CreateLogicalPartitions();
-    bool CreateSnapshotPartitions(android::snapshot::SnapshotManager* sm);
+    bool CreateSnapshotPartitions(SnapshotManager* sm);
     bool MountPartition(const Fstab::iterator& begin, bool erase_same_mounts,
                         Fstab::iterator* end = nullptr);
 
@@ -106,9 +106,10 @@
     // revocation check by DSU installation service.
     void CopyDsuAvbKeys();
 
-    // Pure virtual functions.
-    virtual bool GetDmVerityDevices(std::set<std::string>* devices) = 0;
-    virtual bool SetUpDmVerity(FstabEntry* fstab_entry) = 0;
+    bool GetDmVerityDevices(std::set<std::string>* devices);
+    bool SetUpDmVerity(FstabEntry* fstab_entry);
+
+    bool InitAvbHandle();
 
     bool need_dm_verity_;
     bool dsu_not_on_userdata_ = false;
@@ -122,19 +123,6 @@
     // Reads all AVB keys before chroot into /system, as they might be used
     // later when mounting other partitions, e.g., /vendor and /product.
     std::map<std::string, std::vector<std::string>> preload_avb_key_blobs_;
-};
-
-class FirstStageMountVBootV2 : public FirstStageMount {
-  public:
-    friend void SetInitAvbVersionInRecovery();
-
-    FirstStageMountVBootV2(Fstab fstab);
-    ~FirstStageMountVBootV2() override = default;
-
-  protected:
-    bool GetDmVerityDevices(std::set<std::string>* devices) override;
-    bool SetUpDmVerity(FstabEntry* fstab_entry) override;
-    bool InitAvbHandle();
 
     std::vector<std::string> vbmeta_partitions_;
     AvbUniquePtr avb_handle_;
@@ -220,10 +208,6 @@
 
 // Class Definitions
 // -----------------
-FirstStageMount::FirstStageMount(Fstab fstab) : need_dm_verity_(false), fstab_(std::move(fstab)) {
-    super_partition_name_ = fs_mgr_get_super_partition_name();
-}
-
 Result<std::unique_ptr<FirstStageMount>> FirstStageMount::Create() {
     auto fstab = ReadFirstStageFstab();
     if (!fstab.ok()) {
@@ -233,7 +217,7 @@
     return std::make_unique<FirstStageMountVBootV2>(std::move(*fstab));
 }
 
-bool FirstStageMount::DoCreateDevices() {
+bool FirstStageMountVBootV2::DoCreateDevices() {
     if (!InitDevices()) return false;
 
     // Mount /metadata before creating logical partitions, since we need to
@@ -255,7 +239,7 @@
     return true;
 }
 
-bool FirstStageMount::DoFirstStageMount() {
+bool FirstStageMountVBootV2::DoFirstStageMount() {
     if (!IsDmLinearEnabled() && fstab_.empty()) {
         // Nothing to mount.
         LOG(INFO) << "First stage mount skipped (missing/incompatible/empty fstab in device tree)";
@@ -267,7 +251,7 @@
     return true;
 }
 
-bool FirstStageMount::InitDevices() {
+bool FirstStageMountVBootV2::InitDevices() {
     std::set<std::string> devices;
     GetSuperDeviceName(&devices);
 
@@ -288,14 +272,14 @@
     return true;
 }
 
-bool FirstStageMount::IsDmLinearEnabled() {
+bool FirstStageMountVBootV2::IsDmLinearEnabled() {
     for (const auto& entry : fstab_) {
         if (entry.fs_mgr_flags.logical) return true;
     }
     return false;
 }
 
-void FirstStageMount::GetSuperDeviceName(std::set<std::string>* devices) {
+void FirstStageMountVBootV2::GetSuperDeviceName(std::set<std::string>* devices) {
     // Add any additional devices required for dm-linear mappings.
     if (!IsDmLinearEnabled()) {
         return;
@@ -307,7 +291,7 @@
 // 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) {
+bool FirstStageMountVBootV2::InitRequiredDevices(std::set<std::string> devices) {
     if (!block_dev_init_.InitDeviceMapper()) {
         return false;
     }
@@ -317,7 +301,8 @@
     return block_dev_init_.InitDevices(std::move(devices));
 }
 
-bool FirstStageMount::InitDmLinearBackingDevices(const android::fs_mgr::LpMetadata& metadata) {
+bool FirstStageMountVBootV2::InitDmLinearBackingDevices(
+        const android::fs_mgr::LpMetadata& metadata) {
     std::set<std::string> devices;
 
     auto partition_names = android::fs_mgr::GetBlockDevicePartitionNames(metadata);
@@ -334,7 +319,7 @@
     return InitRequiredDevices(std::move(devices));
 }
 
-bool FirstStageMount::CreateLogicalPartitions() {
+bool FirstStageMountVBootV2::CreateLogicalPartitions() {
     if (!IsDmLinearEnabled()) {
         return true;
     }
@@ -365,7 +350,7 @@
     return android::fs_mgr::CreateLogicalPartitions(*metadata.get(), super_path_);
 }
 
-bool FirstStageMount::CreateSnapshotPartitions(SnapshotManager* sm) {
+bool FirstStageMountVBootV2::CreateSnapshotPartitions(SnapshotManager* sm) {
     // When COW images are present for snapshots, they are stored on
     // the data partition.
     if (!InitRequiredDevices({"userdata"})) {
@@ -400,8 +385,8 @@
     return true;
 }
 
-bool FirstStageMount::MountPartition(const Fstab::iterator& begin, bool erase_same_mounts,
-                                     Fstab::iterator* end) {
+bool FirstStageMountVBootV2::MountPartition(const Fstab::iterator& begin, bool erase_same_mounts,
+                                            Fstab::iterator* end) {
     // Sets end to begin + 1, so we can just return on failure below.
     if (end) {
         *end = begin + 1;
@@ -445,7 +430,7 @@
     return mounted;
 }
 
-void FirstStageMount::PreloadAvbKeys() {
+void FirstStageMountVBootV2::PreloadAvbKeys() {
     for (const auto& entry : fstab_) {
         // No need to cache the key content if it's empty, or is already cached.
         if (entry.avb_keys.empty() || preload_avb_key_blobs_.count(entry.avb_keys)) {
@@ -492,7 +477,7 @@
 // If system is in the fstab then we're not a system-as-root device, and in
 // this case, we mount system first then pivot to it.  From that point on,
 // we are effectively identical to a system-as-root device.
-bool FirstStageMount::TrySwitchSystemAsRoot() {
+bool FirstStageMountVBootV2::TrySwitchSystemAsRoot() {
     UseDsuIfPresent();
     // Preloading all AVB keys from the ramdisk before switching root to /system.
     PreloadAvbKeys();
@@ -521,7 +506,7 @@
     return true;
 }
 
-bool FirstStageMount::MountPartitions() {
+bool FirstStageMountVBootV2::MountPartitions() {
     if (!TrySwitchSystemAsRoot()) return false;
 
     if (!SkipMountingPartitions(&fstab_, true /* verbose */)) return false;
@@ -604,7 +589,7 @@
 // copy files to /metadata is NOT fatal, because it is auxiliary to perform
 // public key matching before booting into DSU images on next boot. The actual
 // public key matching will still be done on next boot to DSU.
-void FirstStageMount::CopyDsuAvbKeys() {
+void FirstStageMountVBootV2::CopyDsuAvbKeys() {
     std::error_code ec;
     // Removing existing keys in gsi::kDsuAvbKeyDir as they might be stale.
     std::filesystem::remove_all(gsi::kDsuAvbKeyDir, ec);
@@ -620,7 +605,7 @@
     }
 }
 
-void FirstStageMount::UseDsuIfPresent() {
+void FirstStageMountVBootV2::UseDsuIfPresent() {
     std::string error;
 
     if (!android::gsi::CanBootIntoGsi(&error)) {
@@ -657,10 +642,10 @@
     TransformFstabForDsu(&fstab_, active_dsu, dsu_partitions);
 }
 
-// First retrieve any vbmeta partitions from device tree (legacy) then read through the fstab
-// for any further vbmeta partitions.
 FirstStageMountVBootV2::FirstStageMountVBootV2(Fstab fstab)
-    : FirstStageMount(std::move(fstab)), avb_handle_(nullptr) {
+    : need_dm_verity_(false), fstab_(std::move(fstab)), avb_handle_(nullptr) {
+    super_partition_name_ = fs_mgr_get_super_partition_name();
+
     std::string device_tree_vbmeta_parts;
     read_android_dt_file("vbmeta/parts", &device_tree_vbmeta_parts);
 
@@ -793,39 +778,6 @@
     return true;
 }
 
-// Public functions
-// ----------------
-// Creates devices and logical partitions from storage devices
-bool DoCreateDevices() {
-    auto fsm = FirstStageMount::Create();
-    if (!fsm.ok()) {
-        LOG(ERROR) << "Failed to create FirstStageMount: " << fsm.error();
-        return false;
-    }
-    return (*fsm)->DoCreateDevices();
-}
-
-// Mounts partitions specified by fstab in device tree.
-bool DoFirstStageMount(bool create_devices) {
-    // Skips first stage mount if we're in recovery mode.
-    if (IsRecoveryMode()) {
-        LOG(INFO) << "First stage mount skipped (recovery mode)";
-        return true;
-    }
-
-    auto fsm = FirstStageMount::Create();
-    if (!fsm.ok()) {
-        LOG(ERROR) << "Failed to create FirstStageMount " << fsm.error();
-        return false;
-    }
-
-    if (create_devices) {
-        if (!(*fsm)->DoCreateDevices()) return false;
-    }
-
-    return (*fsm)->DoFirstStageMount();
-}
-
 void SetInitAvbVersionInRecovery() {
     if (!IsRecoveryMode()) {
         LOG(INFO) << "Skipped setting INIT_AVB_VERSION (not in recovery mode)";
diff --git a/init/first_stage_mount.h b/init/first_stage_mount.h
index 2f4e663..ab14966 100644
--- a/init/first_stage_mount.h
+++ b/init/first_stage_mount.h
@@ -16,11 +16,28 @@
 
 #pragma once
 
+#include <memory>
+
+#include "result.h"
+
 namespace android {
 namespace init {
 
-bool DoCreateDevices();
-bool DoFirstStageMount(bool create_devices);
+class FirstStageMount {
+  public:
+    virtual ~FirstStageMount() = default;
+
+    // The factory method to create a FirstStageMount instance.
+    static Result<std::unique_ptr<FirstStageMount>> Create();
+    // Creates devices and logical partitions from storage devices
+    virtual bool DoCreateDevices() = 0;
+    // Mounts fstab entries read from device tree.
+    virtual bool DoFirstStageMount() = 0;
+
+  protected:
+    FirstStageMount() = default;
+};
+
 void SetInitAvbVersionInRecovery();
 
 }  // namespace init