Merge "Revert "Handle exceptions.""
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 80d0807..f40da7e 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -21,7 +21,7 @@
   ],
   "postsubmit": [
     {
-      "name": "odsign_e2e_tests"
+      "name": "odsign_e2e_tests_full"
     }
   ],
   "imports": [
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index e782bd2..a1e81d2 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -16,22 +16,63 @@
 """sign_virt_apex is a command line tool for sign the Virt APEX file.
 
 Typical usage:
-  sign_virt_apex [-v] [--avbtool path_to_avbtool] [--signing_args args] payload_key payload_dir
+  sign_virt_apex payload_key payload_dir
+    -v, --verbose
+    --verify
+    --avbtool path_to_avbtool
+    --signing_args args
 
 sign_virt_apex uses external tools which are assumed to be available via PATH.
 - avbtool (--avbtool can override the tool)
 - lpmake, lpunpack, simg2img, img2simg
 """
 import argparse
-import glob
 import hashlib
 import os
 import re
 import shlex
-import shutil
 import subprocess
 import sys
 import tempfile
+import traceback
+from concurrent import futures
+
+# pylint: disable=line-too-long,consider-using-with
+
+# Use executor to parallelize the invocation of external tools
+# If a task depends on another, pass the future object of the previous task as wait list.
+# Every future object created by a task should be consumed with AwaitAll()
+# so that exceptions are propagated .
+executor = futures.ThreadPoolExecutor()
+
+# Temporary directory for unpacked super.img.
+# We could put its creation/deletion into the task graph as well, but
+# having it as a global setup is much simpler.
+unpack_dir = tempfile.TemporaryDirectory()
+
+# tasks created with Async() are kept in a list so that they are awaited
+# before exit.
+tasks = []
+
+# create an async task and return a future value of it.
+def Async(fn, *args, wait=None, **kwargs):
+
+    # wrap a function with AwaitAll()
+    def wrapped():
+        AwaitAll(wait)
+        fn(*args, **kwargs)
+
+    task = executor.submit(wrapped)
+    tasks.append(task)
+    return task
+
+
+# waits for task (captured in fs as future values) with future.result()
+# so that any exception raised during task can be raised upward.
+def AwaitAll(fs):
+    if fs:
+        for f in fs:
+            f.result()
 
 
 def ParseArgs(argv):
@@ -71,7 +112,8 @@
     return args
 
 
-def RunCommand(args, cmd, env=None, expected_return_values={0}):
+def RunCommand(args, cmd, env=None, expected_return_values=None):
+    expected_return_values = expected_return_values or {0}
     env = env or {}
     env.update(os.environ.copy())
 
@@ -218,7 +260,7 @@
     if info is None:
         return
 
-    with TempDirectory() as work_dir:
+    with tempfile.TemporaryDirectory() as work_dir:
         algorithm = info['Algorithm']
         rollback_index = info['Rollback Index']
         rollback_index_location = info['Rollback Index Location']
@@ -254,18 +296,14 @@
             f.truncate(65536)
 
 
-class TempDirectory(object):
-
-    def __enter__(self):
-        self.name = tempfile.mkdtemp()
-        return self.name
-
-    def __exit__(self, *unused):
-        shutil.rmtree(self.name)
+def UnpackSuperImg(args, super_img, work_dir):
+    tmp_super_img = os.path.join(work_dir, 'super.img')
+    RunCommand(args, ['simg2img', super_img, tmp_super_img])
+    RunCommand(args, ['lpunpack', tmp_super_img, work_dir])
 
 
 def MakeSuperImage(args, partitions, output):
-    with TempDirectory() as work_dir:
+    with tempfile.TemporaryDirectory() as work_dir:
         cmd = ['lpmake', '--device-size=auto', '--metadata-slots=2',  # A/B
                '--metadata-size=65536', '--sparse', '--output=' + output]
 
@@ -281,6 +319,22 @@
         RunCommand(args, cmd)
 
 
+def SignSuperImg(args, key, super_img, work_dir):
+    # unpack super.img
+    UnpackSuperImg(args, super_img, work_dir)
+
+    system_a_img = os.path.join(work_dir, 'system_a.img')
+    vendor_a_img = os.path.join(work_dir, 'vendor_a.img')
+
+    # re-sign each partition
+    system_a_f = Async(AddHashTreeFooter, args, key, system_a_img)
+    vendor_a_f = Async(AddHashTreeFooter, args, key, vendor_a_img)
+
+    # 3. re-pack super.img
+    partitions = {"system_a": system_a_img, "vendor_a": vendor_a_img}
+    Async(MakeSuperImage, args, partitions, super_img, wait=[system_a_f, vendor_a_f])
+
+
 def ReplaceBootloaderPubkey(args, key, bootloader, bootloader_pubkey):
     if os.path.basename(bootloader) in args.key_overrides:
         key = args.key_overrides[os.path.basename(bootloader)]
@@ -305,134 +359,116 @@
         bl_f.write(new_pubkey)
 
 
+# dict of (key, file) for re-sign/verification. keys are un-versioned for readability.
+virt_apex_files = {
+    'bootloader.pubkey': 'etc/microdroid_bootloader.avbpubkey',
+    'bootloader': 'etc/microdroid_bootloader',
+    'boot.img': 'etc/fs/microdroid_boot-5.10.img',
+    'vendor_boot.img': 'etc/fs/microdroid_vendor_boot-5.10.img',
+    'init_boot.img': 'etc/fs/microdroid_init_boot.img',
+    'super.img': 'etc/fs/microdroid_super.img',
+    'vbmeta.img': 'etc/fs/microdroid_vbmeta.img',
+    'vbmeta_bootconfig.img': 'etc/fs/microdroid_vbmeta_bootconfig.img',
+    'bootconfig.normal': 'etc/microdroid_bootconfig.normal',
+    'bootconfig.app_debuggable': 'etc/microdroid_bootconfig.app_debuggable',
+    'bootconfig.full_debuggable': 'etc/microdroid_bootconfig.full_debuggable',
+    'uboot_env.img': 'etc/uboot_env.img'
+}
+
+
+def TargetFiles(input_dir):
+    return {k: os.path.join(input_dir, v) for k, v in virt_apex_files.items()}
+
+
 def SignVirtApex(args):
     key = args.key
     input_dir = args.input_dir
+    files = TargetFiles(input_dir)
 
-    # target files in the Virt APEX
-    bootloader_pubkey = os.path.join(
-        input_dir, 'etc', 'microdroid_bootloader.avbpubkey')
-    bootloader = os.path.join(input_dir, 'etc', 'microdroid_bootloader')
-    boot_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_boot-5.10.img')
-    vendor_boot_img = os.path.join(
-        input_dir, 'etc', 'fs', 'microdroid_vendor_boot-5.10.img')
-    init_boot_img = os.path.join(
-        input_dir, 'etc', 'fs', 'microdroid_init_boot.img')
-    super_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_super.img')
-    vbmeta_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_vbmeta.img')
-    vbmeta_bootconfig_img = os.path.join(
-        input_dir, 'etc', 'fs', 'microdroid_vbmeta_bootconfig.img')
-    bootconfig_normal = os.path.join(
-        input_dir, 'etc', 'microdroid_bootconfig.normal')
-    bootconfig_app_debuggable = os.path.join(
-        input_dir, 'etc', 'microdroid_bootconfig.app_debuggable')
-    bootconfig_full_debuggable = os.path.join(
-        input_dir, 'etc', 'microdroid_bootconfig.full_debuggable')
-    uboot_env_img = os.path.join(
-        input_dir, 'etc', 'uboot_env.img')
+    # unpacked files (will be unpacked from super.img below)
+    system_a_img = os.path.join(unpack_dir.name, 'system_a.img')
+    vendor_a_img = os.path.join(unpack_dir.name, 'vendor_a.img')
 
-    # Key(pubkey) for bootloader should match with the one used to make VBmeta below
+    # Key(pubkey) embedded in bootloader should match with the one used to make VBmeta below
     # while it's okay to use different keys for other image files.
-    ReplaceBootloaderPubkey(args, key, bootloader, bootloader_pubkey)
+    replace_f = Async(ReplaceBootloaderPubkey, args,
+                      key, files['bootloader'], files['bootloader.pubkey'])
 
     # re-sign bootloader, boot.img, vendor_boot.img, and init_boot.img
-    AddHashFooter(args, key, bootloader)
-    AddHashFooter(args, key, boot_img)
-    AddHashFooter(args, key, vendor_boot_img)
-    AddHashFooter(args, key, init_boot_img)
+    Async(AddHashFooter, args, key, files['bootloader'], wait=[replace_f])
+    boot_img_f = Async(AddHashFooter, args, key, files['boot.img'])
+    vendor_boot_img_f = Async(AddHashFooter, args, key, files['vendor_boot.img'])
+    init_boot_img_f = Async(AddHashFooter, args, key, files['init_boot.img'])
 
     # re-sign super.img
-    with TempDirectory() as work_dir:
-        # unpack super.img
-        tmp_super_img = os.path.join(work_dir, 'super.img')
-        RunCommand(args, ['simg2img', super_img, tmp_super_img])
-        RunCommand(args, ['lpunpack', tmp_super_img, work_dir])
+    super_img_f = Async(SignSuperImg, args, key, files['super.img'], unpack_dir.name)
 
-        system_a_img = os.path.join(work_dir, 'system_a.img')
-        vendor_a_img = os.path.join(work_dir, 'vendor_a.img')
-        partitions = {"system_a": system_a_img, "vendor_a": vendor_a_img}
-
-        # re-sign partitions in super.img
-        for img in partitions.values():
-            AddHashTreeFooter(args, key, img)
-
-        # re-pack super.img
-        MakeSuperImage(args, partitions, super_img)
-
-        # re-generate vbmeta from re-signed {boot, vendor_boot, init_boot, system_a, vendor_a}.img
-        # Ideally, making VBmeta should be done out of TempDirectory block. But doing it here
-        # to avoid unpacking re-signed super.img for system/vendor images which are available
-        # in this block.
-        MakeVbmetaImage(args, key, vbmeta_img, images=[
-                        boot_img, vendor_boot_img, init_boot_img, system_a_img, vendor_a_img])
+    # re-generate vbmeta from re-signed {boot, vendor_boot, init_boot, system_a, vendor_a}.img
+    Async(MakeVbmetaImage, args, key, files['vbmeta.img'],
+          images=[files['boot.img'], files['vendor_boot.img'],
+                  files['init_boot.img'], system_a_img, vendor_a_img],
+          wait=[boot_img_f, vendor_boot_img_f, init_boot_img_f, super_img_f])
 
     # Re-sign bootconfigs and the uboot_env with the same key
     bootconfig_sign_key = key
-    AddHashFooter(args, bootconfig_sign_key, bootconfig_normal)
-    AddHashFooter(args, bootconfig_sign_key, bootconfig_app_debuggable)
-    AddHashFooter(args, bootconfig_sign_key, bootconfig_full_debuggable)
-    AddHashFooter(args, bootconfig_sign_key, uboot_env_img)
+    Async(AddHashFooter, args, bootconfig_sign_key, files['bootconfig.normal'])
+    Async(AddHashFooter, args, bootconfig_sign_key, files['bootconfig.app_debuggable'])
+    Async(AddHashFooter, args, bootconfig_sign_key, files['bootconfig.full_debuggable'])
+    Async(AddHashFooter, args, bootconfig_sign_key, files['uboot_env.img'])
 
     # Re-sign vbmeta_bootconfig with chained_partitions to "bootconfig" and
     # "uboot_env". Note that, for now, `key` and `bootconfig_sign_key` are the
     # same, but technically they can be different. Vbmeta records pubkeys which
     # signed chained partitions.
-    MakeVbmetaImage(args, key, vbmeta_bootconfig_img, chained_partitions={
-                    'bootconfig': bootconfig_sign_key,
-                    'uboot_env': bootconfig_sign_key,
+    Async(MakeVbmetaImage, args, key, files['vbmeta_bootconfig.img'], chained_partitions={
+        'bootconfig': bootconfig_sign_key,
+        'uboot_env': bootconfig_sign_key,
     })
 
 
 def VerifyVirtApex(args):
-    # Generator to emit avbtool-signed items along with its pubkey digest.
-    # This supports lpmake-packed images as well.
-    def Recur(target_dir):
-        for file in glob.glob(os.path.join(target_dir, 'etc', '**', '*'), recursive=True):
-            cur_item = os.path.relpath(file, target_dir)
+    key = args.key
+    input_dir = args.input_dir
+    files = TargetFiles(input_dir)
 
-            if not os.path.isfile(file):
-                continue
+    # unpacked files
+    UnpackSuperImg(args, files['super.img'], unpack_dir.name)
+    system_a_img = os.path.join(unpack_dir.name, 'system_a.img')
+    vendor_a_img = os.path.join(unpack_dir.name, 'vendor_a.img')
 
-            # avbpubkey
-            if cur_item == 'etc/microdroid_bootloader.avbpubkey':
-                with open(file, 'rb') as f:
-                    yield (cur_item, hashlib.sha1(f.read()).hexdigest())
-                continue
+    # Read pubkey digest from the input key
+    with tempfile.NamedTemporaryFile() as pubkey_file:
+        ExtractAvbPubkey(args, key, pubkey_file.name)
+        with open(pubkey_file.name, 'rb') as f:
+            pubkey = f.read()
+            pubkey_digest = hashlib.sha1(pubkey).hexdigest()
 
-            # avbtool signed
-            info, _ = AvbInfo(args, file)
-            if info:
-                yield (cur_item, info['Public key (sha1)'])
-                continue
+    def contents(file):
+        with open(file, 'rb') as f:
+            return f.read()
 
-            # logical partition
-            with TempDirectory() as tmp_dir:
-                unsparsed = os.path.join(tmp_dir, os.path.basename(file))
-                _, rc = RunCommand(
-                    # exit with 255 if it's not sparsed
-                    args, ['simg2img', file, unsparsed], expected_return_values={0, 255})
-                if rc == 0:
-                    with TempDirectory() as unpack_dir:
-                        # exit with 64 if it's not a logical partition.
-                        _, rc = RunCommand(
-                            args, ['lpunpack', unsparsed, unpack_dir], expected_return_values={0, 64})
-                        if rc == 0:
-                            nested_items = list(Recur(unpack_dir))
-                            if len(nested_items) > 0:
-                                for (item, key) in nested_items:
-                                    yield ('%s!/%s' % (cur_item, item), key)
-                                continue
-    # Read pubkey digest
-    with TempDirectory() as tmp_dir:
-        pubkey_file = os.path.join(tmp_dir, 'avbpubkey')
-        ExtractAvbPubkey(args, args.key, pubkey_file)
-        with open(pubkey_file, 'rb') as f:
-            pubkey_digest = hashlib.sha1(f.read()).hexdigest()
+    def check_equals_pubkey(file):
+        assert contents(file) == pubkey, 'pubkey mismatch: %s' % file
 
-    # Check every avbtool-signed item against the input key
-    for (item, pubkey) in Recur(args.input_dir):
-        assert pubkey == pubkey_digest, '%s: key mismatch: %s != %s' % (
-            item, pubkey, pubkey_digest)
+    def check_contains_pubkey(file):
+        assert contents(file).find(pubkey) != -1, 'pubkey missing: %s' % file
+
+    def check_avb_pubkey(file):
+        info, _ = AvbInfo(args, file)
+        assert info is not None, 'no avbinfo: %s' % file
+        assert info['Public key (sha1)'] == pubkey_digest, 'pubkey mismatch: %s' % file
+
+    for f in files.values():
+        if f == files['bootloader.pubkey']:
+            Async(check_equals_pubkey, f)
+        elif f == files['bootloader']:
+            Async(check_contains_pubkey, f)
+        elif f == files['super.img']:
+            Async(check_avb_pubkey, system_a_img)
+            Async(check_avb_pubkey, vendor_a_img)
+        else:
+            # Check pubkey for other files using avbtool
+            Async(check_avb_pubkey, f)
 
 
 def main(argv):
@@ -442,8 +478,10 @@
             VerifyVirtApex(args)
         else:
             SignVirtApex(args)
-    except Exception as e:
-        print(e)
+        # ensure all tasks are completed without exceptions
+        AwaitAll(tasks)
+    except: # pylint: disable=bare-except
+        traceback.print_exc()
         sys.exit(1)
 
 
diff --git a/apex/virtualizationservice.rc b/apex/virtualizationservice.rc
index 7e71105..02b2081 100644
--- a/apex/virtualizationservice.rc
+++ b/apex/virtualizationservice.rc
@@ -14,8 +14,8 @@
 
 service virtualizationservice /apex/com.android.virt/bin/virtualizationservice
     class main
-    user virtualizationservice
-    group virtualizationservice
+    user system
+    group system
     interface aidl android.system.virtualizationservice
     disabled
     oneshot
diff --git a/apkdmverity/TEST_MAPPING b/apkdmverity/TEST_MAPPING
index 1bbec76..997b3f9 100644
--- a/apkdmverity/TEST_MAPPING
+++ b/apkdmverity/TEST_MAPPING
@@ -1,5 +1,5 @@
 {
-  "presubmit" : [
+  "postsubmit" : [
     {
       "name" : "apkdmverity.test"
     }
diff --git a/apkdmverity/src/loopdevice.rs b/apkdmverity/src/loopdevice.rs
index 376abd4..35ae154 100644
--- a/apkdmverity/src/loopdevice.rs
+++ b/apkdmverity/src/loopdevice.rs
@@ -25,8 +25,10 @@
 
 use anyhow::{Context, Result};
 use data_model::DataInit;
+use libc::O_DIRECT;
 use std::fs::{File, OpenOptions};
 use std::mem::size_of;
+use std::os::unix::fs::OpenOptionsExt;
 use std::os::unix::io::AsRawFd;
 use std::path::{Path, PathBuf};
 use std::thread;
@@ -37,8 +39,7 @@
 
 // These are old-style ioctls, thus *_bad.
 nix::ioctl_none_bad!(_loop_ctl_get_free, LOOP_CTL_GET_FREE);
-nix::ioctl_write_int_bad!(_loop_set_fd, LOOP_SET_FD);
-nix::ioctl_write_ptr_bad!(_loop_set_status64, LOOP_SET_STATUS64, loop_info64);
+nix::ioctl_write_ptr_bad!(_loop_configure, LOOP_CONFIGURE, loop_config);
 #[cfg(test)]
 nix::ioctl_none_bad!(_loop_clr_fd, LOOP_CLR_FD);
 
@@ -49,14 +50,9 @@
     Ok(unsafe { _loop_ctl_get_free(ctrl_file.as_raw_fd()) }?)
 }
 
-fn loop_set_fd(device_file: &File, fd: i32) -> Result<i32> {
+fn loop_configure(device_file: &File, config: &loop_config) -> Result<i32> {
     // SAFETY: this ioctl changes the state in kernel, but not the state in this process.
-    Ok(unsafe { _loop_set_fd(device_file.as_raw_fd(), fd) }?)
-}
-
-fn loop_set_status64(device_file: &File, info: &loop_info64) -> Result<i32> {
-    // SAFETY: this ioctl changes the state in kernel, but not the state in this process.
-    Ok(unsafe { _loop_set_status64(device_file.as_raw_fd(), info) }?)
+    Ok(unsafe { _loop_configure(device_file.as_raw_fd(), config) }?)
 }
 
 #[cfg(test)]
@@ -67,7 +63,12 @@
 }
 
 /// Creates a loop device and attach the given file at `path` as the backing store.
-pub fn attach<P: AsRef<Path>>(path: P, offset: u64, size_limit: u64) -> Result<PathBuf> {
+pub fn attach<P: AsRef<Path>>(
+    path: P,
+    offset: u64,
+    size_limit: u64,
+    direct_io: bool,
+) -> Result<PathBuf> {
     // Attaching a file to a loop device can make a race condition; a loop device number obtained
     // from LOOP_CTL_GET_FREE might have been used by another thread or process. In that case the
     // subsequet LOOP_CONFIGURE ioctl returns with EBUSY. Try until it succeeds.
@@ -81,7 +82,7 @@
 
     let begin = Instant::now();
     loop {
-        match try_attach(&path, offset, size_limit) {
+        match try_attach(&path, offset, size_limit, direct_io) {
             Ok(loop_dev) => return Ok(loop_dev),
             Err(e) => {
                 if begin.elapsed() > TIMEOUT {
@@ -99,7 +100,12 @@
 #[cfg(target_os = "android")]
 const LOOP_DEV_PREFIX: &str = "/dev/block/loop";
 
-fn try_attach<P: AsRef<Path>>(path: P, offset: u64, size_limit: u64) -> Result<PathBuf> {
+fn try_attach<P: AsRef<Path>>(
+    path: P,
+    offset: u64,
+    size_limit: u64,
+    direct_io: bool,
+) -> Result<PathBuf> {
     // Get a free loop device
     wait_for_path(LOOP_CONTROL)?;
     let ctrl_file = OpenOptions::new()
@@ -112,21 +118,19 @@
     // Construct the loop_info64 struct
     let backing_file = OpenOptions::new()
         .read(true)
+        .custom_flags(if direct_io { O_DIRECT } else { 0 })
         .open(&path)
         .context(format!("failed to open {:?}", path.as_ref()))?;
     // safe because the size of the array is the same as the size of the struct
-    let mut info: loop_info64 =
-        *DataInit::from_mut_slice(&mut [0; size_of::<loop_info64>()]).unwrap();
-    info.lo_offset = offset;
-    info.lo_sizelimit = size_limit;
-    info.lo_flags |= Flag::LO_FLAGS_DIRECT_IO | Flag::LO_FLAGS_READ_ONLY;
-
-    // Special case: don't use direct IO when the backing file is already a loop device, which
-    // happens only during test. DirectIO-on-loop-over-loop makes the outer loop device
-    // unaccessible.
-    #[cfg(test)]
-    if path.as_ref().to_str().unwrap().starts_with(LOOP_DEV_PREFIX) {
-        info.lo_flags.remove(Flag::LO_FLAGS_DIRECT_IO);
+    let mut config: loop_config =
+        *DataInit::from_mut_slice(&mut [0; size_of::<loop_config>()]).unwrap();
+    config.fd = backing_file.as_raw_fd() as u32;
+    config.block_size = 4096;
+    config.info.lo_offset = offset;
+    config.info.lo_sizelimit = size_limit;
+    config.info.lo_flags = Flag::LO_FLAGS_READ_ONLY;
+    if direct_io {
+        config.info.lo_flags.insert(Flag::LO_FLAGS_DIRECT_IO);
     }
 
     // Configure the loop device to attach the backing file
@@ -137,8 +141,8 @@
         .write(true)
         .open(&device_path)
         .context(format!("failed to open {:?}", &device_path))?;
-    loop_set_fd(&device_file, backing_file.as_raw_fd() as i32)?;
-    loop_set_status64(&device_file, &info)?;
+    loop_configure(&device_file, &config)
+        .context(format!("Failed to configure {:?}", &device_path))?;
 
     Ok(PathBuf::from(device_path))
 }
@@ -150,3 +154,46 @@
     loop_clr_fd(&device_file)?;
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs;
+    use std::path::Path;
+
+    fn create_empty_file(path: &Path, size: u64) {
+        let f = File::create(path).unwrap();
+        f.set_len(size).unwrap();
+    }
+
+    fn is_direct_io(dev: &Path) -> bool {
+        let dio = Path::new("/sys/block").join(dev.file_name().unwrap()).join("loop/dio");
+        "1" == fs::read_to_string(&dio).unwrap().trim()
+    }
+
+    #[test]
+    fn attach_loop_device_with_dio() {
+        let a_dir = tempfile::TempDir::new().unwrap();
+        let a_file = a_dir.path().join("test");
+        let a_size = 4096u64;
+        create_empty_file(&a_file, a_size);
+        let dev = attach(a_file, 0, a_size, /*direct_io*/ true).unwrap();
+        scopeguard::defer! {
+            detach(&dev).unwrap();
+        }
+        assert!(is_direct_io(&dev));
+    }
+
+    #[test]
+    fn attach_loop_device_without_dio() {
+        let a_dir = tempfile::TempDir::new().unwrap();
+        let a_file = a_dir.path().join("test");
+        let a_size = 4096u64;
+        create_empty_file(&a_file, a_size);
+        let dev = attach(a_file, 0, a_size, /*direct_io*/ false).unwrap();
+        scopeguard::defer! {
+            detach(&dev).unwrap();
+        }
+        assert!(!is_direct_io(&dev));
+    }
+}
diff --git a/apkdmverity/src/loopdevice/sys.rs b/apkdmverity/src/loopdevice/sys.rs
index d32987a..fa87548 100644
--- a/apkdmverity/src/loopdevice/sys.rs
+++ b/apkdmverity/src/loopdevice/sys.rs
@@ -24,8 +24,7 @@
 pub const LOOP_CONTROL: &str = "/dev/loop-control";
 
 pub const LOOP_CTL_GET_FREE: libc::c_ulong = 0x4C82;
-pub const LOOP_SET_FD: libc::c_ulong = 0x4C00;
-pub const LOOP_SET_STATUS64: libc::c_ulong = 0x4C04;
+pub const LOOP_CONFIGURE: libc::c_ulong = 0x4C0A;
 #[cfg(test)]
 pub const LOOP_CLR_FD: libc::c_ulong = 0x4C01;
 
diff --git a/apkdmverity/src/main.rs b/apkdmverity/src/main.rs
index dbf3131..16dd480 100644
--- a/apkdmverity/src/main.rs
+++ b/apkdmverity/src/main.rs
@@ -94,7 +94,11 @@
         if apk_size % BLOCK_SIZE != 0 {
             bail!("The size of {:?} is not multiple of {}.", &apk, BLOCK_SIZE)
         }
-        (loopdevice::attach(&apk, 0, apk_size)?, apk_size)
+        (
+            loopdevice::attach(&apk, 0, apk_size, /*direct_io*/ true)
+                .context("Failed to attach APK to a loop device")?,
+            apk_size,
+        )
     };
 
     // Parse the idsig file to locate the merkle tree in it, then attach the file to a loop device
@@ -105,7 +109,11 @@
     )?;
     let offset = sig.merkle_tree_offset;
     let size = sig.merkle_tree_size as u64;
-    let hash_device = loopdevice::attach(&idsig, offset, size)?;
+    // Due to unknown reason(b/191344832), we can't enable "direct IO" for the IDSIG file (backing
+    // the hash). For now we don't use "direct IO" but it seems OK since the IDSIG file is very
+    // small and the benefit of direct-IO would be negliable.
+    let hash_device = loopdevice::attach(&idsig, offset, size, /*direct_io*/ false)
+        .context("Failed to attach idsig to a loop device")?;
 
     // Build a dm-verity target spec from the information from the idsig file. The apk and the
     // idsig files are used as the data device and the hash device, respectively.
@@ -318,11 +326,11 @@
         // already a block device, `enable_verity` uses the block device as it is. The detatching
         // of the data device is done in the scopeguard for the return value of `enable_verity`
         // below. Only the idsig_loop_device needs detatching.
-        let apk_loop_device = loopdevice::attach(&apk_path, 0, apk_size).unwrap();
-        let idsig_loop_device =
-            scopeguard::guard(loopdevice::attach(&idsig_path, 0, idsig_size).unwrap(), |dev| {
-                loopdevice::detach(dev).unwrap()
-            });
+        let apk_loop_device = loopdevice::attach(&apk_path, 0, apk_size, true).unwrap();
+        let idsig_loop_device = scopeguard::guard(
+            loopdevice::attach(&idsig_path, 0, idsig_size, false).unwrap(),
+            |dev| loopdevice::detach(dev).unwrap(),
+        );
 
         let name = "loop_as_input";
         // Run the program WITH the loop devices, not the regular files.
diff --git a/authfs/src/file/remote_file.rs b/authfs/src/file/remote_file.rs
index 039285f..4c112bd 100644
--- a/authfs/src/file/remote_file.rs
+++ b/authfs/src/file/remote_file.rs
@@ -40,7 +40,6 @@
 }
 
 pub struct RemoteFileReader {
-    // This needs to have Sync trait to be used in fuse::worker::start_message_loop.
     service: VirtFdService,
     file_fd: i32,
 }
@@ -81,7 +80,6 @@
 }
 
 pub struct RemoteMerkleTreeReader {
-    // This needs to be a Sync to be used in fuse::worker::start_message_loop.
     service: VirtFdService,
     file_fd: i32,
 }
@@ -108,7 +106,6 @@
 }
 
 pub struct RemoteFileEditor {
-    // This needs to have Sync trait to be used in fuse::worker::start_message_loop.
     service: VirtFdService,
     file_fd: i32,
 }
diff --git a/authfs/src/fusefs.rs b/authfs/src/fusefs.rs
index 511db68..beb6b30 100644
--- a/authfs/src/fusefs.rs
+++ b/authfs/src/fusefs.rs
@@ -184,10 +184,9 @@
 
 type DirHandleTable = BTreeMap<Handle, Arc<DirEntriesSnapshot>>;
 
-// AuthFS needs to be `Sync` to be accepted by fuse::worker::start_message_loop as a `FileSystem`.
+// AuthFS needs to be `Sync` to be used with the `fuse` crate.
 pub struct AuthFs {
-    /// Table for `Inode` to `InodeState` lookup. This needs to be `Sync` to be used in
-    /// `fuse::worker::start_message_loop`.
+    /// Table for `Inode` to `InodeState` lookup.
     inode_table: RwLock<BTreeMap<Inode, InodeState>>,
 
     /// The next available inode number.
diff --git a/authfs/src/fusefs/mount.rs b/authfs/src/fusefs/mount.rs
index 294c6b1..7f8bac1 100644
--- a/authfs/src/fusefs/mount.rs
+++ b/authfs/src/fusefs/mount.rs
@@ -16,16 +16,18 @@
 
 use fuse::mount::MountOption;
 use std::fs::OpenOptions;
+use std::num::NonZeroU8;
 use std::os::unix::io::AsRawFd;
 use std::path::Path;
 
 use super::AuthFs;
 
-/// Maximum bytes in the write transaction to the FUSE device. This limits the maximum buffer
-/// size in a read request (including FUSE protocol overhead) that the filesystem writes to.
+/// Maximum bytes (excluding the FUSE header) `AuthFs` will receive from the kernel for write
+/// operations by another process.
 pub const MAX_WRITE_BYTES: u32 = 65536;
 
-/// Maximum bytes in a read operation.
+/// Maximum bytes (excluding the FUSE header) `AuthFs` will receive from the kernel for read
+/// operations by another process.
 /// TODO(victorhsieh): This option is deprecated by FUSE. Figure out if we can remove this.
 const MAX_READ_BYTES: u32 = 65536;
 
@@ -34,6 +36,7 @@
     authfs: AuthFs,
     mountpoint: &Path,
     extra_options: &Option<String>,
+    threads: Option<NonZeroU8>,
 ) -> Result<(), fuse::Error> {
     let dev_fuse = OpenOptions::new()
         .read(true)
@@ -61,5 +64,10 @@
     )
     .expect("Failed to mount fuse");
 
-    fuse::worker::start_message_loop(dev_fuse, MAX_WRITE_BYTES, MAX_READ_BYTES, authfs)
+    let mut config = fuse::FuseConfig::new();
+    config.dev_fuse(dev_fuse).max_write(MAX_WRITE_BYTES).max_read(MAX_READ_BYTES);
+    if let Some(num) = threads {
+        config.num_threads(u8::from(num).into());
+    }
+    config.enter_message_loop(authfs)
 }
diff --git a/authfs/src/main.rs b/authfs/src/main.rs
index bdca5b4..60318e8 100644
--- a/authfs/src/main.rs
+++ b/authfs/src/main.rs
@@ -37,6 +37,7 @@
 use protobuf::Message;
 use std::convert::TryInto;
 use std::fs::File;
+use std::num::NonZeroU8;
 use std::path::{Path, PathBuf};
 use structopt::StructOpt;
 
@@ -67,6 +68,10 @@
     #[structopt(short = "o")]
     extra_options: Option<String>,
 
+    /// Number of threads to serve FUSE requests.
+    #[structopt(short = "j")]
+    thread_number: Option<NonZeroU8>,
+
     /// A read-only remote file with integrity check. Can be multiple.
     ///
     /// For example, `--remote-ro-file 5:sha256-1234abcd` tells the filesystem to associate the
@@ -312,7 +317,12 @@
     let mut authfs = AuthFs::new(RemoteFsStatsReader::new(service.clone()));
     prepare_root_dir_entries(service, &mut authfs, &args)?;
 
-    fusefs::mount_and_enter_message_loop(authfs, &args.mount_point, &args.extra_options)?;
+    fusefs::mount_and_enter_message_loop(
+        authfs,
+        &args.mount_point,
+        &args.extra_options,
+        args.thread_number,
+    )?;
     bail!("Unexpected exit after the handler loop")
 }
 
diff --git a/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java b/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
index 5d36f16..64658a9 100644
--- a/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
+++ b/authfs/tests/java/src/com/android/fs/AuthFsHostTest.java
@@ -16,21 +16,28 @@
 
 package com.android.virt.fs;
 
+import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestLogData;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeFalse;
 
 import android.platform.test.annotations.RootPermissionTest;
 import android.virt.test.CommandRunner;
 import android.virt.test.VirtualizationTestCaseBase;
 
+import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
 import com.android.compatibility.common.util.PollingCheck;
+import com.android.tradefed.build.IBuildInfo;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.invoker.TestInformation;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
@@ -47,7 +54,8 @@
 import org.junit.rules.TestName;
 import org.junit.runner.RunWith;
 
-import java.util.Optional;
+import java.io.File;
+import java.io.FileNotFoundException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -62,6 +70,12 @@
     /** Output directory where the test can generate output on Android */
     private static final String TEST_OUTPUT_DIR = "/data/local/tmp/authfs/output_dir";
 
+    /** File name of the test APK */
+    private static final String TEST_APK_NAME = "MicrodroidTestApp.apk";
+
+    /** VM config entry path in the test APK */
+    private static final String VM_CONFIG_PATH_IN_APK = "assets/vm_config_extra_apk.json";
+
     /** Path to open_then_run on Android */
     private static final String OPEN_THEN_RUN_BIN = "/data/local/tmp/open_then_run";
 
@@ -75,9 +89,7 @@
     private static final String AUTHFS_BIN = "/system/bin/authfs";
 
     /** Idsig paths to be created for each APK in the "extra_apks" of vm_config_extra_apk.json. */
-    private static final String[] EXTRA_IDSIG_PATHS = new String[] {
-        TEST_DIR + "BuildManifest.apk.idsig",
-    };
+    private static final String EXTRA_IDSIG_PATH = TEST_DIR + "BuildManifest.apk.idsig";
 
     /** Build manifest path in the VM. 0 is the index of extra_apks in vm_config_extra_apk.json. */
     private static final String BUILD_MANIFEST_PATH = "/mnt/extra-apk/0/assets/build_manifest.pb";
@@ -99,7 +111,7 @@
     private static final int VMADDR_CID_HOST = 2;
 
     private static CommandRunner sAndroid;
-    private static String sCid;
+    private static CommandRunner sMicrodroid;
     private static boolean sAssumptionFailed;
 
     private ExecutorService mThreadPool = Executors.newCachedThreadPool();
@@ -127,30 +139,20 @@
             return;
         }
 
-        prepareVirtualizationTestSetup(androidDevice);
-
         // For each test case, boot and adb connect to a new Microdroid
         CLog.i("Starting the shared VM");
-        final String apkName = "MicrodroidTestApp.apk";
-        final String packageName = "com.android.microdroid.test";
-        final String configPath = "assets/vm_config_extra_apk.json"; // path inside the APK
-        sCid =
-                startMicrodroid(
-                        androidDevice,
-                        testInfo.getBuildInfo(),
-                        apkName,
-                        packageName,
-                        EXTRA_IDSIG_PATHS,
-                        configPath,
-                        /* debug */ true,
-                        /* use default memoryMib */ 0,
-                        Optional.empty(),
-                        Optional.empty());
-        adbConnectToMicrodroid(androidDevice, sCid);
+        ITestDevice microdroidDevice =
+                MicrodroidBuilder
+                        .fromFile(findTestApk(testInfo.getBuildInfo()), VM_CONFIG_PATH_IN_APK)
+                        .debugLevel("full")
+                        .addExtraIdsigPath(EXTRA_IDSIG_PATH)
+                        .build((TestDevice) androidDevice);
 
         // Root because authfs (started from shell in this test) currently require root to open
         // /dev/fuse and mount the FUSE.
-        rootMicrodroid();
+        assertThat(microdroidDevice.enableAdbRoot()).isTrue();
+
+        sMicrodroid = new CommandRunner(microdroidDevice);
     }
 
     @AfterClassWithInfo
@@ -158,13 +160,12 @@
             throws DeviceNotAvailableException {
         assertNotNull(sAndroid);
 
-        if (sCid != null) {
+        if (sMicrodroid != null) {
             CLog.i("Shutting down shared VM");
-            shutdownMicrodroid(sAndroid.getDevice(), sCid);
-            sCid = null;
+            ((TestDevice) testInfo.getDevice()).shutdownMicrodroid(sMicrodroid.getDevice());
+            sMicrodroid = null;
         }
 
-        cleanUpVirtualizationTestSetup(sAndroid.getDevice());
         sAndroid = null;
     }
 
@@ -176,9 +177,13 @@
 
     @After
     public void tearDown() throws Exception {
+        if (sMicrodroid != null) {
+            sMicrodroid.tryRun("killall authfs");
+            sMicrodroid.tryRun("umount " + MOUNT_DIR);
+        }
+
+        assertNotNull(sAndroid);
         sAndroid.tryRun("killall fd_server");
-        tryRunOnMicrodroid("killall authfs");
-        tryRunOnMicrodroid("umount " + MOUNT_DIR);
 
         // Even though we only run one VM for the whole class, and could have collect the VM log
         // after all tests are done, TestLogData doesn't seem to work at class level. Hence,
@@ -319,7 +324,7 @@
 
         // Verify
         // Force dropping the page cache, so that the next read can be validated.
-        runOnMicrodroid("echo 1 > /proc/sys/vm/drop_caches");
+        sMicrodroid.run("echo 1 > /proc/sys/vm/drop_caches");
         // A read will fail if the backing data has been tampered.
         assertFalse(checkReadAtFileOffsetOnMicrodroid(
                 destPath, /* offset */ 0, /* number */ 4096));
@@ -419,8 +424,8 @@
 
         // Action
         // Can create nested directories and can create a file in one.
-        runOnMicrodroid("mkdir " + authfsOutputDir + "/new_dir");
-        runOnMicrodroid("mkdir -p " + authfsOutputDir + "/we/need/to/go/deeper");
+        sMicrodroid.run("mkdir " + authfsOutputDir + "/new_dir");
+        sMicrodroid.run("mkdir -p " + authfsOutputDir + "/we/need/to/go/deeper");
         createFileWithOnesOnMicrodroid(authfsOutputDir + "/new_dir/file1", 10000);
         createFileWithOnesOnMicrodroid(authfsOutputDir + "/we/need/file2", 10000);
 
@@ -452,7 +457,7 @@
         runAuthFsOnMicrodroid("--remote-new-rw-dir 3 --cid " + VMADDR_CID_HOST);
 
         // Action & Verify
-        runOnMicrodroid("echo -n foo > " + authfsOutputDir + "/file");
+        sMicrodroid.run("echo -n foo > " + authfsOutputDir + "/file");
         assertEquals(getFileSizeInBytesOnMicrodroid(authfsOutputDir + "/file"), 3);
         // Can override a file and write normally.
         createFileWithOnesOnMicrodroid(authfsOutputDir + "/file", 10000);
@@ -472,13 +477,13 @@
         runFdServerOnAndroid("--open-dir 3:" + androidOutputDir, "--rw-dirs 3");
         runAuthFsOnMicrodroid("--remote-new-rw-dir 3 --cid " + VMADDR_CID_HOST);
 
-        runOnMicrodroid("echo -n foo > " + authfsOutputDir + "/file");
-        runOnMicrodroid("test -f " + authfsOutputDir + "/file");
+        sMicrodroid.run("echo -n foo > " + authfsOutputDir + "/file");
+        sMicrodroid.run("test -f " + authfsOutputDir + "/file");
         sAndroid.run("test -f " + androidOutputDir + "/file");
 
         // Action & Verify
-        runOnMicrodroid("rm " + authfsOutputDir + "/file");
-        runOnMicrodroid("test ! -f " + authfsOutputDir + "/file");
+        sMicrodroid.run("rm " + authfsOutputDir + "/file");
+        sMicrodroid.run("test ! -f " + authfsOutputDir + "/file");
         sAndroid.run("test ! -f " + androidOutputDir + "/file");
     }
 
@@ -491,20 +496,20 @@
         runFdServerOnAndroid("--open-dir 3:" + androidOutputDir, "--rw-dirs 3");
         runAuthFsOnMicrodroid("--remote-new-rw-dir 3 --cid " + VMADDR_CID_HOST);
 
-        runOnMicrodroid("mkdir -p " + authfsOutputDir + "/dir/dir2");
-        runOnMicrodroid("echo -n foo > " + authfsOutputDir + "/dir/file");
+        sMicrodroid.run("mkdir -p " + authfsOutputDir + "/dir/dir2");
+        sMicrodroid.run("echo -n foo > " + authfsOutputDir + "/dir/file");
         sAndroid.run("test -d " + androidOutputDir + "/dir/dir2");
 
         // Action & Verify
-        runOnMicrodroid("rmdir " + authfsOutputDir + "/dir/dir2");
-        runOnMicrodroid("test ! -d " + authfsOutputDir + "/dir/dir2");
+        sMicrodroid.run("rmdir " + authfsOutputDir + "/dir/dir2");
+        sMicrodroid.run("test ! -d " + authfsOutputDir + "/dir/dir2");
         sAndroid.run("test ! -d " + androidOutputDir + "/dir/dir2");
         // Can only delete a directory if empty
         assertFailedOnMicrodroid("rmdir " + authfsOutputDir + "/dir");
-        runOnMicrodroid("test -d " + authfsOutputDir + "/dir");  // still there
-        runOnMicrodroid("rm " + authfsOutputDir + "/dir/file");
-        runOnMicrodroid("rmdir " + authfsOutputDir + "/dir");
-        runOnMicrodroid("test ! -d " + authfsOutputDir + "/dir");
+        sMicrodroid.run("test -d " + authfsOutputDir + "/dir");  // still there
+        sMicrodroid.run("rm " + authfsOutputDir + "/dir/file");
+        sMicrodroid.run("rmdir " + authfsOutputDir + "/dir");
+        sMicrodroid.run("test ! -d " + authfsOutputDir + "/dir");
         sAndroid.run("test ! -d " + androidOutputDir + "/dir");
     }
 
@@ -517,10 +522,10 @@
         runFdServerOnAndroid("--open-dir 3:" + androidOutputDir, "--rw-dirs 3");
         runAuthFsOnMicrodroid("--remote-new-rw-dir 3 --cid " + VMADDR_CID_HOST);
 
-        runOnMicrodroid("touch " + authfsOutputDir + "/some_file");
-        runOnMicrodroid("mkdir " + authfsOutputDir + "/some_dir");
-        runOnMicrodroid("touch " + authfsOutputDir + "/some_dir/file");
-        runOnMicrodroid("mkdir " + authfsOutputDir + "/some_dir/dir");
+        sMicrodroid.run("touch " + authfsOutputDir + "/some_file");
+        sMicrodroid.run("mkdir " + authfsOutputDir + "/some_dir");
+        sMicrodroid.run("touch " + authfsOutputDir + "/some_dir/file");
+        sMicrodroid.run("mkdir " + authfsOutputDir + "/some_dir/dir");
 
         // Action & Verify
         // Cannot create directory if an entry with the same name already exists.
@@ -542,12 +547,12 @@
         // Create a file with some data. Test the existence.
         String outputPath = authfsOutputDir + "/out";
         String androidOutputPath = androidOutputDir + "/out";
-        runOnMicrodroid("echo -n 123 > " + outputPath);
-        runOnMicrodroid("test -f " + outputPath);
+        sMicrodroid.run("echo -n 123 > " + outputPath);
+        sMicrodroid.run("test -f " + outputPath);
         sAndroid.run("test -f " + androidOutputPath);
 
         // Action
-        String output = runOnMicrodroid(
+        String output = sMicrodroid.run(
                 // Open the file for append and read
                 "exec 4>>" + outputPath + " 5<" + outputPath + "; "
                 // Delete the file from the directory
@@ -560,7 +565,7 @@
         // Verify
         // Output contains all written data, while the files are deleted.
         assertEquals("123456", output);
-        runOnMicrodroid("test ! -f " + outputPath);
+        sMicrodroid.run("test ! -f " + outputPath);
         sAndroid.run("test ! -f " + androidOutputDir + "/out");
     }
 
@@ -590,7 +595,7 @@
                 + VMADDR_CID_HOST);
 
         // Verify
-        runOnMicrodroid("test -f " + authfsInputDir + "/system/framework/services.jar");
+        sMicrodroid.run("test -f " + authfsInputDir + "/system/framework/services.jar");
         assertFailedOnMicrodroid("test -f " + authfsInputDir + "/system/bin/sh");
     }
 
@@ -602,14 +607,14 @@
 
         // Action
         String authfsOutputDir = MOUNT_DIR + "/3";
-        runOnMicrodroid("mkdir -p " + authfsOutputDir + "/dir/dir2/dir3");
-        runOnMicrodroid("touch " + authfsOutputDir + "/dir/dir2/dir3/file1");
-        runOnMicrodroid("touch " + authfsOutputDir + "/dir/dir2/dir3/file2");
-        runOnMicrodroid("touch " + authfsOutputDir + "/dir/dir2/dir3/file3");
-        runOnMicrodroid("touch " + authfsOutputDir + "/file");
+        sMicrodroid.run("mkdir -p " + authfsOutputDir + "/dir/dir2/dir3");
+        sMicrodroid.run("touch " + authfsOutputDir + "/dir/dir2/dir3/file1");
+        sMicrodroid.run("touch " + authfsOutputDir + "/dir/dir2/dir3/file2");
+        sMicrodroid.run("touch " + authfsOutputDir + "/dir/dir2/dir3/file3");
+        sMicrodroid.run("touch " + authfsOutputDir + "/file");
 
         // Verify
-        String[] actual = runOnMicrodroid("cd " + authfsOutputDir + "; find |sort").split("\n");
+        String[] actual = sMicrodroid.run("cd " + authfsOutputDir + "; find |sort").split("\n");
         String[] expected = new String[] {
                 ".",
                 "./dir",
@@ -622,14 +627,14 @@
         assertEquals(expected, actual);
 
         // Add more entries.
-        runOnMicrodroid("mkdir -p " + authfsOutputDir + "/dir2");
-        runOnMicrodroid("touch " + authfsOutputDir + "/file2");
+        sMicrodroid.run("mkdir -p " + authfsOutputDir + "/dir2");
+        sMicrodroid.run("touch " + authfsOutputDir + "/file2");
         // Check new entries. Also check that the types are correct.
-        actual = runOnMicrodroid(
+        actual = sMicrodroid.run(
                 "cd " + authfsOutputDir + "; find -maxdepth 1 -type f |sort").split("\n");
         expected = new String[] {"./file", "./file2"};
         assertEquals(expected, actual);
-        actual = runOnMicrodroid(
+        actual = sMicrodroid.run(
                 "cd " + authfsOutputDir + "; find -maxdepth 1 -type d |sort").split("\n");
         expected = new String[] {".", "./dir", "./dir2"};
         assertEquals(expected, actual);
@@ -643,7 +648,7 @@
 
         // Action & Verify
         // Change mode
-        runOnMicrodroid("chmod 321 " + MOUNT_DIR + "/3");
+        sMicrodroid.run("chmod 321 " + MOUNT_DIR + "/3");
         expectFileMode("--wx-w---x", MOUNT_DIR + "/3", TEST_OUTPUT_DIR + "/file");
         // Can't set the disallowed bits
         assertFailedOnMicrodroid("chmod +s " + MOUNT_DIR + "/3");
@@ -659,14 +664,14 @@
         // Action & Verify
         String authfsOutputDir = MOUNT_DIR + "/3";
         // Create with umask
-        runOnMicrodroid("umask 000; mkdir " + authfsOutputDir + "/dir");
-        runOnMicrodroid("umask 022; mkdir " + authfsOutputDir + "/dir/dir2");
+        sMicrodroid.run("umask 000; mkdir " + authfsOutputDir + "/dir");
+        sMicrodroid.run("umask 022; mkdir " + authfsOutputDir + "/dir/dir2");
         expectFileMode("drwxrwxrwx", authfsOutputDir + "/dir", TEST_OUTPUT_DIR + "/dir");
         expectFileMode("drwxr-xr-x", authfsOutputDir + "/dir/dir2", TEST_OUTPUT_DIR + "/dir/dir2");
         // Change mode
-        runOnMicrodroid("chmod -w " + authfsOutputDir + "/dir/dir2");
+        sMicrodroid.run("chmod -w " + authfsOutputDir + "/dir/dir2");
         expectFileMode("dr-xr-xr-x", authfsOutputDir + "/dir/dir2", TEST_OUTPUT_DIR + "/dir/dir2");
-        runOnMicrodroid("chmod 321 " + authfsOutputDir + "/dir");
+        sMicrodroid.run("chmod 321 " + authfsOutputDir + "/dir");
         expectFileMode("d-wx-w---x", authfsOutputDir + "/dir", TEST_OUTPUT_DIR + "/dir");
         // Can't set the disallowed bits
         assertFailedOnMicrodroid("chmod +s " + authfsOutputDir + "/dir/dir2");
@@ -682,14 +687,14 @@
         // Action & Verify
         String authfsOutputDir = MOUNT_DIR + "/3";
         // Create with umask
-        runOnMicrodroid("umask 000; echo -n foo > " + authfsOutputDir + "/file");
-        runOnMicrodroid("umask 022; echo -n foo > " + authfsOutputDir + "/file2");
+        sMicrodroid.run("umask 000; echo -n foo > " + authfsOutputDir + "/file");
+        sMicrodroid.run("umask 022; echo -n foo > " + authfsOutputDir + "/file2");
         expectFileMode("-rw-rw-rw-", authfsOutputDir + "/file", TEST_OUTPUT_DIR + "/file");
         expectFileMode("-rw-r--r--", authfsOutputDir + "/file2", TEST_OUTPUT_DIR + "/file2");
         // Change mode
-        runOnMicrodroid("chmod -w " + authfsOutputDir + "/file");
+        sMicrodroid.run("chmod -w " + authfsOutputDir + "/file");
         expectFileMode("-r--r--r--", authfsOutputDir + "/file", TEST_OUTPUT_DIR + "/file");
-        runOnMicrodroid("chmod 321 " + authfsOutputDir + "/file2");
+        sMicrodroid.run("chmod 321 " + authfsOutputDir + "/file2");
         expectFileMode("--wx-w---x", authfsOutputDir + "/file2", TEST_OUTPUT_DIR + "/file2");
         // Can't set the disallowed bits
         assertFailedOnMicrodroid("chmod +s " + authfsOutputDir + "/file");
@@ -705,7 +710,16 @@
         // Verify
         // Magic matches. Has only 2 inodes (root and "/3").
         assertEquals(
-                FUSE_SUPER_MAGIC_HEX + " 2", runOnMicrodroid("stat -f -c '%t %c' " + MOUNT_DIR));
+                FUSE_SUPER_MAGIC_HEX + " 2", sMicrodroid.run("stat -f -c '%t %c' " + MOUNT_DIR));
+    }
+
+    private static File findTestApk(IBuildInfo buildInfo) {
+        try {
+            return (new CompatibilityBuildHelper(buildInfo)).getTestFile(TEST_APK_NAME);
+        } catch (FileNotFoundException e) {
+            fail("Missing test file: " + TEST_APK_NAME);
+            return null;
+        }
     }
 
     private void expectBackingFileConsistency(
@@ -719,8 +733,8 @@
                 "Inconsistent file hash on the backend storage", hashOnAuthFs, hashOfBackingFile);
     }
 
-    private String computeFileHashOnMicrodroid(String path) {
-        String result = runOnMicrodroid("sha256sum " + path);
+    private String computeFileHashOnMicrodroid(String path) throws DeviceNotAvailableException {
+        String result = sMicrodroid.run("sha256sum " + path);
         String[] tokens = result.split("\\s");
         if (tokens.length > 0) {
             return tokens[0];
@@ -735,7 +749,7 @@
         // TODO(b/182576497): cp returns error because close(2) returns ENOSYS in the current authfs
         // implementation. We should probably fix that since programs can expect close(2) return 0.
         String cmd = "cat " + src + " > " + dest;
-        return tryRunOnMicrodroid(cmd) != null;
+        return sMicrodroid.tryRun(cmd) != null;
     }
 
     private String computeFileHashOnAndroid(String path) throws DeviceNotAvailableException {
@@ -751,38 +765,42 @@
 
     private void expectFileMode(String expected, String microdroidPath, String androidPath)
             throws DeviceNotAvailableException {
-        String actual = runOnMicrodroid("stat -c '%A' " + microdroidPath);
+        String actual = sMicrodroid.run("stat -c '%A' " + microdroidPath);
         assertEquals("Inconsistent mode for " + microdroidPath, expected, actual);
 
         actual = sAndroid.run("stat -c '%A' " + androidPath);
         assertEquals("Inconsistent mode for " + androidPath + " (android)", expected, actual);
     }
 
-    private boolean resizeFileOnMicrodroid(String path, long size) {
-        CommandResult result = runOnMicrodroidForResult("truncate -c -s " + size + " " + path);
+    private boolean resizeFileOnMicrodroid(String path, long size)
+            throws DeviceNotAvailableException {
+        CommandResult result = sMicrodroid.runForResult("truncate -c -s " + size + " " + path);
         return result.getStatus() == CommandStatus.SUCCESS;
     }
 
-    private long getFileSizeInBytesOnMicrodroid(String path) {
-        return Long.parseLong(runOnMicrodroid("stat -c '%s' " + path));
+    private long getFileSizeInBytesOnMicrodroid(String path) throws DeviceNotAvailableException {
+        return Long.parseLong(sMicrodroid.run("stat -c '%s' " + path));
     }
 
-    private void createFileWithOnesOnMicrodroid(String filePath, long numberOfOnes) {
-        runOnMicrodroid(
+    private void createFileWithOnesOnMicrodroid(String filePath, long numberOfOnes)
+            throws DeviceNotAvailableException {
+        sMicrodroid.run(
                 "yes $'\\x01' | tr -d '\\n' | dd bs=1 count=" + numberOfOnes + " of=" + filePath);
     }
 
-    private boolean checkReadAtFileOffsetOnMicrodroid(String filePath, long offset, long size) {
+    private boolean checkReadAtFileOffsetOnMicrodroid(String filePath, long offset, long size)
+            throws DeviceNotAvailableException {
         String cmd = "dd if=" + filePath + " of=/dev/null bs=1 count=" + size;
         if (offset > 0) {
             cmd += " skip=" + offset;
         }
-        CommandResult result = runOnMicrodroidForResult(cmd);
+        CommandResult result = sMicrodroid.runForResult(cmd);
         return result.getStatus() == CommandStatus.SUCCESS;
     }
 
     private boolean writeZerosAtFileOffsetOnMicrodroid(
-            String filePath, long offset, long numberOfZeros, boolean writeThrough) {
+            String filePath, long offset, long numberOfZeros, boolean writeThrough)
+            throws DeviceNotAvailableException {
         String cmd = "dd if=/dev/zero of=" + filePath + " bs=1 count=" + numberOfZeros
                 + " conv=notrunc";
         if (offset > 0) {
@@ -791,7 +809,7 @@
         if (writeThrough) {
             cmd += " direct";
         }
-        CommandResult result = runOnMicrodroidForResult(cmd);
+        CommandResult result = sMicrodroid.runForResult(cmd);
         return result.getStatus() == CommandStatus.SUCCESS;
     }
 
@@ -810,9 +828,14 @@
                     // authfs may fail to start if fd_server is not yet listening on the vsock
                     // ("Error: Invalid raw AIBinder"). Just restart if that happens.
                     while (starting.get()) {
-                        CLog.i("Starting authfs");
-                        CommandResult result = runOnMicrodroidForResult(cmd);
-                        CLog.w("authfs has stopped: " + result);
+                        try {
+                            CLog.i("Starting authfs");
+                            CommandResult result = sMicrodroid.runForResult(cmd);
+                            CLog.w("authfs has stopped: " + result);
+                        } catch (DeviceNotAvailableException e) {
+                            CLog.e("Error running authfs", e);
+                            throw new RuntimeException(e);
+                        }
                     }
                 });
         try {
@@ -855,8 +878,8 @@
                 });
     }
 
-    private boolean isMicrodroidDirectoryOnFuse(String path) {
-        String fs_type = tryRunOnMicrodroid("stat -f -c '%t' " + path);
+    private boolean isMicrodroidDirectoryOnFuse(String path) throws DeviceNotAvailableException {
+        String fs_type = sMicrodroid.tryRun("stat -f -c '%t' " + path);
         return FUSE_SUPER_MAGIC_HEX.equals(fs_type);
     }
 }
diff --git a/compos/apk/assets/vm_config.json b/compos/apk/assets/vm_config.json
index b93c3f7..c4abfd9 100644
--- a/compos/apk/assets/vm_config.json
+++ b/compos/apk/assets/vm_config.json
@@ -25,5 +25,6 @@
     {
       "name": "{CLASSPATH}"
     }
-  ]
+  ],
+  "export_tombstones": true
 }
diff --git a/compos/apk/assets/vm_config_staged.json b/compos/apk/assets/vm_config_staged.json
index 83fa6eb..0be6e78 100644
--- a/compos/apk/assets/vm_config_staged.json
+++ b/compos/apk/assets/vm_config_staged.json
@@ -26,5 +26,6 @@
     {
       "name": "{CLASSPATH}"
     }
-  ]
+  ],
+  "export_tombstones": true
 }
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 9787434..16dc2cf 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -38,12 +38,12 @@
 use compos_aidl_interface::aidl::com::android::compos::ICompOsService::ICompOsService;
 use log::{info, warn};
 use rustutils::system_properties;
-use std::fs::File;
+use std::fs::{self, File};
 use std::io::{BufRead, BufReader};
 use std::num::NonZeroU32;
 use std::os::raw;
 use std::os::unix::io::IntoRawFd;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::sync::{Arc, Condvar, Mutex};
 use std::thread;
 
@@ -95,8 +95,8 @@
         let apex_dir = Path::new(COMPOS_APEX_ROOT);
         let data_dir = Path::new(COMPOS_DATA_ROOT);
 
-        let apk_fd = File::open(apex_dir.join("app/CompOSPayloadApp/CompOSPayloadApp.apk"))
-            .context("Failed to open config APK file")?;
+        let config_apk = Self::locate_config_apk(apex_dir)?;
+        let apk_fd = File::open(config_apk).context("Failed to open config APK file")?;
         let apk_fd = ParcelFileDescriptor::new(apk_fd);
         let idsig_fd = prepare_idsig(service, &apk_fd, idsig)?;
 
@@ -166,6 +166,21 @@
         Ok(VmInstance { vm, cid })
     }
 
+    fn locate_config_apk(apex_dir: &Path) -> Result<PathBuf> {
+        // Our config APK will be in a directory under app, but the name of the directory is at the
+        // discretion of the build system. So just look in each sub-directory until we find it.
+        // (In practice there will be exactly one directory, so this shouldn't take long.)
+        let app_dir = apex_dir.join("app");
+        for dir in fs::read_dir(app_dir).context("Reading app dir")? {
+            let apk_file = dir?.path().join("CompOSPayloadApp.apk");
+            if apk_file.is_file() {
+                return Ok(apk_file);
+            }
+        }
+
+        bail!("Failed to locate CompOSPayloadApp.apk")
+    }
+
     /// Create and return an RPC Binder connection to the Comp OS service in the VM.
     pub fn get_service(&self) -> Result<Strong<dyn ICompOsService>> {
         let mut vsock_factory = VsockFactory::new(&*self.vm);
diff --git a/compos/tests/java/android/compos/test/ComposTestCase.java b/compos/tests/java/android/compos/test/ComposTestCase.java
index 6773eb7..eec9e39 100644
--- a/compos/tests/java/android/compos/test/ComposTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposTestCase.java
@@ -182,15 +182,8 @@
     private void killVmAndReconnectAdb() throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
 
-        // When a VM exits, we tend to see adb disconnecting. So we attempt to reconnect
-        // when we kill it to avoid problems. Of course VirtualizationService may exit anyway
-        // (it's an on-demand service and all its clients have gone), taking the VM with it,
-        // which makes this a bit unpredictable.
-        reconnectHostAdb(getDevice());
         android.tryRun("killall", "crosvm");
-        reconnectHostAdb(getDevice());
         android.tryRun("stop", "virtualizationservice");
-        reconnectHostAdb(getDevice());
 
         // Delete stale data
         android.tryRun("rm", "-rf", "/data/misc/virtualizationservice/*");
diff --git a/demo/assets/vm_config.json b/demo/assets/vm_config.json
index b814394..da420ff 100644
--- a/demo/assets/vm_config.json
+++ b/demo/assets/vm_config.json
@@ -9,5 +9,6 @@
       "hello",
       "microdroid"
     ]
-  }
+  },
+  "export_tombstones": true
 }
