Merge "Increase timeout for Terminal on nested virt" into main
diff --git a/android/virtmgr/src/aidl.rs b/android/virtmgr/src/aidl.rs
index 6a268f9..daa755e 100644
--- a/android/virtmgr/src/aidl.rs
+++ b/android/virtmgr/src/aidl.rs
@@ -20,7 +20,7 @@
 use crate::crosvm::{AudioConfig, CrosvmConfig, DiskFile, SharedPathConfig, DisplayConfig, GpuConfig, InputDeviceOption, PayloadState, UsbConfig, VmContext, VmInstance, VmState};
 use crate::debug_config::{DebugConfig, DebugPolicy};
 use crate::dt_overlay::{create_device_tree_overlay, VM_DT_OVERLAY_MAX_SIZE, VM_DT_OVERLAY_PATH};
-use crate::payload::{add_microdroid_payload_images, add_microdroid_system_images, add_microdroid_vendor_image};
+use crate::payload::{ApexInfoList, add_microdroid_payload_images, add_microdroid_system_images, add_microdroid_vendor_image};
 use crate::selinux::{check_tee_service_permission, getfilecon, getprevcon, SeContext};
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
 use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::{
@@ -137,6 +137,21 @@
 static SUPPORTED_OS_NAMES: LazyLock<HashSet<String>> =
     LazyLock::new(|| get_supported_os_names().expect("Failed to get list of supported os names"));
 
+static CALLING_EXE_PATH: LazyLock<Option<PathBuf>> = LazyLock::new(|| {
+    let calling_exe_link = format!("/proc/{}/exe", get_calling_pid());
+    match fs::read_link(&calling_exe_link) {
+        Ok(calling_exe_path) => Some(calling_exe_path),
+        Err(e) => {
+            // virtmgr forked from apps fails to read /proc probably due to hidepid=2. As we
+            // discourage vendor apps, regarding such cases as system is safer.
+            // TODO(b/383969737): determine if this is okay. Or find a way how to track origins of
+            // apps.
+            warn!("can't read_link '{calling_exe_link}': {e:?}; regarding as system");
+            None
+        }
+    }
+});
+
 fn create_or_update_idsig_file(
     input_fd: &ParcelFileDescriptor,
     idsig_fd: &ParcelFileDescriptor,
@@ -406,10 +421,38 @@
     }
 }
 
