Add flags to start acloud with efi android loader

Example usage:
--android-efi-loader-build-id 11544241
--android-efi-loader-artifact gbl_x86_32.efi

Test: acloud-dev create
      --android-efi-loader-build-id 11544241
      --android-efi-loader-artifact gbl_x86_32.efi
Test: ./run_tests.sh
Bug: 322983262
Change-Id: I1f98b5be707ac81536f6f40e25be30c0b961b3cd
Signed-off-by: Dmitrii Merkurev <dimorinny@google.com>

11 files changed, 92 insertions, 21 deletions

diff --git a/create/avd_spec.py b/create/avd_spec.py
index f5da7fed..8b2f984a 100644
--- a/create/avd_spec.py
+++ b/create/avd_spec.py
@@ -139,6 +139,7 @@ class AVDSpec():
         self._ota_build_info = {}
         self._host_package_build_info = {}
         self._bootloader_build_info = {}
+        self._android_efi_loader_build_info = {}
         self._hw_property = None
         self._hw_customize = False
         self._remote_host = None
@@ -672,6 +673,9 @@ class AVDSpec():
             constants.BUILD_ID: args.bootloader_build_id,
             constants.BUILD_BRANCH: args.bootloader_branch,
             constants.BUILD_TARGET: args.bootloader_build_target}
+        self._android_efi_loader_build_info = {
+            constants.BUILD_ID: args.android_efi_loader_build_id,
+            constants.BUILD_ARTIFACT: args.android_efi_loader_artifact}
         self._host_package_build_info = {
             constants.BUILD_ID: args.host_package_build_id,
             constants.BUILD_BRANCH: args.host_package_branch,
@@ -963,6 +967,11 @@ class AVDSpec():
         return self._bootloader_build_info
 
     @property
+    def android_efi_loader_build_info(self):
+        """Return android efi loader build info."""
+        return self._android_efi_loader_build_info
+
+    @property
     def flavor(self):
         """Return flavor."""
         return self._flavor
diff --git a/create/create_args.py b/create/create_args.py
index cba574d5..fadf528d 100644
--- a/create/create_args.py
+++ b/create/create_args.py
@@ -173,6 +173,18 @@ def AddCommonCreateArgs(parser):
         help="'cuttlefish only' Bootloader build target.",
         required=False)
     parser.add_argument(
+        "--android-efi-loader-build-id",
+        type=str,
+        dest="android_efi_loader_build_id",
+        help="'cuttlefish only' Android EFI loader build id, e.g. P2804227",
+        required=False)
+    parser.add_argument(
+        "--android-efi-loader-artifact",
+        type=str,
+        dest="android_efi_loader_artifact",
+        help="'cuttlefish only' Android EFI loader artifact name, e.g. gbl_aarch64.efi",
+        required=False)
+    parser.add_argument(
         "--kernel-build-id",
         type=str,
         dest="kernel_build_id",
diff --git a/create/remote_image_local_instance.py b/create/remote_image_local_instance.py
index a49b63df..e8feaf39 100644
--- a/create/remote_image_local_instance.py
+++ b/create/remote_image_local_instance.py
@@ -122,6 +122,7 @@ def DownloadAndProcessImageFiles(avd_spec):
         avd_spec.kernel_build_info,
         avd_spec.boot_build_info,
         avd_spec.bootloader_build_info,
+        avd_spec.android_efi_loader_build_info,
         avd_spec.ota_build_info,
         avd_spec.host_package_build_info)
 
diff --git a/internal/lib/android_build_client.py b/internal/lib/android_build_client.py
index b4ad23fd..4960bd0d 100644
--- a/internal/lib/android_build_client.py
+++ b/internal/lib/android_build_client.py
@@ -165,11 +165,12 @@ class AndroidBuildClient(base_cloud_client.BaseCloudApiClient):
         os.chmod(local_dest, fetch_cvd_stat.st_mode | stat.S_IEXEC)
 
     @staticmethod
-    def ProcessBuild(build_info):
+    def ProcessBuild(build_info, ignore_artifact=False):
         """Create a Cuttlefish fetch_cvd build string.
 
         Args:
             build_info: The dictionary that contains build information.
+            ignore_artifact: Avoid adding artifact part to fetch_cvd build string
 
         Returns:
             A string, used in the fetch_cvd cmd or None if all args are None.
@@ -177,17 +178,22 @@ class AndroidBuildClient(base_cloud_client.BaseCloudApiClient):
         build_id = build_info.get(constants.BUILD_ID)
         build_target = build_info.get(constants.BUILD_TARGET)
         branch = build_info.get(constants.BUILD_BRANCH)
-        if not build_target:
-            return build_id or branch
+        artifact = build_info.get(constants.BUILD_ARTIFACT)
 
-        if build_target and not branch:
-            branch = _DEFAULT_BRANCH
-        return (build_id or branch) + "/" + build_target
+        result = build_id or branch
+        if build_target is not None:
+            result = result or _DEFAULT_BRANCH
+            result += "/" + build_target
+
+        if not ignore_artifact and artifact:
+            result += "{" + artifact + "}"
+
+        return result
 
     def GetFetchBuildArgs(self, default_build_info, system_build_info,
                           kernel_build_info, boot_build_info,
-                          bootloader_build_info, ota_build_info,
-                          host_package_build_info):
+                          bootloader_build_info, android_efi_loader_build_info,
+                          ota_build_info, host_package_build_info):
         """Get args from build information for fetch_cvd.
 
         Each build_info is a dictionary that contains 3 items, for example,