diff --git a/libs/avb_bindgen/Android.bp b/libs/avb_bindgen/Android.bp
index 1035498..1e62864 100644
--- a/libs/avb_bindgen/Android.bp
+++ b/libs/avb_bindgen/Android.bp
@@ -4,6 +4,7 @@
 
 rust_bindgen {
     name: "libavb_bindgen",
+    host_supported: true,
     wrapper_src: "bindgen/avb.h",
     crate_name: "avb_bindgen",
     source_stem: "bindings",
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index b265d5c..b7d844f 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -77,7 +77,9 @@
         "linkerconfig",
         "servicemanager.microdroid",
         "tombstoned",
+        "tombstone_transmit.microdroid",
         "cgroups.json",
+        "task_profiles.json",
         "public.libraries.android.txt",
 
         "microdroid_compatibility_matrix",
@@ -88,11 +90,9 @@
         "microdroid_property_contexts",
         "microdroid_service_contexts",
 
-        // TODO(b/195425111) these four should be added automatically
-        "android.hardware.security.secureclock-V1-ndk",
-        "android.hardware.security.sharedsecret-V1-ndk",
-        "libcrypto",
-        "liblzma",
+        // TODO(b/195425111) these should be added automatically
+        "libcrypto", // used by many (init_second_stage, microdroid_manager, toybox, etc)
+        "liblzma", // used by init_second_stage
     ] + microdroid_shell_and_utilities,
     multilib: {
         common: {
@@ -111,9 +111,6 @@
                 "authfs_service",
                 "microdroid_manager",
                 "zipfuse",
-
-                // TODO(b/184872979): Needed by authfs. Remove once the Rust API is created.
-                "libbinder_rpc_unstable",
             ],
         },
     },
