Merge "vmbase: Replace CRYPTO_sysrand() with getentropy()" into main
diff --git a/TEST_MAPPING b/TEST_MAPPING
index e2cf9a5..323b827 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -6,6 +6,9 @@
       "name": "MicrodroidHostTestCases"
     },
     {
+      "name": "ComposHostTestCases"
+    },
+    {
       "name": "MicrodroidTestApp"
     },
     {
@@ -35,9 +38,6 @@
       "name": "ComposBenchmarkApp"
     },
     {
-      "name": "ComposHostTestCases"
-    },
-    {
       "name": "AVFHostTestCases"
     }
   ],
diff --git a/apex/Android.bp b/apex/Android.bp
index fedcfcd..7ef2c79 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -106,6 +106,7 @@
     ],
     host_required: [
         "vm_shell",
+        "prepare_device_vfio",
     ],
     apps: [
         "EmptyPayloadApp",
@@ -136,6 +137,12 @@
     installable: false,
 }
 
+sh_binary_host {
+    name: "prepare_device_vfio",
+    src: "prepare_device_vfio.sh",
+    filename: "prepare_device_vfio.sh",
+}
+
 // Virt apex needs a custom signer for its payload
 python_binary_host {
     name: "sign_virt_apex",
diff --git a/apex/prepare_device_vfio.sh b/apex/prepare_device_vfio.sh
new file mode 100755
index 0000000..de2d502
--- /dev/null
+++ b/apex/prepare_device_vfio.sh
@@ -0,0 +1,176 @@
+#!/bin/bash
+
+# Copyright 2023 Google Inc. All rights reserved.
+#
+# 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.
+
+# prepare_device_vfio.sh: prepares a device for VFIO assignment by binding a VFIO driver to it
+
+adb="${ADB:="adb"}" # ADB command to use
+vfio_dir="/dev/vfio"
+platform_bus="/sys/bus/platform"
+vfio_reset_required="/sys/module/vfio_platform/parameters/reset_required"
+vfio_noiommu_param="/sys/module/vfio/parameters/enable_unsafe_noiommu_mode"
+vfio_unsafe_interrupts_param="/sys/module/vfio_iommu_type1/parameters/allow_unsafe_interrupts"
+
+function print_help() {
+    echo "prepare_device_vfio.sh prepares a device for VFIO assignment"
+    echo ""
+    echo " Usage:"
+    echo "    $0 DEVICE_NAME"
+    echo "      Prepare device DEVICE_NAME for VFIO assignment."
+    echo ""
+    echo "    help - prints this help message"
+}
+
+function cmd() {
+    $adb shell $@
+}
+
+function tcmd() {
+    trap "echo \"Error: adb shell command '$@' failed\" ; exit 1" ERR
+    $adb shell $@
+}
+
+function ensure_root() {
+    # Check user id
+    if [ $(cmd "id -u") != 0 ]; then
+        read -p "Must run as root; restart ADBD? [y/n] " answer
+        case $answer in
+            [Yy]* )
+                $adb root && $adb wait-for-device && sleep 3 || exit 1
+                ;;
+            * )
+                exit 1
+        esac
+    fi
+}
+
+function check_vfio() {
+    cmd "[ -c $vfio_dir/vfio ]"
+    if [ $? -ne 0 ]; then
+        echo "cannot find $vfio_dir/vfio"
+        exit 1
+    fi
+
+    cmd "[ -d $platform_bus/drivers/vfio-platform ]"
+    if [ $? -ne 0 ]; then
+        echo "VFIO-platform is not supported"
+        exit 1
+    fi
+}
+
+function check_device() {
+    cmd "[ -d $device_sys ]"
+    if [ $? -ne 0 ]; then
+        echo "no device $device ($device_sys)"
+        exit 1
+    fi
+}
+
+function get_device_iommu_group() {
+    local group=$(cmd "basename \$(readlink \"$device_sys/iommu_group\")")
+    if [ $? -eq 0 ]; then
+        echo $group
+    else
+        echo ""
+    fi
+}
+
+function misc_setup() {
+    # VFIO NOIOMMU check
+    if [ -z "$group" ]; then
+        echo "$device_sys does not have an IOMMU group - setting $vfio_noiommu_param"
+        tcmd "echo y > \"$vfio_noiommu_param\""
+    fi
+
+    # Disable SELinux to allow virtualizationmanager and crosvm to access sysfs
+    echo "[*WARN*] setenforce=0: SELinux is disabled"
+    tcmd "setenforce 0"
+
+    # Samsung IOMMU does not report interrupt remapping support, so enable unsafe uinterrupts
+    if [ -n "$group" ]; then
+        local iommu_drv=$(cmd "basename \$(readlink \"$device_sys/iommu/device/driver\")")
+        if [ "$iommu_drv" = "samsung-sysmmu-v9" ]; then
+            tcmd "echo y > \"$vfio_unsafe_interrupts_param\""
+        fi
+    fi
+}
+
+function bind_vfio_driver() {
+    # Check if non-VFIO driver is currently bound, ie unbinding is needed
+    cmd "[ -e \"$device_driver\" ] && \
+        [ ! \$(basename \$(readlink \"$device_driver\")) = \"vfio-platform\" ]"
+            if [ $? -eq 0 ]; then
+                # Unbind current driver
+                tcmd "echo \"$device\" > \"$device_driver/unbind\""
+            fi
+
+    # Bind to VFIO driver
+    cmd "[ ! -e \"$device_driver\" ]"
+    if [ $? -eq 0 ]; then
+        # Bind vfio-platform driver
+        tcmd "echo \"vfio-platform\" > \"$device_sys/driver_override\""
+        tcmd "echo \"$device\" > \"$platform_bus/drivers_probe\""
+        sleep 2
+    fi
+}
+
+function verify_vfio_driver() {
+    # Verify new VFIO file structure
+    group=$(get_device_iommu_group)
+    if [ -z "$group" ]; then
+        echo "cannot setup VFIO-NOIOMMU for $device_sys"
+        exit 1
+    fi
+
+    cmd "[ ! -c \"$vfio_dir/$group\" ] || \
+        [ ! -e \"$device_driver\" ] || \
+        [ ! \$(basename \$(readlink \"$device_driver\")) = \"vfio-platform\" ]"
+    if [ $? -eq 0 ]; then
+        echo "could not bind $device to VFIO platform driver"
+
+        if [ $(cmd "cat $vfio_reset_required") = Y ]; then
+            echo "VFIO device reset handler must be registered. Either unset $vfio_reset_required, \
+or register a reset handler for $device_sys"
+        fi
+        exit 1
+    fi
+}
+
+function prepare_device() {
+    device="$1"
+    device_sys="/sys/bus/platform/devices/$device"
+    device_driver="$device_sys/driver"
+
+    ensure_root
+    check_vfio
+    check_device
+    group=$(get_device_iommu_group)
+    misc_setup
+
+    bind_vfio_driver
+    verify_vfio_driver
+
+    echo "Device: $device_sys"
+    echo "IOMMU group: $group"
+    echo "VFIO group file: $vfio_dir/$group"
+    echo "Ready!"
+}
+
+cmd=$1
+
+case $cmd in
+    ""|help) print_help ;;
+    *) prepare_device "$cmd" $@ ;;
+esac
diff --git a/apkdmverity/Android.bp b/apkdmverity/Android.bp
index fae7e99..8429263 100644
--- a/apkdmverity/Android.bp
+++ b/apkdmverity/Android.bp
@@ -40,7 +40,7 @@
     name: "apkdmverity.test",
     defaults: [
         "apkdmverity.defaults",
-        "ignorabletest.defaults",
+        "rdroidtest.defaults",
     ],
     test_suites: ["general-tests"],
     compile_multilib: "first",
diff --git a/apkdmverity/src/main.rs b/apkdmverity/src/main.rs
index 55953a9..d9e9e2b 100644
--- a/apkdmverity/src/main.rs
+++ b/apkdmverity/src/main.rs
@@ -153,12 +153,12 @@
 }
 
 #[cfg(test)]
