Use the injection system for InitBootloaderEnvPartition

Test: launch_cvd --noresume
Bug: 194843194
Change-Id: I22cf825692f134cab27d60f93255e79751c0cd0d
diff --git a/host/commands/assemble_cvd/boot_config.cc b/host/commands/assemble_cvd/boot_config.cc
index 33d24f4..0f74847 100644
--- a/host/commands/assemble_cvd/boot_config.cc
+++ b/host/commands/assemble_cvd/boot_config.cc
@@ -80,63 +80,89 @@
 
 }  // namespace
 
+class InitBootloaderEnvPartitionImpl : public InitBootloaderEnvPartition {
+ public:
+  INJECT(InitBootloaderEnvPartitionImpl(
+      const CuttlefishConfig& config,
+      const CuttlefishConfig::InstanceSpecific& instance))
+      : config_(config), instance_(instance) {}
 
-bool InitBootloaderEnvPartition(const CuttlefishConfig& config,
-                                const CuttlefishConfig::InstanceSpecific& instance) {
-  auto boot_env_image_path = instance.uboot_env_image_path();
-  auto tmp_boot_env_image_path = boot_env_image_path + ".tmp";
-  auto uboot_env_path = instance.PerInstancePath("mkenvimg_input");
-  auto kernel_cmdline =
-      android::base::Join(KernelCommandLineFromConfig(config), " ");
-  // If the bootconfig isn't supported in the guest kernel, the bootconfig args
-  // need to be passed in via the uboot env. This won't be an issue for protect
-  // kvm which is running a kernel with bootconfig support.
-  if (!config.bootconfig_supported()) {
-    auto bootconfig_args =
-        android::base::Join(BootconfigArgsFromConfig(config, instance), " ");
-    // "androidboot.hardware" kernel parameter has changed to "hardware" in
-    // bootconfig and needs to be replaced before being used in the kernel
-    // cmdline.
-    bootconfig_args = android::base::StringReplace(
-        bootconfig_args, " hardware=", " androidboot.hardware=", true);
-    // TODO(b/182417593): Until we pass the module parameters through
-    // modules.options, we pass them through bootconfig using
-    // 'kernel.<key>=<value>' But if we don't support bootconfig, we need to
-    // rename them back to the old cmdline version
-    bootconfig_args =
-        android::base::StringReplace(bootconfig_args, " kernel.", " ", true);
-    kernel_cmdline += " ";
-    kernel_cmdline += bootconfig_args;
-  }
-  if (!WriteEnvironment(config, kernel_cmdline, uboot_env_path)) {
-    LOG(ERROR) << "Unable to write out plaintext env '" << uboot_env_path << ".'";
-    return false;
-  }
+  // Feature
+  std::string Name() const override { return "InitBootloaderEnvPartitionImpl"; }
+  bool Enabled() const override { return !config_.protected_vm(); }
 
-  auto mkimage_path = HostBinaryPath("mkenvimage");
-  Command cmd(mkimage_path);
-  cmd.AddParameter("-s");
-  cmd.AddParameter("4096");
-  cmd.AddParameter("-o");
-  cmd.AddParameter(tmp_boot_env_image_path);
-  cmd.AddParameter(uboot_env_path);
-  int success = cmd.Start().Wait();
-  if (success != 0) {
-    LOG(ERROR) << "Unable to run mkenvimage. Exited with status " << success;
-    return false;
-  }
-
-  if(!FileExists(boot_env_image_path) || ReadFile(boot_env_image_path) != ReadFile(tmp_boot_env_image_path)) {
-    if(!RenameFile(tmp_boot_env_image_path, boot_env_image_path)) {
-      LOG(ERROR) << "Unable to delete the old env image.";
+ private:
+  std::unordered_set<Feature*> Dependencies() const override { return {}; }
+  bool Setup() override {
+    auto boot_env_image_path = instance_.uboot_env_image_path();
+    auto tmp_boot_env_image_path = boot_env_image_path + ".tmp";
+    auto uboot_env_path = instance_.PerInstancePath("mkenvimg_input");
+    auto kernel_cmdline =
+        android::base::Join(KernelCommandLineFromConfig(config_), " ");
+    // If the bootconfig isn't supported in the guest kernel, the bootconfig
+    // args need to be passed in via the uboot env. This won't be an issue for
+    // protect kvm which is running a kernel with bootconfig support.
+    if (!config_.bootconfig_supported()) {
+      auto bootconfig_args = android::base::Join(
+          BootconfigArgsFromConfig(config_, instance_), " ");
+      // "androidboot.hardware" kernel parameter has changed to "hardware" in
+      // bootconfig and needs to be replaced before being used in the kernel
+      // cmdline.
+      bootconfig_args = android::base::StringReplace(
+          bootconfig_args, " hardware=", " androidboot.hardware=", true);
+      // TODO(b/182417593): Until we pass the module parameters through
+      // modules.options, we pass them through bootconfig using
+      // 'kernel.<key>=<value>' But if we don't support bootconfig, we need to
+      // rename them back to the old cmdline version
+      bootconfig_args =
+          android::base::StringReplace(bootconfig_args, " kernel.", " ", true);
+      kernel_cmdline += " ";
+      kernel_cmdline += bootconfig_args;
+    }
+    if (!WriteEnvironment(config_, kernel_cmdline, uboot_env_path)) {
+      LOG(ERROR) << "Unable to write out plaintext env '" << uboot_env_path
+                 << ".'";
       return false;
     }
-    LOG(DEBUG) << "Updated bootloader environment image.";
-  } else {
-    RemoveFile(tmp_boot_env_image_path);
+
+    auto mkimage_path = HostBinaryPath("mkenvimage");
+    Command cmd(mkimage_path);
+    cmd.AddParameter("-s");
+    cmd.AddParameter("4096");
+    cmd.AddParameter("-o");
+    cmd.AddParameter(tmp_boot_env_image_path);
+    cmd.AddParameter(uboot_env_path);
+    int success = cmd.Start().Wait();
+    if (success != 0) {
+      LOG(ERROR) << "Unable to run mkenvimage. Exited with status " << success;
+      return false;
+    }
+
+    if (!FileExists(boot_env_image_path) ||
+        ReadFile(boot_env_image_path) != ReadFile(tmp_boot_env_image_path)) {
+      if (!RenameFile(tmp_boot_env_image_path, boot_env_image_path)) {
+        LOG(ERROR) << "Unable to delete the old env image.";
+        return false;
+      }
+      LOG(DEBUG) << "Updated bootloader environment image.";
+    } else {
+      RemoveFile(tmp_boot_env_image_path);
+    }
+
+    return true;
   }
 
-  return true;
+  const CuttlefishConfig& config_;
+  const CuttlefishConfig::InstanceSpecific& instance_;
+};
+
+fruit::Component<fruit::Required<const CuttlefishConfig,
+                                 const CuttlefishConfig::InstanceSpecific>,
+                 InitBootloaderEnvPartition>
+InitBootloaderEnvPartitionComponent() {
+  return fruit::createComponent()
+      .bind<InitBootloaderEnvPartition, InitBootloaderEnvPartitionImpl>()
+      .addMultibinding<Feature, InitBootloaderEnvPartition>();
 }
 
 } // namespace cuttlefish
diff --git a/host/commands/assemble_cvd/boot_config.h b/host/commands/assemble_cvd/boot_config.h
index 1a81b79..74109dd 100644
--- a/host/commands/assemble_cvd/boot_config.h
+++ b/host/commands/assemble_cvd/boot_config.h
@@ -16,11 +16,19 @@
 #pragma once
 
 #include <string>
-#include <host/libs/config/cuttlefish_config.h>
+
+#include <fruit/fruit.h>
+
+#include "host/libs/config/cuttlefish_config.h"
+#include "host/libs/config/feature.h"
 
 namespace cuttlefish {
 
-bool InitBootloaderEnvPartition(const CuttlefishConfig& config,
-                                const CuttlefishConfig::InstanceSpecific& instance);
+class InitBootloaderEnvPartition : public Feature {};
+
+fruit::Component<fruit::Required<const CuttlefishConfig,
+                                 const CuttlefishConfig::InstanceSpecific>,
+                 InitBootloaderEnvPartition>
+InitBootloaderEnvPartitionComponent();
 
 } // namespace cuttlefish
diff --git a/host/commands/assemble_cvd/disk_flags.cc b/host/commands/assemble_cvd/disk_flags.cc
index 79a645a..e07e6a7 100644
--- a/host/commands/assemble_cvd/disk_flags.cc
+++ b/host/commands/assemble_cvd/disk_flags.cc
@@ -542,6 +542,68 @@
   const CuttlefishConfig::InstanceSpecific& instance_;
 };
 
+class InitializePstore : public Feature {
+ public:
+  INJECT(InitializePstore(const CuttlefishConfig& config,
+                          const CuttlefishConfig::InstanceSpecific& instance))
+      : config_(config), instance_(instance) {}
+
+  // Feature
+  std::string Name() const override { return "InitializePstore"; }
+  bool Enabled() const override { return !config_.protected_vm(); }
+
+ private:
+  std::unordered_set<Feature*> Dependencies() const override { return {}; }
+  bool Setup() {
+    if (FileExists(instance_.pstore_path())) {
+      return true;
+    }
+    bool success =
+        CreateBlankImage(instance_.pstore_path(), 2 /* mb */, "none");
+    if (!success) {
+      LOG(ERROR) << "Failed to create pstore_path \"" << instance_.pstore_path()
+                 << "\"";
+      return false;
+    }
+    return true;
+  }
+
+  const CuttlefishConfig& config_;
+  const CuttlefishConfig::InstanceSpecific& instance_;
+};
+
+class InitializeSdCard : public Feature {
+ public:
+  INJECT(InitializeSdCard(const CuttlefishConfig& config,
+                          const CuttlefishConfig::InstanceSpecific& instance))
+      : config_(config), instance_(instance) {}
+
+  // Feature
+  std::string Name() const override { return "InitializeSdCard"; }
+  bool Enabled() const override {
+    return FLAGS_use_sdcard && !config_.protected_vm();
+  }
+
+ private:
+  std::unordered_set<Feature*> Dependencies() const override { return {}; }
+  bool Setup() {
+    if (FileExists(instance_.sdcard_path())) {
+      return true;
+    }
+    bool success = CreateBlankImage(instance_.sdcard_path(),
+                                    FLAGS_blank_sdcard_image_mb, "sdcard");
+    if (!success) {
+      LOG(ERROR) << "Failed to create sdcard \"" << instance_.sdcard_path()
+                 << "\"";
+      return false;
+    }
+    return true;
+  }
+
+  const CuttlefishConfig& config_;
+  const CuttlefishConfig::InstanceSpecific& instance_;
+};
+
 static fruit::Component<> DiskChangesComponent(const FetcherConfig* fetcher,
                                                const CuttlefishConfig* config) {
   return fruit::createComponent()
@@ -565,7 +627,10 @@
       .bindInstance(*fetcher)
       .bindInstance(*config)
       .bindInstance(*instance)
-      .addMultibinding<Feature, InitializeAccessKregistryImage>();
+      .addMultibinding<Feature, InitializeAccessKregistryImage>()
+      .addMultibinding<Feature, InitializePstore>()
+      .addMultibinding<Feature, InitializeSdCard>()
+      .install(InitBootloaderEnvPartitionComponent);
 }
 
 void CreateDynamicDiskFiles(const FetcherConfig& fetcher_config,
@@ -592,18 +657,6 @@
   // support.
   if (!FLAGS_protected_vm) {
     for (const auto& instance : config.Instances()) {
-      if (!FileExists(instance.pstore_path())) {
-        CreateBlankImage(instance.pstore_path(), 2 /* mb */, "none");
-      }
-
-      if (FLAGS_use_sdcard && !FileExists(instance.sdcard_path())) {
-        CreateBlankImage(instance.sdcard_path(),
-                         FLAGS_blank_sdcard_image_mb, "sdcard");
-      }
-
-      CHECK(InitBootloaderEnvPartition(config, instance))
-          << "Failed to create bootloader environment partition";
-
       const auto frp = instance.factory_reset_protected_path();
       if (!FileExists(frp)) {
         CreateBlankImage(frp, 1 /* mb */, "none");