@@ -532,7 +529,7 @@
 genrule {
     name: "microdroid_uboot_env_gen",
     tools: [
-        "mkenvimage_host",
+        "mkenvimage_slim",
         "avbtool",
     ],
     srcs: [
@@ -540,7 +537,7 @@
         ":microdroid_sign_key",
     ],
     out: ["output.img"],
-    cmd: "$(location mkenvimage_host) -s 4096 -o $(out) $(location uboot-env.txt) && " +
+    cmd: "$(location mkenvimage_slim) -output_path $(out) -input_path $(location uboot-env.txt) && " +
         "$(location avbtool) add_hash_footer " +
         "--algorithm SHA256_RSA4096 " +
         "--partition_name uboot_env " +
@@ -612,4 +609,4 @@
     src: "microdroid_event-log-tags",
     filename: "event-log-tags",
     installable: false,
-}
\ No newline at end of file
+}
diff --git a/microdroid/init.rc b/microdroid/init.rc
index f6d5092..93cced8 100644
--- a/microdroid/init.rc
+++ b/microdroid/init.rc
@@ -173,6 +173,10 @@
     mkdir /data/local 0751 root root
     mkdir /data/local/tmp 0771 shell shell
 
+service tombstone_transmit /system/bin/tombstone_transmit.microdroid -cid 2 -port 2000
+    user root
+    group system
+
 service apexd-vm /system/bin/apexd --vm
     user root
     group system
