Read additional physical partitions from AndroidBuild

base/cvd/cuttlefish/host/commands/assemble_cvd/BUILD.bazel

@@ -163,6 +163,7 @@ cf_cc_library(
         "//cuttlefish/host/commands/assemble_cvd:flags_defaults",
         "//cuttlefish/host/commands/assemble_cvd:instance_image_files",
         "//cuttlefish/host/commands/assemble_cvd:super_image_mixer",
+        "//cuttlefish/host/commands/assemble_cvd/android_build:android_builds",
         "//cuttlefish/host/commands/assemble_cvd/disk:access_kregistry",
         "//cuttlefish/host/commands/assemble_cvd/disk:ap_composite_disk",
         "//cuttlefish/host/commands/assemble_cvd/disk:bootloader_present",

base/cvd/cuttlefish/host/commands/assemble_cvd/assemble_cvd.cc

@@ -299,7 +299,8 @@ Result<const CuttlefishConfig*> InitFilesystemAndCreateConfig(
     const SuperImageFlag& super_image,
     const SystemImageDirFlag& system_image_dir,
     const VendorBootImageFlag& vendor_boot_image,
-    const VmManagerFlag& vm_manager_flag, const Defaults& defaults) {
+    const VmManagerFlag& vm_manager_flag, const Defaults& defaults,
+    AndroidBuilds& android_builds) {
   {
     // The config object is created here, but only exists in memory until the
     // SaveConfig line below. Don't launch cuttlefish subprocesses between these
@@ -360,9 +361,9 @@ Result<const CuttlefishConfig*> InitFilesystemAndCreateConfig(
       Result<std::optional<ChromeOsStateImage>> chrome_os_state =
           CF_EXPECT(ChromeOsStateImage::Reuse(instance));
       if (chrome_os_state.ok()) {
-        Result<DiskBuilder> os_builder =
-            OsCompositeDiskBuilder(config, instance, *chrome_os_state,
-                                   instance_image_files, system_image_dir);
+        Result<DiskBuilder> os_builder = OsCompositeDiskBuilder(
+            config, instance, *chrome_os_state, instance_image_files,
+            android_builds.ForIndex(index), system_image_dir);
         if (!os_builder.ok()) {
           creating_os_disk = true;
         } else {
@@ -469,7 +470,8 @@ Result<const CuttlefishConfig*> InitFilesystemAndCreateConfig(
     CF_EXPECT(Symlink(first_instance, double_legacy_instance_dir));
   }
 
-  CF_EXPECT(CreateDynamicDiskFiles(fetcher_configs, *config, system_image_dir));
+  CF_EXPECT(CreateDynamicDiskFiles(fetcher_configs, *config, android_builds,
+                                   system_image_dir));
 
   return config;
 }
@@ -666,7 +668,7 @@ Result<int> AssembleCvdMain(int argc, char** argv) {
       InitFilesystemAndCreateConfig(
           std::move(fetcher_configs), guest_configs, injector, log, boot_image,
           initramfs_path, kernel_path, super_image, system_image_dir,
-          vendor_boot_image, vm_manager_flag, *defaults),
+          vendor_boot_image, vm_manager_flag, *defaults, android_builds),
       "Failed to create config");
 
   std::cout << GetConfigFilePath(*config) << "\n";

base/cvd/cuttlefish/host/commands/assemble_cvd/create_dynamic_disk_files.cc

@@ -29,6 +29,7 @@
 #include "absl/strings/match.h"
 
 #include "cuttlefish/common/libs/utils/files.h"
+#include "cuttlefish/host/commands/assemble_cvd/android_build/android_builds.h"
 #include "cuttlefish/host/commands/assemble_cvd/boot_config.h"
 #include "cuttlefish/host/commands/assemble_cvd/boot_image_utils.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk/access_kregistry.h"
@@ -98,6 +99,7 @@ Result<BuildArchive> FindImgZip(const FetcherConfig& fetcher_config,
 
 Result<void> CreateDynamicDiskFiles(
     const FetcherConfigs& fetcher_configs, const CuttlefishConfig& config,
+    AndroidBuilds& android_builds,
     const SystemImageDirFlag& system_image_dirs) {
   std::vector<std::vector<std::unique_ptr<ImageFile>>> image_files =
       InstanceImageFiles(config);
@@ -170,9 +172,9 @@ Result<void> CreateDynamicDiskFiles(
       CF_EXPECT(image_file->Generate());
     }
 
-    DiskBuilder os_disk_builder = CF_EXPECT(
-        OsCompositeDiskBuilder(config, instance, chrome_os_state,
-                               instance_image_files, system_image_dirs));
+    DiskBuilder os_disk_builder = CF_EXPECT(OsCompositeDiskBuilder(
+        config, instance, chrome_os_state, instance_image_files,
+        android_builds.ForIndex(instance_index), system_image_dirs));
     const auto os_built_composite =
         CF_EXPECT(os_disk_builder.BuildCompositeDiskIfNecessary());
 

base/cvd/cuttlefish/host/commands/assemble_cvd/create_dynamic_disk_files.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include "cuttlefish/host/commands/assemble_cvd/android_build/android_builds.h"
 #include "cuttlefish/host/commands/assemble_cvd/flags/system_image_dir.h"
 #include "cuttlefish/host/libs/config/cuttlefish_config.h"
 #include "cuttlefish/host/libs/config/fetcher_configs.h"
@@ -25,6 +26,6 @@ namespace cuttlefish {
 
 Result<void> CreateDynamicDiskFiles(const FetcherConfigs& fetcher_config,
                                     const CuttlefishConfig& config,
-                                    const SystemImageDirFlag&);
+                                    AndroidBuilds&, const SystemImageDirFlag&);
 
 }  // namespace cuttlefish

base/cvd/cuttlefish/host/commands/assemble_cvd/disk/BUILD.bazel

@@ -23,6 +23,7 @@ cf_cc_library(
     hdrs = ["android_composite_disk_config.h"],
     deps = [
         "//cuttlefish/common/libs/utils:files",
+        "//cuttlefish/host/commands/assemble_cvd/android_build",
         "//cuttlefish/host/commands/assemble_cvd/disk:image_file",
         "//cuttlefish/host/commands/assemble_cvd/flags:system_image_dir",
         "//cuttlefish/host/libs/config:cuttlefish_config",
@@ -39,6 +40,7 @@ cf_cc_library(
     hdrs = ["android_efi_loader_composite_disk.h"],
     deps = [
         "//cuttlefish/common/libs/utils:files",
+        "//cuttlefish/host/commands/assemble_cvd/android_build",
         "//cuttlefish/host/commands/assemble_cvd/disk:android_composite_disk_config",
         "//cuttlefish/host/commands/assemble_cvd/disk:image_file",
         "//cuttlefish/host/commands/assemble_cvd/flags:system_image_dir",
@@ -280,6 +282,7 @@ cf_cc_library(
     deps = [
         "//cuttlefish/host/commands/assemble_cvd:assemble_cvd_flags",
         "//cuttlefish/host/commands/assemble_cvd:disk_builder",
+        "//cuttlefish/host/commands/assemble_cvd/android_build",
         "//cuttlefish/host/commands/assemble_cvd/disk:android_composite_disk_config",
         "//cuttlefish/host/commands/assemble_cvd/disk:android_efi_loader_composite_disk",
         "//cuttlefish/host/commands/assemble_cvd/disk:chromeos_composite_disk",

base/cvd/cuttlefish/host/commands/assemble_cvd/disk/android_composite_disk_config.cc

@@ -29,6 +29,7 @@
 #include "android-base/strings.h"
 
 #include "cuttlefish/common/libs/utils/files.h"
+#include "cuttlefish/host/commands/assemble_cvd/android_build/android_build.h"
 #include "cuttlefish/host/commands/assemble_cvd/flags/system_image_dir.h"
 #include "cuttlefish/host/libs/config/cuttlefish_config.h"
 #include "cuttlefish/host/libs/image_aggregator/image_aggregator.h"
@@ -71,7 +72,7 @@ std::optional<ImagePartition> HibernationImage(
 Result<std::vector<ImagePartition>> AndroidCompositeDiskConfig(
     const CuttlefishConfig::InstanceSpecific& instance,
     const std::vector<std::unique_ptr<ImageFile>>& image_files,
-    const SystemImageDirFlag& system_image_dir) {
+    AndroidBuild& android_build, const SystemImageDirFlag& system_image_dir) {
   std::vector<ImagePartition> partitions;
 
   const std::set<std::string_view> ab_partitions = {
@@ -90,21 +91,28 @@ Result<std::vector<ImagePartition>> AndroidCompositeDiskConfig(
       kPartitions.vvmtruststore,
   };
 
-  const std::map<std::string_view, std::string> primary_paths = {
-      {kPartitions.boot, instance.new_boot_image()},
-      {kPartitions.init_boot, instance.init_boot_image()},
-      {kPartitions.metadata, ""},
-      {kPartitions.misc, ""},
-      {kPartitions.super, instance.new_super_image()},
-      {kPartitions.userdata, instance.new_data_image()},
-      {kPartitions.vbmeta, instance.new_vbmeta_image()},
-      {kPartitions.vbmeta_system, instance.vbmeta_system_image()},
-      {kPartitions.vbmeta_system_dlkm, instance.new_vbmeta_system_dlkm_image()},
-      {kPartitions.vbmeta_vendor_dlkm, instance.new_vbmeta_vendor_dlkm_image()},
-      {kPartitions.vendor_boot, instance.new_vendor_boot_image()},
-      {kPartitions.vvmtruststore, instance.vvmtruststore_path()},
+  std::map<std::string, std::string> primary_paths = {
+      {std::string(kPartitions.boot), instance.new_boot_image()},
+      {std::string(kPartitions.init_boot), instance.init_boot_image()},
+      {std::string(kPartitions.metadata), ""},
+      {std::string(kPartitions.misc), ""},
+      {std::string(kPartitions.super), instance.new_super_image()},
+      {std::string(kPartitions.userdata), instance.new_data_image()},
+      {std::string(kPartitions.vbmeta), instance.new_vbmeta_image()},
+      {std::string(kPartitions.vbmeta_system), instance.vbmeta_system_image()},
+      {std::string(kPartitions.vbmeta_system_dlkm),
+       instance.new_vbmeta_system_dlkm_image()},
+      {std::string(kPartitions.vbmeta_vendor_dlkm),
+       instance.new_vbmeta_vendor_dlkm_image()},
+      {std::string(kPartitions.vendor_boot), instance.new_vendor_boot_image()},
+      {std::string(kPartitions.vvmtruststore), instance.vvmtruststore_path()},
   };
 
+  for (std::string partition : CF_EXPECT(android_build.PhysicalPartitions())) {
+    std::string path = CF_EXPECT(android_build.ImageFile(partition));
+    primary_paths.try_emplace(std::move(partition), std::move(path));
+  }
+
   const std::map<std::string_view, std::string> fallback_paths = {
       {kPartitions.super, instance.super_image()},
       {kPartitions.vbmeta, instance.vbmeta_image()},

base/cvd/cuttlefish/host/commands/assemble_cvd/disk/android_composite_disk_config.h

@@ -18,6 +18,7 @@
 #include <memory>
 #include <vector>
 
+#include "cuttlefish/host/commands/assemble_cvd/android_build/android_build.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk/image_file.h"
 #include "cuttlefish/host/commands/assemble_cvd/flags/system_image_dir.h"
 #include "cuttlefish/host/libs/config/cuttlefish_config.h"
@@ -28,6 +29,7 @@ namespace cuttlefish {
 
 Result<std::vector<ImagePartition>> AndroidCompositeDiskConfig(
     const CuttlefishConfig::InstanceSpecific& instance,
-    const std::vector<std::unique_ptr<ImageFile>>&, const SystemImageDirFlag&);
+    const std::vector<std::unique_ptr<ImageFile>>&, AndroidBuild&,
+    const SystemImageDirFlag&);
 
 }  // namespace cuttlefish

base/cvd/cuttlefish/host/commands/assemble_cvd/disk/android_efi_loader_composite_disk.cc

@@ -20,6 +20,7 @@
 #include <vector>
 
 #include "cuttlefish/common/libs/utils/files.h"
+#include "cuttlefish/host/commands/assemble_cvd/android_build/android_build.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk/android_composite_disk_config.h"
 #include "cuttlefish/host/commands/assemble_cvd/flags/system_image_dir.h"
 #include "cuttlefish/host/libs/config/cuttlefish_config.h"
@@ -32,9 +33,9 @@ namespace cuttlefish {
 Result<std::vector<ImagePartition>> AndroidEfiLoaderCompositeDiskConfig(
     const CuttlefishConfig::InstanceSpecific& instance,
     const std::vector<std::unique_ptr<ImageFile>>& image_files,
-    const SystemImageDirFlag& system_image_dir) {
-  std::vector<ImagePartition> partitions = CF_EXPECT(
-      AndroidCompositeDiskConfig(instance, image_files, system_image_dir));
+    AndroidBuild& android_build, const SystemImageDirFlag& system_image_dir) {
+  std::vector<ImagePartition> partitions = CF_EXPECT(AndroidCompositeDiskConfig(
+      instance, image_files, android_build, system_image_dir));
   // Cuttlefish uboot EFI bootflow by default looks at the first partition
   // for EFI application. Thus we put "android_esp" at the beginning.
   partitions.insert(

base/cvd/cuttlefish/host/commands/assemble_cvd/disk/android_efi_loader_composite_disk.h

@@ -19,6 +19,7 @@
 #include <memory>
 #include <vector>
 
+#include "cuttlefish/host/commands/assemble_cvd/android_build/android_build.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk/image_file.h"
 #include "cuttlefish/host/commands/assemble_cvd/flags/system_image_dir.h"
 #include "cuttlefish/host/libs/config/cuttlefish_config.h"
@@ -29,6 +30,7 @@ namespace cuttlefish {
 
 Result<std::vector<ImagePartition>> AndroidEfiLoaderCompositeDiskConfig(
     const CuttlefishConfig::InstanceSpecific& instance,
-    const std::vector<std::unique_ptr<ImageFile>>&, const SystemImageDirFlag&);
+    const std::vector<std::unique_ptr<ImageFile>>&, AndroidBuild&,
+    const SystemImageDirFlag&);
 
 }  // namespace cuttlefish

base/cvd/cuttlefish/host/commands/assemble_cvd/disk/os_composite_disk.cc

@@ -22,6 +22,7 @@
 
 #include "absl/log/check.h"
 
+#include "cuttlefish/host/commands/assemble_cvd/android_build/android_build.h"
 #include "cuttlefish/host/commands/assemble_cvd/assemble_cvd_flags.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk/android_composite_disk_config.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk/android_efi_loader_composite_disk.h"
@@ -44,14 +45,14 @@ Result<std::vector<ImagePartition>> GetOsCompositeDiskConfig(
     const CuttlefishConfig::InstanceSpecific& instance,
     const std::optional<ChromeOsStateImage>& chrome_os_state,
     const std::vector<std::unique_ptr<ImageFile>>& image_files,
-    const SystemImageDirFlag& system_image_dir) {
+    AndroidBuild& android_build, const SystemImageDirFlag& system_image_dir) {
   switch (instance.boot_flow()) {
     case BootFlow::Android:
-      return CF_EXPECT(
-          AndroidCompositeDiskConfig(instance, image_files, system_image_dir));
+      return CF_EXPECT(AndroidCompositeDiskConfig(
+          instance, image_files, android_build, system_image_dir));
     case BootFlow::AndroidEfiLoader:
       return CF_EXPECT(AndroidEfiLoaderCompositeDiskConfig(
-          instance, image_files, system_image_dir));
+          instance, image_files, android_build, system_image_dir));
     case BootFlow::ChromeOs:
       CHECK(chrome_os_state.has_value());
       return ChromeOsCompositeDiskConfig(instance, *chrome_os_state);
@@ -71,7 +72,7 @@ Result<DiskBuilder> OsCompositeDiskBuilder(
     const CuttlefishConfig::InstanceSpecific& instance,
     const std::optional<ChromeOsStateImage>& chrome_os_state,
     const std::vector<std::unique_ptr<ImageFile>>& image_files,
-    const SystemImageDirFlag& system_image_dir) {
+    AndroidBuild& android_build, const SystemImageDirFlag& system_image_dir) {
   auto builder =
       DiskBuilder()
           .VmManager(config.vm_manager())
@@ -84,8 +85,9 @@ Result<DiskBuilder> OsCompositeDiskBuilder(
         .CompositeDiskPath(instance.chromeos_disk());
   }
   return builder
-      .Partitions(CF_EXPECT(GetOsCompositeDiskConfig(
-          instance, chrome_os_state, image_files, system_image_dir)))
+      .Partitions(CF_EXPECT(GetOsCompositeDiskConfig(instance, chrome_os_state,
+                                                     image_files, android_build,
+                                                     system_image_dir)))
       .HeaderPath(instance.PerInstancePath("os_composite_gpt_header.img"))
       .FooterPath(instance.PerInstancePath("os_composite_gpt_footer.img"))
       .CompositeDiskPath(instance.os_composite_disk_path());

base/cvd/cuttlefish/host/commands/assemble_cvd/disk/os_composite_disk.h

@@ -20,6 +20,7 @@
 #include <optional>
 #include <vector>
 
+#include "cuttlefish/host/commands/assemble_cvd/android_build/android_build.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk/chromeos_state.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk/image_file.h"
 #include "cuttlefish/host/commands/assemble_cvd/disk_builder.h"
@@ -33,6 +34,7 @@ Result<DiskBuilder> OsCompositeDiskBuilder(
     const CuttlefishConfig& config,
     const CuttlefishConfig::InstanceSpecific& instance,
     const std::optional<ChromeOsStateImage>&,
-    const std::vector<std::unique_ptr<ImageFile>>&, const SystemImageDirFlag&);
+    const std::vector<std::unique_ptr<ImageFile>>&, AndroidBuild&,
+    const SystemImageDirFlag&);
 
 }  // namespace cuttlefish