@@ -206,6 +212,7 @@ class AndroidBuildClient(base_cloud_client.BaseCloudApiClient):
                              constants.BUILD_ARTIFACT which is mapped to the
                              boot image name.
             bootloader_build_info: The build that provides the bootloader.
+            android_efi_loader_build_info: The build that provides the Android EFI loader.
             ota_build_info: The build that provides the OTA tools.
             host_package_build_info: The build that provides the host package.
 
@@ -223,10 +230,13 @@ class AndroidBuildClient(base_cloud_client.BaseCloudApiClient):
         bootloader_build = self.ProcessBuild(bootloader_build_info)
         if bootloader_build:
             fetch_cvd_args.append(f"-bootloader_build={bootloader_build}")
+        android_efi_loader_build = self.ProcessBuild(android_efi_loader_build_info)
+        if android_efi_loader_build:
+            fetch_cvd_args.append(f"-android_efi_loader_build {android_efi_loader_build}")
         kernel_build = self.GetKernelBuild(kernel_build_info)
         if kernel_build:
             fetch_cvd_args.append(f"-kernel_build={kernel_build}")
-        boot_build = self.ProcessBuild(boot_build_info)
+        boot_build = self.ProcessBuild(boot_build_info, ignore_artifact=True)
         if boot_build:
             fetch_cvd_args.append(f"-boot_build={boot_build}")
             boot_artifact = boot_build_info.get(constants.BUILD_ARTIFACT)
diff --git a/internal/lib/android_build_client_test.py b/internal/lib/android_build_client_test.py
index 30e1df4a..247a2f7c 100644
--- a/internal/lib/android_build_client_test.py
+++ b/internal/lib/android_build_client_test.py
@@ -179,6 +179,10 @@ class AndroidBuildClientTest(driver_test_lib.BaseDriverTest):
         ota_build = {constants.BUILD_ID: "4567",
                      constants.BUILD_BRANCH: "ota_branch",
                      constants.BUILD_TARGET: "ota_target"}
+        bootloader_build = {constants.BUILD_ID: "10111213",
+                            constants.BUILD_TARGET: "boot_crosvm_x86_64"}
+        android_efi_loader_build = {constants.BUILD_ID: "6789",
+                                    constants.BUILD_ARTIFACT: "gbl_x86_32.efi"}
         boot_build = {constants.BUILD_ID: "5678",
                       constants.BUILD_BRANCH: "boot_branch",
                       constants.BUILD_TARGET: "boot_target",
@@ -192,7 +196,7 @@ class AndroidBuildClientTest(driver_test_lib.BaseDriverTest):
         self.assertEqual(
             expected_args,
             self.client.GetFetchBuildArgs(
-                default_build, {}, {}, {}, {}, {}, {}))
+                default_build, {}, {}, {}, {}, {}, {}, {}))
 
         # Test base image with system image.
         expected_args = ["-default_build=1234/base_target",
@@ -200,7 +204,7 @@ class AndroidBuildClientTest(driver_test_lib.BaseDriverTest):
         self.assertEqual(
             expected_args,
             self.client.GetFetchBuildArgs(
-                default_build, system_build, {}, {}, {}, {}, {}))
+                default_build, system_build, {}, {}, {}, {}, {}, {}))
 
         # Test base image with kernel image.
         expected_args = ["-default_build=1234/base_target",
@@ -208,7 +212,7 @@ class AndroidBuildClientTest(driver_test_lib.BaseDriverTest):
         self.assertEqual(
             expected_args,
             self.client.GetFetchBuildArgs(
-                default_build, {}, kernel_build, {}, {}, {}, {}))
+                default_build, {}, kernel_build, {}, {}, {}, {}, {}))
 
         # Test base image with boot image.
         expected_args = ["-default_build=1234/base_target",
@@ -217,7 +221,23 @@ class AndroidBuildClientTest(driver_test_lib.BaseDriverTest):
         self.assertEqual(
             expected_args,
             self.client.GetFetchBuildArgs(
-                default_build, {}, {}, boot_build, {}, {}, {}))
+                default_build, {}, {}, boot_build, {}, {}, {}, {}))
+
+        # Test base image with bootloader.
+        expected_args = ["-default_build=1234/base_target",
+                         "-bootloader_build=10111213/boot_crosvm_x86_64"]
+        self.assertEqual(
+            expected_args,
+            self.client.GetFetchBuildArgs(
+                default_build, {}, {}, {}, bootloader_build, {}, {}, {}))
+
+        # Test base image with android efi.
+        expected_args = ["-default_build=1234/base_target",
+                         "-android_efi_loader_build 6789{gbl_x86_32.efi}"]
+        self.assertEqual(
+            expected_args,
+            self.client.GetFetchBuildArgs(
+                default_build, {}, {}, {}, {}, android_efi_loader_build, {}, {}))
 
         # Test base image with otatools.
         expected_args = ["-default_build=1234/base_target",
@@ -225,7 +245,7 @@ class AndroidBuildClientTest(driver_test_lib.BaseDriverTest):
         self.assertEqual(
             expected_args,
             self.client.GetFetchBuildArgs(
-                default_build, {}, {}, {}, {}, ota_build, {}))
+                default_build, {}, {}, {}, {}, {}, ota_build, {}))
 
         # Test base image with host_package.
         expected_args = ["-default_build=1234/base_target",
@@ -233,7 +253,7 @@ class AndroidBuildClientTest(driver_test_lib.BaseDriverTest):
         self.assertEqual(
             expected_args,
             self.client.GetFetchBuildArgs(
-                default_build, {}, {}, {}, {}, {}, host_package_build))
+                default_build, {}, {}, {}, {}, {}, {}, host_package_build))
 
     def testGetFetchCertArg(self):
         """Test GetFetchCertArg."""