diff --git a/microdroid/payload/config/src/lib.rs b/microdroid/payload/config/src/lib.rs
index 67e8feb..b82544f 100644
--- a/microdroid/payload/config/src/lib.rs
+++ b/microdroid/payload/config/src/lib.rs
@@ -38,6 +38,10 @@
     /// Tells VirtualizationService to use staged APEXes if possible
     #[serde(default)]
     pub prefer_staged: bool,
+
+    /// Whether to export the tomsbtones (VM crashes) out of VM to host
+    /// This does not have a default & the value is expected to be in json for deserialization
+    pub export_tombstones: bool,
 }
 
 /// OS config
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 9e159d2..8a638db 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -146,17 +146,17 @@
     }
 }
 
-fn dice_derivation(verified_data: MicrodroidData, payload_config_path: &str) -> Result<()> {
+fn dice_derivation(verified_data: &MicrodroidData, payload_config_path: &str) -> Result<()> {
     // Calculate compound digests of code and authorities
     let mut code_hash_ctx = digest::Context::new(&digest::SHA512);
     let mut authority_hash_ctx = digest::Context::new(&digest::SHA512);
     code_hash_ctx.update(verified_data.apk_data.root_hash.as_ref());
     authority_hash_ctx.update(verified_data.apk_data.pubkey.as_ref());
-    for extra_apk in verified_data.extra_apks_data {
+    for extra_apk in &verified_data.extra_apks_data {
         code_hash_ctx.update(extra_apk.root_hash.as_ref());
         authority_hash_ctx.update(extra_apk.pubkey.as_ref());
     }
-    for apex in verified_data.apex_data {
+    for apex in &verified_data.apex_data {
         code_hash_ctx.update(apex.root_digest.as_ref());
         authority_hash_ctx.update(apex.public_key.as_ref());
     }
@@ -189,7 +189,7 @@
             authorityHash: authority_hash,
             authorityDescriptor: None,
             mode: if app_debuggable { Mode::DEBUG } else { Mode::NORMAL },
-            hidden: verified_data.salt.try_into().unwrap(),
+            hidden: verified_data.salt.clone().try_into().unwrap(),
         }])
         .context("IDiceMaintenance::demoteSelf failed")?;
     Ok(())
@@ -240,6 +240,10 @@
         instance.write_microdroid_data(&verified_data).context("Failed to write identity data")?;
     }
 
