Merge "Allow multiple xml files for early VMs" into main
diff --git a/android/virtmgr/src/aidl.rs b/android/virtmgr/src/aidl.rs
index 87d7a88..e2b2804 100644
--- a/android/virtmgr/src/aidl.rs
+++ b/android/virtmgr/src/aidl.rs
@@ -59,7 +59,7 @@
     Key::Key, PubKey::PubKey, SessionIdSignature::SessionIdSignature, SessionInfo::SessionInfo,
     SessionInitiationInfo::SessionInitiationInfo,
 };
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{anyhow, bail, ensure, Context, Result};
 use apkverify::{HashAlgorithm, V4Signature};
 use avflog::LogResult;
 use binder::{
@@ -76,7 +76,7 @@
 use rustutils::system_properties;
 use semver::VersionReq;
 use serde::Deserialize;
-use std::collections::HashSet;
+use std::collections::{HashSet, HashMap};
 use std::convert::TryInto;
 use std::fs;
 use std::ffi::CStr;
@@ -2007,22 +2007,21 @@
 }
 
 // KEEP IN SYNC WITH early_vms.xsd
-#[derive(Debug, Deserialize, PartialEq)]
+#[derive(Clone, Debug, Deserialize, PartialEq)]
 struct EarlyVm {
-    #[allow(dead_code)]
     name: String,
-    #[allow(dead_code)]
     cid: i32,
-    #[allow(dead_code)]
     path: String,
 }
 
 #[derive(Debug, Default, Deserialize)]
 struct EarlyVms {
-    #[allow(dead_code)]
     early_vm: Vec<EarlyVm>,
 }
 
+static EARLY_VMS_CACHE: LazyLock<Mutex<HashMap<String, Vec<EarlyVm>>>> =
+    LazyLock::new(|| Mutex::new(HashMap::new()));
+
 fn range_for_partition(partition: &str) -> Result<Range<Cid>> {
     match partition {
         "system" => Ok(100..200),
@@ -2031,7 +2030,7 @@
     }
 }
 
-fn find_early_vm(xml_path: &Path, cid_range: &Range<Cid>, name: &str) -> Result<EarlyVm> {
+fn get_early_vms_in_path(xml_path: &Path) -> Result<Vec<EarlyVm>> {
     if !xml_path.exists() {
         bail!("{} doesn't exist", xml_path.display());
     }
@@ -2043,35 +2042,74 @@
     let early_vms: EarlyVms = serde_xml_rs::from_str(&xml)
         .with_context(|| format!("Can't parse {}", xml_path.display()))?;
 
-    let mut found_vm: Option<EarlyVm> = None;
+    Ok(early_vms.early_vm)
+}
 
-    for early_vm in early_vms.early_vm {
+fn validate_cid_range(early_vms: &[EarlyVm], cid_range: &Range<Cid>) -> Result<()> {
+    for early_vm in early_vms {
+        let cid = early_vm
+            .cid
+            .try_into()
+            .with_context(|| format!("VM '{}' uses Invalid CID {}", early_vm.name, early_vm.cid))?;
+
+        ensure!(
+            cid_range.contains(&cid),
+            "VM '{}' uses CID {cid} which is out of range. Available CIDs: {cid_range:?}",
+            early_vm.name
+        );
+    }
+    Ok(())
+}
+
+fn get_early_vms_in_partition(partition: &str) -> Result<Vec<EarlyVm>> {
+    let mut cache = EARLY_VMS_CACHE.lock().unwrap();
+
+    if let Some(result) = cache.get(partition) {
+        return Ok(result.clone());
+    }
+
+    let pattern = format!("/{partition}/etc/avf/early_vms*.xml");
+    let mut early_vms = Vec::new();
+    for entry in glob::glob(&pattern).with_context(|| format!("Failed to glob {}", &pattern))? {
+        match entry {
+            Ok(path) => early_vms.extend(get_early_vms_in_path(&path)?),
+            Err(e) => error!("Error while globbing (but continuing) {}: {}", &pattern, e),
+        }
+    }
+
+    validate_cid_range(&early_vms, &range_for_partition(partition)?)
+        .with_context(|| format!("CID validation for {partition} failed"))?;
+
+    cache.insert(partition.to_owned(), early_vms.clone());
+
+    Ok(early_vms)
+}
+
+fn find_early_vm<'a>(early_vms: &'a [EarlyVm], name: &str) -> Result<&'a EarlyVm> {
+    let mut found_vm: Option<&EarlyVm> = None;
+
+    for early_vm in early_vms {
         if early_vm.name != name {
             continue;
         }
 
-        let cid = early_vm
-            .cid
-            .try_into()
-            .with_context(|| format!("Invalid CID value {}", early_vm.cid))?;
-
-        if !cid_range.contains(&cid) {
-            bail!("VM '{}' uses CID {cid} which is out of range. Available CIDs for '{}': {cid_range:?}", xml_path.display(), early_vm.name);
-        }
-
         if found_vm.is_some() {
-            bail!("Multiple VMs named {name} are found in {}", xml_path.display());
+            bail!("Multiple VMs named '{name}' are found");
         }
 
         found_vm = Some(early_vm);
     }
 
-    found_vm.ok_or_else(|| anyhow!("Can't find {name} in {}", xml_path.display()))
+    found_vm.ok_or_else(|| anyhow!("Can't find a VM named '{name}'"))
 }
 
 fn find_early_vm_for_partition(partition: &str, name: &str) -> Result<EarlyVm> {
-    let cid_range = range_for_partition(partition)?;
-    find_early_vm(Path::new(&format!("/{partition}/etc/avf/early_vms.xml")), &cid_range, name)
+    let early_vms = get_early_vms_in_partition(partition)
+        .with_context(|| format!("Failed to get early VMs from {partition}"))?;
+
+    Ok(find_early_vm(&early_vms, name)
+        .with_context(|| format!("Failed to find early VM '{name}' in {partition}"))?
+        .clone())
 }
 
 #[cfg(test)]
