Merge "libsnapshot: Move ImageManager creation to DeviceInfo."
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;