+    // To minimize the exposure to untrusted data, derive dice profile as soon as possible.
+    info!("DICE derivation for payload");
+    dice_derivation(&verified_data, &metadata.payload_config_path)?;
+
     // Before reading a file from the APK, start zipfuse
     run_zipfuse(
         "fscontext=u:object_r:zipfusefs:s0,context=u:object_r:system_file:s0",
@@ -254,6 +258,13 @@
     );
 
     let config = load_config(Path::new(&metadata.payload_config_path))?;
+
+    // Start tombstone_transmit if enabled
+    if config.export_tombstones {
+        system_properties::write("ctl.start", "tombstone_transmit")
+            .context("Failed to start tombstone_transmit")?;
+    }
+
     if config.extra_apks.len() != verified_data.extra_apks_data.len() {
         return Err(anyhow!(
             "config expects {} extra apks, but found only {}",
@@ -263,9 +274,6 @@
     }
     mount_extra_apks(&config)?;
 
-    info!("DICE derivation for payload");
-    dice_derivation(verified_data, &metadata.payload_config_path)?;
-
     // Wait until apex config is done. (e.g. linker configuration for apexes)
     // TODO(jooyung): wait until sys.boot_completed?
     wait_for_apex_config_done()?;
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index ce096b6..fbdd8d7 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -1,3 +1,7 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
 rust_ffi_static {
     name: "libpvmfw",
     crate_name: "pvmfw",
diff --git a/tests/benchmark/assets/vm_config.json b/tests/benchmark/assets/vm_config.json
index d431877..d8cbb86 100644
--- a/tests/benchmark/assets/vm_config.json
+++ b/tests/benchmark/assets/vm_config.json
@@ -6,6 +6,7 @@
     "type": "microdroid_launcher",
     "command": "empty_payload.so",
     "args": []
-  }
+  },
+  "export_tombstones": true
 }
 
