Merge "Add avf_early_vm_test" into main
diff --git a/tests/early_vm_test/Android.bp b/tests/early_vm_test/Android.bp
new file mode 100644
index 0000000..dbb0c28
--- /dev/null
+++ b/tests/early_vm_test/Android.bp
@@ -0,0 +1,53 @@
+prebuilt_etc {
+    name: "avf_early_vm_test_kernel",
+    filename: "rialto.bin",
+    src: ":empty_file",
+    target: {
+        android_arm64: {
+            src: ":rialto_signed",
+        },
+    },
+    installable: false,
+    system_ext_specific: true,
+    visibility: ["//visibility:private"],
+}
+
+rust_binary {
+    name: "avf_early_vm_test_launcher",
+    crate_name: "avf_early_vm_test_launcher",
+    srcs: ["src/main.rs"],
+    rustlibs: [
+        "android.system.virtualizationservice-rust",
+        "libanyhow",
+        "libclap",
+        "libhypervisor_props",
+        "liblog_rust",
+        "libservice_vm_comm",
+        "libservice_vm_manager",
+        "libvmclient",
+    ],
+    cfgs: select(release_flag("RELEASE_AVF_ENABLE_EARLY_VM"), {
+        true: ["early_vm_enabled"],
+        default: [],
+    }),
+    prefer_rlib: true,
+    system_ext_specific: true,
+    compile_multilib: "first",
+    installable: false,
+}
+
+python_test_host {
+    name: "avf_early_vm_test",
+    main: "avf_early_vm_test.py",
+    srcs: ["avf_early_vm_test.py"],
+    device_first_data: [
+        ":avf_early_vm_test_kernel",
+        ":avf_early_vm_test_launcher",
+    ],
+    data: ["early_vms_rialto_test.xml"],
+    test_suites: ["general-tests"],
+    test_config: "AndroidTest.xml",
+    test_options: {
+        unit_test: false,
+    },
+}
diff --git a/tests/early_vm_test/AndroidTest.xml b/tests/early_vm_test/AndroidTest.xml
new file mode 100644
index 0000000..3eae96d
--- /dev/null
+++ b/tests/early_vm_test/AndroidTest.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright 2025 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<configuration description="Runs avf_early_vm_test.">
+    <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer">
+        <option name="force-root" value="true"/>
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.PushFilePreparer">
+    <option name="abort-on-push-failure" value="true" />
+        <option name="remount-system" value="true" />
+        <option name="push-file" key="avf_early_vm_test_launcher" value="/system_ext/bin/avf_early_vm_test_launcher" />
+        <option name="push-file" key="rialto.bin" value="/system_ext/etc/avf/rialto_test.bin" />
+        <option name="push-file" key="early_vms_rialto_test.xml" value="/system_ext/etc/avf/early_vms_rialto_test.xml" />
+    </target_preparer>
+
+    <test class="com.android.tradefed.testtype.python.PythonBinaryHostTest">
+        <option name="par-file-name" value="avf_early_vm_test" />
+        <option name="test-timeout" value="5m" />
+    </test>
+</configuration>
diff --git a/tests/early_vm_test/TEST_MAPPING b/tests/early_vm_test/TEST_MAPPING
new file mode 100644
index 0000000..1f2335b
--- /dev/null
+++ b/tests/early_vm_test/TEST_MAPPING
@@ -0,0 +1,9 @@
+// When adding or removing tests here, don't forget to amend _all_modules list in
+// wireless/android/busytown/ath_config/configs/prod/avf/tests.gcl
+{
+  "avf-presubmit": [
+    {
+      "name": "avf_early_vm_test"
+    }
+  ]
+}
diff --git a/tests/early_vm_test/avf_early_vm_test.py b/tests/early_vm_test/avf_early_vm_test.py
new file mode 100644
index 0000000..0003351
--- /dev/null
+++ b/tests/early_vm_test/avf_early_vm_test.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+#
+# Copyright 2025 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import logging
+import os
+import subprocess
+import unittest
+
+_DEFAULT_COMMAND_TIMEOUT = 300
+_LAUNCHER_PATH = "/system_ext/bin/avf_early_vm_test_launcher"
+_RIALTO_PATH = "/system_ext/etc/avf/rialto_test.bin"
+
+def _RunCommand(cmd, timeout=_DEFAULT_COMMAND_TIMEOUT):
+    with subprocess.Popen(args=cmd,
+                          stderr=subprocess.PIPE,
+                          stdout=subprocess.PIPE,
+                          universal_newlines=True) as proc:
+        try:
+            out, err = proc.communicate(timeout=timeout)
+            returncode = proc.returncode
+        except subprocess.TimeoutExpired:
+            proc.kill()
+            out, err = proc.communicate()
+            returncode = proc.returncode
+
+    return out, err, returncode
+
+class AvfEarlyVmTest(unittest.TestCase):
+    def setUp(self):
+        self._serial_number = os.environ.get("ANDROID_SERIAL")
+        self.assertTrue(self._serial_number, "$ANDROID_SERIAL is empty.")
+
+    def _TestAvfEarlyVm(self, protected):
+        adb_cmd = ["adb", "-s", self._serial_number, "shell", _LAUNCHER_PATH, "--kernel",
+                   _RIALTO_PATH]
+        if protected:
+            adb_cmd.append("--protected")
+
+        _, err, returncode = _RunCommand(adb_cmd)
+        self.assertEqual(returncode, 0, f"{adb_cmd} failed: {err}")
+
+    def testAvfEarlyVmNonProtected(self):
+        self._TestAvfEarlyVm(False)
+
+    def testAvfEarlyVmProtected(self):
+        self._TestAvfEarlyVm(True)
+
+if __name__ == "__main__":
+    # Setting verbosity is required to generate output that the TradeFed test
+    # runner can parse.
+    unittest.main(verbosity=3)
diff --git a/tests/early_vm_test/early_vms_rialto_test.xml b/tests/early_vm_test/early_vms_rialto_test.xml
new file mode 100644
index 0000000..799fc3f
--- /dev/null
+++ b/tests/early_vm_test/early_vms_rialto_test.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright 2025 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<early_vms>
+    <early_vm>
+        <name>avf_early_vm_test_launcher</name>
+        <cid>299</cid>
+        <path>/system_ext/bin/avf_early_vm_test_launcher</path>
+    </early_vm>
+</early_vms>
diff --git a/tests/early_vm_test/src/main.rs b/tests/early_vm_test/src/main.rs
new file mode 100644
index 0000000..a3c80ca
--- /dev/null
+++ b/tests/early_vm_test/src/main.rs
@@ -0,0 +1,116 @@
+// Copyright 2025 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Tests running an early VM
+
+use android_system_virtualizationservice::{
+    aidl::android::system::virtualizationservice::{
+        IVirtualizationService::IVirtualizationService, VirtualMachineConfig::VirtualMachineConfig,
+        VirtualMachineRawConfig::VirtualMachineRawConfig,
+    },
+    binder::{ParcelFileDescriptor, ProcessState, Strong},
+};
+use anyhow::{Context, Result};
+use clap::Parser;
+use log::info;
+use std::fs::File;
+use std::path::PathBuf;
+
+use service_vm_comm::{Request, Response, VmType};
+use service_vm_manager::ServiceVm;
+use vmclient::VmInstance;
+
+const VM_MEMORY_MB: i32 = 16;
+
+#[derive(Parser)]
+/// Collection of CLI for avf_early_vm_test_rialto
+pub struct Args {
+    /// Path to the Rialto kernel image.
+    #[arg(long)]
+    kernel: PathBuf,
+
+    /// Whether the VM is protected or not.
+    #[arg(long)]
+    protected: bool,
+}
+
+fn get_service() -> Result<Strong<dyn IVirtualizationService>> {
+    let virtmgr = vmclient::VirtualizationService::new_early()
+        .context("Failed to spawn VirtualizationService")?;
+    virtmgr.connect().context("Failed to connect to VirtualizationService")
+}
+
+fn main() -> Result<()> {
+    if std::env::consts::ARCH != "aarch64" {
+        info!("{} not supported. skipping test", std::env::consts::ARCH);
+        return Ok(());
+    }
+
+    if !cfg!(early_vm_enabled) {
+        info!("early VM disabled. skipping test");
+        return Ok(());
+    }
+
+    let args = Args::parse();
+
+    if args.protected {
+        if !hypervisor_props::is_protected_vm_supported()? {
+            info!("pVMs are not supported on device. skipping test");
+            return Ok(());
+        }
+    } else if !hypervisor_props::is_vm_supported()? {
+        info!("non-pVMs are not supported on device. skipping test");
+        return Ok(());
+    }
+
+    let service = get_service()?;
+    let kernel =
+        File::open(&args.kernel).with_context(|| format!("Failed to open {:?}", &args.kernel))?;
+    let kernel = ParcelFileDescriptor::new(kernel);
+
+    let vm_config = VirtualMachineConfig::RawConfig(VirtualMachineRawConfig {
+        name: "avf_early_vm_test_launcher".to_owned(),
+        kernel: Some(kernel),
+        protectedVm: args.protected,
+        memoryMib: VM_MEMORY_MB,
+        platformVersion: "~1.0".to_owned(),
+        ..Default::default()
+    });
+
+    let vm_instance = VmInstance::create(
+        service.as_ref(),
+        &vm_config,
+        // console_in, console_out, and log will be redirected to the kernel log by virtmgr
+        None, // console_in
+        None, // console_out
+        None, // log
+        None, // dump_dt
+        None, // callback
+    )
+    .context("Failed to create VM")?;
+
+    ProcessState::start_thread_pool();
+
+    let vm_type = if args.protected { VmType::ProtectedVm } else { VmType::NonProtectedVm };
+    let mut vm_service = ServiceVm::start_vm(vm_instance, vm_type)?;
+
+    let request_data = vec![1, 2, 3, 4, 5];
+    let reversed_data = vec![5, 4, 3, 2, 1];
+    let response = vm_service
+        .process_request(Request::Reverse(request_data))
+        .context("Failed to process request")?;
+    assert_eq!(Response::Reverse(reversed_data), response);
+
+    Ok(())
+}