libsnapshot: Move ImageManager creation to DeviceInfo.

This refactoring moves ImageManager creation out of SnapshotManager,
where it had always been rather awkward, and into IDeviceInfo so it can
participate in dependency injection.

The "first-stage init" state is now part of IDeviceInfo as well.

Bug: 187396878
Test: vts_libsnapshot_test
      libsnapshot_fuzzer_test
Change-Id: Ic4b3e93527cc074ec18c0c26440dbdb2504d0d5e
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 104484d..6a764e4 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -266,7 +266,7 @@
         "android.hardware.boot@1.1",
         "libbrotli",
         "libc++fs",
-        "libfs_mgr",
+        "libfs_mgr_binder",
         "libgsi",
         "libgmock",
         "liblp",
@@ -421,8 +421,8 @@
         "snapuserd_server.cpp",
         "snapuserd.cpp",
         "snapuserd_daemon.cpp",
-	"snapuserd_worker.cpp",
-	"snapuserd_readahead.cpp",
+        "snapuserd_worker.cpp",
+        "snapuserd_readahead.cpp",
     ],
 
     cflags: [
@@ -564,7 +564,7 @@
     srcs: [
         "cow_snapuserd_test.cpp",
         "snapuserd.cpp",
-	"snapuserd_worker.cpp",
+        "snapuserd_worker.cpp",
     ],
     cflags: [
         "-Wall",
@@ -581,7 +581,7 @@
         "libsnapshot_snapuserd",
         "libcutils_sockets",
         "libz",
-	"libfs_mgr",
+        "libfs_mgr",
         "libdm",
     ],
     header_libs: [
diff --git a/fs_mgr/libsnapshot/corpus/avoid-io-in-fuzzer.txt b/fs_mgr/libsnapshot/corpus/avoid-io-in-fuzzer.txt
new file mode 100644
index 0000000..c474f4c
--- /dev/null
+++ b/fs_mgr/libsnapshot/corpus/avoid-io-in-fuzzer.txt
@@ -0,0 +1,41 @@
+device_info_data {
+  allow_set_slot_as_unbootable: true
+  is_recovery: true
+}
+is_super_metadata_valid: true
+super_data {
+  partitions {
+    partition_name: "sys_a"
+    new_partition_info {
+      size: 3145728
+    }
+  }
+  partitions {
+    partition_name: "vnnd_"
+    new_partition_info {
+      size: 3145728
+    }
+  }
+  partitions {
+    partition_name: "prd_a"
+    new_partition_info {
+    }
+  }
+  dynamic_partition_metadata {
+    groups {
+      name: "group_google_dp_a"
+      size: 34375467008
+      partition_names: "sys_a"
+      partition_names: "vnd_a"
+      partition_names: "prd_a"
+    }
+  }
+}
+has_metadata_snapshots_dir: true
+actions {
+  handle_imminent_data_wipe: true
+}
+actions {
+  begin_update {
+  }
+}
diff --git a/fs_mgr/libsnapshot/device_info.cpp b/fs_mgr/libsnapshot/device_info.cpp
index 0e90100..14ce0ee 100644
--- a/fs_mgr/libsnapshot/device_info.cpp
+++ b/fs_mgr/libsnapshot/device_info.cpp
@@ -17,6 +17,7 @@
 #include <android-base/logging.h>
 #include <fs_mgr.h>
 #include <fs_mgr_overlayfs.h>
+#include <libfiemap/image_manager.h>
 
 namespace android {
 namespace snapshot {
@@ -26,6 +27,7 @@
 using android::hardware::boot::V1_0::CommandResult;
 #endif
 
+using namespace std::chrono_literals;
 using namespace std::string_literals;
 
 #ifdef __ANDROID_RECOVERY__
@@ -34,10 +36,6 @@
 constexpr bool kIsRecovery = false;
 #endif
 
-std::string DeviceInfo::GetGsidDir() const {
-    return "ota"s;
-}
-
 std::string DeviceInfo::GetMetadataDir() const {
     return "/metadata/ota"s;
 }
@@ -100,6 +98,10 @@
     return kIsRecovery;
 }
 
+bool DeviceInfo::IsFirstStageInit() const {
+    return first_stage_init_;
+}
+
 bool DeviceInfo::SetSlotAsUnbootable([[maybe_unused]] unsigned int slot) {
 #ifdef LIBSNAPSHOT_USE_HAL
     if (!EnsureBootHal()) {
@@ -120,5 +122,22 @@
 #endif
 }
 
+std::unique_ptr<android::fiemap::IImageManager> DeviceInfo::OpenImageManager() const {
+    return IDeviceInfo::OpenImageManager("ota");
+}
+
+std::unique_ptr<android::fiemap::IImageManager> ISnapshotManager::IDeviceInfo::OpenImageManager(
+        const std::string& gsid_dir) const {
+    if (IsRecovery() || IsFirstStageInit()) {
+        android::fiemap::ImageManager::DeviceInfo device_info = {
+                .is_recovery = {IsRecovery()},
+        };
+        return android::fiemap::ImageManager::Open(gsid_dir, device_info);
+    } else {
+        // For now, use a preset timeout.
+        return android::fiemap::IImageManager::Open(gsid_dir, 15000ms);
+    }
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/fs_mgr/libsnapshot/device_info.h b/fs_mgr/libsnapshot/device_info.h
index d8d3d91..7999c99 100644
--- a/fs_mgr/libsnapshot/device_info.h
+++ b/fs_mgr/libsnapshot/device_info.h
@@ -29,7 +29,6 @@
     using MergeStatus = android::hardware::boot::V1_1::MergeStatus;
 
   public:
-    std::string GetGsidDir() const override;
     std::string GetMetadataDir() const override;
     std::string GetSlotSuffix() const override;
     std::string GetOtherSlotSuffix() const override;
@@ -39,11 +38,16 @@
     bool SetBootControlMergeStatus(MergeStatus status) override;
     bool SetSlotAsUnbootable(unsigned int slot) override;
     bool IsRecovery() const override;
+    std::unique_ptr<IImageManager> OpenImageManager() const override;
+    bool IsFirstStageInit() const override;
+
+    void set_first_stage_init(bool value) { first_stage_init_ = value; }
 
   private:
     bool EnsureBootHal();
 
     android::fs_mgr::PartitionOpener opener_;
+    bool first_stage_init_ = false;
 #ifdef LIBSNAPSHOT_USE_HAL
     android::sp<android::hardware::boot::V1_1::IBootControl> boot_control_;
 #endif
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/mock_device_info.h b/fs_mgr/libsnapshot/include/libsnapshot/mock_device_info.h
index ef9d648..573a85b 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/mock_device_info.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/mock_device_info.h
@@ -22,7 +22,6 @@
 
 class MockDeviceInfo : public SnapshotManager::IDeviceInfo {
   public:
-    MOCK_METHOD(std::string, GetGsidDir, (), (const, override));
     MOCK_METHOD(std::string, GetMetadataDir, (), (const, override));
     MOCK_METHOD(std::string, GetSlotSuffix, (), (const, override));
     MOCK_METHOD(std::string, GetOtherSlotSuffix, (), (const, override));
@@ -32,6 +31,9 @@
     MOCK_METHOD(bool, SetBootControlMergeStatus, (MergeStatus status), (override));
     MOCK_METHOD(bool, SetSlotAsUnbootable, (unsigned int slot), (override));
     MOCK_METHOD(bool, IsRecovery, (), (const, override));
+    MOCK_METHOD(bool, IsFirstStageInit, (), (const, override));
+    MOCK_METHOD(std::unique_ptr<android::fiemap::IImageManager>, OpenImageManager, (),
+                (const, override));
 };
 
 }  // namespace android::snapshot
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index d2fdfd6..603e896 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -94,8 +94,9 @@
     // Dependency injection for testing.
     class IDeviceInfo {
       public:
+        using IImageManager = android::fiemap::IImageManager;
+
         virtual ~IDeviceInfo() {}
-        virtual std::string GetGsidDir() const = 0;
         virtual std::string GetMetadataDir() const = 0;
         virtual std::string GetSlotSuffix() const = 0;
         virtual std::string GetOtherSlotSuffix() const = 0;
@@ -107,6 +108,11 @@
         virtual bool SetSlotAsUnbootable(unsigned int slot) = 0;
         virtual bool IsRecovery() const = 0;
         virtual bool IsTestDevice() const { return false; }
+        virtual bool IsFirstStageInit() const = 0;
+        virtual std::unique_ptr<IImageManager> OpenImageManager() const = 0;
+
+        // Helper method for implementing OpenImageManager.
+        std::unique_ptr<IImageManager> OpenImageManager(const std::string& gsid_dir) const;
     };
     virtual ~ISnapshotManager() = default;
 
@@ -420,7 +426,6 @@
     bool EnsureSnapuserdConnected();
 
     // Helpers for first-stage init.
-    bool ForceLocalImageManager();
     const std::unique_ptr<IDeviceInfo>& device() const { return device_; }
 
     // Helper functions for tests.
@@ -765,7 +770,6 @@
     std::string metadata_dir_;
     std::unique_ptr<IDeviceInfo> device_;
     std::unique_ptr<IImageManager> images_;
-    bool has_local_image_manager_ = false;
     bool use_first_stage_snapuserd_ = false;
     bool in_factory_data_reset_ = false;
     std::function<bool(const std::string&)> uevent_regen_callback_;
diff --git a/fs_mgr/libsnapshot/include_test/libsnapshot/test_helpers.h b/fs_mgr/libsnapshot/include_test/libsnapshot/test_helpers.h
index b038527..4e7ccf1 100644
--- a/fs_mgr/libsnapshot/include_test/libsnapshot/test_helpers.h
+++ b/fs_mgr/libsnapshot/include_test/libsnapshot/test_helpers.h
@@ -77,7 +77,6 @@
         : TestDeviceInfo(fake_super) {
         set_slot_suffix(slot_suffix);
     }
-    std::string GetGsidDir() const override { return "ota/test"s; }
     std::string GetMetadataDir() const override { return "/metadata/ota/test"s; }
     std::string GetSlotSuffix() const override { return slot_suffix_; }
     std::string GetOtherSlotSuffix() const override { return slot_suffix_ == "_a" ? "_b" : "_a"; }
@@ -96,6 +95,10 @@
         return true;
     }
     bool IsTestDevice() const override { return true; }
+    bool IsFirstStageInit() const override { return first_stage_init_; }
+    std::unique_ptr<IImageManager> OpenImageManager() const override {
+        return IDeviceInfo::OpenImageManager("ota/test");
+    }
 
     bool IsSlotUnbootable(uint32_t slot) { return unbootable_slots_.count(slot) != 0; }
 
@@ -104,6 +107,7 @@
         opener_ = std::make_unique<TestPartitionOpener>(path);
     }
     void set_recovery(bool value) { recovery_ = value; }
+    void set_first_stage_init(bool value) { first_stage_init_ = value; }
     MergeStatus merge_status() const { return merge_status_; }
 
   private:
@@ -111,6 +115,7 @@
     std::unique_ptr<TestPartitionOpener> opener_;
     MergeStatus merge_status_;
     bool recovery_ = false;
+    bool first_stage_init_ = false;
     std::unordered_set<uint32_t> unbootable_slots_;
 };
 
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 245742e..e2c03ae 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -95,18 +95,16 @@
     if (!info) {
         info = new DeviceInfo();
     }
-    auto sm = std::unique_ptr<SnapshotManager>(new SnapshotManager(info));
-    if (info->IsRecovery()) {
-        sm->ForceLocalImageManager();
-    }
-    return sm;
+    return std::unique_ptr<SnapshotManager>(new SnapshotManager(info));
 }
 
 std::unique_ptr<SnapshotManager> SnapshotManager::NewForFirstStageMount(IDeviceInfo* info) {
-    auto sm = New(info);
-    if (!sm || !sm->ForceLocalImageManager()) {
-        return nullptr;
+    if (!info) {
+        DeviceInfo* impl = new DeviceInfo();
+        impl->set_first_stage_init(true);
+        info = impl;
     }
+    auto sm = New(info);
 
     // The first-stage version of snapuserd is explicitly started by init. Do
     // not attempt to using it during tests (which run in normal AOSP).
@@ -117,7 +115,6 @@
 }
 
 SnapshotManager::SnapshotManager(IDeviceInfo* device) : device_(device) {
-    gsid_dir_ = device_->GetGsidDir();
     metadata_dir_ = device_->GetMetadataDir();
 }
 
@@ -538,9 +535,7 @@
 
     bool ok;
     std::string cow_dev;
-    if (has_local_image_manager_) {
-        // If we forced a local image manager, it means we don't have binder,
-        // which means first-stage init. We must use device-mapper.
+    if (device_->IsRecovery() || device_->IsFirstStageInit()) {
         const auto& opener = device_->GetPartitionOpener();
         ok = images_->MapImageWithDeviceMapper(opener, cow_image_name, &cow_dev);
     } else {
@@ -1836,6 +1831,10 @@
         return false;
     }
 
+    if (!EnsureImageManager()) {
+        return false;
+    }
+
     for (const auto& partition : metadata->partitions) {
         if (GetPartitionGroupName(metadata->groups[partition.group_index]) == kCowGroupName) {
             LOG(INFO) << "Skip mapping partition " << GetPartitionName(partition) << " in group "
@@ -2556,8 +2555,7 @@
 bool SnapshotManager::EnsureImageManager() {
     if (images_) return true;
 
-    // For now, use a preset timeout.
-    images_ = android::fiemap::IImageManager::Open(gsid_dir_, 15000ms);
+    images_ = device_->OpenImageManager();
     if (!images_) {
         LOG(ERROR) << "Could not open ImageManager";
         return false;
@@ -2582,20 +2580,6 @@
     return true;
 }
 
-bool SnapshotManager::ForceLocalImageManager() {
-    android::fiemap::ImageManager::DeviceInfo device_info = {
-            .is_recovery = {device_->IsRecovery()},
-    };
-
-    images_ = android::fiemap::ImageManager::Open(gsid_dir_, device_info);
-    if (!images_) {
-        LOG(ERROR) << "Could not open ImageManager";
-        return false;
-    }
-    has_local_image_manager_ = true;
-    return true;
-}
-
 void SnapshotManager::UnmapAndDeleteCowPartition(MetadataBuilder* current_metadata) {
     std::vector<std::string> to_delete;
     for (auto* existing_cow_partition : current_metadata->ListPartitionsInGroup(kCowGroupName)) {
diff --git a/fs_mgr/libsnapshot/snapshot_fuzz_utils.cpp b/fs_mgr/libsnapshot/snapshot_fuzz_utils.cpp
index 8926535..0096f85 100644
--- a/fs_mgr/libsnapshot/snapshot_fuzz_utils.cpp
+++ b/fs_mgr/libsnapshot/snapshot_fuzz_utils.cpp
@@ -381,8 +381,9 @@
     CheckDetachLoopDevices({Basename(fake_super_), Basename(fake_data_block_device_)});
 }
 
-std::unique_ptr<IImageManager> SnapshotFuzzEnv::CheckCreateFakeImageManager(
-        const std::string& metadata_dir, const std::string& data_dir) {
+std::unique_ptr<IImageManager> SnapshotFuzzEnv::CheckCreateFakeImageManager() {
+    auto metadata_dir = fake_root_->tmp_path() + "/images_manager_metadata";
+    auto data_dir = fake_data_mount_point_ + "/image_manager_data";
     PCHECK(Mkdir(metadata_dir));
     PCHECK(Mkdir(data_dir));
     return SnapshotFuzzImageManager::Open(metadata_dir, data_dir);
@@ -428,13 +429,9 @@
         PCHECK(Mkdir(metadata_dir + "/snapshots"));
     }
 
-    ret.device_info = new SnapshotFuzzDeviceInfo(data.device_info_data(),
+    ret.device_info = new SnapshotFuzzDeviceInfo(this, data.device_info_data(),
                                                  std::move(partition_opener), metadata_dir);
     auto snapshot = SnapshotManager::New(ret.device_info /* takes ownership */);
-    snapshot->images_ =
-            CheckCreateFakeImageManager(fake_root_->tmp_path() + "/images_manager_metadata",
-                                        fake_data_mount_point_ + "/image_manager_data");
-    snapshot->has_local_image_manager_ = data.manager_data().is_local_image_manager();
     ret.snapshot = std::move(snapshot);
 
     return ret;
diff --git a/fs_mgr/libsnapshot/snapshot_fuzz_utils.h b/fs_mgr/libsnapshot/snapshot_fuzz_utils.h
index 5319e69..3ed27c8 100644
--- a/fs_mgr/libsnapshot/snapshot_fuzz_utils.h
+++ b/fs_mgr/libsnapshot/snapshot_fuzz_utils.h
@@ -65,6 +65,8 @@
     // ISnapshotManager.
     SnapshotTestModule CheckCreateSnapshotManager(const SnapshotFuzzData& data);
 
+    std::unique_ptr<android::fiemap::IImageManager> CheckCreateFakeImageManager();
+
     // Return path to super partition.
     const std::string& super() const;
 
@@ -79,8 +81,6 @@
     std::string fake_data_block_device_;
     std::unique_ptr<AutoDevice> mounted_data_;
 
-    static std::unique_ptr<android::fiemap::IImageManager> CheckCreateFakeImageManager(
-            const std::string& metadata_dir, const std::string& data_dir);
     static std::unique_ptr<AutoDevice> CheckMapImage(const std::string& fake_persist_path,
                                                      uint64_t size,
                                                      android::dm::LoopControl* control,
@@ -95,15 +95,15 @@
 class SnapshotFuzzDeviceInfo : public ISnapshotManager::IDeviceInfo {
   public:
     // Client is responsible for maintaining the lifetime of |data|.
-    SnapshotFuzzDeviceInfo(const FuzzDeviceInfoData& data,
+    SnapshotFuzzDeviceInfo(SnapshotFuzzEnv* env, const FuzzDeviceInfoData& data,
                            std::unique_ptr<TestPartitionOpener>&& partition_opener,
                            const std::string& metadata_dir)
-        : data_(&data),
+        : env_(env),
+          data_(&data),
           partition_opener_(std::move(partition_opener)),
           metadata_dir_(metadata_dir) {}
 
     // Following APIs are mocked.
-    std::string GetGsidDir() const override { return "fuzz_ota"; }
     std::string GetMetadataDir() const override { return metadata_dir_; }
     std::string GetSuperDevice(uint32_t) const override {
         // TestPartitionOpener can recognize this.
@@ -124,10 +124,15 @@
         return data_->allow_set_slot_as_unbootable();
     }
     bool IsRecovery() const override { return data_->is_recovery(); }
+    bool IsFirstStageInit() const override { return false; }
+    std::unique_ptr<IImageManager> OpenImageManager() const {
+        return env_->CheckCreateFakeImageManager();
+    }
 
     void SwitchSlot() { switched_slot_ = !switched_slot_; }
 
   private:
+    SnapshotFuzzEnv* env_;
     const FuzzDeviceInfoData* data_;
     std::unique_ptr<TestPartitionOpener> partition_opener_;
     std::string metadata_dir_;
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 9226b7f..6018643 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -403,6 +403,7 @@
     }
 
     std::unique_ptr<SnapshotManager> NewManagerForFirstStageMount(TestDeviceInfo* info) {
+        info->set_first_stage_init(true);
         auto init = SnapshotManager::NewForFirstStageMount(info);
         if (!init) {
             return nullptr;