diff --git a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
index 40be248..440ae18 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -82,28 +82,12 @@
         // disconnect from microdroid
         tryRunOnHost("adb", "disconnect", MICRODROID_SERIAL);
 
-        reconnectHostAdb(androidDevice);
-
         // kill stale VMs and directories
         android.tryRun("killall", "crosvm");
         android.tryRun("stop", "virtualizationservice");
         android.tryRun("rm", "-rf", "/data/misc/virtualizationservice/*");
     }
 
-    public static void reconnectHostAdb(ITestDevice androidDevice)
-            throws DeviceNotAvailableException {
-        CommandRunner android = new CommandRunner(androidDevice);
-
-        // Make sure we're connected to the host adb; this connection seems to get dropped when a VM
-        // exits.
-        for (int retry = 0; retry < 10; ++retry) {
-            if (android.tryRun("true") != null) {
-                break;
-            }
-            androidDevice.waitForDeviceOnline(1000);
-        }
-    }
-
     public static void testIfDeviceIsCapable(ITestDevice androidDevice) throws Exception {
         assumeTrue("Need an actual TestDevice", androidDevice instanceof TestDevice);
         TestDevice testDevice = (TestDevice) androidDevice;
@@ -121,7 +105,7 @@
     }
 
     // Run an arbitrary command in the host side and returns the result
-    private static String runOnHost(String... cmd) {
+    public static String runOnHost(String... cmd) {
         return runOnHostWithTimeout(10000, cmd);
     }
 
@@ -149,6 +133,19 @@
         return result.getStdout().trim();
     }
 
+    // Same as runOnMicrodroid, but keeps retrying on error till timeout
+    private static String runOnMicrodroidRetryingOnFailure(String... cmd) {
+        final long timeoutMs = 30000; // 30 sec. Microdroid is extremely slow on GCE-on-CF.
+        int attempts = (int) MICRODROID_ADB_CONNECT_TIMEOUT_MINUTES * 60 * 1000 / 500;
+        CommandResult result = RunUtil.getDefault()
+                .runTimedCmdRetry(timeoutMs, 500, attempts,
+                        "adb", "-s", MICRODROID_SERIAL, "shell", join(cmd));
+        if (result.getStatus() != CommandStatus.SUCCESS) {
+            fail(join(cmd) + " has failed: " + result);
+        }
+        return result.getStdout().trim();
+    }
+
     // Same as runOnMicrodroid, but returns null on error.
     public static String tryRunOnMicrodroid(String... cmd) {
         CommandResult result = runOnMicrodroidForResult(cmd);
@@ -344,34 +341,20 @@
 
         // Shutdown the VM
         android.run(VIRT_APEX + "bin/vm", "stop", cid);
-
-        // TODO(192660485): Figure out why shutting down the VM disconnects adb on cuttlefish
-        // temporarily. Without this wait, the rest of `runOnAndroid/skipIfFail` fails due to the
-        // connection loss, and results in assumption error exception for the rest of the tests.
-        try {
-            Thread.sleep(1000);
-        } catch (InterruptedException e) {
-            Thread.currentThread().interrupt();
-        }
     }
 
     public static void rootMicrodroid() {
         runOnHost("adb", "-s", MICRODROID_SERIAL, "root");
-
-        // TODO(192660959): Figure out the root cause and remove the sleep. For unknown reason,
-        // even though `adb root` actually wait-for-disconnect then wait-for-device, the next
-        // `adb -s $MICRODROID_SERIAL shell ...` often fails with "adb: device offline".
-        try {
-            Thread.sleep(1000);
-            runOnHostWithTimeout(
-                    MICRODROID_ADB_CONNECT_TIMEOUT_MINUTES * 60 * 1000,
-                    "adb",
-                    "-s",
-                    MICRODROID_SERIAL,
-                    "wait-for-device");
-        } catch (InterruptedException e) {
-            Thread.currentThread().interrupt();
-        }
+        runOnHostWithTimeout(
+                MICRODROID_ADB_CONNECT_TIMEOUT_MINUTES * 60 * 1000,
+                "adb",
+                "-s",
+                MICRODROID_SERIAL,
+                "wait-for-device");
+        // There have been tests when adb wait-for-device succeeded but the following command
+        // fails with error: closed. Hence, we run adb shell true in microdroid with retries
+        // before returning.
+        runOnMicrodroidRetryingOnFailure("true");
     }
 
     // Establish an adb connection to microdroid by letting Android forward the connection to
diff --git a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
index 5b71eba..32bdf3b 100644
--- a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
+++ b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
@@ -22,6 +22,7 @@
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
@@ -100,6 +101,13 @@
                 false);
     }
 
+    // Wait until logd-init starts. The service is one of the last services that are started in
+    // the microdroid boot procedure. Therefore, waiting for the service means that we wait for
+    // the boot to complete. TODO: we need a better marker eventually.
+    private void waitForLogdInit() {
+        tryRunOnMicrodroid("watch -e \"getprop init.svc.logd-reinit | grep '^$'\"");
+    }
+
     @Test
     public void testCreateVmRequiresPermission() throws Exception {
         // Revoke the MANAGE_VIRTUAL_MACHINE permission for the test app
@@ -145,7 +153,9 @@
         command.add(virtApexDir.getPath());
 
         CommandResult result = runUtil.runTimedCmd(
-                                    20 * 1000,
+                                    // sign_virt_apex is so slow on CI server that this often times
+                                    // out. Until we can make it fast, use 50s for timeout
+                                    50 * 1000,
                                     "/bin/bash",
                                     "-c",
                                     String.join(" ", command));
@@ -350,6 +360,41 @@
     }
 
     @Test
