Merge "docs: Recommends ferrochrome.sh" into main
diff --git a/apex/virtualizationservice.xml b/apex/virtualizationservice.xml
index 60f466f..97f6b3c 100644
--- a/apex/virtualizationservice.xml
+++ b/apex/virtualizationservice.xml
@@ -1,5 +1,5 @@
 <manifest version="1.0" type="framework">
-    <hal format="aidl">
+    <hal format="aidl" min-level="202404">
         <name>android.hardware.security.keymint</name>
         <version>3</version>
         <fqname>IRemotelyProvisionedComponent/avf</fqname>
diff --git a/docs/custom_vm.md b/docs/custom_vm.md
index adae13c..6c51795 100644
--- a/docs/custom_vm.md
+++ b/docs/custom_vm.md
@@ -238,18 +238,6 @@
 $ adb shell pm clear com.google.android.virtualization.vmlauncher
 ```
 
-### Inside guest OS (for ChromiumOS only)
-
-Go to the network setting and configure as below.
-
-* IP: 192.168.1.2 (other addresses in the 192.168.1.0/24 subnet also works)
-* netmask: 255.255.255.0
-* gateway: 192.168.1.1
-* DNS: 8.8.8.8 (or any DNS server you know)
-
-These settings are persistent; stored in chromiumos_test_image.bin. So you
-don’t have to repeat this next time.
-
 ### Debugging
 
 To open the serial console (interactive terminal):
@@ -258,17 +246,27 @@
 ```
 
 To see console logs only, check
-`/data/data/com.android.virtualization.vmlauncher/files/console.log`
-Or
-`/data/data/com.google.android.virtualization.vmlauncher/files/console.log`
+`/data/data/com{,.google}.android.virtualization.vmlauncher/files/console.log`
+
+For HSUM enabled devices,
+`/data/user/${current_user_id}/com{,.google}.android.virtualization.vmlauncher/files/console.log`
+
+You can monitor console out as follows
 
 ```shell
-$ adb shell su root tail +0 -F /data/data/com{,.google}.android.virtualization.vmlauncher/files/console.log
+$ adb shell 'su root tail +0 -F /data/user/$(am get-current-user)/com{,.google}.android.virtualization.vmlauncher/files/console.log'
 ```
 
-For ChromiumOS, you can ssh-in. Use following commands after network setup.
+For ChromiumOS, you can enter to the console via SSH connection. Check your IP
+address of ChromiumOS VM from the ethernet network setting page and follow
+commands below.
 
 ```shell
-$ adb kill-server ; adb start-server; adb forward tcp:9222 tcp:9222
+$ adb kill-server ; adb start-server
+$ adb shell nc -s localhost -L -p 9222 nc ${CHROMIUMOS_IPV4_ADDR} 22 # This command won't be terminated.
+$ adb forward tcp:9222 tcp:9222
 $ ssh -oProxyCommand=none -o UserKnownHostsFile=/dev/null root@localhost -p 9222
 ```
+
+For ChromiumOS, you would need to login after enthering its console.
+The user ID and the password is `root` and `test0000` respectively.
diff --git a/ferrochrome_app/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java b/ferrochrome_app/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
index d9e5229..7c18537 100644
--- a/ferrochrome_app/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
+++ b/ferrochrome_app/java/com/android/virtualization/ferrochrome/FerrochromeActivity.java
@@ -56,6 +56,7 @@
     private static final Path IMAGE_VERSION_INFO =
             Path.of(EXTERNAL_STORAGE_DIR + "ferrochrome_image_version");
     private static final Path VM_CONFIG_PATH = Path.of(EXTERNAL_STORAGE_DIR + "vm_config.json");
+    private static final int REQUEST_CODE_VMLAUNCHER = 1;
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -109,10 +110,17 @@
                     }
                     updateStatus("Done.");
                     updateStatus("Starting Ferrochrome...");
-                    runOnUiThread(() -> startActivity(intent));
+                    runOnUiThread(() -> startActivityForResult(intent, REQUEST_CODE_VMLAUNCHER));
                 });
     }
 