@@ -2314,6 +2352,87 @@
                 <path>/system/bin/vm_demo_native_early</path>
             </early_vm>
             <early_vm>
+                <name>vm_demo_native_early_2</name>
+                <cid>456</cid>
+                <path>/system/bin/vm_demo_native_early_2</path>
+            </early_vm>
+        </early_vms>
+        "#,
+        )?;
+
+        let cid_range = 100..1000;
+
+        let early_vms = get_early_vms_in_path(&xml_path)?;
+        validate_cid_range(&early_vms, &cid_range)?;
+
+        let test_cases = [
+            (
+                "vm_demo_native_early",
+                EarlyVm {
+                    name: "vm_demo_native_early".to_owned(),
+                    cid: 123,
+                    path: "/system/bin/vm_demo_native_early".to_owned(),
+                },
+            ),
+            (
+                "vm_demo_native_early_2",
+                EarlyVm {
+                    name: "vm_demo_native_early_2".to_owned(),
+                    cid: 456,
+                    path: "/system/bin/vm_demo_native_early_2".to_owned(),
+                },
+            ),
+        ];
+
+        for (name, expected) in test_cases {
+            let result = find_early_vm(&early_vms, name)?;
+            assert_eq!(result, &expected);
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_invalid_cid_validation() -> Result<()> {
+        let tmp_dir = tempfile::TempDir::new()?;
+        let xml_path = tmp_dir.path().join("early_vms.xml");
+
+        let cid_range = 100..1000;
+
+        for cid in [-1, 999999] {
+            std::fs::write(
+                &xml_path,
+                format!(
+                    r#"<?xml version="1.0" encoding="utf-8"?>
+        <early_vms>
+            <early_vm>
+                <name>vm_demo_invalid_cid</name>
+                <cid>{cid}</cid>
+                <path>/system/bin/vm_demo_invalid_cid</path>
+            </early_vm>
+        </early_vms>
+        "#
+                ),
+            )?;
+
+            let early_vms = get_early_vms_in_path(&xml_path)?;
+            assert!(validate_cid_range(&early_vms, &cid_range).is_err(), "should fail");
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_duplicated_early_vms() -> Result<()> {
+        let tmp_dir = tempfile::TempDir::new()?;
+        let tmp_dir_path = tmp_dir.path().to_owned();
+        let xml_path = tmp_dir_path.join("early_vms.xml");
+
+        std::fs::write(
+            &xml_path,
+            br#"<?xml version="1.0" encoding="utf-8"?>
+        <early_vms>
+            <early_vm>
                 <name>vm_demo_duplicated_name</name>
                 <cid>456</cid>
                 <path>/system/bin/vm_demo_duplicated_name_1</path>
@@ -2323,42 +2442,16 @@
                 <cid>789</cid>
                 <path>/system/bin/vm_demo_duplicated_name_2</path>
             </early_vm>
-            <early_vm>
-                <name>vm_demo_invalid_cid_1</name>
-                <cid>-1</cid>
-                <path>/system/bin/vm_demo_invalid_cid_1</path>
-            </early_vm>
-            <early_vm>
-                <name>vm_demo_invalid_cid_2</name>
-                <cid>999999</cid>
-                <path>/system/bin/vm_demo_invalid_cid_2</path>
-            </early_vm>
         </early_vms>
         "#,
         )?;
 
         let cid_range = 100..1000;
 
-        let result = find_early_vm(&xml_path, &cid_range, "vm_demo_native_early")?;
-        let expected = EarlyVm {
-            name: "vm_demo_native_early".to_owned(),
-            cid: 123,
-            path: "/system/bin/vm_demo_native_early".to_owned(),
-        };
-        assert_eq!(result, expected);
+        let early_vms = get_early_vms_in_path(&xml_path)?;
+        validate_cid_range(&early_vms, &cid_range)?;
 
-        assert!(
-            find_early_vm(&xml_path, &cid_range, "vm_demo_duplicated_name").is_err(),
-            "should fail"
-        );
-        assert!(
-            find_early_vm(&xml_path, &cid_range, "vm_demo_invalid_cid_1").is_err(),
-            "should fail"
-        );
-        assert!(
-            find_early_vm(&xml_path, &cid_range, "vm_demo_invalid_cid_2").is_err(),
-            "should fail"
-        );
+        assert!(find_early_vm(&early_vms, "vm_demo_duplicated_name").is_err(), "should fail");
 
         Ok(())
     }
diff --git a/docs/early_vm.md b/docs/early_vm.md
index 44b71ff..3f21f11 100644
--- a/docs/early_vm.md
+++ b/docs/early_vm.md
@@ -8,8 +8,9 @@
 
 To run an early VM, clients must follow these steps.
 
-1) Early VMs need to be defined in `{partition}/etc/avf/early_vms.xml`. The
-schema for this file is defined in [`early_vms.xsd`](../android/virtmgr/early_vms.xsd).
+1) Early VMs must be defined in XML files located at
+`{partition}/etc/avf/early_vms*.xml`. Schema for these files is defined in
+[`early_vms.xsd`](../android/virtmgr/early_vms.xsd).
 
 ```early_vms.xml
 <early_vms>
@@ -25,6 +26,9 @@
 connection with `early_virtmgr` and create a VM named `vm_demo_native_early`,
 which will be assigned the static CID 123.
 
+Multiple XML files matching the glob pattern
+`{partition}/etc/avf/early_vms*.xml` can be used to define early VMs.
+
 2) The client must have the following three or four capabilities.
 
 * `IPC_LOCK`