+    public void testTombstonesAreBeingForwarded() throws Exception {
+        // Note this test relies on logcat values being printed by tombstone_transmit on
+        // and the reeceiver on host (virtualization_service)
+        final String configPath = "assets/vm_config.json"; // path inside the APK
+        final String cid =
+                startMicrodroid(
+                        getDevice(),
+                        getBuild(),
+                        APK_NAME,
+                        PACKAGE_NAME,
+                        configPath,
+                        /* debug */ true,
+                        minMemorySize(),
+                        Optional.of(NUM_VCPUS),
+                        Optional.of(CPU_AFFINITY));
+        adbConnectToMicrodroid(getDevice(), cid);
+        waitForLogdInit();
+        runOnMicrodroid("logcat -c");
+        // We need root permission to write to /data/tombstones/
+        rootMicrodroid();
+        // Write a test tombstone file in /data/tombstones
+        runOnMicrodroid("echo -n \'Test tombstone in VM with 34 bytes\'"
+                    + "> /data/tombstones/transmit.txt");
+        // check if the tombstone have been tranferred from VM
+        assertNotEquals(runOnMicrodroid("timeout 15s logcat | grep -m 1 "
+                            + "'tombstone_transmit.microdroid:.*data/tombstones/transmit.txt'"),
+                "");
+        // Confirm that tombstone is received (from host logcat)
+        assertNotEquals(runOnHost("adb", "-s", getDevice().getSerialNumber(),
+                            "logcat", "-d", "-e",
+                            "Received 34 bytes from guest & wrote to tombstone file.*"),
+                "");
+    }
+
+    @Test
     public void testMicrodroidBoots() throws Exception {
         final String configPath = "assets/vm_config.json"; // path inside the APK
         final String cid =
@@ -364,12 +409,7 @@
                         Optional.of(NUM_VCPUS),
                         Optional.of(CPU_AFFINITY));
         adbConnectToMicrodroid(getDevice(), cid);
-
-        // Wait until logd-init starts. The service is one of the last services that are started in
-        // the microdroid boot procedure. Therefore, waiting for the service means that we wait for
-        // the boot to complete. TODO: we need a better marker eventually.
-        tryRunOnMicrodroid("watch -e \"getprop init.svc.logd-reinit | grep '^$'\"");
-
+        waitForLogdInit();
         // Test writing to /data partition
         runOnMicrodroid("echo MicrodroidTest > /data/local/tmp/test.txt");
         assertThat(runOnMicrodroid("cat /data/local/tmp/test.txt"), is("MicrodroidTest"));
diff --git a/tests/testapk/assets/vm_config.json b/tests/testapk/assets/vm_config.json
index b814394..da420ff 100644
--- a/tests/testapk/assets/vm_config.json
+++ b/tests/testapk/assets/vm_config.json
@@ -9,5 +9,6 @@
       "hello",
       "microdroid"
     ]
-  }
+  },
+  "export_tombstones": true
 }
diff --git a/tests/testapk/assets/vm_config_extra_apk.json b/tests/testapk/assets/vm_config_extra_apk.json
index a5bae63..d7d3dd7 100644
--- a/tests/testapk/assets/vm_config_extra_apk.json
+++ b/tests/testapk/assets/vm_config_extra_apk.json
@@ -14,5 +14,6 @@
     {
       "path": "/system/etc/security/fsverity/BuildManifest.apk"
     }
