diff --git a/authfs/TEST_MAPPING b/authfs/TEST_MAPPING
index 3c84b76..ab6111b 100644
--- a/authfs/TEST_MAPPING
+++ b/authfs/TEST_MAPPING
@@ -6,5 +6,10 @@
     {
       "name": "AuthFsHostTest"
     }
+  ],
+  "avf-postsubmit": [
+    {
+      "name": "AuthFsBenchmarks"
+    }
   ]
 }
diff --git a/authfs/fd_server/src/main.rs b/authfs/fd_server/src/main.rs
index 93a788b..f1fffdd 100644
--- a/authfs/fd_server/src/main.rs
+++ b/authfs/fd_server/src/main.rs
@@ -29,7 +29,7 @@
 use clap::Parser;
 use log::debug;
 use nix::sys::stat::{umask, Mode};
-use rpcbinder::run_rpc_server;
+use rpcbinder::run_vsock_rpc_server;
 use std::collections::BTreeMap;
 use std::fs::File;
 use std::os::unix::io::{FromRawFd, OwnedFd};
@@ -137,7 +137,7 @@
 
     let service = FdService::new_binder(fd_pool).as_binder();
     debug!("fd_server is starting as a rpc service.");
-    let retval = run_rpc_server(service, RPC_SERVICE_PORT, || {
+    let retval = run_vsock_rpc_server(service, RPC_SERVICE_PORT, || {
         debug!("fd_server is ready");
         // Close the ready-fd if we were given one to signal our readiness.
         drop(ready_fd.take());
diff --git a/authfs/tests/Android.bp b/authfs/tests/Android.bp
deleted file mode 100644
index 0177254..0000000
--- a/authfs/tests/Android.bp
+++ /dev/null
@@ -1,55 +0,0 @@
-package {
-    default_applicable_licenses: ["Android-Apache-2.0"],
-}
-
-java_test_host {
-    name: "AuthFsHostTest",
-    srcs: ["java/**/*.java"],
-    libs: [
-        "tradefed",
-        "compatibility-tradefed",
-        "compatibility-host-util",
-    ],
-    static_libs: [
-        "MicrodroidHostTestHelper",
-    ],
-    test_suites: ["general-tests"],
-    data_device_bins_first: [
-        "open_then_run",
-        "fsverity",
-    ],
-    per_testcase_directory: true,
-    data: [
-        ":authfs_test_files",
-        ":CtsApkVerityTestPrebuiltFiles",
-        ":MicrodroidTestApp",
-        ":measure_io",
-    ],
-}
-
-rust_test {
-    name: "open_then_run",
-    crate_name: "open_then_run",
-    srcs: ["open_then_run.rs"],
-    edition: "2021",
-    rustlibs: [
-        "libandroid_logger",
-        "libanyhow",
-        "libclap",
-        "libcommand_fds",
-        "liblibc",
-        "liblog_rust",
-    ],
-    test_suites: ["general-tests"],
-    test_harness: false,
-}
-
-cc_binary {
-    name: "measure_io",
-    srcs: [
-        "measure_io.cpp",
-    ],
-    shared_libs: [
-        "libbase",
-    ],
-}
diff --git a/authfs/tests/benchmarks/Android.bp b/authfs/tests/benchmarks/Android.bp
new file mode 100644
index 0000000..b198328
--- /dev/null
+++ b/authfs/tests/benchmarks/Android.bp
@@ -0,0 +1,37 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+java_test_host {
+    name: "AuthFsBenchmarks",
+    srcs: ["src/java/com/android/fs/benchmarks/*.java"],
+    libs: [
+        "tradefed",
+    ],
+    static_libs: [
+        "AuthFsHostTestCommon",
+        "MicrodroidHostTestHelper",
+    ],
+    test_suites: ["general-tests"],
+    data_device_bins_first: [
+        "open_then_run",
+        "fsverity",
+    ],
+    per_testcase_directory: true,
+    data: [
+        ":authfs_test_files",
+        ":CtsApkVerityTestPrebuiltFiles",
+        ":MicrodroidTestApp",
+        ":measure_io",
+    ],
+}
+
+cc_binary {
+    name: "measure_io",
+    srcs: [
+        "src/measure_io.cpp",
+    ],
+    shared_libs: [
+        "libbase",
+    ],
+}
diff --git a/authfs/tests/benchmarks/AndroidTest.xml b/authfs/tests/benchmarks/AndroidTest.xml
new file mode 100644
index 0000000..7ca3a80
--- /dev/null
+++ b/authfs/tests/benchmarks/AndroidTest.xml
@@ -0,0 +1,64 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2022 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="Config for authfs tests">
+    <!-- Need root to start virtualizationservice -->
+    <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer"/>
+
+    <!-- Still need to define SELinux policy for authfs and fd_server properly. -->
+    <target_preparer class="com.android.tradefed.targetprep.DisableSELinuxTargetPreparer"/>
+
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="throw-if-cmd-fail" value="true" />
+        <!-- Prepare test directories. -->
+        <option name="run-command" value="mkdir -p /data/local/tmp/authfs/mnt" />
+        <option name="teardown-command" value="rm -rf /data/local/tmp/authfs" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.PushFilePreparer">
+        <option name="cleanup" value="true" />
+        <option name="abort-on-push-failure" value="true" />
+
+        <!-- Test executable -->
+        <option name="push-file" key="open_then_run" value="/data/local/tmp/open_then_run" />
+        <option name="push-file" key="fsverity" value="/data/local/tmp/fsverity" />
+
+        <!-- Test data files -->
+        <option name="push-file" key="cert.der" value="/data/local/tmp/authfs/cert.der" />
+        <option name="push-file" key="input.4m" value="/data/local/tmp/authfs/input.4m" />
+        <option name="push-file" key="input.4m.fsv_meta"
+            value="/data/local/tmp/authfs/input.4m.fsv_meta" />
+
+        <!-- Just pick a file with signature that can be trused on the device. -->
+        <option name="push-file" key="CtsApkVerityTestAppPrebuilt.apk"
+            value="/data/local/tmp/authfs/input.apk" />
+        <option name="push-file" key="CtsApkVerityTestAppPrebuilt.apk.fsv_sig"
+            value="/data/local/tmp/authfs/input.apk.fsv_sig" />
+    </target_preparer>
+
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="throw-if-cmd-fail" value="true" />
+        <!-- Now that the files are pushed to the device, enable fs-verity for the targeting file.
+             It works because the signature is trusted on all CTS compatible devices. -->
+        <option name="run-command"
+            value="cd /data/local/tmp/authfs;
+                   ../fsverity enable input.apk --signature=input.apk.fsv_sig" />
+    </target_preparer>
+
+    <test class="com.android.compatibility.common.tradefed.testtype.JarHostTest" >
+        <option name="jar" value="AuthFsBenchmarks.jar" />
+    </test>
+</configuration>
diff --git a/authfs/tests/java/src/com/android/fs/AuthFsBenchmarks.java b/authfs/tests/benchmarks/src/java/com/android/fs/benchmarks/AuthFsBenchmarks.java
similarity index 98%
rename from authfs/tests/java/src/com/android/fs/AuthFsBenchmarks.java
rename to authfs/tests/benchmarks/src/java/com/android/fs/benchmarks/AuthFsBenchmarks.java
index af2c892..5e9073a 100644
--- a/authfs/tests/java/src/com/android/fs/AuthFsBenchmarks.java
+++ b/authfs/tests/benchmarks/src/java/com/android/fs/benchmarks/AuthFsBenchmarks.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.virt.fs;
+package com.android.virt.fs.benchmarks;
 
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 
@@ -22,6 +22,7 @@
 
 import android.platform.test.annotations.RootPermissionTest;
 
+import com.android.fs.common.AuthFsTestRule;
 import com.android.microdroid.test.common.MetricsProcessor;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.invoker.TestInformation;
diff --git a/authfs/tests/measure_io.cpp b/authfs/tests/benchmarks/src/measure_io.cpp
similarity index 100%
rename from authfs/tests/measure_io.cpp
rename to authfs/tests/benchmarks/src/measure_io.cpp
diff --git a/authfs/tests/common/Android.bp b/authfs/tests/common/Android.bp
new file mode 100644
index 0000000..ec426c7
--- /dev/null
+++ b/authfs/tests/common/Android.bp
@@ -0,0 +1,33 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+java_library_host {
+    name: "AuthFsHostTestCommon",
+    srcs: ["src/java/**/*.java"],
+    libs: [
+        "compatibility-host-util",
+        "compatibility-tradefed",
+        "tradefed",
+    ],
+    static_libs: [
+        "MicrodroidHostTestHelper",
+    ],
+}
+
+rust_test {
+    name: "open_then_run",
+    crate_name: "open_then_run",
+    srcs: ["src/open_then_run.rs"],
+    edition: "2021",
+    rustlibs: [
+        "libandroid_logger",
+        "libanyhow",
+        "libclap",
+        "libcommand_fds",
+        "liblibc",
+        "liblog_rust",
+    ],
+    test_suites: ["general-tests"],
+    test_harness: false,
+}
diff --git a/authfs/tests/java/src/com/android/fs/AuthFsTestRule.java b/authfs/tests/common/src/java/com/android/fs/common/AuthFsTestRule.java
similarity index 90%
rename from authfs/tests/java/src/com/android/fs/AuthFsTestRule.java
rename to authfs/tests/common/src/java/com/android/fs/common/AuthFsTestRule.java
index a9fdd10..994f23b 100644
--- a/authfs/tests/java/src/com/android/fs/AuthFsTestRule.java
+++ b/authfs/tests/common/src/java/com/android/fs/common/AuthFsTestRule.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.virt.fs;
+package com.android.fs.common;
 
 import static com.android.microdroid.test.host.LogArchiver.archiveLogThenDelete;
 import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
@@ -50,22 +50,22 @@
 /** Custom TestRule for AuthFs tests. */
 public class AuthFsTestRule extends TestLogData {
     /** FUSE's magic from statfs(2) */
-    static final String FUSE_SUPER_MAGIC_HEX = "65735546";
+    public static final String FUSE_SUPER_MAGIC_HEX = "65735546";
 
     /** VM config entry path in the test APK */
     private static final String VM_CONFIG_PATH_IN_APK = "assets/vm_config.json";
 
     /** Test directory on Android where data are located */
-    static final String TEST_DIR = "/data/local/tmp/authfs";
+    public static final String TEST_DIR = "/data/local/tmp/authfs";
 
     /** File name of the test APK */
     private static final String TEST_APK_NAME = "MicrodroidTestApp.apk";
 
     /** Output directory where the test can generate output on Android */
-    static final String TEST_OUTPUT_DIR = "/data/local/tmp/authfs/output_dir";
+    public static final String TEST_OUTPUT_DIR = "/data/local/tmp/authfs/output_dir";
 
     /** Mount point of authfs on Microdroid during the test */
-    static final String MOUNT_DIR = "/data/local/tmp/mnt";
+    public static final String MOUNT_DIR = "/data/local/tmp/mnt";
 
     /** VM's log file */
     private static final String LOG_PATH = TEST_OUTPUT_DIR + "/log.txt";
@@ -91,7 +91,7 @@
 
     private final ExecutorService mThreadPool = Executors.newCachedThreadPool();
 
-    static void setUpAndroid(TestInformation testInfo) throws Exception {
+    public static void setUpAndroid(TestInformation testInfo) throws Exception {
         assertNotNull(testInfo.getDevice());
         if (!(testInfo.getDevice() instanceof TestDevice)) {
             CLog.w("Unexpected type of ITestDevice. Skipping.");
@@ -110,28 +110,28 @@
         }
     }
 
-    static void tearDownAndroid() {
+    public static void tearDownAndroid() {
         sAndroid = null;
     }
 
     /** This method is supposed to be called after {@link #setUpTest()}. */
-    static CommandRunner getAndroid() {
+    public static CommandRunner getAndroid() {
         assertThat(sAndroid).isNotNull();
         return sAndroid;
     }
 
     /** This method is supposed to be called after {@link #setUpTest()}. */
-    static CommandRunner getMicrodroid() {
+    public static CommandRunner getMicrodroid() {
         assertThat(sMicrodroid).isNotNull();
         return sMicrodroid;
     }
 
-    static ITestDevice getMicrodroidDevice() {
+    public static ITestDevice getMicrodroidDevice() {
         assertThat(sMicrodroidDevice).isNotNull();
         return sMicrodroidDevice;
     }
 
-    static void startMicrodroid() throws DeviceNotAvailableException {
+    public static void startMicrodroid() throws DeviceNotAvailableException {
         CLog.i("Starting the shared VM");
         assertThat(sMicrodroidDevice).isNull();
         sMicrodroidDevice =
@@ -151,7 +151,7 @@
         assertThat(sMicrodroidDevice.enableAdbRoot()).isTrue();
     }
 
-    static void shutdownMicrodroid() throws DeviceNotAvailableException {
+    public static void shutdownMicrodroid() throws DeviceNotAvailableException {
         assertNotNull(sMicrodroidDevice);
         getDevice().shutdownMicrodroid(sMicrodroidDevice);
         sMicrodroidDevice = null;
@@ -172,7 +172,7 @@
                 description);
     }
 
-    void runFdServerOnAndroid(String helperFlags, String fdServerFlags)
+    public void runFdServerOnAndroid(String helperFlags, String fdServerFlags)
             throws DeviceNotAvailableException {
         String cmd =
                 "cd "
@@ -188,7 +188,7 @@
         Future<?> unusedFuture = mThreadPool.submit(() -> runForResult(sAndroid, cmd, "fd_server"));
     }
 
-    void runAuthFsOnMicrodroid(String flags) {
+    public void runAuthFsOnMicrodroid(String flags) {
         String cmd = AUTHFS_BIN + " " + MOUNT_DIR + " " + flags + " --cid " + VMADDR_CID_HOST;
 
         AtomicBoolean starting = new AtomicBoolean(true);
@@ -215,7 +215,7 @@
         }
     }
 
-    static File findTestFile(IBuildInfo buildInfo, String fileName) {
+    public static File findTestFile(IBuildInfo buildInfo, String fileName) {
         try {
             return (new CompatibilityBuildHelper(buildInfo)).getTestFile(fileName);
         } catch (FileNotFoundException e) {
diff --git a/authfs/tests/open_then_run.rs b/authfs/tests/common/src/open_then_run.rs
similarity index 100%
rename from authfs/tests/open_then_run.rs
rename to authfs/tests/common/src/open_then_run.rs
diff --git a/authfs/tests/hosttests/Android.bp b/authfs/tests/hosttests/Android.bp
new file mode 100644
index 0000000..4b8151d
--- /dev/null
+++ b/authfs/tests/hosttests/Android.bp
@@ -0,0 +1,26 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+java_test_host {
+    name: "AuthFsHostTest",
+    srcs: ["java/src/com/android/fs/*.java"],
+    libs: [
+        "tradefed",
+    ],
+    static_libs: [
+        "MicrodroidHostTestHelper",
+        "AuthFsHostTestCommon",
+    ],
+    test_suites: ["general-tests"],
+    data_device_bins_first: [
+        "open_then_run",
+        "fsverity",
+    ],
+    per_testcase_directory: true,
+    data: [
+        ":authfs_test_files",
+        ":CtsApkVerityTestPrebuiltFiles",
+        ":MicrodroidTestApp",
+    ],
+}
diff --git a/authfs/tests/AndroidTest.xml b/authfs/tests/hosttests/AndroidTest.xml
similarity index 100%
rename from authfs/tests/AndroidTest.xml
rename to authfs/tests/hosttests/AndroidTest.xml
diff --git a/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java b/authfs/tests/hosttests/java/src/com/android/fs/AuthFsHostTest.java
similarity index 99%
rename from authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
rename to authfs/tests/hosttests/java/src/com/android/fs/AuthFsHostTest.java
index c4169f6..3157dfd 100644
--- a/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
+++ b/authfs/tests/hosttests/java/src/com/android/fs/AuthFsHostTest.java
@@ -24,6 +24,7 @@
 
 import android.platform.test.annotations.RootPermissionTest;
 
+import com.android.fs.common.AuthFsTestRule;
 import com.android.microdroid.test.host.CommandRunner;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.invoker.TestInformation;
diff --git a/compos/src/compsvc_main.rs b/compos/src/compsvc_main.rs
index 991725d..a4e3903 100644
--- a/compos/src/compsvc_main.rs
+++ b/compos/src/compsvc_main.rs
@@ -25,7 +25,7 @@
 use anyhow::{bail, Result};
 use compos_common::COMPOS_VSOCK_PORT;
 use log::{debug, error};
-use rpcbinder::run_rpc_server;
+use rpcbinder::run_vsock_rpc_server;
 use std::panic;
 use vm_payload_bindgen::AVmPayload_notifyPayloadReady;
 
@@ -48,7 +48,7 @@
     let service = compsvc::new_binder()?.as_binder();
     debug!("compsvc is starting as a rpc service.");
     // SAFETY: Invokes a method from the bindgen library `vm_payload_bindgen`.
-    let retval = run_rpc_server(service, COMPOS_VSOCK_PORT, || unsafe {
+    let retval = run_vsock_rpc_server(service, COMPOS_VSOCK_PORT, || unsafe {
         AVmPayload_notifyPayloadReady();
     });
     if retval {
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 61c41a7..81f97f3 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -124,6 +124,13 @@
             "android.permission.MANAGE_VIRTUAL_MACHINE";
 
     /**
+     * The permission needed to create a virtual machine with more advanced configuration options.
+     */
+    public static final String USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION =
+            "android.permission.USE_CUSTOM_VIRTUAL_MACHINE";
+
+
+    /**
      * Status of a virtual machine
      *
      * @hide
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index 9f62579..d8ff8c6 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -21,6 +21,7 @@
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.RequiresPermission;
 import android.content.Context;
 import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
@@ -352,7 +353,7 @@
             }
 
             if (!mProtectedVmSet) {
-                throw new IllegalStateException("protectedVm must be set explicitly");
+                throw new IllegalStateException("setProtectedVm(t/f) must be called explicitly");
             }
 
             if (mProtectedVm
@@ -376,6 +377,7 @@
          *
          * @hide
          */
+        @RequiresPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION)
         @NonNull
         public Builder setPayloadConfigPath(@NonNull String payloadConfigPath) {
             mPayloadConfigPath = Objects.requireNonNull(payloadConfigPath);
@@ -383,8 +385,8 @@
         }
 
         /**
-         * Sets the path within the APK to the payload binary file that will be executed within
-         * the VM.
+         * Sets the path within the {@code lib/<ABI>} directory of the APK to the payload binary
+         * file that will be executed within the VM.
          *
          * @hide
          */
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 81f94bc..0364896 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -168,7 +168,8 @@
         "grep ro\\.build\\.version\\.security_patch= $(location :buildinfo.prop) && " +
         "grep ro\\.build\\.version\\.known_codenames= $(location :buildinfo.prop) && " +
         "cat $(location build.prop) && " +
-        "echo ro.product.cpu.abilist=x86_64) > $(out)",
+        "echo ro.product.cpu.abilist=x86_64 && " +
+        "echo ro.product.cpu.abi=x86_64) > $(out)",
 }
 
 genrule {
@@ -185,7 +186,8 @@
         "grep ro\\.build\\.version\\.security_patch= $(location :buildinfo.prop) && " +
         "grep ro\\.build\\.version\\.known_codenames= $(location :buildinfo.prop) && " +
         "cat $(location build.prop) && " +
-        "echo ro.product.cpu.abilist=arm64-v8a) > $(out)",
+        "echo ro.product.cpu.abilist=arm64-v8a && " +
+        "echo ro.product.cpu.abi=arm64-v8a) > $(out)",
 }
 
 android_filesystem {
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index 8741fb8..4b06b3e 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -34,7 +34,6 @@
         "libonce_cell",
         "libopenssl",
         "libprotobuf",
-        "libregex",
         "librpcbinder_rs",
         "librustutils",
         "libscopeguard",
diff --git a/microdroid_manager/src/procutil.rs b/microdroid_manager/src/procutil.rs
index b323ca9..f9479c4 100644
--- a/microdroid_manager/src/procutil.rs
+++ b/microdroid_manager/src/procutil.rs
@@ -12,10 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use anyhow::{Context, Result};
+use anyhow::{bail, Result};
 use libc::{sysconf, _SC_CLK_TCK};
-use regex::Regex;
-use std::fs::{self, File};
+use std::fs::File;
 use std::io::{BufRead, BufReader};
 
 const MILLIS_PER_SEC: i64 = 1000;
@@ -57,19 +56,20 @@
 //   sys: 10771070
 //   idle: 10480973587
 pub fn get_cpu_time() -> Result<CpuTime> {
-    let re = Regex::new(r"^cpu\s+([\d]+)\s([\d]+)\s([\d]+)\s([\d]+)").unwrap();
-
     let mut proc_stat = BufReader::new(File::open("/proc/stat")?);
     let mut line = String::new();
     proc_stat.read_line(&mut line)?;
-    let data_list = re.captures(&line).context("Failed to capture values")?;
+    let data_list: Vec<_> = line.split_whitespace().filter_map(|s| s.parse::<i64>().ok()).collect();
+    if data_list.len() < 4 {
+        bail!("Failed to extract numeric values in /proc/stat :\n{}", line);
+    }
 
     let ticks_per_sec = unsafe { sysconf(_SC_CLK_TCK) } as i64;
     let cpu_time = CpuTime {
-        user: data_list.get(1).unwrap().as_str().parse::<i64>()? * MILLIS_PER_SEC / ticks_per_sec,
-        nice: data_list.get(2).unwrap().as_str().parse::<i64>()? * MILLIS_PER_SEC / ticks_per_sec,
-        sys: data_list.get(3).unwrap().as_str().parse::<i64>()? * MILLIS_PER_SEC / ticks_per_sec,
-        idle: data_list.get(4).unwrap().as_str().parse::<i64>()? * MILLIS_PER_SEC / ticks_per_sec,
+        user: data_list[0] * MILLIS_PER_SEC / ticks_per_sec,
+        nice: data_list[1] * MILLIS_PER_SEC / ticks_per_sec,
+        sys: data_list[2] * MILLIS_PER_SEC / ticks_per_sec,
+        idle: data_list[3] * MILLIS_PER_SEC / ticks_per_sec,
     };
     Ok(cpu_time)
 }
@@ -105,21 +105,23 @@
 //   buffer: 10231296
 //   cached: 189502836
 pub fn get_mem_info() -> Result<MemInfo> {
-    let re = Regex::new(r"^.*?:\s+([0-9]+)\skB").unwrap();
-
-    let proc_mem_info = fs::read_to_string("/proc/meminfo")?;
-    let data_list: Vec<_> = proc_mem_info
-        .trim()
-        .splitn(6, '\n')
-        .map(|s| re.captures(s).context("Failed to capture values").ok()?.get(1))
-        .collect();
+    let mut proc_stat = BufReader::new(File::open("/proc/meminfo")?);
+    let mut lines = String::new();
+    for _ in 0..5 {
+        proc_stat.read_line(&mut lines)?;
+    }
+    let data_list: Vec<_> =
+        lines.split_whitespace().filter_map(|s| s.parse::<i64>().ok()).collect();
+    if data_list.len() != 5 {
+        bail!("Failed to extract numeric values in /proc/meminfo :\n{}", lines);
+    }
 
     let mem_info = MemInfo {
-        total: data_list[0].unwrap().as_str().parse::<i64>()?,
-        free: data_list[1].unwrap().as_str().parse::<i64>()?,
-        available: data_list[2].unwrap().as_str().parse::<i64>()?,
-        buffer: data_list[3].unwrap().as_str().parse::<i64>()?,
-        cached: data_list[4].unwrap().as_str().parse::<i64>()?,
+        total: data_list[0],
+        free: data_list[1],
+        available: data_list[2],
+        buffer: data_list[3],
+        cached: data_list[4],
     };
     Ok(mem_info)
 }
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index 84cb18c..296644a 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -9,6 +9,7 @@
     srcs: ["src/main.rs"],
     edition: "2021",
     rustlibs: [
+        "liblog_rust_nostd",
         "libvmbase",
     ],
     apex_available: ["com.android.virt"],
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
new file mode 100644
index 0000000..95dc0b0
--- /dev/null
+++ b/pvmfw/src/entry.rs
@@ -0,0 +1,81 @@
+// Copyright 2022, 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.
+
+//! Low-level entry and exit points of pvmfw.
+
+use crate::helpers::FDT_MAX_SIZE;
+use crate::jump_to_payload;
+use crate::mmio_guard;
+use core::slice;
+use log::{debug, LevelFilter};
+use vmbase::{console, logger, main, power::reboot};
+
+#[derive(Debug, Clone)]
+enum RebootReason {
+    /// An unexpected internal error happened.
+    InternalError,
+}
+
+main!(start);
+
+/// Entry point for pVM firmware.
+pub fn start(fdt_address: u64, payload_start: u64, payload_size: u64, _arg3: u64) {
+    // Limitations in this function:
+    // - can't access non-pvmfw memory (only statically-mapped memory)
+    // - can't access MMIO (therefore, no logging)
+
+    match main_wrapper(fdt_address as usize, payload_start as usize, payload_size as usize) {
+        Ok(_) => jump_to_payload(fdt_address, payload_start),
+        Err(_) => reboot(),
+    }
+
+    // if we reach this point and return, vmbase::entry::rust_entry() will call power::shutdown().
+}
+
+/// Sets up the environment for main() and wraps its result for start().
+///
+/// Provide the abstractions necessary for start() to abort the pVM boot and for main() to run with
+/// the assumption that its environment has been properly configured.
+fn main_wrapper(fdt: usize, payload: usize, payload_size: usize) -> Result<(), RebootReason> {
+    // Limitations in this function:
+    // - only access MMIO once (and while) it has been mapped and configured
+    // - only perform logging once the logger has been initialized
+    // - only access non-pvmfw memory once (and while) it has been mapped
+    logger::init(LevelFilter::Info).map_err(|_| RebootReason::InternalError)?;
+
+    // TODO: Check that the FDT is fully contained in RAM.
+    // SAFETY - We trust the VMM, for now.
+    let fdt = unsafe { slice::from_raw_parts_mut(fdt as *mut u8, FDT_MAX_SIZE) };
+    // TODO: Check that the payload is fully contained in RAM and doesn't overlap with the FDT.
+    // SAFETY - We trust the VMM, for now.
+    let payload = unsafe { slice::from_raw_parts(payload as *const u8, payload_size) };
+
+    // Use debug!() to avoid printing to the UART if we failed to configure it as only local
+    // builds that have tweaked the logger::init() call will actually attempt to log the message.
+
+    mmio_guard::init().map_err(|e| {
+        debug!("{e}");
+        RebootReason::InternalError
+    })?;
+
+    mmio_guard::map(console::BASE_ADDRESS).map_err(|e| {
+        debug!("Failed to configure the UART: {e}");
+        RebootReason::InternalError
+    })?;
+
+    // This wrapper allows main() to be blissfully ignorant of platform details.
+    crate::main(fdt, payload);
+
+    Ok(())
+}
diff --git a/pvmfw/src/exceptions.rs b/pvmfw/src/exceptions.rs
index 0fb2911..03fc220 100644
--- a/pvmfw/src/exceptions.rs
+++ b/pvmfw/src/exceptions.rs
@@ -20,14 +20,14 @@
 use vmbase::{console::emergency_write_str, eprintln, power::reboot};
 
 const ESR_32BIT_EXT_DABT: u64 = 0x96000010;
-const UART_PAGE: u64 = page_4kb_of(console::BASE_ADDRESS as u64);
+const UART_PAGE: usize = page_4kb_of(console::BASE_ADDRESS);
 
 #[no_mangle]
 extern "C" fn sync_exception_current(_elr: u64, _spsr: u64) {
     let esr = read_esr();
     let far = read_far();
     // Don't print to the UART if we're handling the exception it could raise.
-    if esr != ESR_32BIT_EXT_DABT || page_4kb_of(far) != UART_PAGE {
+    if esr != ESR_32BIT_EXT_DABT || page_4kb_of(far as usize) != UART_PAGE {
         emergency_write_str("sync_exception_current\n");
         print_esr(esr);
     }
diff --git a/pvmfw/src/helpers.rs b/pvmfw/src/helpers.rs
index 781c1ac..cade62e 100644
--- a/pvmfw/src/helpers.rs
+++ b/pvmfw/src/helpers.rs
@@ -14,23 +14,15 @@
 
 //! Miscellaneous helper functions.
 
+pub const FDT_MAX_SIZE: usize = 2 << 20;
+pub const SIZE_4KB: usize = 4 << 10;
+
 /// Computes the address of the page containing a given address.
-pub const fn page_of(addr: u64, page_size: u64) -> u64 {
+pub const fn page_of(addr: usize, page_size: usize) -> usize {
     addr & !(page_size - 1)
 }
 
-/// Validates a page size and computes the address of the page containing a given address.
-pub const fn checked_page_of(addr: u64, page_size: u64) -> Option<u64> {
-    if page_size.is_power_of_two() {
-        Some(page_of(addr, page_size))
-    } else {
-        None
-    }
-}
-
 /// Computes the address of the 4KiB page containing a given address.
-pub const fn page_4kb_of(addr: u64) -> u64 {
-    const PAGE_SIZE: u64 = 4 << 10;
-
-    page_of(addr, PAGE_SIZE)
+pub const fn page_4kb_of(addr: usize) -> usize {
+    page_of(addr, SIZE_4KB)
 }
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index eb97961..d1951b3 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -17,59 +17,24 @@
 #![no_main]
 #![no_std]
 
+mod entry;
 mod exceptions;
 mod helpers;
+mod mmio_guard;
 mod smccc;
 
-use core::fmt;
-use helpers::checked_page_of;
+use log::{debug, info};
 
-use vmbase::{console, main, power::reboot, println};
-
-#[derive(Debug, Clone)]
-enum Error {
-    /// Failed to configure the UART; no logs available.
-    FailedUartSetup,
-}
-
-impl fmt::Display for Error {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let msg = match self {
-            Self::FailedUartSetup => "Failed to configure the UART",
-        };
-        write!(f, "{}", msg)
-    }
-}
-
-fn main(fdt_address: u64, payload_start: u64, payload_size: u64, arg3: u64) -> Result<(), Error> {
-    // We need to inform the hypervisor that the MMIO page containing the UART may be shared back.
-    let uart = console::BASE_ADDRESS as u64;
-    let mmio_granule = smccc::mmio_guard_info().map_err(|_| Error::FailedUartSetup)?;
-    let uart_page = checked_page_of(uart, mmio_granule).ok_or(Error::FailedUartSetup)?;
-    smccc::mmio_guard_map(uart_page).map_err(|_| Error::FailedUartSetup)?;
-
-    println!("pVM firmware");
-    println!(
-        "fdt_address={:#018x}, payload_start={:#018x}, payload_size={:#018x}, x3={:#018x}",
-        fdt_address, payload_start, payload_size, arg3,
+fn main(fdt: &mut [u8], payload: &[u8]) {
+    info!("pVM firmware");
+    debug!(
+        "fdt_address={:#018x}, payload_start={:#018x}, payload_size={:#018x}",
+        fdt.as_ptr() as usize,
+        payload.as_ptr() as usize,
+        payload.len(),
     );
 
-    println!("Starting payload...");
-
-    Ok(())
-}
-
-main!(main_wrapper);
-
-/// Entry point for pVM firmware.
-pub fn main_wrapper(fdt_address: u64, payload_start: u64, payload_size: u64, arg3: u64) {
-    match main(fdt_address, payload_start, payload_size, arg3) {
-        Ok(()) => jump_to_payload(fdt_address, payload_start),
-        Err(e) => {
-            println!("Boot rejected: {}", e);
-        }
-    }
-    reboot()
+    info!("Starting payload...");
 }
 
 fn jump_to_payload(fdt_address: u64, payload_start: u64) {
diff --git a/pvmfw/src/mmio_guard.rs b/pvmfw/src/mmio_guard.rs
new file mode 100644
index 0000000..4cde737
--- /dev/null
+++ b/pvmfw/src/mmio_guard.rs
@@ -0,0 +1,53 @@
+// Copyright 2022, 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.
+
+//! Safe MMIO_GUARD support.
+
+use crate::helpers;
+use crate::smccc;
+use core::{fmt, result};
+
+#[derive(Debug, Clone)]
+pub enum Error {
+    /// Failed to obtain the MMIO_GUARD granule size.
+    InfoFailed(smccc::Error),
+    /// Failed to MMIO_GUARD_MAP a page.
+    MapFailed(smccc::Error),
+    /// The MMIO_GUARD granule used by the hypervisor is not supported.
+    UnsupportedGranule(usize),
+}
+
+type Result<T> = result::Result<T, Error>;
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::InfoFailed(e) => write!(f, "Failed to get the MMIO_GUARD granule: {e}"),
+            Self::MapFailed(e) => write!(f, "Failed to MMIO_GUARD map: {e}"),
+            Self::UnsupportedGranule(g) => write!(f, "Unsupported MMIO_GUARD granule: {g}"),
+        }
+    }
+}
+
+pub fn init() -> Result<()> {
+    let mmio_granule = smccc::mmio_guard_info().map_err(Error::InfoFailed)? as usize;
+    if mmio_granule != helpers::SIZE_4KB {
+        return Err(Error::UnsupportedGranule(mmio_granule));
+    }
+    Ok(())
+}
+
+pub fn map(addr: usize) -> Result<()> {
+    smccc::mmio_guard_map(helpers::page_4kb_of(addr) as u64).map_err(Error::MapFailed)
+}
diff --git a/pvmfw/src/smccc.rs b/pvmfw/src/smccc.rs
index e3a2b05..18128e1 100644
--- a/pvmfw/src/smccc.rs
+++ b/pvmfw/src/smccc.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use core::fmt;
+use core::{fmt, result};
 
 // TODO(b/245889995): use psci-0.1.1 crate
 #[inline(always)]
@@ -69,20 +69,22 @@
             Self::NotSupported => write!(f, "SMCCC call not supported"),
             Self::NotRequired => write!(f, "SMCCC call not required"),
             Self::InvalidParameter => write!(f, "SMCCC call received non-supported value"),
-            Self::Unexpected(v) => write!(f, "Unexpected SMCCC return value '{}'", v),
-            Self::Unknown(e) => write!(f, "Unknown SMCCC return value '{}'", e),
+            Self::Unexpected(v) => write!(f, "Unexpected SMCCC return value {} ({0:#x})", v),
+            Self::Unknown(e) => write!(f, "Unknown SMCCC return value {} ({0:#x})", e),
         }
     }
 }
 
-fn check_smccc_err(ret: i64) -> Result<(), Error> {
+type Result<T> = result::Result<T, Error>;
+
+fn check_smccc_err(ret: i64) -> Result<()> {
     match check_smccc_value(ret)? {
         0 => Ok(()),
         v => Err(Error::Unexpected(v)),
     }
 }
 
-fn check_smccc_value(ret: i64) -> Result<u64, Error> {
+fn check_smccc_value(ret: i64) -> Result<u64> {
     match ret {
         x if x >= 0 => Ok(ret as u64),
         -1 => Err(Error::NotSupported),
@@ -96,7 +98,7 @@
 const VENDOR_HYP_KVM_MMIO_GUARD_MAP_FUNC_ID: u32 = 0xc6000007;
 
 /// Issue pKVM-specific MMIO_GUARD_INFO HVC64.
-pub fn mmio_guard_info() -> Result<u64, Error> {
+pub fn mmio_guard_info() -> Result<u64> {
     let args = [0u64; 17];
 
     let res = hvc64(VENDOR_HYP_KVM_MMIO_GUARD_INFO_FUNC_ID, args);
@@ -105,7 +107,7 @@
 }
 
 /// Issue pKVM-specific MMIO_GUARD_MAP HVC64.
-pub fn mmio_guard_map(ipa: u64) -> Result<(), Error> {
+pub fn mmio_guard_map(ipa: u64) -> Result<()> {
     let mut args = [0u64; 17];
     args[0] = ipa;
 
diff --git a/tests/benchmark/AndroidManifest.xml b/tests/benchmark/AndroidManifest.xml
index ff18130..c39b91c 100644
--- a/tests/benchmark/AndroidManifest.xml
+++ b/tests/benchmark/AndroidManifest.xml
@@ -16,6 +16,7 @@
 <manifest xmlns:android="http://schemas.android.com/apk/res/android"
       package="com.android.microdroid.benchmark">
     <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
+    <uses-permission android:name="android.permission.USE_CUSTOM_VIRTUAL_MACHINE" />
     <application>
         <uses-library android:name="android.system.virtualmachine" android:required="false" />
     </application>
diff --git a/tests/benchmark/AndroidTest.xml b/tests/benchmark/AndroidTest.xml
index 4949d22..0214cd9 100644
--- a/tests/benchmark/AndroidTest.xml
+++ b/tests/benchmark/AndroidTest.xml
@@ -13,7 +13,7 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<configuration description="Runs sample instrumentation test.">
+<configuration description="Runs Microdroid benchmarks.">
     <option name="config-descriptor:metadata" key="component" value="security" />
     <option name="config-descriptor:metadata" key="parameter" value="not_instant_app" />
     <option name="config-descriptor:metadata" key="parameter" value="not_multi_abi" />
@@ -25,11 +25,6 @@
     <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer">
         <option name="force-root" value="true" />
     </target_preparer>
-    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
-      <option
-        name="run-command"
-        value="pm grant com.android.microdroid.benchmark android.permission.MANAGE_VIRTUAL_MACHINE" />
-    </target_preparer>
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.microdroid.benchmark" />
         <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index b1a1160..23d546d 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -80,20 +80,23 @@
 
     @Before
     public void setup() {
+        grantPermission(VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION);
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
         prepareTestSetup(mProtectedVm);
         mInstrumentation = getInstrumentation();
     }
 
     private boolean canBootMicrodroidWithMemory(int mem)
             throws VirtualMachineException, InterruptedException, IOException {
-        final int trialCount = 5;
+        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
+                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setDebugLevel(DEBUG_LEVEL_NONE)
+                .setMemoryMib(mem)
+                .build();
 
         // returns true if succeeded at least once.
+        final int trialCount = 5;
         for (int i = 0; i < trialCount; i++) {
-            VirtualMachineConfig.Builder builder =
-                    mInner.newVmConfigBuilder("assets/vm_config.json");
-            VirtualMachineConfig normalConfig =
-                    builder.setDebugLevel(DEBUG_LEVEL_NONE).setMemoryMib(mem).build();
             mInner.forceCreateNewVirtualMachine("test_vm_minimum_memory", normalConfig);
 
             if (tryBootVm(TAG, "test_vm_minimum_memory").payloadStarted) return true;
@@ -144,7 +147,8 @@
         for (int i = 0; i < trialCount; i++) {
 
             // To grab boot events from log, set debug mode to FULL
-            VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder("assets/vm_config.json")
+            VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
+                    .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
                     .setDebugLevel(DEBUG_LEVEL_FULL)
                     .setMemoryMib(256)
                     .build();
@@ -191,7 +195,8 @@
 
     @Test
     public void testVsockTransferFromHostToVM() throws Exception {
-        VirtualMachineConfig config = mInner.newVmConfigBuilder("assets/vm_config_io.json")
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadConfigPath("assets/vm_config_io.json")
                 .setDebugLevel(DEBUG_LEVEL_FULL)
                 .build();
         List<Double> transferRates = new ArrayList<>(IO_TEST_TRIAL_COUNT);
@@ -217,7 +222,8 @@
     }
 
     private void testVirtioBlkReadRate(boolean isRand) throws Exception {
-        VirtualMachineConfig config = mInner.newVmConfigBuilder("assets/vm_config_io.json")
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadConfigPath("assets/vm_config_io.json")
                 .setDebugLevel(DEBUG_LEVEL_FULL)
                 .build();
         List<Double> readRates = new ArrayList<>(IO_TEST_TRIAL_COUNT);
@@ -282,11 +288,11 @@
     @Test
     public void testMemoryUsage() throws Exception {
         final String vmName = "test_vm_mem_usage";
-        VirtualMachineConfig config =
-                mInner.newVmConfigBuilder("assets/vm_config_io.json")
-                        .setDebugLevel(DEBUG_LEVEL_NONE)
-                        .setMemoryMib(256)
-                        .build();
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadConfigPath("assets/vm_config_io.json")
+                .setDebugLevel(DEBUG_LEVEL_NONE)
+                .setMemoryMib(256)
+                .build();
         mInner.forceCreateNewVirtualMachine(vmName, config);
         VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
         MemoryUsageListener listener = new MemoryUsageListener(this::executeCommand);
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 842d713..c394756 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -165,8 +165,8 @@
         }
     };
 
-    if (!RunRpcServerCallback(test_service->asBinder().get(), test_service->SERVICE_PORT, callback,
-                              nullptr)) {
+    if (!RunVsockRpcServerCallback(test_service->asBinder().get(), test_service->SERVICE_PORT,
+                                   callback, nullptr)) {
         return Error() << "RPC Server failed to run";
     }
     return {};
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index f3aa199..17f773e 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -159,6 +159,8 @@
             throws DeviceNotAvailableException, InterruptedException {
         CommandRunner android = new CommandRunner(getDevice());
         unlockScreen(android);
+        // Ensure we are killing the app to get the cold app startup time
+        android.run("am force-stop " + pkgName);
         android.run("echo 3 > /proc/sys/vm/drop_caches");
         String vmStartAppLog = android.run("am", "start -W -S " + pkgName);
         assertNotNull(vmStartAppLog);
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 07705ba..66cd211 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
@@ -19,6 +19,7 @@
 
 import static org.junit.Assume.assumeNoException;
 
+import android.app.Instrumentation;
 import android.app.UiAutomation;
 import android.content.Context;
 import android.os.ParcelFileDescriptor;
@@ -32,6 +33,7 @@
 
 import androidx.annotation.CallSuper;
 import androidx.test.core.app.ApplicationProvider;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.microdroid.test.common.MetricsProcessor;
 import com.android.virt.VirtualizationTestHelper;
@@ -58,6 +60,20 @@
                 SystemProperties.get("debug.hypervisor.metrics_tag"));
     }
 
+    protected final void grantPermission(String permission) {
+        Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
+        UiAutomation uiAutomation = instrumentation.getUiAutomation();
+        uiAutomation.grantRuntimePermission(instrumentation.getContext().getPackageName(),
+                permission);
+    }
+
+    protected final void revokePermission(String permission) {
+        Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
+        UiAutomation uiAutomation = instrumentation.getUiAutomation();
+        uiAutomation.revokeRuntimePermission(instrumentation.getContext().getPackageName(),
+                permission);
+    }
+
     // TODO(b/220920264): remove Inner class; this is a hack to hide virt APEX types
     protected static class Inner {
         private final boolean mProtectedVm;
@@ -82,10 +98,6 @@
             return new VirtualMachineConfig.Builder(mContext).setProtectedVm(mProtectedVm);
         }
 
-        public VirtualMachineConfig.Builder newVmConfigBuilder(String payloadConfigPath) {
-            return newVmConfigBuilder().setPayloadConfigPath(payloadConfigPath);
-        }
-
         /**
          * Creates a new virtual machine, potentially removing an existing virtual machine with
          * given name.
diff --git a/tests/hostside/AndroidTest.xml b/tests/hostside/AndroidTest.xml
index 5c3e5d1..18728ad 100644
--- a/tests/hostside/AndroidTest.xml
+++ b/tests/hostside/AndroidTest.xml
@@ -13,7 +13,7 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<configuration description="Tests for microdroid">
+<configuration description="Host driven tests for Microdroid">
     <option name="test-suite-tag" value="cts" />
     <option name="config-descriptor:metadata" key="component" value="security" />
     <option name="config-descriptor:metadata" key="parameter" value="not_instant_app" />
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
index 6346288..36379af 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
@@ -20,10 +20,9 @@
 import static com.android.microdroid.test.host.CommandResultSubject.command_results;
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestLogData;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
 
 import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
@@ -129,7 +128,7 @@
 
     // Same as runOnHost, but with custom timeout
     private static String runOnHostWithTimeout(long timeoutMillis, String... cmd) {
-        assertTrue(timeoutMillis >= 0);
+        assertThat(timeoutMillis).isAtLeast(0);
         CommandResult result = RunUtil.getDefault().runTimedCmd(timeoutMillis, cmd);
         assertThat(result).isSuccess();
         return result.getStdout().trim();
@@ -213,8 +212,7 @@
         try {
             return (new CompatibilityBuildHelper(buildInfo)).getTestFile(name);
         } catch (FileNotFoundException e) {
-            fail("Missing test file: " + name);
-            return null;
+            throw new AssertionError("Missing test file: " + name, e);
         }
     }
 
@@ -230,7 +228,8 @@
             throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(device);
         String pathLine = android.run("pm", "path", packageName);
-        assertTrue("package not found", pathLine.startsWith("package:"));
+        assertWithMessage("Package " + packageName + " not found")
+                .that(pathLine).startsWith("package:");
         return pathLine.substring("package:".length());
     }
 
@@ -358,7 +357,7 @@
         // Retrieve the CID from the vm tool output
         Pattern pattern = Pattern.compile("with CID (\\d+)");
         Matcher matcher = pattern.matcher(ret);
-        assertTrue(matcher.find());
+        assertWithMessage("Failed to find CID").that(matcher.find()).isTrue();
         return matcher.group(1);
     }
 
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidTestCase.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
similarity index 84%
rename from tests/hostside/java/com/android/microdroid/test/MicrodroidTestCase.java
rename to tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 57882ce..c9df624 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidTestCase.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -23,13 +23,7 @@
 import static com.google.common.truth.Truth.assertWithMessage;
 
 import static org.hamcrest.CoreMatchers.containsString;
-import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.CoreMatchers.nullValue;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
 
 import static java.util.stream.Collectors.toList;
@@ -44,11 +38,8 @@
 import com.android.os.AtomsProto;
 import com.android.os.StatsLog;
 import com.android.tradefed.device.DeviceNotAvailableException;
-import com.android.tradefed.result.TestDescription;
-import com.android.tradefed.result.TestResult;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
-import com.android.tradefed.testtype.junit4.DeviceTestRunOptions;
 import com.android.tradefed.util.CommandResult;
 import com.android.tradefed.util.FileUtil;
 import com.android.tradefed.util.RunUtil;
@@ -79,7 +70,7 @@
 import java.util.regex.Pattern;
 
 @RunWith(DeviceJUnit4ClassRunner.class)
-public class MicrodroidTestCase extends MicrodroidHostTestCaseBase {
+public class MicrodroidHostTests extends MicrodroidHostTestCaseBase {
     private static final String APK_NAME = "MicrodroidTestApp.apk";
     private static final String PACKAGE_NAME = "com.android.microdroid.test";
     private static final String SHELL_PACKAGE_NAME = "com.android.shell";
@@ -99,44 +90,19 @@
     private int minMemorySize() throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(getDevice());
         String abi = android.run("getprop", "ro.product.cpu.abi");
-        assertTrue(abi != null && !abi.isEmpty());
+        assertThat(abi).isNotEmpty();
         if (abi.startsWith("arm64")) {
             return MIN_MEM_ARM64;
         } else if (abi.startsWith("x86_64")) {
             return MIN_MEM_X86_64;
         }
-        fail("Unsupported ABI: " + abi);
-        return 0;
+        throw new AssertionError("Unsupported ABI: " + abi);
     }
 
     private void waitForBootComplete() {
         runOnMicrodroidForResult("watch -e \"getprop dev.bootcomplete | grep '^0$'\"");
     }
 
-    @Test
-    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C-1-4"})
-    public void testCreateVmRequiresPermission() throws Exception {
-        // Revoke the MANAGE_VIRTUAL_MACHINE permission for the test app
-        CommandRunner android = new CommandRunner(getDevice());
-        android.run("pm", "revoke", PACKAGE_NAME, "android.permission.MANAGE_VIRTUAL_MACHINE");
-
-        // Run MicrodroidTests#connectToVmService test, which should fail
-        final DeviceTestRunOptions options =
-                new DeviceTestRunOptions(PACKAGE_NAME)
-                        .setTestClassName(PACKAGE_NAME + ".MicrodroidTests")
-                        .setTestMethodName("connectToVmService[protectedVm=false]")
-                        .setCheckResults(false);
-        assertFalse(runDeviceTests(options));
-
-        Map<TestDescription, TestResult> results = getLastDeviceRunResults().getTestResults();
-        assertThat(results).hasSize(1);
-        TestResult result = results.values().toArray(new TestResult[0])[0];
-        assertTrue(
-                "The test should fail with a permission error",
-                result.getStackTrace()
-                        .contains("android.permission.MANAGE_VIRTUAL_MACHINE permission"));
-    }
-
     private static JSONObject newPartition(String label, String path) {
         return new JSONObject(Map.of("label", label, "path", path));
     }
@@ -165,7 +131,7 @@
         String path = mkPayload.getParentFile().getPath() + separator + System.getenv("PATH");
         runUtil.setEnvVariable("PATH", path);
 
-        List<String> command = new ArrayList<String>();
+        List<String> command = new ArrayList<>();
         command.add("mk_payload");
         command.add("--metadata-only");
         command.add(configFile.toString());
@@ -197,7 +163,7 @@
         String path = signVirtApex.getParentFile().getPath() + separator + System.getenv("PATH");
         runUtil.setEnvVariable("PATH", path);
 
-        List<String> command = new ArrayList<String>();
+        List<String> command = new ArrayList<>();
         command.add("sign_virt_apex");
         for (Map.Entry<String, File> entry : keyOverrides.entrySet()) {
             String filename = entry.getKey();
@@ -310,8 +276,10 @@
         // Pull the virt apex's etc/ directory (which contains images and microdroid.json)
         File virtApexEtcDir = new File(virtApexDir, "etc");
         // We need only etc/ directory for images
-        assertTrue(virtApexEtcDir.mkdirs());
-        assertTrue(getDevice().pullDir(VIRT_APEX + "etc", virtApexEtcDir));
+        assertWithMessage("Failed to mkdir " + virtApexEtcDir)
+                .that(virtApexEtcDir.mkdirs()).isTrue();
+        assertWithMessage("Failed to pull " + VIRT_APEX + "etc")
+                .that(getDevice().pullDir(VIRT_APEX + "etc", virtApexEtcDir)).isTrue();
 
         resignVirtApex(virtApexDir, key, keyOverrides);
 
@@ -370,13 +338,10 @@
         final String initrdPath = TEST_ROOT + "etc/microdroid_initrd_full_debuggable.img";
         config.put("initrd", initrdPath);
         // Add instance image as a partition in disks[1]
-        disks.put(
-            new JSONObject()
+        disks.put(new JSONObject()
                 .put("writable", true)
-                .put(
-                    "partitions",
-                    new JSONArray()
-                        .put(newPartition("vm-instance", instanceImgPath))));
+                .put("partitions",
+                        new JSONArray().put(newPartition("vm-instance", instanceImgPath))));
         // Add payload image disk with partitions:
         // - payload-metadata
         // - apexes: com.android.os.statsd, com.android.adbd, [sharedlib apex](optional)
@@ -417,7 +382,7 @@
                         configPath);
         Pattern pattern = Pattern.compile("with CID (\\d+)");
         Matcher matcher = pattern.matcher(ret);
-        assertTrue(matcher.find());
+        assertWithMessage("Failed to find CID").that(matcher.find()).isTrue();
         return matcher.group(1);
     }
 
@@ -426,7 +391,7 @@
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
     public void testBootFailsWhenProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
-        assumeTrue(isProtectedVmSupported());
+        assumeTrue("Protected VMs are not supported", isProtectedVmSupported());
 
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
@@ -510,12 +475,13 @@
 
     @Test
     public void testTombstonesAreGeneratedUponCrash() throws Exception {
-        assertTrue(isTombstoneGeneratedWithConfig("assets/vm_config_crash.json"));
+        assertThat(isTombstoneGeneratedWithConfig("assets/vm_config_crash.json")).isTrue();
     }
 
     @Test
     public void testTombstonesAreNotGeneratedIfNotExported() throws Exception {
-        assertFalse(isTombstoneGeneratedWithConfig("assets/vm_config_crash_no_tombstone.json"));
+        assertThat(isTombstoneGeneratedWithConfig("assets/vm_config_crash_no_tombstone.json"))
+                .isFalse();
     }
 
     @Test
@@ -548,27 +514,23 @@
         // Check VmCreationRequested atom and clear the statsd report
         List<StatsLog.EventMetricData> data;
         data = ReportUtils.getEventMetricDataList(getDevice());
-        assertEquals(1, data.size());
-        assertEquals(
-                AtomsProto.Atom.VM_CREATION_REQUESTED_FIELD_NUMBER,
-                data.get(0).getAtom().getPushedCase().getNumber());
+        assertThat(data).hasSize(1);
+        assertThat(data.get(0).getAtom().getPushedCase().getNumber()).isEqualTo(
+                AtomsProto.Atom.VM_CREATION_REQUESTED_FIELD_NUMBER);
         AtomsProto.VmCreationRequested atomVmCreationRequested =
                 data.get(0).getAtom().getVmCreationRequested();
-        assertEquals(
-                AtomsProto.VmCreationRequested.Hypervisor.PKVM,
-                atomVmCreationRequested.getHypervisor());
-        assertFalse(atomVmCreationRequested.getIsProtected());
-        assertTrue(atomVmCreationRequested.getCreationSucceeded());
-        assertEquals(0, atomVmCreationRequested.getBinderExceptionCode());
-        assertEquals("VmRunApp", atomVmCreationRequested.getVmIdentifier());
-        assertEquals(
-                AtomsProto.VmCreationRequested.ConfigType.VIRTUAL_MACHINE_APP_CONFIG,
-                atomVmCreationRequested.getConfigType());
-        assertEquals(NUM_VCPUS, atomVmCreationRequested.getNumCpus());
-        assertEquals(minMemorySize(), atomVmCreationRequested.getMemoryMib());
-        assertEquals(
-                "com.android.art:com.android.compos:com.android.sdkext",
-                atomVmCreationRequested.getApexes());
+        assertThat(atomVmCreationRequested.getHypervisor())
+                .isEqualTo(AtomsProto.VmCreationRequested.Hypervisor.PKVM);
+        assertThat(atomVmCreationRequested.getIsProtected()).isFalse();
+        assertThat(atomVmCreationRequested.getCreationSucceeded()).isTrue();
+        assertThat(atomVmCreationRequested.getBinderExceptionCode()).isEqualTo(0);
+        assertThat(atomVmCreationRequested.getVmIdentifier()).isEqualTo("VmRunApp");
+        assertThat(atomVmCreationRequested.getConfigType())
+                .isEqualTo(AtomsProto.VmCreationRequested.ConfigType.VIRTUAL_MACHINE_APP_CONFIG);
+        assertThat(atomVmCreationRequested.getNumCpus()).isEqualTo(NUM_VCPUS);
+        assertThat(atomVmCreationRequested.getMemoryMib()).isEqualTo(minMemorySize());
+        assertThat(atomVmCreationRequested.getApexes())
+                .isEqualTo("com.android.art:com.android.compos:com.android.sdkext");
 
         // Boot VM with microdroid
         adbConnectToMicrodroid(getDevice(), cid);
@@ -576,12 +538,11 @@
 
         // Check VmBooted atom and clear the statsd report
         data = ReportUtils.getEventMetricDataList(getDevice());
-        assertEquals(1, data.size());
-        assertEquals(
-                AtomsProto.Atom.VM_BOOTED_FIELD_NUMBER,
-                data.get(0).getAtom().getPushedCase().getNumber());
+        assertThat(data).hasSize(1);
+        assertThat(data.get(0).getAtom().getPushedCase().getNumber())
+                .isEqualTo(AtomsProto.Atom.VM_BOOTED_FIELD_NUMBER);
         AtomsProto.VmBooted atomVmBooted = data.get(0).getAtom().getVmBooted();
-        assertEquals("VmRunApp", atomVmBooted.getVmIdentifier());
+        assertThat(atomVmBooted.getVmIdentifier()).isEqualTo("VmRunApp");
 
         // Shutdown VM with microdroid
         shutdownMicrodroid(getDevice(), cid);
@@ -590,18 +551,18 @@
 
         // Check VmExited atom and clear the statsd report
         data = ReportUtils.getEventMetricDataList(getDevice());
-        assertEquals(1, data.size());
-        assertEquals(
-                AtomsProto.Atom.VM_EXITED_FIELD_NUMBER,
-                data.get(0).getAtom().getPushedCase().getNumber());
+        assertThat(data).hasSize(1);
+        assertThat(data.get(0).getAtom().getPushedCase().getNumber())
+                .isEqualTo(AtomsProto.Atom.VM_EXITED_FIELD_NUMBER);
         AtomsProto.VmExited atomVmExited = data.get(0).getAtom().getVmExited();
-        assertEquals("VmRunApp", atomVmExited.getVmIdentifier());
-        assertEquals(AtomsProto.VmExited.DeathReason.KILLED, atomVmExited.getDeathReason());
+        assertThat(atomVmExited.getVmIdentifier()).isEqualTo("VmRunApp");
+        assertThat(atomVmExited.getDeathReason()).isEqualTo(AtomsProto.VmExited.DeathReason.KILLED);
 
         // Check UID and elapsed_time by comparing each other.
-        assertEquals(atomVmCreationRequested.getUid(), atomVmBooted.getUid());
-        assertEquals(atomVmCreationRequested.getUid(), atomVmExited.getUid());
-        assertTrue(atomVmBooted.getElapsedTimeMillis() < atomVmExited.getElapsedTimeMillis());
+        assertThat(atomVmBooted.getUid()).isEqualTo(atomVmCreationRequested.getUid());
+        assertThat(atomVmExited.getUid()).isEqualTo(atomVmCreationRequested.getUid());
+        assertThat(atomVmBooted.getElapsedTimeMillis())
+                .isLessThan(atomVmExited.getElapsedTimeMillis());
     }
 
     @Test
@@ -622,35 +583,33 @@
         waitForBootComplete();
         // Test writing to /data partition
         runOnMicrodroid("echo MicrodroidTest > /data/local/tmp/test.txt");
-        assertThat(runOnMicrodroid("cat /data/local/tmp/test.txt"), is("MicrodroidTest"));
+        assertThat(runOnMicrodroid("cat /data/local/tmp/test.txt")).isEqualTo("MicrodroidTest");
 
         // Check if the APK & its idsig partitions exist
         final String apkPartition = "/dev/block/by-name/microdroid-apk";
-        assertThat(runOnMicrodroid("ls", apkPartition), is(apkPartition));
+        assertThat(runOnMicrodroid("ls", apkPartition)).isEqualTo(apkPartition);
         final String apkIdsigPartition = "/dev/block/by-name/microdroid-apk-idsig";
-        assertThat(runOnMicrodroid("ls", apkIdsigPartition), is(apkIdsigPartition));
+        assertThat(runOnMicrodroid("ls", apkIdsigPartition)).isEqualTo(apkIdsigPartition);
         // Check the vm-instance partition as well
         final String vmInstancePartition = "/dev/block/by-name/vm-instance";
-        assertThat(runOnMicrodroid("ls", vmInstancePartition), is(vmInstancePartition));
+        assertThat(runOnMicrodroid("ls", vmInstancePartition)).isEqualTo(vmInstancePartition);
 
         // Check if the native library in the APK is has correct filesystem info
         final String[] abis = runOnMicrodroid("getprop", "ro.product.cpu.abilist").split(",");
-        assertThat(abis.length, is(1));
+        assertThat(abis).hasLength(1);
         final String testLib = "/mnt/apk/lib/" + abis[0] + "/MicrodroidTestNativeLib.so";
         final String label = "u:object_r:system_file:s0";
-        assertThat(runOnMicrodroid("ls", "-Z", testLib), is(label + " " + testLib));
+        assertThat(runOnMicrodroid("ls", "-Z", testLib)).isEqualTo(label + " " + testLib);
 
         // Check that no denials have happened so far
         CommandRunner android = new CommandRunner(getDevice());
-        assertThat(
-                android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", LOG_PATH), is(nullValue()));
+        assertThat(android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", LOG_PATH)).isNull();
 
-        assertThat(
-                runOnMicrodroid("cat /proc/cpuinfo | grep processor | wc -l"),
-                is(Integer.toString(NUM_VCPUS)));
+        assertThat(runOnMicrodroid("cat /proc/cpuinfo | grep processor | wc -l"))
+                .isEqualTo(Integer.toString(NUM_VCPUS));
 
         // Check that selinux is enabled
-        assertThat(runOnMicrodroid("getenforce"), is("Enforcing"));
+        assertThat(runOnMicrodroid("getenforce")).isEqualTo("Enforcing");
 
         // TODO(b/176805428): adb is broken for nested VM
         if (!isCuttlefish()) {
@@ -722,13 +681,14 @@
     @Test
     public void testCustomVirtualMachinePermission()
             throws DeviceNotAvailableException, IOException, JSONException {
-        assumeTrue(isProtectedVmSupported());
+        assumeTrue("Protected VMs are not supported", isProtectedVmSupported());
         CommandRunner android = new CommandRunner(getDevice());
 
         // Pull etc/microdroid.json
         File virtApexDir = FileUtil.createTempDir("virt_apex");
         File microdroidConfigFile = new File(virtApexDir, "microdroid.json");
-        assertTrue(getDevice().pullFile(VIRT_APEX + "etc/microdroid.json", microdroidConfigFile));
+        assertThat(getDevice().pullFile(VIRT_APEX + "etc/microdroid.json", microdroidConfigFile))
+                .isTrue();
         JSONObject config = new JSONObject(FileUtil.readStringFromFile(microdroidConfigFile));
 
         // USE_CUSTOM_VIRTUAL_MACHINE is enforced only on protected mode
diff --git a/tests/testapk/AndroidManifest.xml b/tests/testapk/AndroidManifest.xml
index 9c8b2d5..ab22546 100644
--- a/tests/testapk/AndroidManifest.xml
+++ b/tests/testapk/AndroidManifest.xml
@@ -16,6 +16,7 @@
 <manifest xmlns:android="http://schemas.android.com/apk/res/android"
       package="com.android.microdroid.test">
     <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
+    <uses-permission android:name="android.permission.USE_CUSTOM_VIRTUAL_MACHINE" />
     <uses-sdk android:minSdkVersion="33" android:targetSdkVersion="33" />
     <application>
         <uses-library android:name="android.system.virtualmachine" android:required="false" />
diff --git a/tests/testapk/AndroidTest.xml b/tests/testapk/AndroidTest.xml
index e8bb1aa..787ebd4 100644
--- a/tests/testapk/AndroidTest.xml
+++ b/tests/testapk/AndroidTest.xml
@@ -13,7 +13,7 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<configuration description="Runs sample instrumentation test.">
+<configuration description="Runs Microdroid device-side tests.">
     <option name="test-suite-tag" value="cts" />
     <option name="config-descriptor:metadata" key="component" value="security" />
     <option name="config-descriptor:metadata" key="parameter" value="not_instant_app" />
@@ -22,11 +22,6 @@
     <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
         <option name="test-file-name" value="MicrodroidTestApp.apk" />
     </target_preparer>
-    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
-      <option
-        name="run-command"
-        value="pm grant com.android.microdroid.test android.permission.MANAGE_VIRTUAL_MACHINE" />
-    </target_preparer>
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.microdroid.test" />
         <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 466acc7..297341b 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -21,6 +21,8 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.TruthJUnit.assume;
 
+import static org.junit.Assert.assertThrows;
+
 import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 
 import android.os.Build;
@@ -35,6 +37,7 @@
 import com.android.microdroid.test.device.MicrodroidDeviceTestBase;
 import com.android.microdroid.testservice.ITestService;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Rule;
@@ -76,9 +79,16 @@
 
     @Before
     public void setup() {
+        grantPermission(VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION);
         prepareTestSetup(mProtectedVm);
     }
 
+    @After
+    public void tearDown() {
+        revokePermission(VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION);
+        revokePermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
+    }
+
     private static final int MIN_MEM_ARM64 = 150;
     private static final int MIN_MEM_X86_64 = 196;
 
@@ -106,12 +116,59 @@
     @Test
     @CddTest(requirements = {
             "9.17/C-1-1",
+            "9.17/C-1-2",
+            "9.17/C-1-4",
+    })
+    public void createVmRequiresPermission() throws Exception {
+        assumeSupportedKernel();
+
+        revokePermission(VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION);
+
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setMemoryMib(minMemoryRequired())
+                .build();
+
+        SecurityException e = assertThrows(SecurityException.class,
+                () -> mInner.forceCreateNewVirtualMachine("test_vm_requires_permission", config));
+        assertThat(e).hasMessageThat()
+                .contains("android.permission.MANAGE_VIRTUAL_MACHINE permission");
+    }
+
+    @Test
+    @CddTest(requirements = {
+            "9.17/C-1-1",
+            "9.17/C-1-2",
+            "9.17/C-1-4",
+    })
+    public void createVmWithConfigRequiresPermission() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadConfigPath("assets/vm_config.json")
+                .setMemoryMib(minMemoryRequired())
+                .build();
+
+        VirtualMachine vm = mInner.forceCreateNewVirtualMachine(
+                "test_vm_config_requires_permission", config);
+
+        SecurityException e = assertThrows(SecurityException.class, () -> runVmTestService(vm));
+        assertThat(e).hasMessageThat()
+                .contains("android.permission.USE_CUSTOM_VIRTUAL_MACHINE permission");
+    }
+
+
+    @Test
+    @CddTest(requirements = {
+            "9.17/C-1-1",
             "9.17/C-2-1"
     })
     public void extraApk() throws Exception {
         assumeSupportedKernel();
 
-        VirtualMachineConfig config = mInner.newVmConfigBuilder("assets/vm_config_extra_apk.json")
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadConfigPath("assets/vm_config_extra_apk.json")
                 .setMemoryMib(minMemoryRequired())
                 .build();
         VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm_extra_apk", config);
@@ -406,9 +463,11 @@
 
     @Test
     public void bootFailsWhenConfigIsInvalid() throws Exception {
-        VirtualMachineConfig.Builder builder =
-                mInner.newVmConfigBuilder("assets/vm_config_no_task.json");
-        VirtualMachineConfig normalConfig = builder.setDebugLevel(DEBUG_LEVEL_FULL).build();
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
+        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
+                .setPayloadConfigPath("assets/vm_config_no_task.json")
+                .setDebugLevel(DEBUG_LEVEL_FULL)
+                .build();
         mInner.forceCreateNewVirtualMachine("test_vm_invalid_config", normalConfig);
 
         BootResult bootResult = tryBootVm(TAG, "test_vm_invalid_config");
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index d810e3c..5d6ca8b 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -122,8 +122,8 @@
             abort();
         }
     };
-    if (!RunRpcServerCallback(testService->asBinder().get(), testService->SERVICE_PORT, callback,
-                              nullptr)) {
+    if (!RunVsockRpcServerCallback(testService->asBinder().get(), testService->SERVICE_PORT,
+                                   callback, nullptr)) {
         return Error() << "RPC Server failed to run";
     }
 
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 532ddc7..4e3bed2 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -57,7 +57,7 @@
 use disk::QcowFile;
 use log::{debug, error, info, warn};
 use microdroid_payload_config::{OsConfig, Task, TaskType, VmPayloadConfig};
-use rpcbinder::run_rpc_server_with_factory;
+use rpcbinder::run_vsock_rpc_server_with_factory;
 use rustutils::system_properties;
 use semver::VersionReq;
 use std::convert::TryInto;
@@ -320,7 +320,7 @@
         let state = service.state.clone();
         std::thread::spawn(move || {
             debug!("VirtualMachineService is starting as an RPC service.");
-            if run_rpc_server_with_factory(VM_BINDER_SERVICE_PORT as u32, |cid| {
+            if run_vsock_rpc_server_with_factory(VM_BINDER_SERVICE_PORT as u32, |cid| {
                 VirtualMachineService::factory(cid, &state)
             }) {
                 debug!("RPC server has shut down gracefully");
@@ -340,10 +340,18 @@
     ) -> binder::Result<Strong<dyn IVirtualMachine>> {
         check_manage_access()?;
 
-        if let VirtualMachineConfig::RawConfig(config) = config {
-            if config.protectedVm {
-                check_use_custom_virtual_machine()?;
+        let is_custom = match config {
+            VirtualMachineConfig::RawConfig(_) => true,
+            VirtualMachineConfig::AppConfig(config) => {
+                // Some features are reserved for platform apps only, even when using
+                // VirtualMachineAppConfig:
+                // - controlling CPUs;
+                // - specifying a config file in the APK.
+                !config.taskProfiles.is_empty() || matches!(config.payload, Payload::ConfigPath(_))
             }
+        };
+        if is_custom {
+            check_use_custom_virtual_machine()?;
         }
 
         let state = &mut *self.state.lock().unwrap();
@@ -377,18 +385,17 @@
             )
         })?;
 
-        let is_app_config = matches!(config, VirtualMachineConfig::AppConfig(_));
-
-        let config = match config {
-            VirtualMachineConfig::AppConfig(config) => BorrowedOrOwned::Owned(
-                load_app_config(config, &temporary_directory).map_err(|e| {
+        let (is_app_config, config) = match config {
+            VirtualMachineConfig::RawConfig(config) => (false, BorrowedOrOwned::Borrowed(config)),
+            VirtualMachineConfig::AppConfig(config) => {
+                let config = load_app_config(config, &temporary_directory).map_err(|e| {
                     *is_protected = config.protectedVm;
                     let message = format!("Failed to load app config: {:?}", e);
                     error!("{}", message);
                     Status::new_service_specific_error_str(-1, Some(message))
-                })?,
-            ),
-            VirtualMachineConfig::RawConfig(config) => BorrowedOrOwned::Borrowed(config),
+                })?;
+                (true, BorrowedOrOwned::Owned(config))
+            }
         };
         let config = config.as_ref();
         *is_protected = config.protectedVm;
@@ -590,12 +597,6 @@
     config: &VirtualMachineAppConfig,
     temporary_directory: &Path,
 ) -> Result<VirtualMachineRawConfig> {
-    // Controlling CPUs is reserved for platform apps only, even when using
-    // VirtualMachineAppConfig.
-    if !config.taskProfiles.is_empty() {
-        check_use_custom_virtual_machine()?
-    }
-
     let apk_file = clone_file(config.apk.as_ref().unwrap())?;
     let idsig_file = clone_file(config.idsig.as_ref().unwrap())?;
     let instance_file = clone_file(config.instanceImage.as_ref().unwrap())?;
diff --git a/vm/vm_shell.sh b/vm/vm_shell.sh
index ec9243b..c0dd38f 100755
--- a/vm/vm_shell.sh
+++ b/vm/vm_shell.sh
@@ -27,7 +27,7 @@
     adb forward tcp:8000 vsock:${cid}:5555
     adb connect localhost:8000
     adb -s localhost:8000 root
-    sleep 2
+    adb -s localhost:8000 wait-for-device
     adb -s localhost:8000 shell
     exit 0
 }
@@ -40,12 +40,7 @@
     exit 1
 fi
 
-if [ -n "${selected_cid}" ]; then
-    if [[ ! " ${available_cids[*]} " =~ " ${selected_cid} " ]]; then
-        echo VM of CID $selected_cid does not exist. Available CIDs: ${available_cids}
-        exit 1
-    fi
-else
+if [ ! -n "${selected_cid}" ]; then
     PS3="Select CID of VM to adb-shell into: "
     select cid in ${available_cids}
     do
@@ -54,4 +49,9 @@
     done
 fi
 
+if [[ ! " ${available_cids[*]} " =~ " ${selected_cid} " ]]; then
+    echo VM of CID $selected_cid does not exist. Available CIDs: ${available_cids}
+    exit 1
+fi
+
 connect_vm ${selected_cid}