virtualizationservice: "prefer_staged" for apexes

App can tell virtualizationservice to use "staged" apexes for its
payload.

  // vm_config.json
  {
    ..
    "perfer_staged": true
  }

Bug: 199146189
Test: MicrodroidHostTestCases
Test: adb install --staged new-adbd.apex
      vm run-app with prefer_staged: true, started with staged adbd apex
Change-Id: I6320247362e07519e120aed19b290618270e1335
diff --git a/microdroid/payload/config/src/lib.rs b/microdroid/payload/config/src/lib.rs
index 6dc127b..2547f3d 100644
--- a/microdroid/payload/config/src/lib.rs
+++ b/microdroid/payload/config/src/lib.rs
@@ -30,6 +30,10 @@
     /// APEXes to activate in a VM
     #[serde(default)]
     pub apexes: Vec<ApexConfig>,
+
+    /// Tells VirtualizationService to use staged APEXes if possible
+    #[serde(default)]
+    pub prefer_staged: bool,
 }
 
 /// OS config
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 443436d..18d8ade 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -43,6 +43,8 @@
         "libvmconfig",
         "libzip",
         "libvsock",
+        // TODO(b/202115393) stabilize the interface
+        "packagemanager_aidl-rust",
     ],
     shared_libs: [
         "libbinder_rpc_unstable",
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 76c3a16..b3d54ef 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -462,14 +462,13 @@
 
     // Microdroid requires an additional payload disk image and the bootconfig partition.
     if os_name == "microdroid" {
-        let apexes = vm_payload_config.apexes.clone();
         add_microdroid_images(
             config,
             temporary_directory,
             apk_file,
             idsig_file,
             instance_file,
-            apexes,
+            &vm_payload_config,
             &mut vm_config,
         )?;
     }
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index 9662fa3..59ad81c 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -20,9 +20,12 @@
 };
 use android_system_virtualizationservice::binder::ParcelFileDescriptor;
 use anyhow::{anyhow, Context, Result};
+use binder::{wait_for_interface, Strong};
+use log::info;
 use microdroid_metadata::{ApexPayload, ApkPayload, Metadata};
-use microdroid_payload_config::ApexConfig;
+use microdroid_payload_config::{ApexConfig, VmPayloadConfig};
 use once_cell::sync::OnceCell;
+use packagemanager_aidl::aidl::android::content::pm::IPackageManagerNative::IPackageManagerNative;
 use serde::Deserialize;
 use serde_xml_rs::from_reader;
 use std::fs::{File, OpenOptions};
@@ -35,6 +38,8 @@
 
 const APEX_INFO_LIST_PATH: &str = "/apex/apex-info-list.xml";
 
+const PACKAGE_MANAGER_NATIVE_SERVICE: &str = "package_native";
+
 /// Represents the list of APEXes
 #[derive(Debug, Deserialize)]
 struct ApexInfoList {
@@ -74,6 +79,32 @@
     }
 }
 
+struct PackageManager {
+    service: Strong<dyn IPackageManagerNative>,
+    // TODO(b/199146189) use IPackageManagerNative
+    apex_info_list: &'static ApexInfoList,
+}
+
+impl PackageManager {
+    fn new() -> Result<Self> {
+        let service = wait_for_interface(PACKAGE_MANAGER_NATIVE_SERVICE)
+            .context("Failed to find PackageManager")?;
+        let apex_info_list = ApexInfoList::load()?;
+        Ok(Self { service, apex_info_list })
+    }
+
+    fn get_apex_path(&self, name: &str, prefer_staged: bool) -> Result<PathBuf> {
+        if prefer_staged {
+            let apex_info = self.service.getStagedApexInfo(name)?;
+            if let Some(apex_info) = apex_info {
+                info!("prefer_staged: use {} for {}", apex_info.diskImagePath, name);
+                return Ok(PathBuf::from(apex_info.diskImagePath));
+            }
+        }
+        self.apex_info_list.get_path_for(name)
+    }
+}
+
 fn make_metadata_file(
     config_path: &str,
     apexes: &[ApexConfig],
@@ -127,6 +158,7 @@
     idsig_file: File,
     config_path: &str,
     apexes: &[ApexConfig],
+    prefer_staged: bool,
     temporary_directory: &Path,
 ) -> Result<DiskImage> {
     let metadata_file = make_metadata_file(config_path, apexes, temporary_directory)?;
@@ -137,9 +169,9 @@
         writable: false,
     }];
 
-    let apex_info_list = ApexInfoList::load()?;
+    let pm = PackageManager::new()?;
     for (i, apex) in apexes.iter().enumerate() {
-        let apex_path = apex_info_list.get_path_for(&apex.name)?;
+        let apex_path = pm.get_apex_path(&apex.name, prefer_staged)?;
         let apex_file = open_parcel_file(&apex_path, false)?;
         partitions.push(Partition {
             label: format!("microdroid-apex-{}", i),
@@ -167,9 +199,10 @@
     apk_file: File,
     idsig_file: File,
     instance_file: File,
-    mut apexes: Vec<ApexConfig>,
+    vm_payload_config: &VmPayloadConfig,
     vm_config: &mut VirtualMachineRawConfig,
 ) -> Result<()> {
+    let mut apexes = vm_payload_config.apexes.clone();
     apexes.extend(
         MICRODROID_REQUIRED_APEXES.iter().map(|name| ApexConfig { name: name.to_string() }),
     );
@@ -180,6 +213,7 @@
         idsig_file,
         &config.configPath,
         &apexes,
+        vm_payload_config.prefer_staged,
         temporary_directory,
     )?);