+    @Override
+    protected void onActivityResult(int requestCode, int resultCode, Intent data) {
+        if (requestCode == REQUEST_CODE_VMLAUNCHER) {
+            finishAndRemoveTask();
+        }
+    }
+
     private void updateStatus(String line) {
         Log.d(TAG, line);
         runOnUiThread(
diff --git a/pvmfw/platform.dts b/pvmfw/platform.dts
index 2df0768..44834ed 100644
--- a/pvmfw/platform.dts
+++ b/pvmfw/platform.dts
@@ -411,6 +411,7 @@
 		reg = <0x00 0x3000 0x00 0x1000>;
 		clock-frequency = <10>;
 		timeout-sec = <8>;
+		interrupts = <GIC_PPI 0xf PLACEHOLDER>;
 	};
 
 	cpufreq {
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index 84dc14d..939a4ea 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -759,6 +759,84 @@
     Ok(SerialInfo { addrs })
 }
 
+#[derive(Default, Debug, PartialEq)]
+struct WdtInfo {
+    addr: u64,
+    size: u64,
+    irq: [u32; WdtInfo::IRQ_CELLS],
+}
+
+impl WdtInfo {
+    const IRQ_CELLS: usize = 3;
+    const IRQ_NR: u32 = 0xf;
+    const ADDR: u64 = 0x3000;
+    const SIZE: u64 = 0x1000;
+    const GIC_PPI: u32 = 1;
+    const IRQ_TYPE_EDGE_RISING: u32 = 1;
+    const GIC_FDT_IRQ_PPI_CPU_SHIFT: u32 = 8;
+    // TODO(b/350498812): Rework this for >8 vCPUs.
+    const GIC_FDT_IRQ_PPI_CPU_MASK: u32 = 0xff << Self::GIC_FDT_IRQ_PPI_CPU_SHIFT;
+
+    const fn get_expected(num_cpus: usize) -> Self {
+        Self {
+            addr: Self::ADDR,
+            size: Self::SIZE,
+            irq: [
+                Self::GIC_PPI,
+                Self::IRQ_NR,
+                ((((1 << num_cpus) - 1) << Self::GIC_FDT_IRQ_PPI_CPU_SHIFT)
+                    & Self::GIC_FDT_IRQ_PPI_CPU_MASK)
+                    | Self::IRQ_TYPE_EDGE_RISING,
+            ],
+        }
+    }
+}
+
+fn read_wdt_info_from(fdt: &Fdt) -> libfdt::Result<WdtInfo> {
+    let mut node_iter = fdt.compatible_nodes(cstr!("qemu,vcpu-stall-detector"))?;
+    let node = node_iter.next().ok_or(FdtError::NotFound)?;
+    let mut ranges = node.reg()?.ok_or(FdtError::NotFound)?;
+
+    let reg = ranges.next().ok_or(FdtError::NotFound)?;
+    let size = reg.size.ok_or(FdtError::NotFound)?;
+    if ranges.next().is_some() {
+        warn!("Discarding extra vmwdt <reg> entries.");
+    }
+
+    let interrupts = node.getprop_cells(cstr!("interrupts"))?.ok_or(FdtError::NotFound)?;
+    let mut chunks = CellChunkIterator::<{ WdtInfo::IRQ_CELLS }>::new(interrupts);
+    let irq = chunks.next().ok_or(FdtError::NotFound)?;
+
+    if chunks.next().is_some() {
+        warn!("Discarding extra vmwdt <interrupts> entries.");
+    }
+
+    Ok(WdtInfo { addr: reg.addr, size, irq })
+}
+
+fn validate_wdt_info(wdt: &WdtInfo, num_cpus: usize) -> Result<(), RebootReason> {
+    if *wdt != WdtInfo::get_expected(num_cpus) {
+        error!("Invalid watchdog timer: {wdt:?}");
+        return Err(RebootReason::InvalidFdt);
+    }
+
+    Ok(())
+}
+
+fn patch_wdt_info(fdt: &mut Fdt, num_cpus: usize) -> libfdt::Result<()> {
+    let mut interrupts = WdtInfo::get_expected(num_cpus).irq;
+    for v in interrupts.iter_mut() {
+        *v = v.to_be();
+    }
+
+    let mut node = fdt
+        .root_mut()
+        .next_compatible(cstr!("qemu,vcpu-stall-detector"))?
+        .ok_or(libfdt::FdtError::NotFound)?;
+    node.setprop_inplace(cstr!("interrupts"), interrupts.as_bytes())?;
+    Ok(())
+}
+
 /// Patch the DT by deleting the ns16550a compatible nodes whose address are unknown
 fn patch_serial_info(fdt: &mut Fdt, serial_info: &SerialInfo) -> libfdt::Result<()> {
     let name = cstr!("ns16550a");
@@ -862,7 +940,9 @@
         interrupts.take(NUM_INTERRUPTS * CELLS_PER_INTERRUPT).collect();
 
     let num_cpus: u32 = num_cpus.try_into().unwrap();
+    // TODO(b/350498812): Rework this for >8 vCPUs.
     let cpu_mask: u32 = (((0x1 << num_cpus) - 1) & 0xff) << 8;
+
     for v in value.iter_mut().skip(2).step_by(CELLS_PER_INTERRUPT) {
         *v |= cpu_mask;
     }
@@ -1053,6 +1133,12 @@
     })?;
     validate_pci_info(&pci_info, &memory_range)?;
 