-  ]
+  ],
+  "export_tombstones": true
 }
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 36bea72..f7261d3 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -262,11 +262,6 @@
     public void changingDebugLevelInvalidatesVmIdentity()
             throws VirtualMachineException, InterruptedException, IOException {
         assume()
-            .withMessage("Skip on Cuttlefish. b/195765441")
-            .that(android.os.Build.DEVICE)
-            .isNotEqualTo("vsoc_x86_64");
-
-        assume()
             .withMessage("SKip on 5.4 kernel. b/218303240")
             .that(KERNEL_VERSION)
             .isNotEqualTo("5.4");
@@ -344,11 +339,6 @@
     public void instancesOfSameVmHaveDifferentCdis()
             throws VirtualMachineException, InterruptedException {
         assume()
-            .withMessage("Skip on Cuttlefish. b/195765441")
-            .that(android.os.Build.DEVICE)
-            .isNotEqualTo("vsoc_x86_64");
-
-        assume()
             .withMessage("SKip on 5.4 kernel. b/218303240")
             .that(KERNEL_VERSION)
             .isNotEqualTo("5.4");
@@ -368,11 +358,6 @@
     public void sameInstanceKeepsSameCdis()
             throws VirtualMachineException, InterruptedException {
         assume()
-            .withMessage("Skip on Cuttlefish. b/195765441")
-            .that(android.os.Build.DEVICE)
-            .isNotEqualTo("vsoc_x86_64");
-
-        assume()
             .withMessage("SKip on 5.4 kernel. b/218303240")
             .that(KERNEL_VERSION)
             .isNotEqualTo("5.4");
@@ -389,11 +374,6 @@
     public void bccIsSuperficiallyWellFormed()
             throws VirtualMachineException, InterruptedException, CborException {
         assume()
-            .withMessage("Skip on Cuttlefish. b/195765441")
-            .that(android.os.Build.DEVICE)
-            .isNotEqualTo("vsoc_x86_64");
-
-        assume()
             .withMessage("SKip on 5.4 kernel. b/218303240")
             .that(KERNEL_VERSION)
             .isNotEqualTo("5.4");
@@ -529,20 +509,12 @@
     @Test
     public void bootFailsWhenMicrodroidDataIsCompromised()
             throws VirtualMachineException, InterruptedException, IOException {
-        assume().withMessage("Skip on Cuttlefish. b/195765441")
-                .that(android.os.Build.DEVICE)
-                .isNotEqualTo("vsoc_x86_64");
-
         assertThatBootFailsAfterCompromisingPartition(MICRODROID_PARTITION_UUID);
     }
 
     @Test
     public void bootFailsWhenUBootAvbDataIsCompromised()
             throws VirtualMachineException, InterruptedException, IOException {
-        assume().withMessage("Skip on Cuttlefish. b/195765441")
-                .that(android.os.Build.DEVICE)
-                .isNotEqualTo("vsoc_x86_64");
-
         if (mProtectedVm) {
             assertThatBootFailsAfterCompromisingPartition(U_BOOT_AVB_PARTITION_UUID);
         } else {
@@ -554,10 +526,6 @@
     @Test
     public void bootFailsWhenUBootEnvDataIsCompromised()
             throws VirtualMachineException, InterruptedException, IOException {
-        assume().withMessage("Skip on Cuttlefish. b/195765441")
-                .that(android.os.Build.DEVICE)
-                .isNotEqualTo("vsoc_x86_64");
-
         if (mProtectedVm) {
             assertThatBootFailsAfterCompromisingPartition(U_BOOT_ENV_PARTITION_UUID);
         } else {
@@ -569,10 +537,6 @@
     @Test
     public void bootFailsWhenPvmFwDataIsCompromised()
             throws VirtualMachineException, InterruptedException, IOException {
-        assume().withMessage("Skip on Cuttlefish. b/195765441")
-                .that(android.os.Build.DEVICE)
-                .isNotEqualTo("vsoc_x86_64");
-
         if (mProtectedVm) {
             assertThatBootFailsAfterCompromisingPartition(PVM_FW_PARTITION_UUID);
         } else {
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 9b2b740..7a8da96 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -45,6 +45,7 @@
         "libserde_xml_rs",
         "libshared_child",
         "libstatslog_virtualization_rust",
+        "libtombstoned_client_rust",
         "libvmconfig",
         "libzip",
         "libvsock",
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
index 1a16f2a..dff5d46 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -30,6 +30,12 @@
     const int VM_BINDER_SERVICE_PORT = 5000;
 
     /**
+     * Port number that VirtualMachineService listens on connections from the guest VMs for the
+     * tombtones
+     */
+    const int VM_TOMBSTONES_SERVICE_PORT = 2000;
+
+    /**
      * Notifies that the payload has started.
      */
     void notifyPayloadStarted();
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index d9825dc..a2e856c 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -42,7 +42,7 @@
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
     IVirtualMachineService::{
         BnVirtualMachineService, IVirtualMachineService, VM_BINDER_SERVICE_PORT,
-        VM_STREAM_SERVICE_PORT,
+        VM_STREAM_SERVICE_PORT, VM_TOMBSTONES_SERVICE_PORT,
     },
 };
 use anyhow::{anyhow, bail, Context, Result};
@@ -57,13 +57,14 @@
 use std::convert::TryInto;
 use std::ffi::CStr;
 use std::fs::{create_dir, File, OpenOptions};
-use std::io::{Error, ErrorKind, Write};
+use std::io::{Error, ErrorKind, Write, Read};
 use std::num::NonZeroU32;
 use std::os::raw;
 use std::os::unix::io::{FromRawFd, IntoRawFd};
 use std::path::{Path, PathBuf};
 use std::ptr::null_mut;
 use std::sync::{Arc, Mutex, Weak};
+use tombstoned_client::{TombstonedConnection, DebuggerdDumpType};
 use vmconfig::VmConfig;
 use vsock::{SockAddr, VsockListener, VsockStream};
 use zip::ZipArchive;
@@ -86,6 +87,8 @@
 /// Version of the instance image format
 const ANDROID_VM_INSTANCE_VERSION: u16 = 1;
 
+const CHUNK_RECV_MAX_LEN: usize = 1024;
+
 /// Implementation of `IVirtualizationService`, the entry point of the AIDL service.
 #[derive(Debug, Default)]
 pub struct VirtualizationService {
@@ -129,139 +132,26 @@
         console_fd: Option<&ParcelFileDescriptor>,
         log_fd: Option<&ParcelFileDescriptor>,
     ) -> binder::Result<Strong<dyn IVirtualMachine>> {
-        check_manage_access()?;
-        let state = &mut *self.state.lock().unwrap();
-        let console_fd = console_fd.map(clone_file).transpose()?;
-        let log_fd = log_fd.map(clone_file).transpose()?;
-        let requester_uid = ThreadState::get_calling_uid();
-        let requester_sid = get_calling_sid()?;
-        let requester_debug_pid = ThreadState::get_calling_pid();
-        let cid = next_cid().or(Err(ExceptionCode::ILLEGAL_STATE))?;
-
-        // Counter to generate unique IDs for temporary image files.
-        let mut next_temporary_image_id = 0;
-        // Files which are referred to from composite images. These must be mapped to the crosvm
-        // child process, and not closed before it is started.
-        let mut indirect_files = vec![];
-
-        // Make directory for temporary files.
-        let temporary_directory: PathBuf = format!("{}/{}", TEMPORARY_DIRECTORY, cid).into();
-        create_dir(&temporary_directory).map_err(|e| {
-            // At this point, we do not know the protected status of Vm
-            // setting it to false, though this may not be correct.
-            write_vm_creation_stats(false, false);
-            error!(
-                "Failed to create temporary directory {:?} for VM files: {}",
-                temporary_directory, e
-            );
-            new_binder_exception(
-                ExceptionCode::SERVICE_SPECIFIC,
-                format!(
-                    "Failed to create temporary directory {:?} for VM files: {}",
-                    temporary_directory, e
-                ),
-            )
-        })?;
-
-        let is_app_config = matches!(config, VirtualMachineConfig::AppConfig(_));
-
-        let config = match config {
-            VirtualMachineConfig::AppConfig(config) => BorrowedOrOwned::Owned(
-                load_app_config(config, &temporary_directory).map_err(|e| {
-                    error!("Failed to load app config from {}: {}", &config.configPath, e);
-                    write_vm_creation_stats(config.protectedVm, false);
-                    new_binder_exception(
-                        ExceptionCode::SERVICE_SPECIFIC,
-                        format!("Failed to load app config from {}: {}", &config.configPath, e),
-                    )
-                })?,
-            ),
-            VirtualMachineConfig::RawConfig(config) => BorrowedOrOwned::Borrowed(config),
-        };
-        let config = config.as_ref();
-        let protected = config.protectedVm;
-
-        // Check if partition images are labeled incorrectly. This is to prevent random images
-        // which are not protected by the Android Verified Boot (e.g. bits downloaded by apps) from
-        // being loaded in a pVM.  Specifically, for images in the raw config, nothing is allowed
-        // to be labeled as app_data_file. For images in the app config, nothing but the instance
-        // partition is allowed to be labeled as such.
-        config
-            .disks
-            .iter()
-            .flat_map(|disk| disk.partitions.iter())
-            .filter(|partition| {
-                if is_app_config {
-                    partition.label != "vm-instance"
-                } else {
-                    true // all partitions are checked
-                }
-            })
-            .try_for_each(check_label_for_partition)
-            .map_err(|e| new_binder_exception(ExceptionCode::SERVICE_SPECIFIC, e.to_string()))?;
-
-        let zero_filler_path = temporary_directory.join("zero.img");
-        write_zero_filler(&zero_filler_path).map_err(|e| {
-            error!("Failed to make composite image: {}", e);
-            write_vm_creation_stats(protected, false);
-            new_binder_exception(
-                ExceptionCode::SERVICE_SPECIFIC,
-                format!("Failed to make composite image: {}", e),
-            )
-        })?;
-
-        // Assemble disk images if needed.
-        let disks = config
-            .disks
-            .iter()
-            .map(|disk| {
-                assemble_disk_image(
-                    disk,
-                    &zero_filler_path,
-                    &temporary_directory,
-                    &mut next_temporary_image_id,
-                    &mut indirect_files,
-                )
-            })
-            .collect::<Result<Vec<DiskFile>, _>>()?;
-
-        // Actually start the VM.
-        let crosvm_config = CrosvmConfig {
-            cid,
-            bootloader: maybe_clone_file(&config.bootloader)?,
-            kernel: maybe_clone_file(&config.kernel)?,
-            initrd: maybe_clone_file(&config.initrd)?,
-            disks,
-            params: config.params.to_owned(),
-            protected,
-            memory_mib: config.memoryMib.try_into().ok().and_then(NonZeroU32::new),
-            cpus: config.numCpus.try_into().ok().and_then(NonZeroU32::new),
-            cpu_affinity: config.cpuAffinity.clone(),
-            console_fd,
-            log_fd,
-            indirect_files,
-            platform_version: parse_platform_version_req(&config.platformVersion)?,
-        };
-        let instance = Arc::new(
-            VmInstance::new(
-                crosvm_config,
-                temporary_directory,
-                requester_uid,
-                requester_sid,
-                requester_debug_pid,
-            )
-            .map_err(|e| {
-                error!("Failed to create VM with config {:?}: {}", config, e);
-                write_vm_creation_stats(protected, false);
-                new_binder_exception(
-                    ExceptionCode::SERVICE_SPECIFIC,
-                    format!("Failed to create VM: {}", e),
-                )
-            })?,
-        );
-        state.add_vm(Arc::downgrade(&instance));
-        write_vm_creation_stats(protected, true);
-        Ok(VirtualMachine::create(instance))
+        let mut is_protected = false;
+        let ret = self.create_vm_internal(config, console_fd, log_fd, &mut is_protected);
+        match ret {
+            Ok(_) => {
+                let ok_status = Status::ok();
+                write_vm_creation_stats(
+                    is_protected,
+                    /*creation_succeeded*/ true,
+                    ok_status.exception_code() as i32,
+                );
+            }
+            Err(ref e) => {
+                write_vm_creation_stats(
+                    is_protected,
+                    /*creation_succeeded*/ false,
+                    e.exception_code() as i32,
+                );
+            }
+        }
+        ret
     }
 
     /// Initialise an empty partition image of the given size to be used as a writable partition.
@@ -371,6 +261,55 @@
     }
 }
 
+fn handle_stream_connection_tombstoned() -> Result<()> {
+    let listener =
+        VsockListener::bind_with_cid_port(VMADDR_CID_HOST, VM_TOMBSTONES_SERVICE_PORT as u32)?;
+    info!("Listening to tombstones from guests ...");
+    for incoming_stream in listener.incoming() {
+        let mut incoming_stream = match incoming_stream {
+            Err(e) => {
+                warn!("invalid incoming connection: {}", e);
+                continue;
+            }
+            Ok(s) => s,
+        };
+        std::thread::spawn(move || {
+            if let Err(e) = handle_tombstone(&mut incoming_stream) {
+                error!("Failed to write tombstone- {:?}", e);
+            }
+        });
+    }
+    Ok(())
+}
+
+fn handle_tombstone(stream: &mut VsockStream) -> Result<()> {
+    if let Ok(SockAddr::Vsock(addr)) = stream.peer_addr() {
+        info!("Vsock Stream connected to cid={} for tombstones", addr.cid());
+    }
+    let tb_connection =
+        TombstonedConnection::connect(std::process::id() as i32, DebuggerdDumpType::Tombstone)
+            .context("Failed to connect to tombstoned")?;
+    let mut text_output = tb_connection
+        .text_output
+        .as_ref()
+        .ok_or_else(|| anyhow!("Could not get file to write the tombstones on"))?;
+    let mut num_bytes_read = 0;
+    loop {
+        let mut chunk_recv = [0; CHUNK_RECV_MAX_LEN];
+        let n = stream
+            .read(&mut chunk_recv)
+            .context("Failed to read tombstone data from Vsock stream")?;
+        if n == 0 {
+            break;
+        }
+        num_bytes_read += n;
+        text_output.write_all(&chunk_recv[0..n]).context("Failed to write guests tombstones")?;
+    }
+    info!("Received {} bytes from guest & wrote to tombstone file", num_bytes_read);
+    tb_connection.notify_completion()?;
+    Ok(())
+}
+
 impl VirtualizationService {
     pub fn init() -> VirtualizationService {
         let service = VirtualizationService::default();
@@ -381,8 +320,15 @@
             handle_stream_connection_from_vm(state).unwrap();
         });
 
+        std::thread::spawn(|| {
+            if let Err(e) = handle_stream_connection_tombstoned() {
+                warn!("Error receiving tombstone from guest or writing them. Error: {}", e);
+            }
+        });
+
         // binder server for vm
-        let mut state = service.state.clone(); // reference to state (not the state itself) is copied
+        // reference to state (not the state itself) is copied
+        let mut state = service.state.clone();
         std::thread::spawn(move || {
             let state_ptr = &mut state as *mut _ as *mut raw::c_void;
 
@@ -407,11 +353,149 @@
         });
         service
     }
+
+    fn create_vm_internal(
+        &self,
+        config: &VirtualMachineConfig,
+        console_fd: Option<&ParcelFileDescriptor>,
+        log_fd: Option<&ParcelFileDescriptor>,
+        is_protected: &mut bool,
+    ) -> binder::Result<Strong<dyn IVirtualMachine>> {
+        check_manage_access()?;
+        let state = &mut *self.state.lock().unwrap();
+        let console_fd = console_fd.map(clone_file).transpose()?;
+        let log_fd = log_fd.map(clone_file).transpose()?;
+        let requester_uid = ThreadState::get_calling_uid();
+        let requester_sid = get_calling_sid()?;
+        let requester_debug_pid = ThreadState::get_calling_pid();
+        let cid = next_cid().or(Err(ExceptionCode::ILLEGAL_STATE))?;
+
+        // Counter to generate unique IDs for temporary image files.
+        let mut next_temporary_image_id = 0;
+        // Files which are referred to from composite images. These must be mapped to the crosvm
+        // child process, and not closed before it is started.
+        let mut indirect_files = vec![];
+
+        // Make directory for temporary files.
+        let temporary_directory: PathBuf = format!("{}/{}", TEMPORARY_DIRECTORY, cid).into();
+        create_dir(&temporary_directory).map_err(|e| {
+            // At this point, we do not know the protected status of Vm
+            // setting it to false, though this may not be correct.
+            error!(
+                "Failed to create temporary directory {:?} for VM files: {}",
+                temporary_directory, e
+            );
+            new_binder_exception(
+                ExceptionCode::SERVICE_SPECIFIC,
+                format!(
+                    "Failed to create temporary directory {:?} for VM files: {}",
+                    temporary_directory, e
+                ),
+            )
+        })?;
+
+        let is_app_config = matches!(config, VirtualMachineConfig::AppConfig(_));
+
+        let config = match config {
+            VirtualMachineConfig::AppConfig(config) => BorrowedOrOwned::Owned(
+                load_app_config(config, &temporary_directory).map_err(|e| {
+                    error!("Failed to load app config from {}: {}", &config.configPath, e);
+                    *is_protected = config.protectedVm;
+                    new_binder_exception(
+                        ExceptionCode::SERVICE_SPECIFIC,
+                        format!("Failed to load app config from {}: {}", &config.configPath, e),
+                    )
+                })?,
+            ),
+            VirtualMachineConfig::RawConfig(config) => BorrowedOrOwned::Borrowed(config),
+        };
+        let config = config.as_ref();
+        *is_protected = config.protectedVm;
+
+        // Check if partition images are labeled incorrectly. This is to prevent random images
+        // which are not protected by the Android Verified Boot (e.g. bits downloaded by apps) from
+        // being loaded in a pVM.  Specifically, for images in the raw config, nothing is allowed
+        // to be labeled as app_data_file. For images in the app config, nothing but the instance
+        // partition is allowed to be labeled as such.
+        config
+            .disks
+            .iter()
+            .flat_map(|disk| disk.partitions.iter())
+            .filter(|partition| {
+                if is_app_config {
+                    partition.label != "vm-instance"
+                } else {
+                    true // all partitions are checked
+                }
+            })
+            .try_for_each(check_label_for_partition)
+            .map_err(|e| new_binder_exception(ExceptionCode::SERVICE_SPECIFIC, e.to_string()))?;
+
+        let zero_filler_path = temporary_directory.join("zero.img");
+        write_zero_filler(&zero_filler_path).map_err(|e| {
+            error!("Failed to make composite image: {}", e);
+            new_binder_exception(
+                ExceptionCode::SERVICE_SPECIFIC,
+                format!("Failed to make composite image: {}", e),
+            )
+        })?;
+
+        // Assemble disk images if needed.
+        let disks = config
+            .disks
+            .iter()
+            .map(|disk| {
+                assemble_disk_image(
+                    disk,
+                    &zero_filler_path,
+                    &temporary_directory,
+                    &mut next_temporary_image_id,
+                    &mut indirect_files,
+                )
+            })
+            .collect::<Result<Vec<DiskFile>, _>>()?;
+
+        // Actually start the VM.
+        let crosvm_config = CrosvmConfig {
+            cid,
+            bootloader: maybe_clone_file(&config.bootloader)?,
+            kernel: maybe_clone_file(&config.kernel)?,
+            initrd: maybe_clone_file(&config.initrd)?,
+            disks,
+            params: config.params.to_owned(),
+            protected: *is_protected,
+            memory_mib: config.memoryMib.try_into().ok().and_then(NonZeroU32::new),
+            cpus: config.numCpus.try_into().ok().and_then(NonZeroU32::new),
+            cpu_affinity: config.cpuAffinity.clone(),
+            console_fd,
+            log_fd,
+            indirect_files,
+            platform_version: parse_platform_version_req(&config.platformVersion)?,
+        };
+        let instance = Arc::new(
+            VmInstance::new(
+                crosvm_config,
+                temporary_directory,
+                requester_uid,
+                requester_sid,
+                requester_debug_pid,
+            )
+            .map_err(|e| {
+                error!("Failed to create VM with config {:?}: {}", config, e);
+                new_binder_exception(
+                    ExceptionCode::SERVICE_SPECIFIC,
+                    format!("Failed to create VM: {}", e),
+                )
+            })?,
+        );
+        state.add_vm(Arc::downgrade(&instance));
+        Ok(VirtualMachine::create(instance))
+    }
 }
 
 /// Write the stats of VMCreation to statsd
-fn write_vm_creation_stats(protected: bool, success: bool) {
-    match stats_write(Hypervisor::Pkvm, protected, success) {
+fn write_vm_creation_stats(is_protected: bool, creation_succeeded: bool, exception_code: i32) {
+    match stats_write(Hypervisor::Pkvm, is_protected, creation_succeeded, exception_code) {
         Err(e) => {
             warn!("statslog_rust failed with error: {}", e);
         }
diff --git a/zipfuse/src/main.rs b/zipfuse/src/main.rs
index a91642c..c3fae69 100644
--- a/zipfuse/src/main.rs
+++ b/zipfuse/src/main.rs
@@ -82,7 +82,9 @@
         libc::MS_NOSUID | libc::MS_NODEV | libc::MS_RDONLY,
         &mount_options,
     )?;
-    Ok(fuse::worker::start_message_loop(dev_fuse, MAX_READ, MAX_WRITE, ZipFuse::new(zip_file)?)?)
+    let mut config = fuse::FuseConfig::new();
+    config.dev_fuse(dev_fuse).max_write(MAX_WRITE).max_read(MAX_READ);
+    Ok(config.enter_message_loop(ZipFuse::new(zip_file)?)?)
 }
 
 struct ZipFuse {