diff --git a/internal/lib/cvd_compute_client_multi_stage.py b/internal/lib/cvd_compute_client_multi_stage.py
index 6e72667a..c6ba1b9a 100644
--- a/internal/lib/cvd_compute_client_multi_stage.py
+++ b/internal/lib/cvd_compute_client_multi_stage.py
@@ -325,7 +325,7 @@ class CvdComputeClient(android_compute_client.AndroidComputeClient):
     @utils.TimeExecute(function_description="Downloading build on instance")
     def FetchBuild(self, default_build_info, system_build_info,
                    kernel_build_info, boot_build_info, bootloader_build_info,
-                   ota_build_info, host_package_build_info):
+                   android_efi_loader_build_info, ota_build_info, host_package_build_info):
         """Execute fetch_cvd on the remote instance to get Cuttlefish runtime files.
 
         Args:
@@ -334,6 +334,7 @@ class CvdComputeClient(android_compute_client.AndroidComputeClient):
             kernel_build_info: The build that provides the kernel.
             boot_build_info: The build that provides the boot image.
             bootloader_build_info: The build that provides the bootloader.
+            android_efi_loader_build_info: The build that provides the Android EFI app.
             ota_build_info: The build that provides the OTA tools.
             host_package_build_info: The build that provides the host package.
 
@@ -344,8 +345,8 @@ class CvdComputeClient(android_compute_client.AndroidComputeClient):
         fetch_cvd_args = ["-credential_source=gce"]
         fetch_cvd_build_args = self._build_api.GetFetchBuildArgs(
             default_build_info, system_build_info, kernel_build_info,
-            boot_build_info, bootloader_build_info, ota_build_info,
-            host_package_build_info)
+            boot_build_info, bootloader_build_info, android_efi_loader_build_info,
+            ota_build_info, host_package_build_info)
         fetch_cvd_args.extend(fetch_cvd_build_args)
 
         self._ssh.Run("./fetch_cvd " + " ".join(fetch_cvd_args),
diff --git a/internal/lib/cvd_utils.py b/internal/lib/cvd_utils.py
index 14f51178..ac05a078 100644
--- a/internal/lib/cvd_utils.py
+++ b/internal/lib/cvd_utils.py
@@ -993,6 +993,10 @@ def GetRemoteBuildInfoDict(avd_spec):
         {"bootloader_" + key: val
          for key, val in avd_spec.bootloader_build_info.items() if val}
     )
+    build_info_dict.update(
+        {"android_efi_loader_" + key: val
+         for key, val in avd_spec.android_efi_loader_build_info.items() if val}
+    )
     return build_info_dict
 
 
diff --git a/internal/lib/cvd_utils_test.py b/internal/lib/cvd_utils_test.py
index edb0ea04..8580473c 100644
--- a/internal/lib/cvd_utils_test.py
+++ b/internal/lib/cvd_utils_test.py
@@ -617,7 +617,8 @@ class CvdUtilsTest(driver_test_lib.BaseDriverTest):
             remote_image=remote_image,
             kernel_build_info={"build_target": "kernel"},
             system_build_info={},
-            bootloader_build_info={})
+            bootloader_build_info={},
+            android_efi_loader_build_info = {})
         self.assertEqual(remote_image,
                          cvd_utils.GetRemoteBuildInfoDict(mock_avd_spec))
 
@@ -633,6 +634,10 @@ class CvdUtilsTest(driver_test_lib.BaseDriverTest):
             "branch": "aosp_u-boot-mainline",
             "build_id": "400000",
             "build_target": "u-boot_crosvm_x86_64"}
+        android_efi_loader_build_info = {
+            "build_id": "500000",
+            "artifact": "gbl_aarch64.efi"
+        }
         all_build_info = {
             "kernel_branch": "aosp_kernel-common-android12-5.10",
             "kernel_build_id": "200000",
@@ -642,14 +647,18 @@ class CvdUtilsTest(driver_test_lib.BaseDriverTest):
             "system_build_target": "aosp_x86_64-userdebug",
             "bootloader_branch": "aosp_u-boot-mainline",
             "bootloader_build_id": "400000",
-            "bootloader_build_target": "u-boot_crosvm_x86_64"}
+            "bootloader_build_target": "u-boot_crosvm_x86_64",
+            "android_efi_loader_build_id": "500000",
+            "android_efi_loader_artifact": "gbl_aarch64.efi"
+        }
         all_build_info.update(remote_image)
         mock_avd_spec = mock.Mock(
             spec=[],
             remote_image=remote_image,
             kernel_build_info=kernel_build_info,
             system_build_info=system_build_info,
-            bootloader_build_info=bootloader_build_info)
+            bootloader_build_info=bootloader_build_info,
+            android_efi_loader_build_info=android_efi_loader_build_info)
         self.assertEqual(all_build_info,
                          cvd_utils.GetRemoteBuildInfoDict(mock_avd_spec))
 
diff --git a/public/actions/remote_host_cf_device_factory.py b/public/actions/remote_host_cf_device_factory.py
index bfb43c84..4980fb18 100644
--- a/public/actions/remote_host_cf_device_factory.py
+++ b/public/actions/remote_host_cf_device_factory.py
@@ -313,6 +313,7 @@ class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
             self._avd_spec.kernel_build_info,
             self._avd_spec.boot_build_info,
             self._avd_spec.bootloader_build_info,
+            self._avd_spec.android_efi_loader_build_info,
             self._avd_spec.ota_build_info,
             self._avd_spec.host_package_build_info)
 
@@ -339,6 +340,7 @@ class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
             self._avd_spec.kernel_build_info,
             self._avd_spec.boot_build_info,
             self._avd_spec.bootloader_build_info,
+            self._avd_spec.android_efi_loader_build_info,
             self._avd_spec.ota_build_info,
             self._avd_spec.host_package_build_info)
 
@@ -398,6 +400,7 @@ class RemoteHostDeviceFactory(base_device_factory.BaseDeviceFactory):
             self._avd_spec.kernel_build_info,
             self._avd_spec.boot_build_info,
             self._avd_spec.bootloader_build_info,
+            self._avd_spec.android_efi_loader_build_info,
             self._avd_spec.ota_build_info,
             self._avd_spec.host_package_build_info)
         creds_cache_file = os.path.join(_HOME_FOLDER, cfg.creds_cache_file)
diff --git a/public/actions/remote_host_cf_device_factory_test.py b/public/actions/remote_host_cf_device_factory_test.py
index 9ee010de..79c159a3 100644
--- a/public/actions/remote_host_cf_device_factory_test.py
+++ b/public/actions/remote_host_cf_device_factory_test.py
@@ -56,6 +56,7 @@ class RemoteHostDeviceFactoryTest(driver_test_lib.BaseDriverTest):
                          kernel_build_info={},
                          boot_build_info={},
                          bootloader_build_info={},
+                         android_efi_loader_build_info={},
                          ota_build_info={},
                          host_package_build_info={},
                          remote_host="192.0.2.100",
diff --git a/public/actions/remote_instance_cf_device_factory.py b/public/actions/remote_instance_cf_device_factory.py
index b6259997..e5a2f849 100644
--- a/public/actions/remote_instance_cf_device_factory.py
+++ b/public/actions/remote_instance_cf_device_factory.py
@@ -111,6 +111,7 @@ class RemoteInstanceDeviceFactory(gce_device_factory.GCEDeviceFactory):
                     avd_spec.kernel_build_info,
                     avd_spec.boot_build_info,
                     avd_spec.bootloader_build_info,
+                    avd_spec.android_efi_loader_build_info,
                     avd_spec.ota_build_info,
                     avd_spec.host_package_build_info)