-ignorabletest::test_main!();
+rdroidtest::test_main!();
 
 #[cfg(test)]
 mod tests {
     use crate::*;
-    use ignorabletest::test;
+    use rdroidtest::test;
     use std::fs::{File, OpenOptions};
     use std::io::Write;
     use std::ops::Deref;
diff --git a/authfs/src/fusefs.rs b/authfs/src/fusefs.rs
index 87bdffc..64b340a 100644
--- a/authfs/src/fusefs.rs
+++ b/authfs/src/fusefs.rs
@@ -386,7 +386,8 @@
 }
 
 cfg_if::cfg_if! {
-    if #[cfg(all(target_arch = "aarch64", target_pointer_width = "64"))] {
+    if #[cfg(all(any(target_arch = "aarch64", target_arch = "riscv64"),
+                 target_pointer_width = "64"))] {
         fn blk_size() -> libc::c_int { CHUNK_SIZE as libc::c_int }
     } else {
         fn blk_size() -> libc::c_long { CHUNK_SIZE as libc::c_long }
@@ -992,8 +993,8 @@
     fn statfs(&self, _ctx: Context, _inode: Self::Inode) -> io::Result<libc::statvfs64> {
         let remote_stat = self.remote_fs_stats_reader.statfs()?;
 
-        // Safe because we are zero-initializing a struct with only POD fields. Not all fields
-        // matter to FUSE. See also:
+        // SAFETY: We are zero-initializing a struct with only POD fields. Not all fields matter to
+        // FUSE. See also:
         // https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/fs/fuse/inode.c?h=v5.15#n460
         let mut st: libc::statvfs64 = unsafe { zeroed() };
 
diff --git a/authfs/tests/benchmarks/Android.bp b/authfs/tests/benchmarks/Android.bp
index 110d000..cea5a81 100644
--- a/authfs/tests/benchmarks/Android.bp
+++ b/authfs/tests/benchmarks/Android.bp
@@ -17,12 +17,10 @@
     test_suites: ["general-tests"],
     data_device_bins_first: [
         "open_then_run",
-        "fsverity",
     ],
     per_testcase_directory: true,
     data: [
         ":authfs_test_files",
-        ":CtsApkVerityTestPrebuiltFiles",
         ":MicrodroidTestApp",
     ],
     required: ["MicrodroidTestPreparer"],
diff --git a/authfs/tests/benchmarks/AndroidTest.xml b/authfs/tests/benchmarks/AndroidTest.xml
index 9216006..715f352 100644
--- a/authfs/tests/benchmarks/AndroidTest.xml
+++ b/authfs/tests/benchmarks/AndroidTest.xml
@@ -34,32 +34,16 @@
 
         <!-- 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.microdroid.test.preparer.DisableMicrodroidDebugPolicyPreparer" />
 
-    <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>
diff --git a/authfs/tests/hosttests/Android.bp b/authfs/tests/hosttests/Android.bp
index 4b8151d..83ef853 100644
--- a/authfs/tests/hosttests/Android.bp
+++ b/authfs/tests/hosttests/Android.bp
@@ -20,7 +20,6 @@
     per_testcase_directory: true,
     data: [
         ":authfs_test_files",
-        ":CtsApkVerityTestPrebuiltFiles",
         ":MicrodroidTestApp",
     ],
 }
diff --git a/authfs/tests/hosttests/AndroidTest.xml b/authfs/tests/hosttests/AndroidTest.xml
index 2ccc45f..5920630 100644
--- a/authfs/tests/hosttests/AndroidTest.xml
+++ b/authfs/tests/hosttests/AndroidTest.xml
@@ -50,18 +50,13 @@
         <option name="push-file" key="input.4m.fsv_meta.bad_merkle"
             value="/data/local/tmp/authfs/input.4m.fsv_meta.bad_merkle" />
 
-        <!-- 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" />
+        <option name="push-file" key="input.4m" value="/data/local/tmp/authfs/input.file" />
     </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" />
+        <!-- Now that the files are pushed to the device, enable fs-verity for the targeting file. -->
+        <option name="run-command" value="cd /data/local/tmp/authfs; ../fsverity enable input.file" />
     </target_preparer>
 
     <test class="com.android.compatibility.common.tradefed.testtype.JarHostTest" >
diff --git a/authfs/tests/hosttests/java/src/com/android/fs/AuthFsHostTest.java b/authfs/tests/hosttests/java/src/com/android/fs/AuthFsHostTest.java
index 440f5ca..d0a7c66 100644
--- a/authfs/tests/hosttests/java/src/com/android/fs/AuthFsHostTest.java
+++ b/authfs/tests/hosttests/java/src/com/android/fs/AuthFsHostTest.java
@@ -145,17 +145,17 @@
 
     @Test
     public void testReadWithFsverityVerification_FdServerUsesRealFsverityData() throws Exception {
-        // Setup (fs-verity is enabled for input.apk in AndroidTest.xml)
-        runFdServerOnAndroid("--open-ro 3:input.apk", "--ro-fds 3");
-        String expectedDigest = sAndroid.run(
-                FSVERITY_BIN + " digest --compact " + TEST_DIR + "/input.apk");
+        // Setup (fs-verity is enabled for input.file in AndroidTest.xml)
+        runFdServerOnAndroid("--open-ro 3:input.file", "--ro-fds 3");
+        String expectedDigest =
+                sAndroid.run(FSVERITY_BIN + " digest --compact " + TEST_DIR + "/input.file");
         runAuthFsOnMicrodroid("--remote-ro-file 3:sha256-" + expectedDigest);
 
         // Action
         String actualHash = computeFileHash(sMicrodroid, MOUNT_DIR + "/3");
 
         // Verify
-        String expectedHash = computeFileHash(sAndroid, TEST_DIR + "/input.apk");
+        String expectedHash = computeFileHash(sAndroid, TEST_DIR + "/input.file");
         assertEquals("Inconsistent hash from /authfs/3: ", expectedHash, actualHash);
     }
 
diff --git a/compos/composd/src/fd_server_helper.rs b/compos/composd/src/fd_server_helper.rs
index 777ec27..24371b5 100644
--- a/compos/composd/src/fd_server_helper.rs
+++ b/compos/composd/src/fd_server_helper.rs
@@ -107,8 +107,9 @@
 
 fn create_pipe() -> Result<(File, File)> {
     let (raw_read, raw_write) = pipe2(OFlag::O_CLOEXEC)?;
-    // SAFETY: We are the sole owners of these fds as they were just created.
+    // SAFETY: We are the sole owner of raw_read and it is valid as it was just created.
     let read_fd = unsafe { File::from_raw_fd(raw_read) };
+    // SAFETY: We are the sole owner of raw_write and it is valid as it was just created.
     let write_fd = unsafe { File::from_raw_fd(raw_write) };
     Ok((read_fd, write_fd))
 }
diff --git a/compos/src/compsvc_main.rs b/compos/src/compsvc_main.rs
index 77e2daa..b0fc323 100644
--- a/compos/src/compsvc_main.rs
+++ b/compos/src/compsvc_main.rs
@@ -50,11 +50,11 @@
     debug!("compsvc is starting as a rpc service.");
     let param = ptr::null_mut();
     let mut service = compsvc::new_binder()?.as_binder();
+    let service = service.as_native_mut() as *mut AIBinder;
+    // SAFETY: We hold a strong pointer, so the raw pointer remains valid. The bindgen AIBinder
+    // is the same type as sys::AIBinder. It is safe for on_ready to be invoked at any time, with
+    // any parameter.
     unsafe {
-        // SAFETY: We hold a strong pointer, so the raw pointer remains valid. The bindgen AIBinder
-        // is the same type as sys::AIBinder.
-        let service = service.as_native_mut() as *mut AIBinder;
-        // SAFETY: It is safe for on_ready to be invoked at any time, with any parameter.
         AVmPayload_runVsockRpcServer(service, COMPOS_VSOCK_PORT, Some(on_ready), param);
     }
     Ok(())
diff --git a/compos/tests/AndroidTest.xml b/compos/tests/AndroidTest.xml
index 4b414f1..e35b874 100644
--- a/compos/tests/AndroidTest.xml
+++ b/compos/tests/AndroidTest.xml
@@ -15,6 +15,7 @@
 -->
 <configuration description="Tests for CompOS">
     <option name="config-descriptor:metadata" key="mainline-param" value="com.google.android.art.apex" />
+    <option name="config-descriptor:metadata" key="mainline-param" value="com.android.art.apex" />
 
     <!-- Only run tests if the device under test is SDK version 33 (Android 13) or above. -->
     <object type="module_controller"
diff --git a/compos/tests/java/android/compos/test/ComposTestCase.java b/compos/tests/java/android/compos/test/ComposTestCase.java
index 1cebd1a..244d34e 100644
--- a/compos/tests/java/android/compos/test/ComposTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposTestCase.java
@@ -23,6 +23,8 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
+import static org.junit.Assume.assumeFalse;
+
 import android.platform.test.annotations.RootPermissionTest;
 
 import com.android.microdroid.test.host.CommandRunner;
@@ -81,6 +83,8 @@
     @Before
     public void setUp() throws Exception {
         assumeDeviceIsCapable(getDevice());
+        // Test takes too long to run on Cuttlefish (b/292824951).
+        assumeFalse("Skipping test on Cuttlefish", isCuttlefish());
 
         String value = getDevice().getProperty(SYSTEM_SERVER_COMPILER_FILTER_PROP_NAME);
         if (value == null) {
diff --git a/docs/debug/README.md b/docs/debug/README.md
index b7729a4..1e5f096 100644
--- a/docs/debug/README.md
+++ b/docs/debug/README.md
@@ -74,12 +74,14 @@
 
         __overlay__ {
             avf {
-                common {
-                    log = <1>; // Enable kernel log and logcat
-                    ramdump = <1>; // Enable ramdump
-                }
-                microdroid {
-                    adb = <1>; // Enable ADB connection
+                guest {
+                    common {
+                        log = <1>; // Enable kernel log and logcat
+                        ramdump = <1>; // Enable ramdump
+                    }
+                    microdroid {
+                        adb = <1>; // Enable ADB connection
+                    }
                 }
             };
         };
diff --git a/encryptedstore/src/main.rs b/encryptedstore/src/main.rs
index 86fa6da..1a16f49 100644
--- a/encryptedstore/src/main.rs
+++ b/encryptedstore/src/main.rs
@@ -153,6 +153,9 @@
     let mountpoint = CString::new(mountpoint.as_os_str().as_bytes())?;
     let fstype = CString::new("ext4").unwrap();
 
+    // SAFETY: The source, target and filesystemtype are valid C strings. For ext4, data is expected
+    // to be a C string as well, which it is. None of these pointers are retained after mount
+    // returns.
     let ret = unsafe {
         libc::mount(
             source.as_ptr(),
diff --git a/libs/capabilities/src/caps.rs b/libs/capabilities/src/caps.rs
index 1f44a34..bc17fa8 100644
--- a/libs/capabilities/src/caps.rs
+++ b/libs/capabilities/src/caps.rs
@@ -26,8 +26,8 @@
 /// Removes inheritable capabilities set for this process.
 /// See: https://man7.org/linux/man-pages/man7/capabilities.7.html
 pub fn drop_inheritable_caps() -> Result<()> {
+    // SAFETY: we do not manipulate memory handled by libcap.
     unsafe {
-        // SAFETY: we do not manipulate memory handled by libcap.
         let caps = cap_get_proc();
         scopeguard::defer! {
             cap_free(caps as *mut std::os::raw::c_void);
@@ -49,8 +49,8 @@
 pub fn drop_bounding_set() -> Result<()> {
     let mut cap_id: cap_value_t = 0;
     while cap_id <= CAP_LAST_CAP.try_into().unwrap() {
+        // SAFETY: we do not manipulate memory handled by libcap.
         unsafe {
-            // SAFETY: we do not manipulate memory handled by libcap.
             if cap_drop_bound(cap_id) == -1 {
                 let e = Errno::last();
                 bail!("cap_drop_bound failed for {}: {:?}", cap_id, e);
diff --git a/libs/devicemapper/Android.bp b/libs/devicemapper/Android.bp
index b7cdedc..29f2f5f 100644
--- a/libs/devicemapper/Android.bp
+++ b/libs/devicemapper/Android.bp
@@ -33,7 +33,7 @@
     name: "libdm_rust.test",
     defaults: [
         "libdm_rust.defaults",
-        "ignorabletest.defaults",
+        "rdroidtest.defaults",
     ],
     test_suites: ["general-tests"],
     rustlibs: [
diff --git a/libs/devicemapper/src/lib.rs b/libs/devicemapper/src/lib.rs
index 0170795..868ac5a 100644
--- a/libs/devicemapper/src/lib.rs
+++ b/libs/devicemapper/src/lib.rs
@@ -233,13 +233,13 @@
 }
 
 #[cfg(test)]
-ignorabletest::test_main!();
+rdroidtest::test_main!();
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crypt::{CipherType, DmCryptTargetBuilder};
-    use ignorabletest::test;
+    use rdroidtest::test;
     use rustutils::system_properties;
     use std::fs::{read, File, OpenOptions};
     use std::io::Write;
diff --git a/libs/ignorabletest/Android.bp b/libs/ignorabletest/Android.bp
deleted file mode 100644
index 4ae89af..0000000
--- a/libs/ignorabletest/Android.bp
+++ /dev/null
@@ -1,36 +0,0 @@
-rust_library {
-    name: "libignorabletest",
-    host_supported: true,
-    crate_name: "ignorabletest",
-    cargo_env_compat: true,
-    cargo_pkg_version: "0.1.0",
-    srcs: ["src/lib.rs"],
-    edition: "2021",
-    rustlibs: [
-        "liblibtest_mimic",
-        "liblinkme",
-        "liblog_rust",
-        "liblogger",
-    ],
-    proc_macros: ["libpaste"],
-    apex_available: [
-        "//apex_available:platform",
-        "//apex_available:anyapex",
-    ],
-}
-
-rust_defaults {
-    name: "ignorabletest.defaults",
-    test_harness: false,
-    cfgs: ["test"],
-    rustlibs: [
-        "libignorabletest",
-        "liblinkme",
-    ],
-    // Without this flag we get linker errors saying to add it. See
-    // https://github.com/dtolnay/linkme/issues/49 and related issues.
-    ld_flags: [
-        "-z",
-        "nostart-stop-gc",
-    ],
-}
diff --git a/libs/ignorabletest/README.md b/libs/ignorabletest/README.md
deleted file mode 100644
index 77140bd..0000000
--- a/libs/ignorabletest/README.md
+++ /dev/null
@@ -1,66 +0,0 @@
-# ignorabletest
-
-This is a custom Rust test harness which allows tests to be ignored at runtime based on arbitrary
-criteria. The built-in Rust test harness only allows tests to be ignored at compile time, but this
-is often not enough on Android, where we want to ignore tests based on system properties or other
-characteristics of the device on which the test is being run, which are not known at build time.
-
-## Usage
-
-Unfortunately without the built-in support that rustc provides to the standard test harness, this
-one is slightly more cumbersome to use. Firstly, add it to the `rust_test` build rule in your
-`Android.bp` by adding the defaults provided:
-
-```soong
-rust_test {
-    name: "mycrate.test",
-    defaults: ["ignorabletest.defaults"],
-    // ...
-}
-```
-
-If you are testing a binary that has a `main` function, you'll need to remove it from the test
-build:
-
-```rust
-#[cfg(not(test))]
-fn main() {
-    // ...
-}
-```
-
-(If you're testing a library or anything else which doesn't have a `main` function, you don't need
-to worry about this.)
-
-Each test case should be marked with the `ignorabletest::test!` macro, rather than the standard
-`#[test]` attribute:
-
-```rust
-use ignorabletest::test;
-
-test!(one_plus_one);
-fn one_plus_one {
-    assert_eq!(1 + 1, 2);
-}
-```
-
-To ignore a test, you can add an `ignore_if` clause with a boolean expression:
-
-```rust
-use ignorabletest::test;
-
-test!(clap_hands, ignore_if: !feeling_happy());
-fn clap_hands {
-    assert!(HANDS.clap().is_ok());
-}
-```
-
-Somewhere in your main module, you need to use the `test_main` macro to generate an entry point for
-the test harness:
-
-```rust
-#[cfg(test)]
-ignorabletest::test_main!();
-```
-
-You can then run your tests as usual with `atest`.
diff --git a/libs/ignorabletest/src/lib.rs b/libs/ignorabletest/src/lib.rs
deleted file mode 100644
index c7243e6..0000000
--- a/libs/ignorabletest/src/lib.rs
+++ /dev/null
@@ -1,57 +0,0 @@
-//! Test harness which supports ignoring tests at runtime.
-
-pub mod runner;
-
-#[doc(hidden)]
-pub use libtest_mimic as _libtest_mimic;
-#[doc(hidden)]
-pub use linkme as _linkme;
-#[doc(hidden)]
-pub use paste as _paste;
-
-/// Macro to generate the main function for the test harness.
-#[macro_export]
-macro_rules! test_main {
-    () => {
-        #[cfg(test)]
-        fn main() {
-            ignorabletest::runner::main()
-        }
-    };
-}
-
-/// Macro to generate a wrapper function for a single test.
-///
-/// # Usage
-///
-/// ```
-/// test!(test_string_equality);
-/// fn test_string_equality() {
-///   assert_eq!("", "");
-/// }
-/// ```
-#[macro_export]
-macro_rules! test {
-    ($test_name:ident) => {
-        $crate::_paste::paste!(
-            #[$crate::_linkme::distributed_slice($crate::runner::IGNORABLETEST_TESTS)]
-            fn [< __test_ $test_name >]() -> $crate::_libtest_mimic::Trial {
-                $crate::_libtest_mimic::Trial::test(
-                    ::std::stringify!($test_name),
-                    move || ignorabletest::runner::run($test_name),
-                )
-            }
-        );
-    };
-    ($test_name:ident, ignore_if: $ignore_expr:expr) => {
-        $crate::_paste::paste!(
-            #[$crate::_linkme::distributed_slice($crate::runner::IGNORABLETEST_TESTS)]
-            fn [< __test_ $test_name >]() -> $crate::_libtest_mimic::Trial {
-                $crate::_libtest_mimic::Trial::test(
-                    ::std::stringify!($test_name),
-                    move || ignorabletest::runner::run($test_name),
-                ).with_ignored_flag($ignore_expr)
-            }
-        );
-    };
-}
diff --git a/libs/ignorabletest/src/runner.rs b/libs/ignorabletest/src/runner.rs
deleted file mode 100644
index fdac406..0000000
--- a/libs/ignorabletest/src/runner.rs
+++ /dev/null
@@ -1,29 +0,0 @@
-//! Test runner.
-
-use core::ops::{Deref, FnOnce};
-use libtest_mimic::{Arguments, Failed, Trial};
-use linkme::distributed_slice;
-use log::Level;
-use std::env;
-
-/// Command-line arguments to ignore, because they are not supported by libtest-mimic.
-const IGNORED_ARGS: [&str; 2] = ["-Zunstable-options", "--report-time"];
-
-/// The collection of all tests to run.
-#[doc(hidden)]
-#[distributed_slice]
-pub static IGNORABLETEST_TESTS: [fn() -> Trial] = [..];
-
-/// Runs all tests.
-pub fn main() {
-    logger::init(logger::Config::default().with_min_level(Level::Debug));
-    let args = Arguments::from_iter(env::args().filter(|arg| !IGNORED_ARGS.contains(&arg.deref())));
-    let tests = IGNORABLETEST_TESTS.iter().map(|test| test()).collect();
-    libtest_mimic::run(&args, tests).exit();
-}
-
-/// Runs the given test.
-pub fn run(test: impl FnOnce()) -> Result<(), Failed> {
-    test();
-    Ok(())
-}
diff --git a/microdroid/kdump/Android.bp b/microdroid/kdump/Android.bp
index cc681a7..b9a18fe 100644
--- a/microdroid/kdump/Android.bp
+++ b/microdroid/kdump/Android.bp
@@ -18,6 +18,9 @@
     static_executable: true,
     installable: false,
     compile_multilib: "64",
+    sanitize: {
+        hwaddress: false, // HWASAN setup fails when run as init process
+    },
 }
 
 android_filesystem {
diff --git a/microdroid_manager/src/dice.rs b/microdroid_manager/src/dice.rs
index bacefcd..8e078ea 100644
--- a/microdroid_manager/src/dice.rs
+++ b/microdroid_manager/src/dice.rs
@@ -78,7 +78,7 @@
         let mmap_size =
             file.read_u64::<NativeEndian>()
                 .map_err(|error| Error::new(error).context("Reading driver"))? as usize;
-        // It's safe to map the driver as the service will only create a single
+        // SAFETY: It's safe to map the driver as the service will only create a single
         // mapping per process.
         let mmap_addr = unsafe {
             let fd = file.as_raw_fd();
@@ -87,10 +87,10 @@
         if mmap_addr == MAP_FAILED {
             bail!("Failed to mmap {:?}", driver_path);
         }
-        // The slice is created for the region of memory that was just
+        let mmap_buf =
+        // SAFETY: The slice is created for the region of memory that was just
         // successfully mapped into the process address space so it will be
         // accessible and not referenced from anywhere else.
-        let mmap_buf =
             unsafe { slice::from_raw_parts((mmap_addr as *const u8).as_ref().unwrap(), mmap_size) };
         let bcc_handover =
             bcc_handover_parse(mmap_buf).map_err(|_| anyhow!("Failed to parse Bcc Handover"))?;
@@ -149,9 +149,9 @@
 impl Drop for DiceDriver<'_> {
     fn drop(&mut self) {
         if let &mut Self::Real { mmap_addr, mmap_size, .. } = self {
-            // All references to the mapped region have the same lifetime as self. Since self is
-            // being dropped, so are all the references to the mapped region meaning its safe to
-            // unmap.
+            // SAFETY: All references to the mapped region have the same lifetime as self. Since
+            // self is being dropped, so are all the references to the mapped region meaning it's
+            // safe to unmap.
             let ret = unsafe { munmap(mmap_addr, mmap_size) };
             if ret != 0 {
                 log::warn!("Failed to munmap ({})", ret);
diff --git a/microdroid_manager/src/ioutil.rs b/microdroid_manager/src/ioutil.rs
index d36e349..772941d 100644
--- a/microdroid_manager/src/ioutil.rs
+++ b/microdroid_manager/src/ioutil.rs
@@ -28,6 +28,10 @@
 const SLEEP_DURATION: Duration = Duration::from_millis(5);
 
 /// waits for a file with a timeout and returns it
+///
+/// WARNING: This only guarantees file creation. When there's another thread
+///   writing a file and you're waiting for the file, reading the file should be
+///   synchronized with other mechanism than just waiting for the creation.
 pub fn wait_for_file<P: AsRef<Path> + Debug>(path: P, timeout: Duration) -> Result<File> {
     debug!("waiting for {:?}...", path);
     let begin = Instant::now();
@@ -64,15 +68,21 @@
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::fs::rename;
     use std::io::{Read, Write};
 
     #[test]
     fn test_wait_for_file() -> Result<()> {
         let test_dir = tempfile::TempDir::new().unwrap();
         let test_file = test_dir.path().join("test.txt");
+        let temp_file = test_dir.path().join("test.txt~");
         thread::spawn(move || -> io::Result<()> {
+            // Sleep to ensure that `wait_for_file` actually waits
             thread::sleep(Duration::from_secs(1));
-            File::create(test_file)?.write_all(b"test")
+            // Write to a temp file and then rename it to avoid the race between
+            // write and read.
+            File::create(&temp_file)?.write_all(b"test")?;
+            rename(temp_file, test_file)
         });
 
         let test_file = test_dir.path().join("test.txt");
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 1cdcde1..319d2df 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -193,7 +193,7 @@
 
 /// Prepares a socket file descriptor for the vm payload service.
 ///
-/// # Safety requirement
+/// # Safety
 ///
 /// The caller must ensure that this function is the only place that claims ownership
 /// of the file descriptor and it is called only once.
@@ -267,6 +267,8 @@
     if Path::new(ENCRYPTEDSTORE_BACKING_DEVICE).exists() {
         let mountpoint = CString::new(ENCRYPTEDSTORE_MOUNTPOINT).unwrap();
 
+        // SAFETY: `mountpoint` is a valid C string. `syncfs` and `close` are safe for any parameter
+        // values.
         let ret = unsafe {
             let dirfd = libc::open(
                 mountpoint.as_ptr(),
@@ -856,19 +858,15 @@
         }
     };
 
+    // SAFETY: We are not accessing any resource of the parent process. This means we can't make any
+    // log calls inside the closure.
     unsafe {
-        // SAFETY: we are not accessing any resource of the parent process.
         command.pre_exec(|| {
-            info!("dropping capabilities before executing payload");
             // It is OK to continue with payload execution even if the calls below fail, since
             // whether process can use a capability is controlled by the SELinux. Dropping the
             // capabilities here is just another defense-in-depth layer.
-            if let Err(e) = cap::drop_inheritable_caps() {
-                error!("failed to drop inheritable capabilities: {:?}", e);
-            }
-            if let Err(e) = cap::drop_bounding_set() {
-                error!("failed to drop bounding set: {:?}", e);
-            }
+            let _ = cap::drop_inheritable_caps();
+            let _ = cap::drop_bounding_set();
             Ok(())
         });
     }
diff --git a/microdroid_manager/src/swap.rs b/microdroid_manager/src/swap.rs
index 2f4d176..c2b20ac 100644
--- a/microdroid_manager/src/swap.rs
+++ b/microdroid_manager/src/swap.rs
@@ -48,7 +48,7 @@
         .checked_mul(512)
         .ok_or_else(|| anyhow!("sysfs_size too large"))?;
 
-    // safe because we give a constant and known-valid sysconf parameter
+    // SAFETY: We give a constant and known-valid sysconf parameter.
     let pagesize = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) as u64 };
 
     let mut f = OpenOptions::new().read(false).write(true).open(format!("/dev/{}", dev))?;
@@ -75,7 +75,7 @@
 /// Simple "swapon", using libc:: wrapper.
 fn swapon(dev: &str) -> Result<()> {
     let swapon_arg = std::ffi::CString::new(format!("/dev/{}", dev))?;
-    // safe because we give a nul-terminated string and check the result
+    // SAFETY: We give a nul-terminated string and check the result.
     let res = unsafe { libc::swapon(swapon_arg.as_ptr(), 0) };
     if res != 0 {
         return Err(anyhow!("Failed to swapon: {}", Error::last_os_error()));
diff --git a/pvmfw/avb/tests/api_test.rs b/pvmfw/avb/tests/api_test.rs
index aa9ed36..2f45d77 100644
--- a/pvmfw/avb/tests/api_test.rs
+++ b/pvmfw/avb/tests/api_test.rs
@@ -243,10 +243,15 @@
     let total_len = kernel.len() as u64;
     let footer = extract_avb_footer(&kernel)?;
     assert!(footer.vbmeta_offset < total_len);
+    // TODO: use core::mem::offset_of once stable.
+    let footer_addr = ptr::addr_of!(footer) as *const u8;
     let vbmeta_offset_addr = ptr::addr_of!(footer.vbmeta_offset) as *const u8;
-    // SAFETY: It is safe as both raw pointers `vbmeta_offset_addr` and `footer` are not null.
     let vbmeta_offset_start =
-        unsafe { vbmeta_offset_addr.offset_from(ptr::addr_of!(footer) as *const u8) };
+        // SAFETY:
+        // - both raw pointers `vbmeta_offset_addr` and `footer_addr` are not null;
+        // - they are both derived from the `footer` object;
+        // - the offset is known from the struct definition to be a small positive number of bytes.
+        unsafe { vbmeta_offset_addr.offset_from(footer_addr) };
     let footer_start = kernel.len() - size_of::<AvbFooter>();
     let vbmeta_offset_start = footer_start + usize::try_from(vbmeta_offset_start)?;
 
diff --git a/pvmfw/platform.dts b/pvmfw/platform.dts
index 74439d9..cb8e30d 100644
--- a/pvmfw/platform.dts
+++ b/pvmfw/platform.dts
@@ -225,6 +225,8 @@
 			0x3000 0x0 0x0 1 &intc 0 0 GIC_SPI (IRQ_BASE + 5) IRQ_TYPE_LEVEL_HIGH
 			0x3800 0x0 0x0 1 &intc 0 0 GIC_SPI (IRQ_BASE + 6) IRQ_TYPE_LEVEL_HIGH
 			0x4000 0x0 0x0 1 &intc 0 0 GIC_SPI (IRQ_BASE + 7) IRQ_TYPE_LEVEL_HIGH
+			0x4800 0x0 0x0 1 &intc 0 0 GIC_SPI (IRQ_BASE + 8) IRQ_TYPE_LEVEL_HIGH
+			0x5000 0x0 0x0 1 &intc 0 0 GIC_SPI (IRQ_BASE + 9) IRQ_TYPE_LEVEL_HIGH
 		>;
 		interrupt-map-mask = <0xf800 0x0 0x0 0x7
 				      0xf800 0x0 0x0 0x7
@@ -233,6 +235,8 @@
 				      0xf800 0x0 0x0 0x7
 				      0xf800 0x0 0x0 0x7
 				      0xf800 0x0 0x0 0x7
+				      0xf800 0x0 0x0 0x7
+				      0xf800 0x0 0x0 0x7
 				      0xf800 0x0 0x0 0x7>;
 	};
 
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index b3d0402..244b192 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -209,7 +209,7 @@
 impl PciInfo {
     const IRQ_MASK_CELLS: usize = 4;
     const IRQ_MAP_CELLS: usize = 10;
-    const MAX_IRQS: usize = 8;
+    const MAX_IRQS: usize = 10;
 }
 
 type PciAddrRange = AddressRange<(u32, u64), u64, u64>;
diff --git a/pvmfw/src/gpt.rs b/pvmfw/src/gpt.rs
index 1060460..06bf994 100644
--- a/pvmfw/src/gpt.rs
+++ b/pvmfw/src/gpt.rs
@@ -103,7 +103,7 @@
 
     fn new(mut device: VirtIOBlk) -> Result<Self> {
         let mut blk = [0; Self::LBA_SIZE];
-        device.read_block(Header::LBA, &mut blk).map_err(Error::FailedRead)?;
+        device.read_blocks(Header::LBA, &mut blk).map_err(Error::FailedRead)?;
         let header = Header::read_from_prefix(blk.as_slice()).unwrap();
         if !header.is_valid() {
             return Err(Error::InvalidHeader);
@@ -145,11 +145,11 @@
     }
 
     fn read_block(&mut self, index: usize, blk: &mut [u8]) -> Result<()> {
-        self.device.read_block(index, blk).map_err(Error::FailedRead)
+        self.device.read_blocks(index, blk).map_err(Error::FailedRead)
     }
 
     fn write_block(&mut self, index: usize, blk: &[u8]) -> Result<()> {
-        self.device.write_block(index, blk).map_err(Error::FailedWrite)
+        self.device.write_blocks(index, blk).map_err(Error::FailedWrite)
     }
 }
 
diff --git a/rialto/Android.bp b/rialto/Android.bp
index 1840278..ed9a284 100644
--- a/rialto/Android.bp
+++ b/rialto/Android.bp
@@ -101,6 +101,7 @@
         "liblog_rust",
         "libnix",
         "libvmclient",
+        "libvsock",
     ],
     data: [
         ":rialto_bin",
diff --git a/rialto/AndroidTest.xml b/rialto/AndroidTest.xml
new file mode 100644
index 0000000..43c4c90
--- /dev/null
+++ b/rialto/AndroidTest.xml
@@ -0,0 +1,35 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2023 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 rialto_test">
+  <!--
+    We need root privilege to bypass selinux because shell cannot create socket.
+    Otherwise, we hit the following errors:
+
+    avc:  denied  { create } for  scontext=u:r:shell:s0 tcontext=u:r:shell:s0
+     tclass=vsock_socket permissive=0
+  -->
+  <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer"/>
+
+  <target_preparer class="com.android.tradefed.targetprep.PushFilePreparer">
+    <option name="push-file" key="rialto_test" value="/data/local/tmp/rialto_test" />
+  </target_preparer>
+
+  <test class="com.android.tradefed.testtype.rust.RustBinaryTest" >
+    <option name="test-device-path" value="/data/local/tmp" />
+    <option name="module-name" value="rialto_test" />
+  </test>
+</configuration>
\ No newline at end of file
diff --git a/rialto/src/communication.rs b/rialto/src/communication.rs
new file mode 100644
index 0000000..f00393d
--- /dev/null
+++ b/rialto/src/communication.rs
@@ -0,0 +1,85 @@
+// Copyright 2023, 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.
+
+//! Supports for the communication between rialto and host.
+
+use crate::error::{Error, Result};
+use log::info;
+use virtio_drivers::{
+    self,
+    device::socket::{
+        SingleConnectionManager, SocketError, VirtIOSocket, VsockAddr, VsockEventType,
+    },
+    transport::Transport,
+    Hal,
+};
+
+const MAX_RECV_BUFFER_SIZE_BYTES: usize = 64;
+
+pub struct DataChannel<H: Hal, T: Transport> {
+    connection_manager: SingleConnectionManager<H, T>,
+}
+
+impl<H: Hal, T: Transport> From<VirtIOSocket<H, T>> for DataChannel<H, T> {
+    fn from(socket_device_driver: VirtIOSocket<H, T>) -> Self {
+        Self { connection_manager: SingleConnectionManager::new(socket_device_driver) }
+    }
+}
+
+impl<H: Hal, T: Transport> DataChannel<H, T> {
+    /// Connects to the given destination.
+    pub fn connect(&mut self, destination: VsockAddr) -> virtio_drivers::Result {
+        // Use the same port on rialto and host for convenience.
+        self.connection_manager.connect(destination, destination.port)?;
+        self.connection_manager.wait_for_connect()?;
+        info!("Connected to the destination {destination:?}");
+        Ok(())
+    }
+
+    /// Processes the received requests and sends back a reply.
+    pub fn handle_incoming_request(&mut self) -> Result<()> {
+        let mut buffer = [0u8; MAX_RECV_BUFFER_SIZE_BYTES];
+
+        // TODO(b/274441673): Handle the scenario when the given buffer is too short.
+        let len = self.wait_for_recv(&mut buffer).map_err(Error::ReceivingDataFailed)?;
+
+        // TODO(b/291732060): Implement the communication protocol.
+        // Just reverse the received message for now.
+        buffer[..len].reverse();
+        self.connection_manager.send(&buffer[..len])?;
+        Ok(())
+    }
+
+    fn wait_for_recv(&mut self, buffer: &mut [u8]) -> virtio_drivers::Result<usize> {
+        loop {
+            match self.connection_manager.wait_for_recv(buffer)?.event_type {
+                VsockEventType::Disconnected { .. } => {
+                    return Err(SocketError::ConnectionFailed.into())
+                }
+                VsockEventType::Received { length, .. } => return Ok(length),
+                VsockEventType::Connected
+                | VsockEventType::ConnectionRequest
+                | VsockEventType::CreditRequest
+                | VsockEventType::CreditUpdate => {}
+            }
+        }
+    }
+
+    /// Shuts down the data channel.
+    pub fn force_close(&mut self) -> virtio_drivers::Result {
+        self.connection_manager.force_close()?;
+        info!("Connection shutdown.");
+        Ok(())
+    }
+}
diff --git a/rialto/src/error.rs b/rialto/src/error.rs
index 0c1e25d..461870b 100644
--- a/rialto/src/error.rs
+++ b/rialto/src/error.rs
@@ -41,6 +41,10 @@
     VirtIOSocketCreationFailed(virtio_drivers::Error),
     /// Missing socket device.
     MissingVirtIOSocketDevice,
+    /// Failed VirtIO driver operation.
+    VirtIODriverOperationFailed(virtio_drivers::Error),
+    /// Failed to receive data.
+    ReceivingDataFailed(virtio_drivers::Error),
 }
 
 impl fmt::Display for Error {
@@ -58,6 +62,10 @@
                 write!(f, "Failed to create VirtIO Socket device: {e}")
             }
             Self::MissingVirtIOSocketDevice => write!(f, "Missing VirtIO Socket device."),
+            Self::VirtIODriverOperationFailed(e) => {
+                write!(f, "Failed VirtIO driver operation: {e}")
+            }
+            Self::ReceivingDataFailed(e) => write!(f, "Failed to receive data: {e}"),
         }
     }
 }
@@ -91,3 +99,9 @@
         Self::MemoryOperationFailed(e)
     }
 }
+
+impl From<virtio_drivers::Error> for Error {
+    fn from(e: virtio_drivers::Error) -> Self {
+        Self::VirtIODriverOperationFailed(e)
+    }
+}
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index bbc9997..5c6649a 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -17,11 +17,13 @@
 #![no_main]
 #![no_std]
 
+mod communication;
 mod error;
 mod exceptions;
 
 extern crate alloc;
 
+use crate::communication::DataChannel;
 use crate::error::{Error, Result};
 use core::num::NonZeroUsize;
 use core::slice;
@@ -30,6 +32,7 @@
 use libfdt::FdtError;
 use log::{debug, error, info};
 use virtio_drivers::{
+    device::socket::VsockAddr,
     transport::{pci::bus::PciRoot, DeviceType, Transport},
     Hal,
 };
@@ -46,6 +49,20 @@
     },
 };
 
+fn host_addr() -> VsockAddr {
+    const PROTECTED_VM_PORT: u32 = 5679;
+    const NON_PROTECTED_VM_PORT: u32 = 5680;
+    const VMADDR_CID_HOST: u64 = 2;
+
+    let port = if is_protected_vm() { PROTECTED_VM_PORT } else { NON_PROTECTED_VM_PORT };
+    VsockAddr { cid: VMADDR_CID_HOST, port }
+}
+
+fn is_protected_vm() -> bool {
+    // Use MMIO support to determine whether the VM is protected.
+    get_mmio_guard().is_some()
+}
+
 fn new_page_table() -> Result<PageTable> {
     let mut page_table = PageTable::default();
 
@@ -119,6 +136,12 @@
     debug!("PCI root: {pci_root:#x?}");
     let socket_device = find_socket_device::<HalImpl>(&mut pci_root)?;
     debug!("Found socket device: guest cid = {:?}", socket_device.guest_cid());
+
+    let mut data_channel = DataChannel::from(socket_device);
+    data_channel.connect(host_addr())?;
+    data_channel.handle_incoming_request()?;
+    data_channel.force_close()?;
+
     Ok(())
 }
 
diff --git a/rialto/tests/test.rs b/rialto/tests/test.rs
index 98c935d..4ad8eb8 100644
--- a/rialto/tests/test.rs
+++ b/rialto/tests/test.rs
@@ -22,15 +22,21 @@
     },
     binder::{ParcelFileDescriptor, ProcessState},
 };
-use anyhow::{anyhow, Context, Error};
+use anyhow::{anyhow, bail, Context, Error};
 use log::info;
 use std::fs::File;
-use std::io::{self, BufRead, BufReader};
+use std::io::{self, BufRead, BufReader, Read, Write};
 use std::os::unix::io::FromRawFd;
 use std::panic;
 use std::thread;
 use std::time::Duration;
 use vmclient::{DeathReason, VmInstance};
+use vsock::{VsockListener, VMADDR_CID_HOST};
+
+// TODO(b/291732060): Move the port numbers to the common library shared between the host
+// and rialto.
+const PROTECTED_VM_PORT: u32 = 5679;
+const NON_PROTECTED_VM_PORT: u32 = 5680;
 
 const SIGNED_RIALTO_PATH: &str = "/data/local/tmp/rialto_test/arm64/rialto.bin";
 const UNSIGNED_RIALTO_PATH: &str = "/data/local/tmp/rialto_test/arm64/rialto_unsigned.bin";
@@ -124,6 +130,9 @@
     )
     .context("Failed to create VM")?;
 
+    let port = if protected_vm { PROTECTED_VM_PORT } else { NON_PROTECTED_VM_PORT };
+    let check_socket_handle = thread::spawn(move || try_check_socket_connection(port).unwrap());
+
     vm.start().context("Failed to start VM")?;
 
     // Wait for VM to finish, and check that it shut down cleanly.
@@ -132,6 +141,15 @@
         .ok_or_else(|| anyhow!("Timed out waiting for VM exit"))?;
     assert_eq!(death_reason, DeathReason::Shutdown);
 
+    match check_socket_handle.join() {
+        Ok(_) => {
+            info!(
+                "Received the echoed message. \
+                   The socket connection between the host and the service VM works correctly."
+            )
+        }
+        Err(_) => bail!("The socket connection check failed."),
+    }
     Ok(())
 }
 
@@ -140,6 +158,7 @@
 
     // SAFETY: These are new FDs with no previous owner.
     let reader = unsafe { File::from_raw_fd(reader_fd) };
+    // SAFETY: These are new FDs with no previous owner.
     let writer = unsafe { File::from_raw_fd(writer_fd) };
 
     thread::spawn(|| {
@@ -149,3 +168,28 @@
     });
     Ok(writer)
 }
+
+fn try_check_socket_connection(port: u32) -> Result<(), Error> {
+    info!("Setting up the listening socket on port {port}...");
+    let listener = VsockListener::bind_with_cid_port(VMADDR_CID_HOST, port)?;
+    info!("Listening on port {port}...");
+
+    let Some(Ok(mut vsock_stream)) = listener.incoming().next() else {
+        bail!("Failed to get vsock_stream");
+    };
+    info!("Accepted connection {:?}", vsock_stream);
+
+    let message = "Hello from host";
+    vsock_stream.write_all(message.as_bytes())?;
+    vsock_stream.flush()?;
+    info!("Sent message: {:?}.", message);
+
+    let mut buffer = vec![0u8; 30];
+    vsock_stream.set_read_timeout(Some(Duration::from_millis(1_000)))?;
+    let len = vsock_stream.read(&mut buffer)?;
+
+    assert_eq!(message.len(), len);
+    buffer[..len].reverse();
+    assert_eq!(message.as_bytes(), &buffer[..len]);
+    Ok(())
+}
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
index 9c512bf..90ba575 100644
--- a/tests/benchmark/Android.bp
+++ b/tests/benchmark/Android.bp
@@ -26,6 +26,7 @@
     sdk_version: "test_current",
     use_embedded_native_libs: true,
     compile_multilib: "64",
+    required: ["perf-setup"],
     host_required: ["MicrodroidTestPreparer"],
 }
 
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 aed28a8..625f26a 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -78,6 +78,7 @@
     private static final String TAG = "MicrodroidBenchmarks";
     private static final String METRIC_NAME_PREFIX = getMetricPrefix() + "microdroid/";
     private static final int IO_TEST_TRIAL_COUNT = 5;
+    private static final int TEST_TRIAL_COUNT = 5;
     private static final long ONE_MEBI = 1024 * 1024;
 
     @Rule public Timeout globalTimeout = Timeout.seconds(300);
@@ -767,4 +768,35 @@
         }
         reportMetrics(requestLatencies, "latency/vsock", "us");
     }
+
+    @Test
+    public void testVmKillTime() throws Exception {
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config_io.json")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .build();
+        List<Double> vmKillTime = new ArrayList<>(TEST_TRIAL_COUNT);
+
+        for (int i = 0; i < TEST_TRIAL_COUNT; ++i) {
+            VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_kill_time" + i, config);
+            VmEventListener listener =
+                    new VmEventListener() {
+                        @Override
+                        public void onPayloadReady(VirtualMachine vm) {
+                            long start = System.nanoTime();
+                            try {
+                                vm.stop();
+                            } catch (Exception e) {
+                                Log.e(TAG, "Error in vm.stop():" + e);
+                                throw new RuntimeException(e);
+                            }
+                            vmKillTime.add((double) (System.nanoTime() - start) / NANO_TO_MICRO);
+                            super.onPayloadReady(vm);
+                        }
+                    };
+            listener.runToFinish(TAG, vm);
+        }
+        reportMetrics(vmKillTime, "vm_kill_time", "microsecond");
+    }
 }
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index 8c6218c..f98d1d9 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -60,9 +60,6 @@
     // Files that define the "test" instance of CompOS
     private static final String COMPOS_TEST_ROOT = "/data/misc/apexdata/com.android.compos/test/";
 
-    private static final String SYSTEM_SERVER_COMPILER_FILTER_PROP_NAME =
-            "dalvik.vm.systemservercompilerfilter";
-
     private static final String BOOTLOADER_TIME_PROP_NAME = "ro.boot.boottime";
     private static final String BOOTLOADER_PREFIX = "bootloader-";
     private static final String BOOTLOADER_TIME = "bootloader_time";
@@ -74,7 +71,6 @@
     private static final int COMPILE_STAGED_APEX_RETRY_INTERVAL_MS = 10 * 1000;
     private static final int COMPILE_STAGED_APEX_TIMEOUT_SEC = 540;
     private static final int BOOT_COMPLETE_TIMEOUT_MS = 10 * 60 * 1000;
-    private static final double NANOS_IN_SEC = 1_000_000_000.0;
     private static final int ROUND_COUNT = 5;
     private static final int ROUND_IGNORE_STARTUP_TIME = 3;
     private static final String APK_NAME = "MicrodroidTestApp.apk";
@@ -262,32 +258,34 @@
                         .memoryMib(vm_mem_mb)
                         .cpuTopology("match_host")
                         .build(device);
-        microdroidDevice.waitForBootComplete(30000);
-        microdroidDevice.enableAdbRoot();
-
-        CommandRunner microdroid = new CommandRunner(microdroidDevice);
-
-        microdroid.run("mkdir -p /mnt/ramdisk && chmod 777 /mnt/ramdisk");
-        microdroid.run("mount -t tmpfs -o size=32G tmpfs /mnt/ramdisk");
-
-        // Allocate memory for the VM until it fails and make sure that we touch
-        // the allocated memory in the guest to be able to create stage2 fragmentation.
         try {
-            microdroid.tryRun(
-                    String.format(
-                            "cd /mnt/ramdisk && truncate -s %dM sprayMemory"
-                                    + " && dd if=/dev/zero of=sprayMemory bs=1MB count=%d",
-                            vm_mem_mb, vm_mem_mb));
-        } catch (Exception ex) {
-        }
+            microdroidDevice.waitForBootComplete(30000);
+            microdroidDevice.enableAdbRoot();
 
-        // Run the app during the VM run and collect cold startup time.
-        for (int i = 0; i < ROUND_COUNT; i++) {
-            AmStartupTimeCmdParser duringVmStartApp = getColdRunStartupTimes(android, pkgName);
-            metricColector.addStartupTimeMetricDuringVmRun(duringVmStartApp);
-        }
+            CommandRunner microdroid = new CommandRunner(microdroidDevice);
 
-        device.shutdownMicrodroid(microdroidDevice);
+            microdroid.run("mkdir -p /mnt/ramdisk && chmod 777 /mnt/ramdisk");
+            microdroid.run("mount -t tmpfs -o size=32G tmpfs /mnt/ramdisk");
+
+            // Allocate memory for the VM until it fails and make sure that we touch
+            // the allocated memory in the guest to be able to create stage2 fragmentation.
+            try {
+                microdroid.tryRun(
+                        String.format(
+                                "cd /mnt/ramdisk && truncate -s %dM sprayMemory"
+                                        + " && dd if=/dev/zero of=sprayMemory bs=1MB count=%d",
+                                vm_mem_mb, vm_mem_mb));
+            } catch (Exception expected) {
+            }
+
+            // Run the app during the VM run and collect cold startup time.
+            for (int i = 0; i < ROUND_COUNT; i++) {
+                AmStartupTimeCmdParser duringVmStartApp = getColdRunStartupTimes(android, pkgName);
+                metricColector.addStartupTimeMetricDuringVmRun(duringVmStartApp);
+            }
+        } finally {
+            device.shutdownMicrodroid(microdroidDevice);
+        }
 
         // Run the app after the VM run and collect cold startup time.
         for (int i = 0; i < ROUND_COUNT; i++) {
@@ -304,12 +302,12 @@
             String[] lines = startAppLog.split("[\r\n]+");
             mTotalTime = mWaitTime = 0;
 
-            for (int i = 0; i < lines.length; i++) {
-                if (lines[i].contains("TotalTime:")) {
-                    mTotalTime = Integer.parseInt(lines[i].replaceAll("\\D+", ""));
+            for (String line : lines) {
+                if (line.contains("TotalTime:")) {
+                    mTotalTime = Integer.parseInt(line.replaceAll("\\D+", ""));
                 }
-                if (lines[i].contains("WaitTime:")) {
-                    mWaitTime = Integer.parseInt(lines[i].replaceAll("\\D+", ""));
+                if (line.contains("WaitTime:")) {
+                    mWaitTime = Integer.parseInt(line.replaceAll("\\D+", ""));
                 }
             }
         }
@@ -365,9 +363,9 @@
         String content = android.runForResult("cat /proc/meminfo").getStdout().trim();
         String[] lines = content.split("[\r\n]+");
 
-        for (int i = 0; i < lines.length; i++) {
-            if (lines[i].contains("MemFree:")) {
-                freeMemory = Integer.parseInt(lines[i].replaceAll("\\D+", "")) / 1024;
+        for (String line : lines) {
+            if (line.contains("MemFree:")) {
+                freeMemory = Integer.parseInt(line.replaceAll("\\D+", "")) / 1024;
                 return freeMemory;
             }
         }
@@ -416,7 +414,7 @@
 
         CommandRunner android = new CommandRunner(getDevice());
         String result = android.run("dmesg");
-        Pattern pattern = Pattern.compile("\\[(.*)\\].*sys.boot_completed=1.*");
+        Pattern pattern = Pattern.compile("\\[(.*)].*sys.boot_completed=1.*");
         for (String line : result.split("[\r\n]+")) {
             Matcher matcher = pattern.matcher(line);
             if (matcher.find()) {
@@ -568,7 +566,7 @@
     private void compileStagedApex(int timeoutSec) throws Exception {
 
         long timeStart = System.currentTimeMillis();
-        long timeEnd = timeStart + timeoutSec * 1000;
+        long timeEnd = timeStart + timeoutSec * 1000L;
 
         while (true) {
 
@@ -599,7 +597,7 @@
     private void reInstallApex(int timeoutSec) throws Exception {
 
         long timeStart = System.currentTimeMillis();
-        long timeEnd = timeStart + timeoutSec * 1000;
+        long timeEnd = timeStart + timeoutSec * 1000L;
 
         while (true) {
 
diff --git a/tests/helper/src/java/com/android/microdroid/test/common/DeviceProperties.java b/tests/helper/src/java/com/android/microdroid/test/common/DeviceProperties.java
index 23f8ca6..2ea748b 100644
--- a/tests/helper/src/java/com/android/microdroid/test/common/DeviceProperties.java
+++ b/tests/helper/src/java/com/android/microdroid/test/common/DeviceProperties.java
@@ -28,10 +28,12 @@
 
     private static final String KEY_VENDOR_DEVICE = "ro.product.vendor.device";
     private static final String KEY_BUILD_TYPE = "ro.build.type";
+    private static final String KEY_PRODUCT_NAME = "ro.product.name";
     private static final String KEY_METRICS_TAG = "debug.hypervisor.metrics_tag";
 
     private static final String CUTTLEFISH_DEVICE_PREFIX = "vsoc_";
     private static final String USER_BUILD_TYPE = "user";
+    private static final String HWASAN_SUFFIX = "_hwasan";
 
     private final PropertyGetter mPropertyGetter;
 
@@ -53,6 +55,14 @@
     }
 
     /**
+     * @return whether the build is HWASAN.
+     */
+    public boolean isHwasan() {
+        String productName = getProperty(KEY_PRODUCT_NAME);
+        return productName != null && productName.contains(HWASAN_SUFFIX);
+    }
+
+    /**
      * @return whether the device is user build.
      */
     public boolean isUserBuild() {
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 81ccec7..98327a9 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
@@ -89,6 +89,10 @@
         return DeviceProperties.create(getDevice()::getProperty).isCuttlefish();
     }
 
+    protected boolean isHwasan() {
+        return DeviceProperties.create(getDevice()::getProperty).isHwasan();
+    }
+
     protected String getMetricPrefix() {
         return MetricsProcessor.getMetricPrefix(
                 DeviceProperties.create(getDevice()::getProperty).getMetricsTag());
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index deaf08a..82d8571 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -42,7 +42,6 @@
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.TestDevice;
-import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 import com.android.tradefed.util.CommandResult;
@@ -62,12 +61,8 @@
 import org.xml.sax.Attributes;
 import org.xml.sax.helpers.DefaultHandler;
 
-import java.io.BufferedReader;
 import java.io.ByteArrayInputStream;
 import java.io.File;
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.InputStreamReader;
 import java.io.PipedInputStream;
 import java.io.PipedOutputStream;
 import java.util.ArrayList;
@@ -75,7 +70,6 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
@@ -95,15 +89,11 @@
 
     private static final int BOOT_COMPLETE_TIMEOUT = 30000; // 30 seconds
 
-    private static final Pattern sCIDPattern = Pattern.compile("with CID (\\d+)");
-
     private static class VmInfo {
         final Process mProcess;
-        final String mCid;
 
-        VmInfo(Process process, String cid) {
+        VmInfo(Process process) {
             mProcess = process;
-            mCid = cid;
         }
     }
 
@@ -376,44 +366,14 @@
 
         PipedInputStream pis = new PipedInputStream();
         Process process = RunUtil.getDefault().runCmdInBackground(args, new PipedOutputStream(pis));
-        return new VmInfo(process, extractCidFrom(pis));
-    }
-
-    private static Optional<String> tryExtractCidFrom(String str) {
-        Matcher matcher = sCIDPattern.matcher(str);
-        if (matcher.find()) {
-            return Optional.of(matcher.group(1));
-        }
-        return Optional.empty();
-    }
-
-    private static String extractCidFrom(InputStream input) throws IOException {
-        String cid = null;
-        String line;
-        try (BufferedReader out = new BufferedReader(new InputStreamReader(input))) {
-            while ((line = out.readLine()) != null) {
-                CLog.i("VM output: " + line);
-                Optional<String> result = tryExtractCidFrom(line);
-                if (result.isPresent()) {
-                    cid = result.get();
-                    break;
-                }
-            }
-        }
-        assertWithMessage("The output does not contain the expected pattern for CID.")
-                .that(cid)
-                .isNotNull();
-        return cid;
+        return new VmInfo(process);
     }
 
     @Test
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
     public void protectedVmRunsPvmfw() throws Exception {
         // Arrange
-        boolean protectedVm = true;
-        assumeTrue(
-                "Skip if protected VMs are not supported",
-                getAndroidDevice().supportsMicrodroid(protectedVm));
+        assumeProtectedVmSupported();
         final String configPath = "assets/vm_config_apex.json";
 
         // Act
@@ -422,7 +382,7 @@
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(protectedVm)
+                        .protectedVm(true)
                         .build(getAndroidDevice());
 
         // Assert
@@ -441,16 +401,16 @@
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
     public void protectedVmWithImageSignedWithDifferentKeyRunsPvmfw() throws Exception {
         // Arrange
-        boolean protectedVm = true;
-        assumeTrue(
-                "Skip if protected VMs are not supported",
-                getAndroidDevice().supportsMicrodroid(protectedVm));
+        assumeProtectedVmSupported();
         File key = findTestFile("test.com.android.virt.pem");
 
         // Act
         VmInfo vmInfo =
                 runMicrodroidWithResignedImages(
-                        key, /*keyOverrides=*/ Map.of(), protectedVm, /*updateBootconfigs=*/ true);
+                        key,
+                        /*keyOverrides=*/ Map.of(),
+                        /*isProtected=*/ true,
+                        /*updateBootconfigs=*/ true);
 
         // Assert
         vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
@@ -465,6 +425,7 @@
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
     public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
+        assumeNonProtectedVmSupported();
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
         VmInfo vmInfo =
@@ -481,6 +442,8 @@
     @Test
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
     public void testBootFailsWhenVbMetaDigestDoesNotMatchBootconfig() throws Exception {
+        // protectedVmWithImageSignedWithDifferentKeyRunsPvmfw() is the protected case.
+        assumeNonProtectedVmSupported();
         // Sign everything with key1 except vbmeta
         File key = findTestFile("test.com.android.virt.pem");
         // To be able to stop it, it should be a daemon.
@@ -558,10 +521,26 @@
     }
 
     @Test
-    public void testTombstonesAreGeneratedUponUserspaceCrash() throws Exception {
+    public void testTombstonesAreGeneratedUponUserspaceCrashOnNonPvm() throws Exception {
+        assumeNonProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        testTombstonesAreGeneratedUponUserspaceCrash(false);
+    }
+
+    @Test
+    public void testTombstonesAreGeneratedUponUserspaceCrashOnPvm() throws Exception {
+        assumeProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        testTombstonesAreGeneratedUponUserspaceCrash(true);
+    }
+
+    private void testTombstonesAreGeneratedUponUserspaceCrash(boolean protectedVm)
+            throws Exception {
         assertThat(
                         isTombstoneGeneratedWithCmd(
-                                false,
+                                protectedVm,
                                 "assets/vm_config.json",
                                 "kill",
                                 "-SIGSEGV",
@@ -570,10 +549,28 @@
     }
 
     @Test
-    public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash() throws Exception {
+    public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrashOnNonPvm()
+            throws Exception {
+        assumeNonProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash(false);
+    }
+
+    @Test
+    public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrashOnPvm()
+            throws Exception {
+        assumeProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash(true);
+    }
+
+    private void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash(boolean protectedVm)
+            throws Exception {
         assertThat(
                         isTombstoneGeneratedWithCmd(
-                                false,
+                                protectedVm,
                                 "assets/vm_config_no_tombstone.json",
                                 "kill",
                                 "-SIGSEGV",
@@ -597,19 +594,18 @@
 
     @Test
     public void testTombstonesAreGeneratedUponKernelCrashOnNonPvm() throws Exception {
-        testTombstonesAreGeneratedUponKernelCrash(false);
+        assumeNonProtectedVmSupported();
+        testTombstonesAreGeneratedUponKernelCrash(/* protectedVm=*/ false);
     }
 
     @Test
     public void testTombstonesAreGeneratedUponKernelCrashOnPvm() throws Exception {
-        assumeTrue(
-                "Protected VMs are not supported",
-                getAndroidDevice().supportsMicrodroid(/*protectedVm=*/ true));
-        testTombstonesAreGeneratedUponKernelCrash(true);
+        assumeProtectedVmSupported();
+        testTombstonesAreGeneratedUponKernelCrash(/* protectedVm=*/ true);
     }
 
-    private boolean isTombstoneGeneratedWithVmRunApp(boolean debuggable, String... additionalArgs)
-            throws Exception {
+    private boolean isTombstoneGeneratedWithVmRunApp(
+            boolean protectedVm, boolean debuggable, String... additionalArgs) throws Exception {
         // we can't use microdroid builder as it wants ADB connection (debuggable)
         CommandRunner android = new CommandRunner(getDevice());
         String testStartTime = android.runWithTimeout(1000, "date", "'+%Y-%m-%d %H:%M:%S.%N'");
@@ -630,44 +626,130 @@
                                 apkPath,
                                 idsigPath,
                                 instanceImgPath));
+        if (protectedVm) {
+            cmd.add("--protected");
+        }
         Collections.addAll(cmd, additionalArgs);
 
         android.run(cmd.toArray(new String[0]));
         return isTombstoneReceivedFromHostLogcat(testStartTime);
     }
 
-    private boolean isTombstoneGeneratedWithCrashPayload(boolean debuggable) throws Exception {
+    private boolean isTombstoneGeneratedWithCrashPayload(boolean protectedVm, boolean debuggable)
+            throws Exception {
         return isTombstoneGeneratedWithVmRunApp(
-                debuggable, "--payload-binary-name", "MicrodroidCrashNativeLib.so");
+                protectedVm, debuggable, "--payload-binary-name", "MicrodroidCrashNativeLib.so");
     }
 
     @Test
-    public void testTombstonesAreGeneratedWithCrashPayload() throws Exception {
-        assertThat(isTombstoneGeneratedWithCrashPayload(true /* debuggable */)).isTrue();
+    public void testTombstonesAreGeneratedWithCrashPayloadOnPvm() throws Exception {
+        assumeProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        assertThat(
+                        isTombstoneGeneratedWithCrashPayload(
+                                /*protectedVm=*/ true, /*debuggable=*/ true))
+                .isTrue();
     }
 
     @Test
-    public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggable() throws Exception {
-        assertThat(isTombstoneGeneratedWithCrashPayload(false /* debuggable */)).isFalse();
+    public void testTombstonesAreGeneratedWithCrashPayloadOnNonPvm() throws Exception {
+        assumeNonProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        assertThat(
+                        isTombstoneGeneratedWithCrashPayload(
+                                /*protectedVm=*/ false, /*debuggable=*/ true))
+                .isTrue();
     }
 
-    private boolean isTombstoneGeneratedWithCrashConfig(boolean debuggable) throws Exception {
+    @Test
+    public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggableOnPvm()
+            throws Exception {
+        assumeProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        assertThat(
+                        isTombstoneGeneratedWithCrashPayload(
+                                /*protectedVm=*/ true, /*debuggable=*/ false))
+                .isFalse();
+    }
+
+    @Test
+    public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggableOnNonPvm()
+            throws Exception {
+        assumeNonProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        assertThat(
+                        isTombstoneGeneratedWithCrashPayload(
+                                /*protectedVm=*/ false, /*debuggable=*/ false))
+                .isFalse();
+    }
+
+    private boolean isTombstoneGeneratedWithCrashConfig(boolean protectedVm, boolean debuggable)
+            throws Exception {
         return isTombstoneGeneratedWithVmRunApp(
-                debuggable, "--config-path", "assets/vm_config_crash.json");
+                protectedVm, debuggable, "--config-path", "assets/vm_config_crash.json");
     }
 
     @Test
-    public void testTombstonesAreGeneratedWithCrashConfig() throws Exception {
-        assertThat(isTombstoneGeneratedWithCrashConfig(true /* debuggable */)).isTrue();
+    public void testTombstonesAreGeneratedWithCrashConfigOnPvm() throws Exception {
+        assumeProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        assertThat(isTombstoneGeneratedWithCrashConfig(/*protectedVm=*/ true, /*debuggable=*/ true))
+                .isTrue();
     }
 
     @Test
-    public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggable() throws Exception {
-        assertThat(isTombstoneGeneratedWithCrashConfig(false /* debuggable */)).isFalse();
+    public void testTombstonesAreGeneratedWithCrashConfigOnNonPvm() throws Exception {
+        assumeNonProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        assertThat(
+                        isTombstoneGeneratedWithCrashConfig(
+                                /*protectedVm=*/ false, /*debuggable=*/ true))
+                .isTrue();
     }
 
     @Test
-    public void testTelemetryPushedAtoms() throws Exception {
+    public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggableOnPvm()
+            throws Exception {
+        assumeProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        assertThat(
+                        isTombstoneGeneratedWithCrashConfig(
+                                /*protectedVm=*/ true, /*debuggable=*/ false))
+                .isFalse();
+    }
+
+    @Test
+    public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggableOnNonPvm()
+            throws Exception {
+        assumeNonProtectedVmSupported();
+        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
+        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
+        assertThat(
+                        isTombstoneGeneratedWithCrashConfig(
+                                /*protectedVm=*/ false, /*debuggable=*/ false))
+                .isFalse();
+    }
+
+    @Test
+    public void testTelemetryPushedAtomsOnNonPvm() throws Exception {
+        assumeNonProtectedVmSupported();
+        testTelemetryPushedAtoms(false);
+    }
+
+    @Test
+    public void testTelemetryPushedAtomsOnPvm() throws Exception {
+        assumeProtectedVmSupported();
+        testTelemetryPushedAtoms(true);
+    }
+
+    private void testTelemetryPushedAtoms(boolean protectedVm) throws Exception {
         // Reset statsd config and report before the test
         ConfigUtils.removeConfig(getDevice());
         ReportUtils.clearReports(getDevice());
@@ -688,6 +770,7 @@
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
+                        .protectedVm(protectedVm)
                         .build(device);
         microdroid.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         device.shutdownMicrodroid(microdroid);
@@ -714,7 +797,7 @@
                 data.get(0).getAtom().getVmCreationRequested();
         assertThat(atomVmCreationRequested.getHypervisor())
                 .isEqualTo(AtomsProto.VmCreationRequested.Hypervisor.PKVM);
-        assertThat(atomVmCreationRequested.getIsProtected()).isFalse();
+        assertThat(atomVmCreationRequested.getIsProtected()).isEqualTo(protectedVm);
         assertThat(atomVmCreationRequested.getCreationSucceeded()).isTrue();
         assertThat(atomVmCreationRequested.getBinderExceptionCode()).isEqualTo(0);
         assertThat(atomVmCreationRequested.getVmIdentifier()).isEqualTo("VmRunApp");
@@ -743,7 +826,19 @@
 
     @Test
     @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
-    public void testMicrodroidBoots() throws Exception {
+    public void testMicrodroidBootsOnPvm() throws Exception {
+        assumeProtectedVmSupported();
+        testMicrodroidBoots(true);
+    }
+
+    @Test
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
+    public void testMicrodroidBootsOnNonPvm() throws Exception {
+        assumeNonProtectedVmSupported();
+        testMicrodroidBoots(false);
+    }
+
+    private void testMicrodroidBoots(boolean protectedVm) throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
 
         final String configPath = "assets/vm_config.json"; // path inside the APK
@@ -752,6 +847,7 @@
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
+                        .protectedVm(protectedVm)
                         .build(getAndroidDevice());
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
@@ -810,13 +906,25 @@
     }
 
     @Test
-    public void testMicrodroidRamUsage() throws Exception {
+    public void testMicrodroidRamUsageOnPvm() throws Exception {
+        assumeProtectedVmSupported();
+        testMicrodroidRamUsage(true);
+    }
+
+    @Test
+    public void testMicrodroidRamUsageOnNonPvm() throws Exception {
+        assumeNonProtectedVmSupported();
+        testMicrodroidRamUsage(false);
+    }
+
+    private void testMicrodroidRamUsage(boolean protectedVm) throws Exception {
         final String configPath = "assets/vm_config.json";
         mMicrodroidDevice =
                 MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
+                        .protectedVm(protectedVm)
                         .build(getAndroidDevice());
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         mMicrodroidDevice.enableAdbRoot();
@@ -990,6 +1098,18 @@
                         "android.permission.USE_CUSTOM_VIRTUAL_MACHINE");
     }
 
+    private void assumeProtectedVmSupported() throws Exception {
+        assumeTrue(
+                "Test skipped because protected VMs are not supported",
+                getAndroidDevice().supportsMicrodroid(true));
+    }
+
+    private void assumeNonProtectedVmSupported() throws Exception {
+        assumeTrue(
+                "Test skipped because non-protected VMs are not supported",
+                getAndroidDevice().supportsMicrodroid(false));
+    }
+
     private TestDevice getAndroidDevice() {
         TestDevice androidDevice = (TestDevice) getDevice();
         assertThat(androidDevice).isNotNull();
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index b2497b1..daee7c5 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -28,6 +28,7 @@
     ErrorCode::ErrorCode,
 };
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
+    AssignableDevice::AssignableDevice,
     CpuTopology::CpuTopology,
     DiskImage::DiskImage,
     IVirtualMachine::{BnVirtualMachine, IVirtualMachine},
@@ -270,6 +271,12 @@
         // Delegate to the global service, including checking the debug permission.
         GLOBAL_SERVICE.debugListVms()
     }
+
+    /// Get a list of assignable device types.
+    fn getAssignableDevices(&self) -> binder::Result<Vec<AssignableDevice>> {
+        // Delegate to the global service, including checking the permission.
+        GLOBAL_SERVICE.getAssignableDevices()
+    }
 }
 
 impl VirtualizationService {
@@ -1236,7 +1243,7 @@
             return Err(Status::new_service_specific_error_str(
                 -1,
                 Some(format!("cannot find a VM with CID {}", cid)),
-            ))
+            ));
         };
         let instance_img_path = vm.temporary_directory.join("rkpvm_instance.img");
         let instance_img = OpenOptions::new()
diff --git a/virtualizationmanager/src/payload.rs b/virtualizationmanager/src/payload.rs
index ab6f31c..343f3cf 100644
--- a/virtualizationmanager/src/payload.rs
+++ b/virtualizationmanager/src/payload.rs
@@ -81,6 +81,9 @@
 
     #[serde(rename = "provideSharedApexLibs")]
     provide_shared_apex_libs: bool,
+
+    #[serde(rename = "preinstalledModulePath")]
+    preinstalled_path: PathBuf,
 }
 
 impl ApexInfoList {
@@ -275,7 +278,7 @@
     let apex_list = pm.get_apex_list(vm_payload_config.prefer_staged)?;
 
     // collect APEXes from config
-    let mut apex_infos = collect_apex_infos(&apex_list, &vm_payload_config.apexes, debug_config);
+    let mut apex_infos = collect_apex_infos(&apex_list, &vm_payload_config.apexes, debug_config)?;
 
     // Pass sorted list of apexes. Sorting key shouldn't use `path` because it will change after
     // reboot with prefer_staged. `last_update_seconds` is added to distinguish "samegrade"
@@ -376,18 +379,28 @@
     Ok(apexes)
 }
 
+fn check_apexes_are_from_allowed_partitions(requested_apexes: &Vec<&ApexInfo>) -> Result<()> {
+    const ALLOWED_PARTITIONS: [&str; 2] = ["/system", "/system_ext"];
+    for apex in requested_apexes {
+        if !ALLOWED_PARTITIONS.iter().any(|p| apex.preinstalled_path.starts_with(p)) {
+            bail!("Non-system APEX {} is not supported in Microdroid", apex.name);
+        }
+    }
+    Ok(())
+}
+
 // Collect ApexInfos from VM config
 fn collect_apex_infos<'a>(
     apex_list: &'a ApexInfoList,
     apex_configs: &[ApexConfig],
     debug_config: &DebugConfig,
-) -> Vec<&'a ApexInfo> {
+) -> Result<Vec<&'a ApexInfo>> {
     let mut additional_apexes: Vec<&str> = MICRODROID_REQUIRED_APEXES.to_vec();
     if debug_config.should_include_debug_apexes() {
         additional_apexes.extend(MICRODROID_REQUIRED_APEXES_DEBUG.to_vec());
     }
 
-    apex_list
+    let apex_infos = apex_list
         .list
         .iter()
         .filter(|ai| {
@@ -395,7 +408,10 @@
                 || additional_apexes.iter().any(|name| name == &ai.name && ai.is_active)
                 || ai.provide_shared_apex_libs
         })
-        .collect()
+        .collect();
+
+    check_apexes_are_from_allowed_partitions(&apex_infos)?;
+    Ok(apex_infos)
 }
 
 pub fn add_microdroid_vendor_image(vendor_image: File, vm_config: &mut VirtualMachineRawConfig) {
@@ -488,13 +504,14 @@
     }
 
     #[test]
-    fn test_collect_apexes() {
+    fn test_collect_apexes() -> Result<()> {
         let apex_info_list = ApexInfoList {
             list: vec![
                 ApexInfo {
                     // 0
                     name: "com.android.adbd".to_string(),
                     path: PathBuf::from("adbd"),
+                    preinstalled_path: PathBuf::from("/system/adbd"),
                     has_classpath_jar: false,
                     last_update_seconds: 12345678,
                     is_factory: true,
@@ -505,6 +522,7 @@
                     // 1
                     name: "com.android.os.statsd".to_string(),
                     path: PathBuf::from("statsd"),
+                    preinstalled_path: PathBuf::from("/system/statsd"),
                     has_classpath_jar: false,
                     last_update_seconds: 12345678,
                     is_factory: true,
@@ -515,6 +533,7 @@
                     // 2
                     name: "com.android.os.statsd".to_string(),
                     path: PathBuf::from("statsd/updated"),
+                    preinstalled_path: PathBuf::from("/system/statsd"),
                     has_classpath_jar: false,
                     last_update_seconds: 12345678 + 1,
                     is_factory: false,
@@ -545,6 +564,7 @@
                     // 5
                     name: "has_classpath".to_string(),
                     path: PathBuf::from("has_classpath/updated"),
+                    preinstalled_path: PathBuf::from("/system/has_classpath"),
                     has_classpath_jar: true,
                     last_update_seconds: 87654321 + 1,
                     is_factory: false,
@@ -555,6 +575,7 @@
                     // 6
                     name: "apex-foo".to_string(),
                     path: PathBuf::from("apex-foo"),
+                    preinstalled_path: PathBuf::from("/system/apex-foo"),
                     has_classpath_jar: false,
                     last_update_seconds: 87654321,
                     is_factory: true,
@@ -565,6 +586,7 @@
                     // 7
                     name: "apex-foo".to_string(),
                     path: PathBuf::from("apex-foo/updated"),
+                    preinstalled_path: PathBuf::from("/system/apex-foo"),
                     has_classpath_jar: false,
                     last_update_seconds: 87654321 + 1,
                     is_factory: false,
@@ -575,6 +597,7 @@
                     // 8
                     name: "sharedlibs".to_string(),
                     path: PathBuf::from("apex-foo"),
+                    preinstalled_path: PathBuf::from("/system/apex-foo"),
                     last_update_seconds: 87654321,
                     is_factory: true,
                     provide_shared_apex_libs: true,
@@ -584,6 +607,7 @@
                     // 9
                     name: "sharedlibs".to_string(),
                     path: PathBuf::from("apex-foo/updated"),
+                    preinstalled_path: PathBuf::from("/system/apex-foo"),
                     last_update_seconds: 87654321 + 1,
                     is_active: true,
                     provide_shared_apex_libs: true,
@@ -596,7 +620,11 @@
             ApexConfig { name: "{CLASSPATH}".to_string() },
         ];
         assert_eq!(
-            collect_apex_infos(&apex_info_list, &apex_configs, &DebugConfig::new(DebugLevel::FULL)),
+            collect_apex_infos(
+                &apex_info_list,
+                &apex_configs,
+                &DebugConfig::new(DebugLevel::FULL)
+            )?,
             vec![
                 // Pass active/required APEXes
                 &apex_info_list.list[0],
@@ -609,6 +637,55 @@
                 &apex_info_list.list[9],
             ]
         );
+        Ok(())
+    }
+
+    #[test]
+    fn test_check_allowed_partitions_vendor_not_allowed() -> Result<()> {
+        let apex_info_list = ApexInfoList {
+            list: vec![ApexInfo {
+                name: "apex-vendor".to_string(),
+                path: PathBuf::from("apex-vendor"),
+                preinstalled_path: PathBuf::from("/vendor/apex-vendor"),
+                is_active: true,
+                ..Default::default()
+            }],
+        };
+        let apex_configs = vec![ApexConfig { name: "apex-vendor".to_string() }];
+
+        let ret =
+            collect_apex_infos(&apex_info_list, &apex_configs, &DebugConfig::new(DebugLevel::NONE));
+        assert!(ret
+            .is_err_and(|ret| ret.to_string()
+                == "Non-system APEX apex-vendor is not supported in Microdroid"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_check_allowed_partitions_system_ext_allowed() -> Result<()> {
+        let apex_info_list = ApexInfoList {
+            list: vec![ApexInfo {
+                name: "apex-system_ext".to_string(),
+                path: PathBuf::from("apex-system_ext"),
+                preinstalled_path: PathBuf::from("/system_ext/apex-system_ext"),
+                is_active: true,
+                ..Default::default()
+            }],
+        };
+
+        let apex_configs = vec![ApexConfig { name: "apex-system_ext".to_string() }];
+
+        assert_eq!(
+            collect_apex_infos(
+                &apex_info_list,
+                &apex_configs,
+                &DebugConfig::new(DebugLevel::NONE)
+            )?,
+            vec![&apex_info_list.list[0]]
+        );
+
+        Ok(())
     }
 
     #[test]
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 6b39ff9..0732c04 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -36,6 +36,7 @@
         "libstatslog_virtualization_rust",
         "libtombstoned_client_rust",
         "libvsock",
+        "liblazy_static",
     ],
     apex_available: ["com.android.virt"],
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/AssignableDevice.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/AssignableDevice.aidl
new file mode 100644
index 0000000..014d78c
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/AssignableDevice.aidl
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2023 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.
+ */
+package android.system.virtualizationservice;
+
+/** A specification of device to be assigned to the virtual machine. */
+@RustDerive(Clone=true)
+parcelable AssignableDevice {
+    /** Path to SysFS node of the device. */
+    String node;
+
+    /** Kind of the device. */
+    String kind;
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
index 99bc076..df72e49 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
@@ -15,6 +15,7 @@
  */
 package android.system.virtualizationservice;
 
+import android.system.virtualizationservice.AssignableDevice;
 import android.system.virtualizationservice.IVirtualMachine;
 import android.system.virtualizationservice.PartitionType;
 import android.system.virtualizationservice.VirtualMachineConfig;
@@ -55,4 +56,9 @@
      * and as such is only permitted from the shell user.
      */
     VirtualMachineDebugInfo[] debugListVms();
+
+    /**
+     * Get a list of assignable device types.
+     */
+    AssignableDevice[] getAssignableDevices();
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
index cc59b3f..3546355 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -15,6 +15,7 @@
  */
 package android.system.virtualizationservice_internal;
 
+import android.system.virtualizationservice.AssignableDevice;
 import android.system.virtualizationservice.VirtualMachineDebugInfo;
 import android.system.virtualizationservice_internal.AtomVmBooted;
 import android.system.virtualizationservice_internal.AtomVmCreationRequested;
@@ -59,4 +60,17 @@
      * @return the X.509 encoded certificate.
      */
     byte[] requestCertificate(in byte[] csr, in ParcelFileDescriptor instanceImgFd);
+
+    /**
+     * Get a list of assignable devices.
+     */
+    AssignableDevice[] getAssignableDevices();
+
+    /**
+     * Bind given devices to vfio driver.
+     *
+     * @param devices paths of sysfs nodes of devices to assign.
+     * @return a file descriptor containing DTBO for VM.
+     */
+    ParcelFileDescriptor bindDevicesToVfioDriver(in String[] devices);
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 7dfabb0..2dbb6e2 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -19,6 +19,7 @@
 use crate::rkpvm::request_certificate;
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
 use android_system_virtualizationservice::{
+    aidl::android::system::virtualizationservice::AssignableDevice::AssignableDevice,
     aidl::android::system::virtualizationservice::VirtualMachineDebugInfo::VirtualMachineDebugInfo,
     binder::ParcelFileDescriptor,
 };
@@ -32,19 +33,22 @@
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::VM_TOMBSTONES_SERVICE_PORT;
 use anyhow::{anyhow, ensure, Context, Result};
 use binder::{self, BinderFeatures, ExceptionCode, Interface, LazyServiceGuard, Status, Strong};
+use lazy_static::lazy_static;
 use libc::VMADDR_CID_HOST;
 use log::{error, info, warn};
 use rustutils::system_properties;
-use std::collections::HashMap;
-use std::fs::{create_dir, remove_dir_all, set_permissions, Permissions};
+use std::collections::{HashMap, HashSet};
+use std::fs::{canonicalize, create_dir, remove_dir_all, set_permissions, File, Permissions};
 use std::io::{Read, Write};
+use std::os::fd::FromRawFd;
 use std::os::unix::fs::PermissionsExt;
 use std::os::unix::raw::{pid_t, uid_t};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex, Weak};
 use tombstoned_client::{DebuggerdDumpType, TombstonedConnection};
 use vsock::{VsockListener, VsockStream};
-use nix::unistd::{chown, Uid};
+use nix::fcntl::OFlag;
+use nix::unistd::{chown, pipe2, Uid};
 
 /// The unique ID of a VM used (together with a port number) for vsock communication.
 pub type Cid = u32;
@@ -170,6 +174,85 @@
             Status::new_exception_str(ExceptionCode::SERVICE_SPECIFIC, Some(e.to_string()))
         })
     }
+
+    fn getAssignableDevices(&self) -> binder::Result<Vec<AssignableDevice>> {
+        check_use_custom_virtual_machine()?;
+
+        // TODO(b/291191362): read VM DTBO to find assignable devices.
+        Ok(vec![AssignableDevice {
+            kind: "eh".to_owned(),
+            node: "/sys/bus/platform/devices/16d00000.eh".to_owned(),
+        }])
+    }
+
+    fn bindDevicesToVfioDriver(&self, devices: &[String]) -> binder::Result<ParcelFileDescriptor> {
+        check_use_custom_virtual_machine()?;
+
+        let mut set = HashSet::new();
+        for device in devices.iter() {
+            if !set.insert(device) {
+                return Err(Status::new_exception_str(
+                    ExceptionCode::ILLEGAL_ARGUMENT,
+                    Some(format!("duplicated device {device}")),
+                ));
+            }
+            bind_device(device)?;
+        }
+
+        // TODO(b/278008182): create a file descriptor containing DTBO for devices.
+        let (raw_read, raw_write) = pipe2(OFlag::O_CLOEXEC).map_err(|e| {
+            Status::new_exception_str(
+                ExceptionCode::SERVICE_SPECIFIC,
+                Some(format!("can't create fd for DTBO: {e:?}")),
+            )
+        })?;
+        // SAFETY: We are the sole owner of this FD as we just created it, and it is valid and open.
+        let read_fd = unsafe { File::from_raw_fd(raw_read) };
+        // SAFETY: We are the sole owner of this FD as we just created it, and it is valid and open.
+        let _write_fd = unsafe { File::from_raw_fd(raw_write) };
+
+        Ok(ParcelFileDescriptor::new(read_fd))
+    }
+}
+
+lazy_static! {
+    static ref SYSFS_PLATFORM_DEVICES: &'static Path = Path::new("/sys/devices/platform/");
+    static ref VFIO_PLATFORM_DRIVER: &'static Path =
+        Path::new("/sys/bus/platform/drivers/vfio-platform");
+}
+
+fn bind_device(device: &str) -> binder::Result<()> {
+    // Check platform device exists
+    let dev_sysfs_path = canonicalize(device).map_err(|e| {
+        Status::new_exception_str(
+            ExceptionCode::SERVICE_SPECIFIC,
+            Some(format!("can't canonicalize: {e:?}")),
+        )
+    })?;
+    if !dev_sysfs_path.starts_with(*SYSFS_PLATFORM_DEVICES) {
+        return Err(Status::new_exception_str(
+            ExceptionCode::ILLEGAL_ARGUMENT,
+            Some(format!("{device} is not a platform device")),
+        ));
+    }
+
+    // Check platform device is bound to VFIO driver
+    let dev_driver_path = canonicalize(dev_sysfs_path.join("driver")).map_err(|e| {
+        Status::new_exception_str(
+            ExceptionCode::SERVICE_SPECIFIC,
+            Some(format!("can't canonicalize: {e:?}")),
+        )
+    })?;
+    if dev_driver_path != *VFIO_PLATFORM_DRIVER {
+        // TODO(b/278008182): unbind driver and bind to VFIO
+        return Err(Status::new_exception_str(
+            ExceptionCode::UNSUPPORTED_OPERATION,
+            Some("not implemented".to_owned()),
+        ));
+    }
+
+    // already bound to VFIO driver
+    Ok(())
 }
 
 #[derive(Debug, Default)]
@@ -393,3 +476,8 @@
 fn check_manage_access() -> binder::Result<()> {
     check_permission("android.permission.MANAGE_VIRTUAL_MACHINE")
 }
+
+/// Check whether the caller of the current Binder method is allowed to use custom VMs
+fn check_use_custom_virtual_machine() -> binder::Result<()> {
+    check_permission("android.permission.USE_CUSTOM_VIRTUAL_MACHINE")
+}
diff --git a/vm/src/main.rs b/vm/src/main.rs
index d7c2c4d..0c99acb 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -82,8 +82,8 @@
         #[clap(long)]
         log: Option<PathBuf>,
 
-        /// Debug level of the VM. Supported values: "none" (default), and "full".
-        #[clap(long, default_value = "none", value_parser = parse_debug_level)]
+        /// Debug level of the VM. Supported values: "full" (default), and "none".
+        #[clap(long, default_value = "full", value_parser = parse_debug_level)]
         debug: DebugLevel,
 
         /// Run VM in protected mode.
@@ -154,7 +154,7 @@
         #[clap(long)]
         log: Option<PathBuf>,
 
-        /// Debug level of the VM. Supported values: "none" (default), and "full".
+        /// Debug level of the VM. Supported values: "full" (default), and "none".
         #[clap(long, default_value = "full", value_parser = parse_debug_level)]
         debug: DebugLevel,
 
diff --git a/vmbase/example/src/pci.rs b/vmbase/example/src/pci.rs
index 26bc29b..b838539 100644
--- a/vmbase/example/src/pci.rs
+++ b/vmbase/example/src/pci.rs
@@ -78,7 +78,7 @@
             assert_eq!(blk.capacity(), EXPECTED_SECTOR_COUNT as u64);
             let mut data = [0; SECTOR_SIZE_BYTES * EXPECTED_SECTOR_COUNT];
             for i in 0..EXPECTED_SECTOR_COUNT {
-                blk.read_block(i, &mut data[i * SECTOR_SIZE_BYTES..(i + 1) * SECTOR_SIZE_BYTES])
+                blk.read_blocks(i, &mut data[i * SECTOR_SIZE_BYTES..(i + 1) * SECTOR_SIZE_BYTES])
                     .expect("Failed to read block device.");
             }
             for (i, chunk) in data.chunks(size_of::<u32>()).enumerate() {
@@ -89,7 +89,7 @@
         1 => {
             assert_eq!(blk.capacity(), 0);
             let mut data = [0; SECTOR_SIZE_BYTES];
-            assert_eq!(blk.read_block(0, &mut data), Err(Error::IoError));
+            assert_eq!(blk.read_blocks(0, &mut data), Err(Error::IoError));
         }
         _ => panic!("Unexpected VirtIO block device index {}.", index),
     }
diff --git a/vmbase/example/tests/test.rs b/vmbase/example/tests/test.rs
index 3594523..2b6dfbc 100644
--- a/vmbase/example/tests/test.rs
+++ b/vmbase/example/tests/test.rs
@@ -143,6 +143,7 @@
 
     // SAFETY: These are new FDs with no previous owner.
     let reader = unsafe { File::from_raw_fd(reader_fd) };
+    // SAFETY: These are new FDs with no previous owner.
     let writer = unsafe { File::from_raw_fd(writer_fd) };
 
     Ok((reader, writer))
diff --git a/vmbase/src/memory/mod.rs b/vmbase/src/memory/mod.rs
index 898aa10..2f72fc4 100644
--- a/vmbase/src/memory/mod.rs
+++ b/vmbase/src/memory/mod.rs
@@ -23,10 +23,12 @@
 pub use error::MemoryTrackerError;
 pub use page_table::PageTable;
 pub use shared::{
-    alloc_shared, dealloc_shared, handle_permission_fault, handle_translation_fault, MemoryRange,
-    MemoryTracker, MEMORY,
+    handle_permission_fault, handle_translation_fault, MemoryRange, MemoryTracker, MEMORY,
 };
 pub use util::{
-    flush, flushed_zeroize, min_dcache_line_size, page_4kb_of, phys_to_virt, virt_to_phys,
-    PAGE_SIZE, SIZE_128KB, SIZE_2MB, SIZE_4KB, SIZE_4MB, SIZE_64KB,
+    flush, flushed_zeroize, min_dcache_line_size, page_4kb_of, PAGE_SIZE, SIZE_128KB, SIZE_2MB,
+    SIZE_4KB, SIZE_4MB, SIZE_64KB,
 };
+
+pub(crate) use shared::{alloc_shared, dealloc_shared};
+pub(crate) use util::{phys_to_virt, virt_to_phys};
diff --git a/vmbase/src/memory/shared.rs b/vmbase/src/memory/shared.rs
index 173c0ec..dfa29e4 100644
--- a/vmbase/src/memory/shared.rs
+++ b/vmbase/src/memory/shared.rs
@@ -27,6 +27,7 @@
 use alloc::vec::Vec;
 use buddy_system_allocator::{FrameAllocator, LockedFrameAllocator};
 use core::alloc::Layout;
+use core::cmp::max;
 use core::mem::size_of;
 use core::num::NonZeroUsize;
 use core::ops::Range;
@@ -76,10 +77,6 @@
     payload_range: Option<MemoryRange>,
 }
 
-// TODO: Remove this once aarch64-paging crate is updated.
-// SAFETY: Only `PageTable` doesn't implement Send, but it should.
-unsafe impl Send for MemoryTracker {}
-
 impl MemoryTracker {
     const CAPACITY: usize = 5;
     const MMIO_CAPACITY: usize = 5;
@@ -343,7 +340,7 @@
 
 /// Allocates a memory range of at least the given size and alignment that is shared with the host.
 /// Returns a pointer to the buffer.
-pub fn alloc_shared(layout: Layout) -> hyp::Result<NonNull<u8>> {
+pub(crate) fn alloc_shared(layout: Layout) -> hyp::Result<NonNull<u8>> {
     assert_ne!(layout.size(), 0);
     let Some(buffer) = try_shared_alloc(layout) else {
         handle_alloc_error(layout);
@@ -359,7 +356,11 @@
     if let Some(buffer) = shared_pool.alloc_aligned(layout) {
         Some(NonNull::new(buffer as _).unwrap())
     } else if let Some(shared_memory) = SHARED_MEMORY.lock().as_mut() {
-        shared_memory.refill(&mut shared_pool, layout);
+        // Adjusts the layout size to the max of the next power of two and the alignment,
+        // as this is the actual size of the memory allocated in `alloc_aligned()`.
+        let size = max(layout.size().next_power_of_two(), layout.align());
+        let refill_layout = Layout::from_size_align(size, layout.align()).unwrap();
+        shared_memory.refill(&mut shared_pool, refill_layout);
         shared_pool.alloc_aligned(layout).map(|buffer| NonNull::new(buffer as _).unwrap())
     } else {
         None
@@ -374,7 +375,7 @@
 ///
 /// The memory must have been allocated by `alloc_shared` with the same layout, and not yet
 /// deallocated.
-pub unsafe fn dealloc_shared(vaddr: NonNull<u8>, layout: Layout) -> hyp::Result<()> {
+pub(crate) unsafe fn dealloc_shared(vaddr: NonNull<u8>, layout: Layout) -> hyp::Result<()> {
     SHARED_POOL.get().unwrap().lock().dealloc_aligned(vaddr.as_ptr() as usize, layout);
 
     trace!("Deallocated shared buffer at {vaddr:?} with {layout:?}");
diff --git a/vmbase/src/memory/util.rs b/vmbase/src/memory/util.rs
index 48d4c55..2b75414 100644
--- a/vmbase/src/memory/util.rs
+++ b/vmbase/src/memory/util.rs
@@ -88,7 +88,7 @@
 ///
 /// As we use identity mapping for everything, this is just a cast, but it's useful to use it to be
 /// explicit about where we are converting from virtual to physical address.
-pub fn virt_to_phys(vaddr: NonNull<u8>) -> usize {
+pub(crate) fn virt_to_phys(vaddr: NonNull<u8>) -> usize {
     vaddr.as_ptr() as _
 }
 
@@ -96,6 +96,6 @@
 /// physical address.
 ///
 /// Panics if `paddr` is 0.
-pub fn phys_to_virt(paddr: usize) -> NonNull<u8> {
+pub(crate) fn phys_to_virt(paddr: usize) -> NonNull<u8> {
     NonNull::new(paddr as _).unwrap()
 }
diff --git a/zipfuse/src/inode.rs b/zipfuse/src/inode.rs
index ea63422..ef48389 100644
--- a/zipfuse/src/inode.rs
+++ b/zipfuse/src/inode.rs
@@ -210,7 +210,7 @@
                     parent = found;
                     // Update the mode if this is a directory leaf.
                     if !is_file && is_leaf {
-                        let mut inode = table.get_mut(parent).unwrap();
+                        let inode = table.get_mut(parent).unwrap();
                         inode.mode = file.unix_mode().unwrap_or(DEFAULT_DIR_MODE);
                     }
                     continue;
diff --git a/zipfuse/src/main.rs b/zipfuse/src/main.rs
index 20d6fd6..e8be42c 100644
--- a/zipfuse/src/main.rs
+++ b/zipfuse/src/main.rs
@@ -31,7 +31,7 @@
 use std::fs::{File, OpenOptions};
 use std::io;
 use std::io::Read;
-use std::mem::size_of;
+use std::mem::{size_of, MaybeUninit};
 use std::os::unix::io::AsRawFd;
 use std::path::Path;
 use std::path::PathBuf;
@@ -192,7 +192,8 @@
     #[allow(clippy::useless_conversion)]
     fn stat_from(&self, inode: Inode) -> io::Result<libc::stat64> {
         let inode_data = self.find_inode(inode)?;
-        let mut st = unsafe { std::mem::MaybeUninit::<libc::stat64>::zeroed().assume_init() };
+        // SAFETY: All fields of stat64 are valid for zero byte patterns.
+        let mut st = unsafe { MaybeUninit::<libc::stat64>::zeroed().assume_init() };
         st.st_dev = 0;
         st.st_nlink = if let Some(directory) = inode_data.get_directory() {
             (2 + directory.len() as libc::nlink_t).into()