+    let wdt_info = read_wdt_info_from(fdt).map_err(|e| {
+        error!("Failed to read vCPU stall detector info from DT: {e}");
+        RebootReason::InvalidFdt
+    })?;
+    validate_wdt_info(&wdt_info, cpus.len())?;
+
     let serial_info = read_serial_info_from(fdt).map_err(|e| {
         error!("Failed to read serial info from DT: {e}");
         RebootReason::InvalidFdt
@@ -1141,6 +1227,10 @@
         error!("Failed to patch pci info to DT: {e}");
         RebootReason::InvalidFdt
     })?;
+    patch_wdt_info(fdt, info.cpus.len()).map_err(|e| {
+        error!("Failed to patch wdt info to DT: {e}");
+        RebootReason::InvalidFdt
+    })?;
     patch_serial_info(fdt, &info.serial_info).map_err(|e| {
         error!("Failed to patch serial info to DT: {e}");
         RebootReason::InvalidFdt
diff --git a/pvmfw/testdata/test_crosvm_dt_base.dtsi b/pvmfw/testdata/test_crosvm_dt_base.dtsi
index 7d1161a..55f0a14 100644
--- a/pvmfw/testdata/test_crosvm_dt_base.dtsi
+++ b/pvmfw/testdata/test_crosvm_dt_base.dtsi
@@ -141,6 +141,7 @@
 		reg = <0x00 0x3000 0x00 0x1000>;
 		clock-frequency = <0x0a>;
 		timeout-sec = <0x08>;
+		interrupts = <0x01 0xf 0x101>; // <GIC_PPI 0xf IRQ_TYPE_EDGE_RISING>
 	};
 
 	__symbols__ {
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index 0280652..4a61016 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -47,6 +47,7 @@
 import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.regex.Matcher;
@@ -121,9 +122,14 @@
 
     @Test
     public void testNoLongHypSections() throws Exception {
-        assumeTrue("Skip without hypervisor tracing", KvmHypTracer.isSupported(getDevice()));
+        String[] hypEvents = {
+            "hyp_enter", "hyp_exit"
+        };
 
-        KvmHypTracer tracer = new KvmHypTracer(getDevice());
+        assumeTrue("Skip without hypervisor tracing",
+            KvmHypTracer.isSupported(getDevice(), hypEvents));
+
+        KvmHypTracer tracer = new KvmHypTracer(getDevice(), hypEvents);
         String result = tracer.run(COMPOSD_CMD_BIN + " test-compile");
         assertWithMessage("Failed to test compilation VM.")
                 .that(result).ignoringCase().contains("all ok");
@@ -134,6 +140,36 @@
     }
 
     @Test
+    public void testPsciMemProtect() throws Exception {
+        String[] hypEvents = {
+            "psci_mem_protect"
+        };
+
+        assumeTrue("Skip without hypervisor tracing",
+            KvmHypTracer.isSupported(getDevice(), hypEvents));
+        KvmHypTracer tracer = new KvmHypTracer(getDevice(), hypEvents);
+
+        /* We need to wait for crosvm to die so all the VM pages are reclaimed */
+        String result = tracer.run(COMPOSD_CMD_BIN + " test-compile && killall -w crosvm || true");
+        assertWithMessage("Failed to test compilation VM.")
+                .that(result).ignoringCase().contains("all ok");
+
+        List<Integer> values = tracer.getPsciMemProtect();
+
+        assertWithMessage("PSCI MEM_PROTECT events not recorded")
+            .that(values.size()).isGreaterThan(2);
+
+        assertWithMessage("PSCI MEM_PROTECT counter not starting from 0")
+            .that(values.get(0)).isEqualTo(0);
+
+        assertWithMessage("PSCI MEM_PROTECT counter not ending with 0")
+            .that(values.get(values.size() - 1)).isEqualTo(0);
+
+        assertWithMessage("PSCI MEM_PROTECT counter didn't increment")
+            .that(Collections.max(values)).isGreaterThan(0);
+    }
+
+    @Test
     public void testCameraAppStartupTime() throws Exception {
         String[] launchIntentPackages = {
             "com.android.camera2",
diff --git a/tests/ferrochrome/assets/vm_config.json b/tests/ferrochrome/assets/vm_config.json
index 1d32463..358df95 100644
--- a/tests/ferrochrome/assets/vm_config.json
+++ b/tests/ferrochrome/assets/vm_config.json
@@ -2,11 +2,15 @@
     "name": "cros",
     "disks": [
         {
-            "image": "/data/local/tmp/ferrochrome/chromiumos_test_image.bin",
+            "image": "/data/local/tmp/chromiumos_test_image.bin",
             "partitions": [],
             "writable": true
         }
     ],
+    "gpu": {
+        "backend": "virglrenderer",
+        "context_types": ["virgl2"]
+    },
     "params": "root=/dev/vda3 rootwait noinitrd ro enforcing=0 cros_debug cros_secure",
     "protected": false,
     "cpu_topology": "match_host",
@@ -14,4 +18,3 @@
     "memory_mib" : 8096,
     "console_input_device": "ttyS0"
 }
-
diff --git a/tests/ferrochrome/ferrochrome.sh b/tests/ferrochrome/ferrochrome.sh
index d72e882..6814ac5 100755
--- a/tests/ferrochrome/ferrochrome.sh
+++ b/tests/ferrochrome/ferrochrome.sh
@@ -135,7 +135,12 @@
 echo "Starting ferrochrome"
 adb shell am start-activity -a ${ACTION_NAME} > /dev/null
 
-log_path="/data/data/${pkg_name}/files/console.log"
+if [[ $(adb shell getprop ro.fw.mu.headless_system_user) == "true" ]]; then
+  current_user=$(adb shell am get-current-user)
+  log_path="/data/user/${current_user}/${pkg_name}/files/console.log"
+else
+  log_path="/data/data/${pkg_name}/files/console.log"
+fi
 fecr_start_time=${EPOCHSECONDS}
 
 while [[ $((EPOCHSECONDS - fecr_start_time)) -lt ${FECR_BOOT_TIMEOUT} ]]; do
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/KvmHypTracer.java b/tests/hostside/helper/java/com/android/microdroid/test/host/KvmHypTracer.java
index 0d8ee96..5c72358 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/KvmHypTracer.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/KvmHypTracer.java
@@ -30,24 +30,63 @@
 import java.text.ParseException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import javax.annotation.Nonnull;
 
+class KvmHypEvent {
+    public final int cpu;
+    public final double timestamp;
+    public final String name;
+    public final String args;
+    public final boolean valid;
+
+    private static final Pattern LOST_EVENT_PATTERN = Pattern.compile(
+            "^CPU:[0-9]* \\[LOST ([0-9]*) EVENTS\\]");
+
+    public KvmHypEvent(String str) {
+        Matcher matcher = LOST_EVENT_PATTERN.matcher(str);
+        if (matcher.find())
+            throw new OutOfMemoryError("Lost " + matcher.group(1) + " events");
+
+        Pattern pattern = Pattern.compile(
+                "^\\[([0-9]*)\\][ \t]*([0-9]*\\.[0-9]*): (\\S+) (.*)");
+
+        matcher = pattern.matcher(str);
+        if (!matcher.find()) {
+            valid = false;
+            cpu = 0;
+            timestamp = 0;
+            name = "";
+            args = "";
+            CLog.w("Failed to parse hyp event: " + str);
+            return;
+        }
+
+        cpu = Integer.parseInt(matcher.group(1));
+        timestamp = Double.parseDouble(matcher.group(2));
+        name = matcher.group(3);
+        args = matcher.group(4);
+        valid = true;
+    }
+
+    public String toString() {
+        return String.format(
+                "[%03d]\t%f: %s %s", cpu, timestamp, name, args);
+    }
+}
+
 /** This class provides utilities to interact with the hyp tracing subsystem */
 public final class KvmHypTracer {
 
     private static final String HYP_TRACING_ROOT = "/sys/kernel/tracing/hyp/";
-    private static final String HYP_EVENTS[] = { "hyp_enter", "hyp_exit" };
     private static final int DEFAULT_BUF_SIZE_KB = 4 * 1024;
-    private static final Pattern LOST_EVENT_PATTERN = Pattern.compile(
-            "^CPU:[0-9]* \\[LOST ([0-9]*) EVENTS\\]");
-    private static final Pattern EVENT_PATTERN = Pattern.compile(
-            "^\\[([0-9]*)\\][ \t]*([0-9]*\\.[0-9]*): (" + String.join("|", HYP_EVENTS) + ") (.*)");
 
     private final CommandRunner mRunner;
     private final ITestDevice mDevice;
     private final int mNrCpus;
+    private final String mHypEvents[];
 
     private final ArrayList<File> mTraces;
 
@@ -59,22 +98,23 @@
         return "events/hyp/" + event + "/";
     }
 
-    public static boolean isSupported(ITestDevice device) throws Exception {
-        for (String event: HYP_EVENTS) {
+    public static boolean isSupported(ITestDevice device, String[] events) throws Exception {
+        for (String event: events) {
             if (!device.doesFileExist(HYP_TRACING_ROOT + eventDir(event) + "/enable"))
                 return false;
         }
         return true;
     }
 
-    public KvmHypTracer(@Nonnull ITestDevice device) throws Exception {
-        assertWithMessage("Hypervisor tracing not supported")
-                .that(isSupported(device)).isTrue();
+    public KvmHypTracer(@Nonnull ITestDevice device, String[] events) throws Exception {
+        assertWithMessage("Hypervisor events " + String.join(",", events) + " not supported")
+            .that(isSupported(device, events)).isTrue();
 
         mDevice = device;
         mRunner = new CommandRunner(mDevice);
         mTraces = new ArrayList<File>();
         mNrCpus = Integer.parseInt(mRunner.run("nproc"));
+        mHypEvents = events;
     }
 
     public String run(String payload_cmd) throws Exception {
@@ -83,7 +123,7 @@
         setNode("tracing_on", 0);
         mRunner.run("echo 0 | tee " + HYP_TRACING_ROOT + "events/*/*/enable");
         setNode("buffer_size_kb", DEFAULT_BUF_SIZE_KB);
-        for (String event: HYP_EVENTS)
+        for (String event: mHypEvents)
             setNode(eventDir(event) + "/enable", 1);
         setNode("trace", 0);
 
@@ -96,15 +136,21 @@
             cmd += "CPU" + i + "_TRACE_PIPE_PID=$!;";
         }
 
+        String cmd_script = mRunner.run("mktemp -t cmd_script.XXXXXXXXXX");
+        mRunner.run("echo '" + payload_cmd + "' > " + cmd_script);
+
         /* Run the payload with tracing enabled */
         cmd += "echo 1 > tracing_on;";
         String cmd_stdout = mRunner.run("mktemp -t cmd_stdout.XXXXXXXXXX");
-        cmd += payload_cmd + " > " + cmd_stdout + ";";
+        cmd += "sh " + cmd_script + " > " + cmd_stdout + ";";
         cmd += "echo 0 > tracing_on;";
 
-        /* Actively kill the cat subprocesses as trace_pipe is blocking */
-        for (int i = 0; i < mNrCpus; i++)
+        /* Wait for cat to finish reading the pipe interface before killing it */
+        for (int i = 0; i < mNrCpus; i++) {
+            cmd += "while $(test '$(ps -o S -p $CPU" + i
+                + "_TRACE_PIPE_PID | tail -n 1)' = 'R'); do sleep 1; done;";
             cmd += "kill -9 $CPU" + i + "_TRACE_PIPE_PID;";
+        }
         cmd += "wait";
 
         /*
@@ -116,6 +162,8 @@
          */
         mRunner.run(cmd);
 
+        mRunner.run("rm -f " + cmd_script);
+
         for (String t: trace_pipes) {
             File trace = mDevice.pullFile(t);
             assertNotNull(trace);
@@ -128,37 +176,57 @@
         return res;
     }
 
+    private boolean hasEvents(String[] events) {
+        for (String event : events) {
+            if (!Arrays.asList(mHypEvents).contains(event)) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    private KvmHypEvent getNextEvent(BufferedReader br) throws Exception {
+        KvmHypEvent event;
+        String l;
+
+        if ((l = br.readLine()) == null)
+            return null;
+
+        event = new KvmHypEvent(l);
+        if (!event.valid)
+            return null;
+
+        return event;
+    }
+
     public SimpleStats getDurationStats() throws Exception {
+        String[] reqEvents = {"hyp_enter", "hyp_exit"};
         SimpleStats stats = new SimpleStats();
 
+        assertWithMessage("KvmHypTracer() is missing events " + String.join(",", reqEvents))
+            .that(hasEvents(reqEvents)).isTrue();
+
         for (File trace: mTraces) {
             BufferedReader br = new BufferedReader(new FileReader(trace));
             double last = 0.0, hyp_enter = 0.0;
-            String l, prev_event = "";
-            while ((l = br.readLine()) != null) {
-                Matcher matcher = LOST_EVENT_PATTERN.matcher(l);
-                if (matcher.find())
-                    throw new OutOfMemoryError("Lost " + matcher.group(1) + " events");
+            String prev_event = "";
+            KvmHypEvent hypEvent;
 
-                matcher = EVENT_PATTERN.matcher(l);
-                if (!matcher.find()) {
-                    CLog.w("Failed to parse hyp event: " + l);
-                    continue;
-                }
-
-                int cpu = Integer.parseInt(matcher.group(1));
+            while ((hypEvent = getNextEvent(br)) != null) {
+                int cpu = hypEvent.cpu;
                 if (cpu < 0 || cpu >= mNrCpus)
                     throw new ParseException("Incorrect CPU number: " + cpu, 0);
 
-                double cur = Double.parseDouble(matcher.group(2));
+                double cur = hypEvent.timestamp;
                 if (cur < last)
                     throw new ParseException("Time must not go backward: " + cur, 0);
                 last = cur;
 
-                String event = matcher.group(3);
+                String event = hypEvent.name;
                 if (event.equals(prev_event)) {
-                    throw new ParseException("Hyp event found twice in a row: " + trace + " - " + l,
-                                             0);
+                    throw new ParseException("Hyp event found twice in a row: " +
+                                             trace + " - " + hypEvent, 0);
                 }
 
                 switch (event) {
@@ -170,7 +238,7 @@
                         hyp_enter = cur;
                         break;
                     default:
-                        throw new ParseException("Unexpected line in trace" + l, 0);
+                        throw new ParseException("Unexpected line in trace " + hypEvent, 0);
                 }
                 prev_event = event;
             }
@@ -178,4 +246,55 @@
 
         return stats;
     }
+
+    public List<Integer> getPsciMemProtect() throws Exception {
+        String[] reqEvents = {"psci_mem_protect"};
+        List<Integer> psciMemProtect = new ArrayList<>();
+
+        assertWithMessage("KvmHypTracer() is missing events " + String.join(",", reqEvents))
+            .that(hasEvents(reqEvents)).isTrue();
+
+        BufferedReader[] brs = new BufferedReader[mTraces.size()];
+        KvmHypEvent[] next = new KvmHypEvent[mTraces.size()];
+
+        for (int i = 0; i < mTraces.size(); i++) {
+            brs[i] = new BufferedReader(new FileReader(mTraces.get(i)));
+            next[i] = getNextEvent(brs[i]);
+        }
+
+        while (true) {
+            double oldest = Double.MAX_VALUE;
+            int oldestIdx = -1;
+
+            for (int i = 0; i < mTraces.size(); i ++) {
+                if ((next[i] != null) && (next[i].timestamp < oldest)) {
+                    oldest = next[i].timestamp;
+                    oldestIdx = i;
+                }
+            }
+
+            if (oldestIdx < 0)
+                break;
+
+            Pattern pattern = Pattern.compile(
+                "count=([0-9]*) was=([0-9]*)");
+            Matcher matcher = pattern.matcher(next[oldestIdx].args);
+            if (!matcher.find()) {
+                throw new ParseException("Unexpected psci_mem_protect event: " +
+                                         next[oldestIdx], 0);
+            }
+
+            int count = Integer.parseInt(matcher.group(1));
+            int was = Integer.parseInt(matcher.group(2));
+
+            if (psciMemProtect.isEmpty()) {
+                psciMemProtect.add(was);
+            }
+
+            psciMemProtect.add(count);
+            next[oldestIdx] = getNextEvent(brs[oldestIdx]);
+        }
+
+        return psciMemProtect;
+    }
 }
diff --git a/vmbase/src/bionic.rs b/vmbase/src/bionic.rs
index a38fd25..9fe3a4a 100644
--- a/vmbase/src/bionic.rs
+++ b/vmbase/src/bionic.rs
@@ -111,7 +111,7 @@
 ///
 /// # Note
 ///
-/// This Rust functions is missing the last argument of its C/C++ counterpart, a va_list.
+/// This Rust function is missing the last argument of its C/C++ counterpart, a va_list.
 #[no_mangle]
 unsafe extern "C" fn async_safe_fatal_va_list(prefix: *const c_char, format: *const c_char) {
     // SAFETY: The caller guaranteed that both strings were valid and NUL-terminated.
diff --git a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
index c2f218a..d23ee4c 100644
--- a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
+++ b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
@@ -198,12 +198,18 @@
         return configBuilder.build();
     }
 
+    private static boolean isVolumeKey(int keyCode) {
+        return keyCode == KeyEvent.KEYCODE_VOLUME_UP
+                || keyCode == KeyEvent.KEYCODE_VOLUME_DOWN
+                || keyCode == KeyEvent.KEYCODE_VOLUME_MUTE;
+    }
+
     @Override
     public boolean onKeyDown(int keyCode, KeyEvent event) {
         if (mVirtualMachine == null) {
             return false;
         }
-        return mVirtualMachine.sendKeyEvent(event);
+        return !isVolumeKey(keyCode) && mVirtualMachine.sendKeyEvent(event);
     }
 
     @Override
@@ -211,7 +217,7 @@
         if (mVirtualMachine == null) {
             return false;
         }
-        return mVirtualMachine.sendKeyEvent(event);
+        return !isVolumeKey(keyCode) && mVirtualMachine.sendKeyEvent(event);
     }
 
     @Override
@@ -235,41 +241,34 @@
 
                     @Override
                     public void onPayloadStarted(VirtualMachine vm) {
-                        Log.e(TAG, "payload start");
+                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
+                        // this.
                     }
 
                     @Override
                     public void onPayloadReady(VirtualMachine vm) {
-                        // This check doesn't 100% prevent race condition or UI hang.
-                        // However, it's fine for demo.
-                        if (mService.isShutdown()) {
-                            return;
-                        }
-                        Log.d(TAG, "(Payload is ready. Testing VM service...)");
+                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
+                        // this.
                     }
 
                     @Override
                     public void onPayloadFinished(VirtualMachine vm, int exitCode) {
-                        // This check doesn't 100% prevent race condition, but is fine for demo.
-                        if (!mService.isShutdown()) {
-                            Log.d(
-                                    TAG,
-                                    String.format("(Payload finished. exit code: %d)", exitCode));
-                        }
+                        // This event is only from Microdroid-based VM. Custom VM shouldn't emit
+                        // this.
                     }
 
                     @Override
                     public void onError(VirtualMachine vm, int errorCode, String message) {
-                        Log.d(
-                                TAG,
-                                String.format(
-                                        "(Error occurred. code: %d, message: %s)",
-                                        errorCode, message));
+                        Log.e(TAG, "Error from VM. code: " + errorCode + " (" + message + ")");
+                        setResult(RESULT_CANCELED);
+                        finish();
                     }
 
                     @Override
                     public void onStopped(VirtualMachine vm, int reason) {
-                        Log.e(TAG, "vm stop");
+                        Log.d(TAG, "VM stopped. Reason: " + reason);
+                        setResult(RESULT_OK);
+                        finish();
                     }
                 };