-fn find_partition(path: &Path) -> binder::Result<String> {
-    match path.components().nth(1) {
+fn find_partition(path: Option<&Path>) -> binder::Result<String> {
+    let path = match path {
+        Some(path) => path,
+        None => return Ok("system".to_owned()),
+    };
+    let mut components = path.components();
+    match components.nth(1) {
         Some(std::path::Component::Normal(partition)) => {
-            Ok(partition.to_string_lossy().into_owned())
+            if partition != "apex" {
+                return Ok(partition.to_string_lossy().into_owned());
+            }
+
+            // If path is under /apex, find a partition of the preinstalled .apex path
+            let apex_name = match components.next() {
+                Some(std::path::Component::Normal(name)) => name.to_string_lossy(),
+                _ => {
+                    return Err(anyhow!("Can't find apex name for '{}'", path.display()))
+                        .or_service_specific_exception(-1)
+                }
+            };
+
+            let apex_info_list = ApexInfoList::load()
+                .context("Failed to get apex info list")
+                .or_service_specific_exception(-1)?;
+
+            for apex_info in apex_info_list.list.iter() {
+                if apex_info.name == apex_name {
+                    return Ok(apex_info.partition.to_lowercase());
+                }
+            }
+
+            Err(anyhow!("Can't find apex info for '{apex_name}'")).or_service_specific_exception(-1)
         }
         _ => Err(anyhow!("Can't find partition in '{}'", path.display()))
             .or_service_specific_exception(-1),
@@ -424,24 +467,27 @@
     fn create_early_vm_context(
         &self,
         config: &VirtualMachineConfig,
+        calling_exe_path: Option<&Path>,
     ) -> binder::Result<(VmContext, Cid, PathBuf)> {
-        let calling_exe_path = format!("/proc/{}/exe", get_calling_pid());
-        let link = fs::read_link(&calling_exe_path)
-            .context(format!("can't read_link '{calling_exe_path}'"))
-            .or_service_specific_exception(-1)?;
-        let partition = find_partition(&link)?;
-
         let name = match config {
             VirtualMachineConfig::RawConfig(config) => &config.name,
             VirtualMachineConfig::AppConfig(config) => &config.name,
         };
-        let early_vm =
-            find_early_vm_for_partition(&partition, name).or_service_specific_exception(-1)?;
-        if Path::new(&early_vm.path) != link {
+        let calling_partition = find_partition(calling_exe_path)?;
+        let early_vm = find_early_vm_for_partition(&calling_partition, name)
+            .or_service_specific_exception(-1)?;
+        let calling_exe_path = match calling_exe_path {
+            Some(path) => path,
+            None => {
+                return Err(anyhow!("Can't verify the path of PID {}", get_calling_pid()))
+                    .or_service_specific_exception(-1)
+            }
+        };
+        if Path::new(&early_vm.path) != calling_exe_path {
             return Err(anyhow!(
-                "VM '{name}' in partition '{partition}' must be created with '{}', not '{}'",
+                "VM '{name}' in partition '{calling_partition}' must be created with '{}', not '{}'",
                 &early_vm.path,
-                link.display()
+                calling_exe_path.display()
             ))
             .or_service_specific_exception(-1);
         }
@@ -525,7 +571,7 @@
 
         // Allocating VM context checks the MANAGE_VIRTUAL_MACHINE permission.
         let (vm_context, cid, temporary_directory) = if cfg!(early) {
-            self.create_early_vm_context(config)?
+            self.create_early_vm_context(config, CALLING_EXE_PATH.as_deref())?
         } else {
             self.create_vm_context(requester_debug_pid, config)?
         };
@@ -625,7 +671,8 @@
         // Check if files for payloads and bases are NOT coming from /vendor and /odm, as they may
         // have unstable interfaces.
         // TODO(b/316431494): remove once Treble interfaces are stabilized.
-        check_partitions_for_files(config).or_service_specific_exception(-1)?;
+        check_partitions_for_files(config, &find_partition(CALLING_EXE_PATH.as_deref())?)
+            .or_service_specific_exception(-1)?;
 
         let zero_filler_path = temporary_directory.join("zero.img");
         write_zero_filler(&zero_filler_path)
@@ -1261,7 +1308,7 @@
     Ok(vm_config)
 }
 
-fn check_partition_for_file(fd: &ParcelFileDescriptor) -> Result<()> {
+fn check_partition_for_file(fd: &ParcelFileDescriptor, calling_partition: &str) -> Result<()> {
     let path = format!("/proc/self/fd/{}", fd.as_raw_fd());
     let link = fs::read_link(&path).context(format!("can't read_link {path}"))?;
 
@@ -1271,24 +1318,39 @@
         return Ok(());
     }
 
-    if link.starts_with("/vendor") || link.starts_with("/odm") {
-        bail!("vendor or odm file {} can't be used for VM", link.display());
+    let is_fd_vendor = link.starts_with("/vendor") || link.starts_with("/odm");
+    let is_caller_vendor = calling_partition == "vendor" || calling_partition == "odm";
+
+    if is_fd_vendor != is_caller_vendor {
+        bail!("{} can't be used for VM client in {calling_partition}", link.display());
     }
 
     Ok(())
 }
 
-fn check_partitions_for_files(config: &VirtualMachineRawConfig) -> Result<()> {
+fn check_partitions_for_files(
+    config: &VirtualMachineRawConfig,
+    calling_partition: &str,
+) -> Result<()> {
     config
         .disks
         .iter()
         .flat_map(|disk| disk.partitions.iter())
         .filter_map(|partition| partition.image.as_ref())
-        .try_for_each(check_partition_for_file)?;
+        .try_for_each(|fd| check_partition_for_file(fd, calling_partition))?;
 
-    config.kernel.as_ref().map_or(Ok(()), check_partition_for_file)?;
-    config.initrd.as_ref().map_or(Ok(()), check_partition_for_file)?;
-    config.bootloader.as_ref().map_or(Ok(()), check_partition_for_file)?;
+    config
+        .disks
+        .iter()
+        .filter_map(|disk| disk.image.as_ref())
+        .try_for_each(|fd| check_partition_for_file(fd, calling_partition))?;
+
+    config.kernel.as_ref().map_or(Ok(()), |fd| check_partition_for_file(fd, calling_partition))?;
+    config.initrd.as_ref().map_or(Ok(()), |fd| check_partition_for_file(fd, calling_partition))?;
+    config
+        .bootloader
+        .as_ref()
+        .map_or(Ok(()), |fd| check_partition_for_file(fd, calling_partition))?;
 
     Ok(())
 }
@@ -2139,6 +2201,7 @@
     match partition {
         "system" => Ok(100..200),
         "system_ext" | "product" => Ok(200..300),
+        "vendor" | "odm" => Ok(300..400),
         _ => Err(anyhow!("Early VMs are not supported for {partition}")),
     }
 }
diff --git a/android/virtmgr/src/payload.rs b/android/virtmgr/src/payload.rs
index 5811314..bd6bf10 100644
--- a/android/virtmgr/src/payload.rs
+++ b/android/virtmgr/src/payload.rs
@@ -48,15 +48,19 @@
 
 /// Represents the list of APEXes
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
-struct ApexInfoList {
+pub(crate) struct ApexInfoList {
+    /// The list of APEXes
     #[serde(rename = "apex-info")]
-    list: Vec<ApexInfo>,
+    pub(crate) list: Vec<ApexInfo>,
 }
 
+/// Represents info of an APEX
 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq)]
-struct ApexInfo {
+pub(crate) struct ApexInfo {
+    /// Name of APEX
     #[serde(rename = "moduleName")]
-    name: String,
+    pub(crate) name: String,
+
     #[serde(rename = "versionCode")]
     version: u64,
     #[serde(rename = "modulePath")]
@@ -80,11 +84,15 @@
 
     #[serde(rename = "preinstalledModulePath")]
     preinstalled_path: PathBuf,
+
+    /// Partition of APEX
+    #[serde(default)]
+    pub(crate) partition: String,
 }
 
 impl ApexInfoList {
     /// Loads ApexInfoList
-    fn load() -> Result<&'static ApexInfoList> {
+    pub(crate) fn load() -> Result<&'static ApexInfoList> {
         static INSTANCE: OnceCell<ApexInfoList> = OnceCell::new();
         INSTANCE.get_or_try_init(|| {
             let apex_info_list = File::open(APEX_INFO_LIST_PATH)
diff --git a/libs/libavf/Android.bp b/libs/libavf/Android.bp
index e143709..079f4ae 100644
--- a/libs/libavf/Android.bp
+++ b/libs/libavf/Android.bp
@@ -34,7 +34,6 @@
 rust_ffi_static {
     name: "libavf_impl",
     defaults: ["libavf.default"],
-    export_include_dirs: ["include"],
 }
 
 cc_library {
@@ -49,7 +48,7 @@
         "libbinder_rpc_unstable",
         "liblog",
     ],
-    export_static_lib_headers: ["libavf_impl"],
+    export_include_dirs: ["include"],
     apex_available: ["com.android.virt"],
     version_script: "libavf.map.txt",
     stubs: {
diff --git a/libs/libavf/include/android/virtualization.h b/libs/libavf/include/android/virtualization.h
index 3a710f3..af39e9b 100644
--- a/libs/libavf/include/android/virtualization.h
+++ b/libs/libavf/include/android/virtualization.h
@@ -44,7 +44,7 @@
  *
  * \return A new virtual machine raw config object. On failure (such as out of memory), it aborts.
  */
-AVirtualMachineRawConfig* _Nonnull AVirtualMachineRawConfig_create() __INTRODUCED_IN(36);
+AVirtualMachineRawConfig* _Nonnull AVirtualMachineRawConfig_create(void) __INTRODUCED_IN(36);
 
 /**
  * Destroy a virtual machine config object.
@@ -147,6 +147,15 @@
                                             int32_t swiotlbMiB) __INTRODUCED_IN(36);
 
 /**
+ * Set vCPU count. The default is 1.
+ *
+ * \param config a virtual machine config object.
+ * \param n number of vCPUs. Must be positive.
+ */
+void AVirtualMachineRawConfig_setVCpuCount(AVirtualMachineRawConfig* _Nonnull config, int32_t n)
+        __INTRODUCED_IN(36);
+
+/**
  * Set whether the virtual machine's memory will be protected from the host, so the host can't
  * access its memory.
  *
diff --git a/libs/libavf/libavf.map.txt b/libs/libavf/libavf.map.txt
index 1590753..a505d85 100644
--- a/libs/libavf/libavf.map.txt
+++ b/libs/libavf/libavf.map.txt
@@ -9,6 +9,7 @@
     AVirtualMachineRawConfig_addDisk; # apex llndk
     AVirtualMachineRawConfig_setMemoryMiB; # apex llndk
     AVirtualMachineRawConfig_setSwiotlbMiB; # apex llndk
+    AVirtualMachineRawConfig_setVCpuCount; # apex llndk
     AVirtualMachineRawConfig_setProtectedVm; # apex llndk
     AVirtualMachineRawConfig_setHypervisorSpecificAuthMethod; # apex llndk
     AVirtualMachineRawConfig_addCustomMemoryBackingFile; # apex llndk
diff --git a/libs/libavf/src/lib.rs b/libs/libavf/src/lib.rs
index 98db246..9ae2635 100644
--- a/libs/libavf/src/lib.rs
+++ b/libs/libavf/src/lib.rs
@@ -23,14 +23,15 @@
 
 use android_system_virtualizationservice::{
     aidl::android::system::virtualizationservice::{
-        DiskImage::DiskImage, IVirtualizationService::IVirtualizationService,
-        VirtualMachineConfig::VirtualMachineConfig,
+        CpuTopology::CpuTopology, DiskImage::DiskImage,
+        IVirtualizationService::IVirtualizationService, VirtualMachineConfig::VirtualMachineConfig,
         VirtualMachineRawConfig::VirtualMachineRawConfig,
     },
     binder::{ParcelFileDescriptor, Strong},
 };
 use avf_bindgen::AVirtualMachineStopReason;
 use libc::timespec;
+use log::error;
 use vmclient::{DeathReason, VirtualizationService, VmInstance};
 
 /// Create a new virtual machine config object with no properties.
@@ -205,6 +206,22 @@
     config.swiotlbMib = swiotlb_mib;
 }
 
+/// Set vCPU count.
+///
+/// # Safety
+/// `config` must be a pointer returned by `AVirtualMachineRawConfig_create`.
+#[no_mangle]
+pub unsafe extern "C" fn AVirtualMachineRawConfig_setVCpuCount(
+    config: *mut VirtualMachineRawConfig,
+    n: i32,
+) {
+    // SAFETY: `config` is assumed to be a valid, non-null pointer returned by
+    // AVirtualMachineRawConfig_create. It's the only reference to the object.
+    let config = unsafe { &mut *config };
+    config.cpuTopology = CpuTopology::CUSTOM;
+    config.customVcpuCount = n;
+}
+
 /// Set whether a virtual machine is protected or not.
 ///
 /// # Safety
@@ -334,7 +351,10 @@
             }
             0
         }
-        Err(_) => -libc::EIO,
+        Err(e) => {
+            error!("AVirtualMachine_createRaw failed: {e:?}");
+            -libc::EIO
+        }
     }
 }
 
@@ -349,7 +369,10 @@
     let vm = unsafe { &*vm };
     match vm.start() {
         Ok(_) => 0,
-        Err(_) => -libc::EIO,
+        Err(e) => {
+            error!("AVirtualMachine_start failed: {e:?}");
+            -libc::EIO
+        }
     }
 }
 
@@ -364,7 +387,10 @@
     let vm = unsafe { &*vm };
     match vm.stop() {
         Ok(_) => 0,
-        Err(_) => -libc::EIO,
+        Err(e) => {
+            error!("AVirtualMachine_stop failed: {e:?}");
+            -libc::EIO
+        }
     }
 }
 
@@ -379,7 +405,10 @@
     let vm = unsafe { &*vm };
     match vm.connect_vsock(port) {
         Ok(pfd) => pfd.into_raw_fd(),
-        Err(_) => -libc::EIO,
+        Err(e) => {
+            error!("AVirtualMachine_connectVsock failed: {e:?}");
+            -libc::EIO
+        }
     }
 }
 
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index e6f15ff..83dc58e 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -207,9 +207,6 @@
             assume().withMessage("Skip where protected VMs aren't supported")
                     .that(capabilities & VirtualMachineManager.CAPABILITY_PROTECTED_VM)
                     .isNotEqualTo(0);
-            assume().withMessage("Testing protected VMs on GSI isn't supported. b/272443823")
-                    .that(isGsi())
-                    .isFalse();
             // TODO(b/376870129): remove this
             assume().withMessage("pVMs with 16k kernel are not supported yet :(")
                     .that(mOs)