Merge "pvmfw: Document asm! & use u64 for SCTLR fields"
diff --git a/OWNERS b/OWNERS
index ecd24ed..310add7 100644
--- a/OWNERS
+++ b/OWNERS
@@ -1,5 +1,7 @@
 # Welcome to Android KVM!
 #
+# Bug component: 867125
+#
 # If you are not a member of the project please send review requests
 # to one of those listed below.
 dbrazdil@google.com
diff --git a/PREUPLOAD.cfg b/PREUPLOAD.cfg
index a6b1f95..43c89d4 100644
--- a/PREUPLOAD.cfg
+++ b/PREUPLOAD.cfg
@@ -4,10 +4,15 @@
 bpfmt = true
 clang_format = true
 jsonlint = true
+google_java_format = true
 pylint3 = true
 rustfmt = true
 xmllint = true
 
+[Tool Paths]
+google-java-format = ${REPO_ROOT}/prebuilts/tools/common/google-java-format/google-java-format
+google-java-format-diff = ${REPO_ROOT}/prebuilts/tools/common/google-java-format/google-java-format-diff.py
+
 [Builtin Hooks Options]
 clang_format = --commit ${PREUPLOAD_COMMIT} --style file --extensions c,h,cc,cpp
 rustfmt = --config-path=rustfmt.toml
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 8376c1f..dd81738 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -30,6 +30,11 @@
       "name": "AVFHostTestCases"
     }
   ],
+  "postsubmit": [
+    {
+      "name": "CtsMicrodroidDisabledTestCases"
+    }
+  ],
   "imports": [
     {
       "path": "packages/modules/Virtualization/apkdmverity"
diff --git a/apex/Android.bp b/apex/Android.bp
index 52f4384..596493a 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -15,20 +15,58 @@
     "microdroid_vendor_boot",
 ]
 
-apex {
-    name: "com.android.virt",
+soong_config_module_type {
+    name: "virt_apex",
+    module_type: "apex",
+    config_namespace: "ANDROID",
+    bool_variables: ["avf_enabled"],
+    properties: ["defaults"],
+}
 
+virt_apex {
+    name: "com.android.virt",
+    soong_config_variables: {
+        avf_enabled: {
+            defaults: ["com.android.virt_avf_enabled"],
+            conditions_default: {
+                defaults: ["com.android.virt_avf_disabled"],
+            },
+        },
+    },
+}
+
+apex_defaults {
+    name: "com.android.virt_common",
     // TODO(jiyong): make it updatable
     updatable: false,
     future_updatable: true,
     platform_apis: true,
 
-    system_ext_specific: true,
-
     manifest: "manifest.json",
 
     key: "com.android.virt.key",
     certificate: ":com.android.virt.certificate",
+
+    apps: [
+        "android.system.virtualmachine.res",
+    ],
+
+    file_contexts: ":com.android.virt-file_contexts",
+    canned_fs_config: "canned_fs_config",
+
+    bootclasspath_fragments: [
+        "com.android.virt-bootclasspath-fragment",
+    ],
+    jni_libs: [
+        "libvirtualmachine_jni",
+    ],
+}
+
+apex_defaults {
+    name: "com.android.virt_avf_enabled",
+
+    defaults: ["com.android.virt_common"],
+
     custom_sign_tool: "sign_virt_apex",
 
     // crosvm and virtualizationservice are only enabled for 64-bit targets on device
@@ -52,17 +90,9 @@
         "fd_server",
         "vm",
     ],
-    java_libs: [
-        "android.system.virtualmachine",
-    ],
-    jni_libs: [
-        "libvirtualmachine_jni",
-    ],
-    apps: [
-        "android.system.virtualmachine.res",
-    ],
     prebuilts: [
         "com.android.virt.init.rc",
+        "features_com.android.virt.xml",
         "microdroid_initrd_app_debuggable",
         "microdroid_initrd_full_debuggable",
         "microdroid_initrd_normal",
@@ -71,11 +101,18 @@
         "microdroid_bootloader.avbpubkey",
         "microdroid_kernel",
     ],
-    file_contexts: ":com.android.virt-file_contexts",
-    canned_fs_config: "canned_fs_config",
     host_required: [
         "vm_shell",
     ],
+    apps: [
+        "EmptyPayloadApp",
+    ],
+}
+
+apex_defaults {
+    name: "com.android.virt_avf_disabled",
+
+    defaults: ["com.android.virt_common"],
 }
 
 apex_key {
@@ -123,6 +160,8 @@
         // deapexer
         "deapexer",
         "debugfs_static",
+        "blkid",
+        "fsck.erofs",
 
         // sign_virt_apex
         "avbtool",
@@ -172,3 +211,43 @@
         },
     },
 }
+
+// Encapsulate the contributions made by the com.android.virt to the bootclasspath.
+bootclasspath_fragment {
+    name: "com.android.virt-bootclasspath-fragment",
+    contents: ["framework-virtualization"],
+    apex_available: ["com.android.virt"],
+
+    // The bootclasspath_fragments that provide APIs on which this depends.
+    fragments: [
+        {
+            apex: "com.android.art",
+            module: "art-bootclasspath-fragment",
+        },
+    ],
+
+    // Additional stubs libraries that this fragment's contents use which are
+    // not provided by another bootclasspath_fragment.
+    additional_stubs: [
+        "android-non-updatable",
+    ],
+
+    hidden_api: {
+
+        // This module does not contain any split packages.
+        split_packages: [],
+
+        // The following packages and all their subpackages currently only
+        // contain classes from this bootclasspath_fragment. Listing a package
+        // here won't prevent other bootclasspath modules from adding classes in
+        // any of those packages but it will prevent them from adding those
+        // classes into an API surface, e.g. public, system, etc.. Doing so will
+        // result in a build failure due to inconsistent flags.
+        package_prefixes: [
+            "android.system.virtualmachine",
+            "android.system.virtualizationservice",
+            // android.sysprop.*, renamed by jarjar
+            "com.android.system.virtualmachine.sysprop",
+        ],
+    },
+}
diff --git a/apex/canned_fs_config b/apex/canned_fs_config
index 1cf63b6..ce942d3 100644
--- a/apex/canned_fs_config
+++ b/apex/canned_fs_config
@@ -1 +1 @@
-/bin/crosvm 0 2000 0755 capabilities=0x4000
+/bin/virtualizationservice 0 2000 0755 capabilities=0x1000000  # CAP_SYS_RESOURCE
diff --git a/apex/empty-payload-apk/Android.bp b/apex/empty-payload-apk/Android.bp
new file mode 100644
index 0000000..70e6754
--- /dev/null
+++ b/apex/empty-payload-apk/Android.bp
@@ -0,0 +1,26 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_app {
+    name: "EmptyPayloadApp",
+    installable: true,
+    jni_libs: ["MicrodroidEmptyPayloadJniLib"],
+    apex_available: ["com.android.virt"],
+    sdk_version: "system_current",
+    jni_uses_platform_apis: true,
+    min_sdk_version: "UpsideDownCake",
+    target_sdk_version: "UpsideDownCake",
+    compile_multilib: "first",
+    stl: "none",
+}
+
+cc_library {
+    name: "MicrodroidEmptyPayloadJniLib",
+    srcs: ["empty_binary.cpp"],
+    shared_libs: ["libvm_payload#current"],
+    installable: true,
+    apex_available: ["com.android.virt"],
+    compile_multilib: "first",
+    stl: "none",
+}
diff --git a/apex/empty-payload-apk/AndroidManifest.xml b/apex/empty-payload-apk/AndroidManifest.xml
new file mode 100644
index 0000000..e649744
--- /dev/null
+++ b/apex/empty-payload-apk/AndroidManifest.xml
@@ -0,0 +1,9 @@
+<?xml version="1.0" encoding="utf-8"?>
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.microdroid.empty_payload">
+
+    <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
+    <uses-feature android:name="android.software.virtualization_framework" android:required="true" />
+    <application android:testOnly="true" android:hasCode="false" />
+
+</manifest>
diff --git a/apex/empty-payload-apk/empty_binary.cpp b/apex/empty-payload-apk/empty_binary.cpp
new file mode 100644
index 0000000..4308954
--- /dev/null
+++ b/apex/empty-payload-apk/empty_binary.cpp
@@ -0,0 +1,22 @@
+#include <stdio.h>
+#include <unistd.h>
+#include <vm_main.h>
+#include <vm_payload.h>
+
+extern "C" int AVmPayload_main() {
+    // disable buffering to communicate seamlessly
+    setvbuf(stdin, nullptr, _IONBF, 0);
+    setvbuf(stdout, nullptr, _IONBF, 0);
+    setvbuf(stderr, nullptr, _IONBF, 0);
+
+    printf("Hello Microdroid\n");
+
+    AVmPayload_notifyPayloadReady();
+
+    // Wait forever to allow developer to interact with Microdroid shell
+    for (;;) {
+        pause();
+    }
+
+    return 0;
+}
diff --git a/apex/permissions/Android.bp b/apex/permissions/Android.bp
new file mode 100644
index 0000000..0c925ce
--- /dev/null
+++ b/apex/permissions/Android.bp
@@ -0,0 +1,24 @@
+//
+// Copyright (C) 2022 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+prebuilt_etc {
+    name: "features_com.android.virt.xml",
+    sub_dir: "permissions",
+    src: "features_com.android.virt.xml",
+}
diff --git a/apex/permissions/features_com.android.virt.xml b/apex/permissions/features_com.android.virt.xml
new file mode 100644
index 0000000..d2b32e6
--- /dev/null
+++ b/apex/permissions/features_com.android.virt.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2022 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<permissions>
+    <feature name="android.software.virtualization_framework" />
+</permissions>
diff --git a/apex/product_packages.mk b/apex/product_packages.mk
index 5bc0044..4293c80 100644
--- a/apex/product_packages.mk
+++ b/apex/product_packages.mk
@@ -21,7 +21,6 @@
 
 PRODUCT_PACKAGES += \
     com.android.compos \
-    com.android.virt \
 
 # TODO(b/207336449): Figure out how to get these off /system
 PRODUCT_ARTIFACT_PATH_REQUIREMENT_ALLOWED_LIST := \
@@ -33,3 +32,5 @@
 PRODUCT_SYSTEM_EXT_PROPERTIES := ro.config.isolated_compilation_enabled=true
 
 PRODUCT_FSVERITY_GENERATE_METADATA := true
+
+PRODUCT_AVF_ENABLED := true
diff --git a/apex/sign_virt_apex_test.sh b/apex/sign_virt_apex_test.sh
index 640a3d4..03a56ca 100644
--- a/apex/sign_virt_apex_test.sh
+++ b/apex/sign_virt_apex_test.sh
@@ -23,8 +23,11 @@
 # To access host tools
 PATH=$TEST_DIR:$PATH
 DEBUGFS=$TEST_DIR/debugfs_static
+BLKID=$TEST_DIR/blkid
+FSCKEROFS=$TEST_DIR/fsck.erofs
 
-deapexer --debugfs_path $DEBUGFS extract $TEST_DIR/com.android.virt.apex $TMP_ROOT
+deapexer --debugfs_path $DEBUGFS --blkid_path $BLKID --fsckerofs_path $FSCKEROFS \
+  extract $TEST_DIR/com.android.virt.apex $TMP_ROOT
 
 if [ "$(ls -A $TMP_ROOT/etc/fs/)" ]; then
   sign_virt_apex $TEST_DIR/test.com.android.virt.pem $TMP_ROOT
diff --git a/apkdmverity/src/main.rs b/apkdmverity/src/main.rs
index a69b583..6e12e38 100644
--- a/apkdmverity/src/main.rs
+++ b/apkdmverity/src/main.rs
@@ -162,7 +162,7 @@
     }
 
     fn create_block_aligned_file(path: &Path, data: &[u8]) {
-        let mut f = File::create(&path).unwrap();
+        let mut f = File::create(path).unwrap();
         f.write_all(data).unwrap();
 
         // Add padding so that the size of the file is multiple of 4096.
diff --git a/authfs/aidl/com/android/virt/fs/IAuthFsService.aidl b/authfs/aidl/com/android/virt/fs/IAuthFsService.aidl
index b349db2..30cc281 100644
--- a/authfs/aidl/com/android/virt/fs/IAuthFsService.aidl
+++ b/authfs/aidl/com/android/virt/fs/IAuthFsService.aidl
@@ -21,6 +21,8 @@
 
 /** @hide */
 interface IAuthFsService {
+    const String AUTHFS_SERVICE_SOCKET_NAME = "authfs_service";
+
     /**
      * Creates an AuthFS mount given the config. Returns the binder object that represent the AuthFS
      * instance. The AuthFS setup is deleted once the lifetime of the returned binder object ends.
diff --git a/authfs/fd_server/src/main.rs b/authfs/fd_server/src/main.rs
index f1fffdd..21d0e64 100644
--- a/authfs/fd_server/src/main.rs
+++ b/authfs/fd_server/src/main.rs
@@ -29,7 +29,7 @@
 use clap::Parser;
 use log::debug;
 use nix::sys::stat::{umask, Mode};
-use rpcbinder::run_vsock_rpc_server;
+use rpcbinder::RpcServer;
 use std::collections::BTreeMap;
 use std::fs::File;
 use std::os::unix::io::{FromRawFd, OwnedFd};
@@ -135,18 +135,14 @@
     let old_umask = umask(Mode::empty());
     debug!("Setting umask to 0 (old: {:03o})", old_umask.bits());
 
-    let service = FdService::new_binder(fd_pool).as_binder();
     debug!("fd_server is starting as a rpc service.");
-    let retval = run_vsock_rpc_server(service, RPC_SERVICE_PORT, || {
-        debug!("fd_server is ready");
-        // Close the ready-fd if we were given one to signal our readiness.
-        drop(ready_fd.take());
-    });
+    let service = FdService::new_binder(fd_pool).as_binder();
+    let server = RpcServer::new_vsock(service, RPC_SERVICE_PORT)?;
+    debug!("fd_server is ready");
 
-    if retval {
-        debug!("RPC server has shut down gracefully");
-        Ok(())
-    } else {
-        bail!("Premature termination of RPC server");
-    }
+    // Close the ready-fd if we were given one to signal our readiness.
+    drop(ready_fd.take());
+
+    server.join();
+    Ok(())
 }
diff --git a/authfs/service/Android.bp b/authfs/service/Android.bp
index e9eec1e..de6326d 100644
--- a/authfs/service/Android.bp
+++ b/authfs/service/Android.bp
@@ -16,6 +16,7 @@
         "liblibc",
         "liblog_rust",
         "libnix",
+        "librpcbinder_rs",
         "libshared_child",
     ],
     prefer_rlib: true,
diff --git a/authfs/service/authfs_service.rc b/authfs/service/authfs_service.rc
index 9ad0ce6..7edb1ca 100644
--- a/authfs/service/authfs_service.rc
+++ b/authfs/service/authfs_service.rc
@@ -1,2 +1,3 @@
 service authfs_service /system/bin/authfs_service
     disabled
+    socket authfs_service stream 0666 root system
diff --git a/authfs/service/src/main.rs b/authfs/service/src/main.rs
index 77cac9a..e710f07 100644
--- a/authfs/service/src/main.rs
+++ b/authfs/service/src/main.rs
@@ -22,8 +22,9 @@
 
 mod authfs;
 
-use anyhow::{bail, Context, Result};
+use anyhow::{bail, Result};
 use log::*;
+use rpcbinder::RpcServer;
 use std::ffi::OsString;
 use std::fs::{create_dir, read_dir, remove_dir_all, remove_file};
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -31,13 +32,10 @@
 use authfs_aidl_interface::aidl::com::android::virt::fs::AuthFsConfig::AuthFsConfig;
 use authfs_aidl_interface::aidl::com::android::virt::fs::IAuthFs::IAuthFs;
 use authfs_aidl_interface::aidl::com::android::virt::fs::IAuthFsService::{
-    BnAuthFsService, IAuthFsService,
+    BnAuthFsService, IAuthFsService, AUTHFS_SERVICE_SOCKET_NAME,
 };
-use binder::{
-    self, add_service, BinderFeatures, ExceptionCode, Interface, ProcessState, Status, Strong,
-};
+use binder::{self, BinderFeatures, ExceptionCode, Interface, Status, Strong};
 
-const SERVICE_NAME: &str = "authfs_service";
 const SERVICE_ROOT: &str = "/data/misc/authfs";
 
 /// Implementation of `IAuthFsService`.
@@ -117,15 +115,13 @@
 
     clean_up_working_directory()?;
 
-    ProcessState::start_thread_pool();
-
     let service = AuthFsService::new_binder(debuggable).as_binder();
-    add_service(SERVICE_NAME, service)
-        .with_context(|| format!("Failed to register service {}", SERVICE_NAME))?;
-    debug!("{} is running", SERVICE_NAME);
-
-    ProcessState::join_thread_pool();
-    bail!("Unexpected exit after join_thread_pool")
+    debug!("{} is starting as a rpc service.", AUTHFS_SERVICE_SOCKET_NAME);
+    let server = RpcServer::new_init_unix_domain(service, AUTHFS_SERVICE_SOCKET_NAME)?;
+    info!("The RPC server '{}' is running.", AUTHFS_SERVICE_SOCKET_NAME);
+    server.join();
+    info!("The RPC server at '{}' has shut down gracefully.", AUTHFS_SERVICE_SOCKET_NAME);
+    Ok(())
 }
 
 fn main() {
diff --git a/authfs/tests/benchmarks/src/java/com/android/fs/benchmarks/AuthFsBenchmarks.java b/authfs/tests/benchmarks/src/java/com/android/fs/benchmarks/AuthFsBenchmarks.java
index 428c816..32eafb8 100644
--- a/authfs/tests/benchmarks/src/java/com/android/fs/benchmarks/AuthFsBenchmarks.java
+++ b/authfs/tests/benchmarks/src/java/com/android/fs/benchmarks/AuthFsBenchmarks.java
@@ -20,6 +20,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
 import android.cts.host.utils.DeviceJUnit4ClassRunnerWithParameters;
@@ -27,6 +28,7 @@
 import android.platform.test.annotations.RootPermissionTest;
 
 import com.android.fs.common.AuthFsTestRule;
+import com.android.microdroid.test.common.DeviceProperties;
 import com.android.microdroid.test.common.MetricsProcessor;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.metrics.proto.MetricMeasurement.DataType;
@@ -82,9 +84,10 @@
         AuthFsTestRule.setUpAndroid(getTestInformation());
         mAuthFsTestRule.setUpTest();
         assumeTrue(AuthFsTestRule.getDevice().supportsMicrodroid(mProtectedVm));
-        String metricsPrefix =
-                MetricsProcessor.getMetricPrefix(
-                        getDevice().getProperty("debug.hypervisor.metrics_tag"));
+        DeviceProperties deviceProperties = DeviceProperties.create(getDevice()::getProperty);
+        assumeFalse(
+                "Skip on CF; no need to collect metrics on CF", deviceProperties.isCuttlefish());
+        String metricsPrefix = MetricsProcessor.getMetricPrefix(deviceProperties.getMetricsTag());
         mMetricsProcessor = new MetricsProcessor(metricsPrefix + "authfs/");
         AuthFsTestRule.startMicrodroid(mProtectedVm);
     }
@@ -137,6 +140,7 @@
 
             String rate = mAuthFsTestRule.getMicrodroid().run(cmd);
             rates.add(Double.parseDouble(rate));
+            mAuthFsTestRule.killFdServerOnAndroid();
         }
         reportMetrics(rates, mode + "_read", "mb_per_sec");
     }
@@ -149,11 +153,14 @@
         List<Double> rates = new ArrayList<>(TRIAL_COUNT);
         for (int i = 0; i < TRIAL_COUNT + 1; ++i) {
             mAuthFsTestRule.runFdServerOnAndroid(
-                    "--open-rw 5:" + mAuthFsTestRule.TEST_OUTPUT_DIR + "/out.file", "--rw-fds 5");
+                    "--open-rw 5:" + AuthFsTestRule.TEST_OUTPUT_DIR + "/out.file", "--rw-fds 5");
             mAuthFsTestRule.runAuthFsOnMicrodroid("--remote-new-rw-file 5");
 
             String rate = mAuthFsTestRule.getMicrodroid().run(cmd);
             rates.add(Double.parseDouble(rate));
+            mAuthFsTestRule.killFdServerOnAndroid();
+            AuthFsTestRule.getAndroid()
+                    .runForResult("rm", "-rf", AuthFsTestRule.TEST_OUTPUT_DIR + "/out.file");
         }
         reportMetrics(rates, mode + "_write", "mb_per_sec");
     }
diff --git a/authfs/tests/common/src/java/com/android/fs/common/AuthFsTestRule.java b/authfs/tests/common/src/java/com/android/fs/common/AuthFsTestRule.java
index 6087eef..357edea 100644
--- a/authfs/tests/common/src/java/com/android/fs/common/AuthFsTestRule.java
+++ b/authfs/tests/common/src/java/com/android/fs/common/AuthFsTestRule.java
@@ -181,6 +181,10 @@
         Future<?> unusedFuture = mThreadPool.submit(() -> runForResult(sAndroid, cmd, "fd_server"));
     }
 
+    public void killFdServerOnAndroid() throws DeviceNotAvailableException {
+        sAndroid.tryRun("killall fd_server");
+    }
+
     public void runAuthFsOnMicrodroid(String flags) {
         String cmd = AUTHFS_BIN + " " + MOUNT_DIR + " " + flags + " --cid " + VMADDR_CID_HOST;
 
@@ -250,7 +254,7 @@
         }
 
         assertNotNull(sAndroid);
-        sAndroid.tryRun("killall fd_server");
+        killFdServerOnAndroid();
 
         // 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,
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 68e1948..601c6fc 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -32,10 +32,8 @@
 use log::{info, warn};
 use rustutils::system_properties;
 use std::fs::{self, File};
-use std::io::{BufRead, BufReader};
 use std::num::NonZeroU32;
 use std::path::{Path, PathBuf};
-use std::thread;
 use vmclient::{DeathReason, ErrorCode, VmInstance, VmWaitError};
 
 /// This owns an instance of the CompOS VM.
@@ -240,12 +238,7 @@
 
 struct Callback {}
 impl vmclient::VmCallback for Callback {
-    fn on_payload_started(&self, cid: i32, stream: Option<&File>) {
-        if let Some(file) = stream {
-            if let Err(e) = start_logging(file) {
-                warn!("Can't log vm output: {}", e);
-            };
-        }
+    fn on_payload_started(&self, cid: i32) {
         log::info!("VM payload started, cid = {}", cid);
     }
 
@@ -265,19 +258,3 @@
         log::warn!("VM died, cid = {}, reason = {:?}", cid, death_reason);
     }
 }
-
-fn start_logging(file: &File) -> Result<()> {
-    let reader = BufReader::new(file.try_clone().context("Cloning file failed")?);
-    thread::spawn(move || {
-        for line in reader.lines() {
-            match line {
-                Ok(line) => info!("VM: {}", line),
-                Err(e) => {
-                    warn!("Reading VM output failed: {}", e);
-                    break;
-                }
-            }
-        }
-    });
-    Ok(())
-}
diff --git a/compos/compos_key_helper/Android.bp b/compos/compos_key_helper/Android.bp
index c9480fc..f8dc783 100644
--- a/compos/compos_key_helper/Android.bp
+++ b/compos/compos_key_helper/Android.bp
@@ -24,11 +24,12 @@
     defaults: ["compos_key_defaults"],
     srcs: ["compos_key_main.cpp"],
 
+    header_libs: ["vm_payload_restricted_headers"],
     static_libs: [
         "libcompos_key",
     ],
     shared_libs: [
-        "libvm_payload",
+        "libvm_payload#current",
         "libbinder_ndk",
     ],
 }
diff --git a/compos/compos_key_helper/compos_key_main.cpp b/compos/compos_key_helper/compos_key_main.cpp
index 4fb0762..f1637e0 100644
--- a/compos/compos_key_helper/compos_key_main.cpp
+++ b/compos/compos_key_helper/compos_key_main.cpp
@@ -17,7 +17,7 @@
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <unistd.h>
-#include <vm_payload.h>
+#include <vm_payload_restricted.h>
 
 #include <string_view>
 #include <vector>
@@ -38,11 +38,8 @@
 
 Result<Ed25519KeyPair> getSigningKey() {
     Seed seed;
-    if (!AVmPayload_getVmInstanceSecret(kSigningKeySeedIdentifier,
-                                        strlen(kSigningKeySeedIdentifier), seed.data(),
-                                        seed.size())) {
-        return Error() << "Failed to get signing key seed";
-    }
+    AVmPayload_getVmInstanceSecret(kSigningKeySeedIdentifier, strlen(kSigningKeySeedIdentifier),
+                                   seed.data(), seed.size());
     return compos_key::keyFromSeed(seed);
 }
 
@@ -60,16 +57,9 @@
 }
 
 int write_bcc() {
-    size_t bcc_size;
-    if (!AVmPayload_getDiceAttestationChain(nullptr, 0, &bcc_size)) {
-        LOG(ERROR) << "Failed to measure attestation chain";
-        return 1;
-    }
+    size_t bcc_size = AVmPayload_getDiceAttestationChain(nullptr, 0);
     std::vector<uint8_t> bcc(bcc_size);
-    if (!AVmPayload_getDiceAttestationChain(bcc.data(), bcc.size(), &bcc_size)) {
-        LOG(ERROR) << "Failed to get attestation chain";
-        return 1;
-    }
+    AVmPayload_getDiceAttestationChain(bcc.data(), bcc.size());
 
     if (!WriteFully(STDOUT_FILENO, bcc.data(), bcc.size())) {
         PLOG(ERROR) << "Write failed";
diff --git a/compos/src/artifact_signer.rs b/compos/src/artifact_signer.rs
index e51b8dd..d3843fc 100644
--- a/compos/src/artifact_signer.rs
+++ b/compos/src/artifact_signer.rs
@@ -46,7 +46,7 @@
     pub fn add_artifact(&mut self, path: &Path) -> Result<()> {
         // The path we store is where the file will be when it is verified, not where it is now.
         let suffix = path
-            .strip_prefix(&self.base_directory)
+            .strip_prefix(self.base_directory)
             .context("Artifacts must be under base directory")?;
         let target_path = Path::new(TARGET_DIRECTORY).join(suffix);
         let target_path = target_path.to_str().ok_or_else(|| anyhow!("Invalid path"))?;
diff --git a/compos/src/compsvc.rs b/compos/src/compsvc.rs
index 4330bbf..40d14d8 100644
--- a/compos/src/compsvc.rs
+++ b/compos/src/compsvc.rs
@@ -30,14 +30,16 @@
 use crate::artifact_signer::ArtifactSigner;
 use crate::compilation::odrefresh;
 use crate::compos_key;
+use authfs_aidl_interface::aidl::com::android::virt::fs::IAuthFsService::{
+    IAuthFsService, AUTHFS_SERVICE_SOCKET_NAME,
+};
 use binder::{BinderFeatures, ExceptionCode, Interface, Result as BinderResult, Status, Strong};
 use compos_aidl_interface::aidl::com::android::compos::ICompOsService::{
     BnCompOsService, ICompOsService, OdrefreshArgs::OdrefreshArgs,
 };
 use compos_common::binder::to_binder_result;
 use compos_common::odrefresh::{is_system_property_interesting, ODREFRESH_PATH};
-
-const AUTHFS_SERVICE_NAME: &str = "authfs_service";
+use rpcbinder::get_unix_domain_rpc_interface;
 
 /// Constructs a binder object that implements ICompOsService.
 pub fn new_binder() -> Result<Strong<dyn ICompOsService>> {
@@ -127,8 +129,10 @@
 
 impl CompOsService {
     fn do_odrefresh(&self, args: &OdrefreshArgs) -> Result<i8> {
-        let authfs_service = binder::get_interface(AUTHFS_SERVICE_NAME)
-            .context("Unable to connect to AuthFS service")?;
+        log::debug!("Prepare to connect to {}", AUTHFS_SERVICE_SOCKET_NAME);
+        let authfs_service: Strong<dyn IAuthFsService> =
+            get_unix_domain_rpc_interface(AUTHFS_SERVICE_SOCKET_NAME)
+                .with_context(|| format!("Failed to connect to {}", AUTHFS_SERVICE_SOCKET_NAME))?;
         let exit_code = odrefresh(&self.odrefresh_path, args, authfs_service, |output_dir| {
             // authfs only shows us the files we created, so it's ok to just sign everything
             // under the output directory.
@@ -144,7 +148,7 @@
 
 fn add_artifacts(target_dir: &Path, artifact_signer: &mut ArtifactSigner) -> Result<()> {
     for entry in
-        read_dir(&target_dir).with_context(|| format!("Traversing {}", target_dir.display()))?
+        read_dir(target_dir).with_context(|| format!("Traversing {}", target_dir.display()))?
     {
         let entry = entry?;
         let file_type = entry.file_type()?;
diff --git a/compos/src/compsvc_main.rs b/compos/src/compsvc_main.rs
index a4e3903..77e2daa 100644
--- a/compos/src/compsvc_main.rs
+++ b/compos/src/compsvc_main.rs
@@ -22,12 +22,14 @@
 mod compsvc;
 mod fsverity;
 
-use anyhow::{bail, Result};
+use anyhow::Result;
+use binder::unstable_api::AsNative;
 use compos_common::COMPOS_VSOCK_PORT;
 use log::{debug, error};
-use rpcbinder::run_vsock_rpc_server;
+use std::os::raw::c_void;
 use std::panic;
-use vm_payload_bindgen::AVmPayload_notifyPayloadReady;
+use std::ptr;
+use vm_payload_bindgen::{AIBinder, AVmPayload_notifyPayloadReady, AVmPayload_runVsockRpcServer};
 
 fn main() {
     if let Err(e) = try_main() {
@@ -45,16 +47,21 @@
         error!("{}", panic_info);
     }));
 
-    let service = compsvc::new_binder()?.as_binder();
     debug!("compsvc is starting as a rpc service.");
-    // SAFETY: Invokes a method from the bindgen library `vm_payload_bindgen`.
-    let retval = run_vsock_rpc_server(service, COMPOS_VSOCK_PORT, || unsafe {
-        AVmPayload_notifyPayloadReady();
-    });
-    if retval {
-        debug!("RPC server has shut down gracefully");
-        Ok(())
-    } else {
-        bail!("Premature termination of RPC server");
+    let param = ptr::null_mut();
+    let mut service = compsvc::new_binder()?.as_binder();
+    unsafe {
+        // SAFETY: We hold a strong pointer, so the raw pointer remains valid. The bindgen AIBinder
+        // is the same type as sys::AIBinder.
+        let service = service.as_native_mut() as *mut AIBinder;
+        // SAFETY: It is safe for on_ready to be invoked at any time, with any parameter.
+        AVmPayload_runVsockRpcServer(service, COMPOS_VSOCK_PORT, Some(on_ready), param);
     }
+    Ok(())
+}
+
+extern "C" fn on_ready(_param: *mut c_void) {
+    // SAFETY: Invokes a method from the bindgen library `vm_payload_bindgen` which is safe to
+    // call at any time.
+    unsafe { AVmPayload_notifyPayloadReady() };
 }
diff --git a/compos/tests/AndroidTest.xml b/compos/tests/AndroidTest.xml
index 2a84291..f9e6837 100644
--- a/compos/tests/AndroidTest.xml
+++ b/compos/tests/AndroidTest.xml
@@ -14,6 +14,8 @@
      limitations under the License.
 -->
 <configuration description="Tests for CompOS">
+    <option name="config-descriptor:metadata" key="mainline-param" value="com.google.android.art.apex" />
+
     <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer">
         <option name="force-root" value="true" />
     </target_preparer>
diff --git a/demo/Android.bp b/demo/Android.bp
index 8613166..a291ee1 100644
--- a/demo/Android.bp
+++ b/demo/Android.bp
@@ -12,11 +12,9 @@
         "com.android.microdroid.testservice-java",
         "com.google.android.material_material",
     ],
-    libs: [
-        "android.system.virtualmachine",
-    ],
+    sdk_version: "system_current",
     jni_libs: ["MicrodroidTestNativeLib"],
-    platform_apis: true,
+    jni_uses_platform_apis: true,
     use_embedded_native_libs: true,
     v4_signature: true,
     min_sdk_version: "33",
diff --git a/demo/AndroidManifest.xml b/demo/AndroidManifest.xml
index 6669adb..17a7680 100644
--- a/demo/AndroidManifest.xml
+++ b/demo/AndroidManifest.xml
@@ -4,11 +4,11 @@
 
     <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
     <uses-sdk android:minSdkVersion="33" android:targetSdkVersion="33"/>
+    <uses-feature android:name="android.software.virtualization_framework" android:required="true" />
     <application
         android:label="MicrodroidDemo"
         android:theme="@style/Theme.MicrodroidDemo"
         android:testOnly="true">
-        <uses-library android:name="android.system.virtualmachine" android:required="true" />
         <activity android:name=".MainActivity" android:exported="true">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
diff --git a/demo/java/com/android/microdroid/demo/MainActivity.java b/demo/java/com/android/microdroid/demo/MainActivity.java
index b52ef40..77f2ee7 100644
--- a/demo/java/com/android/microdroid/demo/MainActivity.java
+++ b/demo/java/com/android/microdroid/demo/MainActivity.java
@@ -42,7 +42,6 @@
 import com.android.microdroid.testservice.ITestService;
 
 import java.io.BufferedReader;
-import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
@@ -145,6 +144,7 @@
 
     /** Models a virtual machine and outputs from it. */
     public static class VirtualMachineModel extends AndroidViewModel {
+        private static final String VM_NAME = "demo_vm";
         private VirtualMachine mVirtualMachine;
         private final MutableLiveData<String> mConsoleOutput = new MutableLiveData<>();
         private final MutableLiveData<String> mLogOutput = new MutableLiveData<>();
@@ -168,16 +168,7 @@
                         private final ExecutorService mService = mExecutorService;
 
                         @Override
-                        public void onPayloadStarted(VirtualMachine vm,
-                                ParcelFileDescriptor stream) {
-                            if (stream == null) {
-                                mPayloadOutput.postValue("(no output available)");
-                                return;
-                            }
-
-                            InputStream input = new FileInputStream(stream.getFileDescriptor());
-                            mService.execute(new Reader("payload", mPayloadOutput, input));
-                        }
+                        public void onPayloadStarted(VirtualMachine vm) {}
 
                         @Override
                         public void onPayloadReady(VirtualMachine vm) {
@@ -265,20 +256,21 @@
                     builder.setDebugLevel(VirtualMachineConfig.DEBUG_LEVEL_FULL);
                 }
                 VirtualMachineConfig config = builder.build();
-                VirtualMachineManager vmm = VirtualMachineManager.getInstance(getApplication());
-                mVirtualMachine = vmm.getOrCreate("demo_vm", config);
+                VirtualMachineManager vmm =
+                        getApplication().getSystemService(VirtualMachineManager.class);
+                mVirtualMachine = vmm.getOrCreate(VM_NAME, config);
                 try {
                     mVirtualMachine.setConfig(config);
                 } catch (VirtualMachineException e) {
-                    mVirtualMachine.delete();
-                    mVirtualMachine = vmm.create("demo_vm", config);
+                    vmm.delete(VM_NAME);
+                    mVirtualMachine = vmm.create(VM_NAME, config);
                 }
                 mVirtualMachine.run();
                 mVirtualMachine.setCallback(Executors.newSingleThreadExecutor(), callback);
                 mStatus.postValue(mVirtualMachine.getStatus());
 
-                InputStream console = mVirtualMachine.getConsoleOutputStream();
-                InputStream log = mVirtualMachine.getLogOutputStream();
+                InputStream console = mVirtualMachine.getConsoleOutput();
+                InputStream log = mVirtualMachine.getLogOutput();
                 mExecutorService.execute(new Reader("console", mConsoleOutput, console));
                 mExecutorService.execute(new Reader("log", mLogOutput, log));
             } catch (VirtualMachineException e) {
diff --git a/docs/getting_started/index.md b/docs/getting_started/index.md
index 245aba6..f184862 100644
--- a/docs/getting_started/index.md
+++ b/docs/getting_started/index.md
@@ -97,6 +97,15 @@
 If you run into problems, inspect the logs produced by `atest`. Their location is printed at the
 end. The `host_log_*.zip` file should contain the output of individual commands as well as VM logs.
 
+### Custom pvmfw
+
+Hostside tests, which run on the PC and extends `MicrodroidHostTestCaseBase`, can be run with
+a custom `pvmfw`. Use `--module-arg` to push `pvmfw` for individual test methods.
+
+```shell
+atest com.android.microdroid.test.MicrodroidHostTests -- --module-arg MicrodroidHostTestCases:set-option:pvmfw:pvmfw.img
+```
+
 ## Spawning your own VMs with custom kernel
 
 You can spawn your own VMs by passing a JSON config file to the VirtualizationService via the `vm`
@@ -120,10 +129,30 @@
 The `vm` command also has other subcommands for debugging; run `/apex/com.android.virt/bin/vm help`
 for details.
 
+## Spawning your own VMs with custom pvmfw
+
+Set system property `hypervisor.pvmfw.path` to custom `pvmfw` on the device before using `vm` tool.
+`virtualizationservice` will pass the specified `pvmfw` to `crosvm` for protected VMs.
+
+```shell
+adb push pvmfw.img /data/local/tmp/pvmfw.img
+adb root  # required for setprop
+adb shell setprop hypervisor.pvmfw.path /data/local/tmp/pvmfw.img
+```
+
 ## Spawning your own VMs with Microdroid
 
 [Microdroid](../../microdroid/README.md) is a lightweight version of Android that is intended to run
-on pVM. You can manually run the demo app on top of Microdroid as follows:
+on pVM. You can run a Microdroid with empty payload using the following command:
+
+```shell
+adb shell /apex/com.android.virt/bin/vm run-microdroid --debug full
+```
+
+The `instance.img` and `apk.idsig` files will be stored in a subdirectory under
+`/data/local/tmp/microdroid`, that `vm` will create.
+
+Atlernatively, you can manually run the demo app on top of Microdroid as follows:
 
 ```shell
 UNBUNDLED_BUILD_SDKS_FROM_SOURCE=true TARGET_BUILD_APPS=MicrodroidDemoApp m apps_only dist
@@ -133,7 +162,8 @@
   --debug full \
   /data/local/tmp/virt/MicrodroidDemoApp.apk \
   /data/local/tmp/virt/MicrodroidDemoApp.apk.idsig \
-  /data/local/tmp/virt/instance.img assets/vm_config.json
+  /data/local/tmp/virt/instance.img \
+  --payload-path MicrodroidTestNativeLib.so
 ```
 
 ## Building and updating CrosVM and VirtualizationService {#building-and-updating}
diff --git a/encryptedstore/Android.bp b/encryptedstore/Android.bp
new file mode 100644
index 0000000..13ef1b9
--- /dev/null
+++ b/encryptedstore/Android.bp
@@ -0,0 +1,31 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_defaults {
+    name: "encryptedstore.defaults",
+    srcs: ["src/main.rs"],
+    edition: "2021",
+    prefer_rlib: true,
+    rustlibs: [
+        "libandroid_logger",
+        "libanyhow",
+        "liblibc",
+        "libclap",
+        "libhex",
+        "liblog_rust",
+        "libnix",
+        "libdm_rust",
+    ],
+    multilib: {
+        lib32: {
+            enabled: false,
+        },
+    },
+}
+
+rust_binary {
+    name: "encryptedstore",
+    defaults: ["encryptedstore.defaults"],
+    bootstrap: true,
+}
diff --git a/encryptedstore/src/main.rs b/encryptedstore/src/main.rs
new file mode 100644
index 0000000..9c8311d
--- /dev/null
+++ b/encryptedstore/src/main.rs
@@ -0,0 +1,159 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//! `encryptedstore` is a program that (as the name indicates) provides encrypted storage
+//! solution in a VM. This is based on dm-crypt & requires the (64 bytes') key & the backing device.
+//! It uses dm_rust lib.
+
+use anyhow::{ensure, Context, Result};
+use clap::{arg, App};
+use dm::crypt::CipherType;
+use dm::util;
+use log::info;
+use std::ffi::CString;
+use std::fs::{create_dir_all, OpenOptions};
+use std::io::{Error, Read, Write};
+use std::os::unix::ffi::OsStrExt;
+use std::os::unix::fs::FileTypeExt;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+
+const MK2FS_BIN: &str = "/system/bin/mke2fs";
+const UNFORMATTED_STORAGE_MAGIC: &str = "UNFORMATTED-STORAGE";
+
+fn main() -> Result<()> {
+    android_logger::init_once(
+        android_logger::Config::default()
+            .with_tag("encryptedstore")
+            .with_min_level(log::Level::Info),
+    );
+    info!("Starting encryptedstore binary");
+
+    let matches = App::new("encryptedstore")
+        .args(&[
+            arg!(--blkdevice <FILE> "the block device backing the encrypted storage")
+                .required(true),
+            arg!(--key <KEY> "key (in hex) equivalent to 32 bytes)").required(true),
+            arg!(--mountpoint <MOUNTPOINT> "mount point for the storage").required(true),
+        ])
+        .get_matches();
+
+    let blkdevice = Path::new(matches.value_of("blkdevice").unwrap());
+    let key = matches.value_of("key").unwrap();
+    let mountpoint = Path::new(matches.value_of("mountpoint").unwrap());
+    encryptedstore_init(blkdevice, key, mountpoint).context(format!(
+        "Unable to initialize encryptedstore on {:?} & mount at {:?}",
+        blkdevice, mountpoint
+    ))?;
+    Ok(())
+}
+
+fn encryptedstore_init(blkdevice: &Path, key: &str, mountpoint: &Path) -> Result<()> {
+    ensure!(
+        std::fs::metadata(&blkdevice)
+            .context(format!("Failed to get metadata of {:?}", blkdevice))?
+            .file_type()
+            .is_block_device(),
+        "The path:{:?} is not of a block device",
+        blkdevice
+    );
+
+    let needs_formatting =
+        needs_formatting(blkdevice).context("Unable to check if formatting is required")?;
+    let crypt_device =
+        enable_crypt(blkdevice, key, "cryptdev").context("Unable to map crypt device")?;
+
+    // We might need to format it with filesystem if this is a "seen-for-the-first-time" device.
+    if needs_formatting {
+        info!("Freshly formatting the crypt device");
+        format_ext4(&crypt_device)?;
+    }
+    mount(&crypt_device, mountpoint).context(format!("Unable to mount {:?}", crypt_device))?;
+    Ok(())
+}
+
+fn enable_crypt(data_device: &Path, key: &str, name: &str) -> Result<PathBuf> {
+    let dev_size = util::blkgetsize64(data_device)?;
+    let key = hex::decode(key).context("Unable to decode hex key")?;
+
+    // Create the dm-crypt spec
+    let target = dm::crypt::DmCryptTargetBuilder::default()
+        .data_device(data_device, dev_size)
+        .cipher(CipherType::AES256HCTR2)
+        .key(&key)
+        .build()
+        .context("Couldn't build the DMCrypt target")?;
+    let dm = dm::DeviceMapper::new()?;
+    dm.create_crypt_device(name, &target).context("Failed to create dm-crypt device")
+}
+
+// The disk contains UNFORMATTED_STORAGE_MAGIC to indicate we need to format the crypt device.
+// This function looks for it, zeroing it, if present.
+fn needs_formatting(data_device: &Path) -> Result<bool> {
+    let mut file = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .open(data_device)
+        .with_context(|| format!("Failed to open {:?}", data_device))?;
+
+    let mut buf = [0; UNFORMATTED_STORAGE_MAGIC.len()];
+    file.read_exact(&mut buf)?;
+
+    if buf == UNFORMATTED_STORAGE_MAGIC.as_bytes() {
+        buf.fill(0);
+        file.write_all(&buf)?;
+        return Ok(true);
+    }
+    Ok(false)
+}
+
+fn format_ext4(device: &Path) -> Result<()> {
+    let mkfs_options = [
+        "-j",               // Create appropriate sized journal
+        "-O metadata_csum", // Metadata checksum for filesystem integrity
+    ];
+    let mut cmd = Command::new(MK2FS_BIN);
+    let status = cmd
+        .args(mkfs_options)
+        .arg(device)
+        .status()
+        .context(format!("failed to execute {}", MK2FS_BIN))?;
+    ensure!(status.success(), "mkfs failed with {:?}", status);
+    Ok(())
+}
+
+fn mount(source: &Path, mountpoint: &Path) -> Result<()> {
+    create_dir_all(mountpoint).context(format!("Failed to create {:?}", &mountpoint))?;
+    let mount_options = CString::new("").unwrap();
+    let source = CString::new(source.as_os_str().as_bytes())?;
+    let mountpoint = CString::new(mountpoint.as_os_str().as_bytes())?;
+    let fstype = CString::new("ext4").unwrap();
+
+    let ret = unsafe {
+        libc::mount(
+            source.as_ptr(),
+            mountpoint.as_ptr(),
+            fstype.as_ptr(),
+            libc::MS_NOSUID | libc::MS_NODEV | libc::MS_NOEXEC,
+            mount_options.as_ptr() as *const std::ffi::c_void,
+        )
+    };
+    if ret < 0 {
+        Err(Error::last_os_error()).context("mount failed")
+    } else {
+        Ok(())
+    }
+}
diff --git a/javalib/Android.bp b/javalib/Android.bp
index 51dd381..71287f2 100644
--- a/javalib/Android.bp
+++ b/javalib/Android.bp
@@ -11,9 +11,10 @@
 }
 
 java_sdk_library {
-    name: "android.system.virtualmachine",
-    installable: true,
-    compile_dex: true,
+    name: "framework-virtualization",
+
+    // TODO(b/243512044): introduce non-updatable-framework-module-defaults
+    defaults: ["framework-module-defaults"],
 
     jarjar_rules: "jarjar-rules.txt",
 
@@ -25,12 +26,36 @@
     ],
 
     apex_available: ["com.android.virt"],
+
     permitted_packages: [
         "android.system.virtualmachine",
         "android.system.virtualizationservice",
         // android.sysprop.*, renamed by jarjar
         "com.android.system.virtualmachine.sysprop",
     ],
+    errorprone: {
+        enabled: true,
+        javacflags: [
+            // We use @GuardedBy and we want a test failure if our locking isn't consistent with it.
+            "-Xep:GuardedBy:ERROR",
+        ],
+    },
+
+    test: {
+        enabled: true,
+        sdk_version: "module_current",
+    },
+
+    sdk_version: "core_platform",
+    stub_only_libs: [
+        "android_module_lib_stubs_current",
+    ],
+    impl_only_libs: [
+        "framework",
+    ],
+    impl_library_visibility: [
+        "//packages/modules/Virtualization:__subpackages__",
+    ],
 }
 
 prebuilt_apis {
diff --git a/javalib/api/module-lib-current.txt b/javalib/api/module-lib-current.txt
new file mode 100644
index 0000000..4d59764
--- /dev/null
+++ b/javalib/api/module-lib-current.txt
@@ -0,0 +1,9 @@
+// Signature format: 2.0
+package android.system.virtualmachine {
+
+  public class VirtualizationFrameworkInitializer {
+    method public static void registerServiceWrappers();
+  }
+
+}
+
diff --git a/javalib/api/module-lib-removed.txt b/javalib/api/module-lib-removed.txt
new file mode 100644
index 0000000..d802177
--- /dev/null
+++ b/javalib/api/module-lib-removed.txt
@@ -0,0 +1 @@
+// Signature format: 2.0
diff --git a/javalib/api/system-current.txt b/javalib/api/system-current.txt
index d802177..16995c5 100644
--- a/javalib/api/system-current.txt
+++ b/javalib/api/system-current.txt
@@ -1 +1,108 @@
 // Signature format: 2.0
+package android.system.virtualmachine {
+
+  public class VirtualMachine implements java.lang.AutoCloseable {
+    method public void clearCallback();
+    method public void close();
+    method @NonNull public android.os.IBinder connectToVsockServer(int) throws android.system.virtualmachine.VirtualMachineException;
+    method @NonNull public android.os.ParcelFileDescriptor connectVsock(int) throws android.system.virtualmachine.VirtualMachineException;
+    method public int getCid() throws android.system.virtualmachine.VirtualMachineException;
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig getConfig();
+    method @NonNull public java.io.InputStream getConsoleOutput() throws android.system.virtualmachine.VirtualMachineException;
+    method @NonNull public java.io.InputStream getLogOutput() throws android.system.virtualmachine.VirtualMachineException;
+    method @NonNull public String getName();
+    method public int getStatus();
+    method @RequiresPermission(android.system.virtualmachine.VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION) public void run() throws android.system.virtualmachine.VirtualMachineException;
+    method public void setCallback(@NonNull java.util.concurrent.Executor, @NonNull android.system.virtualmachine.VirtualMachineCallback);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig setConfig(@NonNull android.system.virtualmachine.VirtualMachineConfig) throws android.system.virtualmachine.VirtualMachineException;
+    method public void stop() throws android.system.virtualmachine.VirtualMachineException;
+    method @NonNull public android.system.virtualmachine.VirtualMachineDescriptor toDescriptor() throws android.system.virtualmachine.VirtualMachineException;
+    field public static final String MANAGE_VIRTUAL_MACHINE_PERMISSION = "android.permission.MANAGE_VIRTUAL_MACHINE";
+    field public static final int STATUS_DELETED = 2; // 0x2
+    field public static final int STATUS_RUNNING = 1; // 0x1
+    field public static final int STATUS_STOPPED = 0; // 0x0
+    field public static final String USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION = "android.permission.USE_CUSTOM_VIRTUAL_MACHINE";
+  }
+
+  public interface VirtualMachineCallback {
+    method public void onError(@NonNull android.system.virtualmachine.VirtualMachine, int, @NonNull String);
+    method public void onPayloadFinished(@NonNull android.system.virtualmachine.VirtualMachine, int);
+    method public void onPayloadReady(@NonNull android.system.virtualmachine.VirtualMachine);
+    method public void onPayloadStarted(@NonNull android.system.virtualmachine.VirtualMachine);
+    method public void onRamdump(@NonNull android.system.virtualmachine.VirtualMachine, @NonNull android.os.ParcelFileDescriptor);
+    method public void onStopped(@NonNull android.system.virtualmachine.VirtualMachine, int);
+    field public static final int ERROR_PAYLOAD_CHANGED = 2; // 0x2
+    field public static final int ERROR_PAYLOAD_INVALID_CONFIG = 3; // 0x3
+    field public static final int ERROR_PAYLOAD_VERIFICATION_FAILED = 1; // 0x1
+    field public static final int ERROR_UNKNOWN = 0; // 0x0
+    field public static final int STOP_REASON_BOOTLOADER_INSTANCE_IMAGE_CHANGED = 10; // 0xa
+    field public static final int STOP_REASON_BOOTLOADER_PUBLIC_KEY_MISMATCH = 9; // 0x9
+    field public static final int STOP_REASON_CRASH = 6; // 0x6
+    field public static final int STOP_REASON_ERROR = 4; // 0x4
+    field public static final int STOP_REASON_HANGUP = 16; // 0x10
+    field public static final int STOP_REASON_INFRASTRUCTURE_ERROR = 0; // 0x0
+    field public static final int STOP_REASON_KILLED = 1; // 0x1
+    field public static final int STOP_REASON_MICRODROID_FAILED_TO_CONNECT_TO_VIRTUALIZATION_SERVICE = 11; // 0xb
+    field public static final int STOP_REASON_MICRODROID_INVALID_PAYLOAD_CONFIG = 14; // 0xe
+    field public static final int STOP_REASON_MICRODROID_PAYLOAD_HAS_CHANGED = 12; // 0xc
+    field public static final int STOP_REASON_MICRODROID_PAYLOAD_VERIFICATION_FAILED = 13; // 0xd
+    field public static final int STOP_REASON_MICRODROID_UNKNOWN_RUNTIME_ERROR = 15; // 0xf
+    field public static final int STOP_REASON_PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED = 8; // 0x8
+    field public static final int STOP_REASON_PVM_FIRMWARE_PUBLIC_KEY_MISMATCH = 7; // 0x7
+    field public static final int STOP_REASON_REBOOT = 5; // 0x5
+    field public static final int STOP_REASON_SHUTDOWN = 3; // 0x3
+    field public static final int STOP_REASON_UNKNOWN = 2; // 0x2
+    field public static final int STOP_REASON_VIRTUALIZATION_SERVICE_DIED = -1; // 0xffffffff
+  }
+
+  public final class VirtualMachineConfig {
+    method @NonNull public String getApkPath();
+    method @NonNull public int getDebugLevel();
+    method @IntRange(from=0) public int getMemoryMib();
+    method @IntRange(from=1) public int getNumCpus();
+    method @Nullable public String getPayloadBinaryPath();
+    method @Nullable public String getPayloadConfigPath();
+    method public boolean isCompatibleWith(@NonNull android.system.virtualmachine.VirtualMachineConfig);
+    method public boolean isProtectedVm();
+    field public static final int DEBUG_LEVEL_APP_ONLY = 1; // 0x1
+    field public static final int DEBUG_LEVEL_FULL = 2; // 0x2
+    field public static final int DEBUG_LEVEL_NONE = 0; // 0x0
+  }
+
+  public static final class VirtualMachineConfig.Builder {
+    ctor public VirtualMachineConfig.Builder(@NonNull android.content.Context);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig build();
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setApkPath(@NonNull String);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setDebugLevel(int);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setMemoryMib(@IntRange(from=0) int);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setNumCpus(@IntRange(from=1) int);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setPayloadBinaryPath(@NonNull String);
+    method @NonNull @RequiresPermission(android.system.virtualmachine.VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION) public android.system.virtualmachine.VirtualMachineConfig.Builder setPayloadConfigPath(@NonNull String);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setProtectedVm(boolean);
+  }
+
+  public final class VirtualMachineDescriptor implements android.os.Parcelable {
+    method public int describeContents();
+    method public void writeToParcel(@NonNull android.os.Parcel, int);
+    field @NonNull public static final android.os.Parcelable.Creator<android.system.virtualmachine.VirtualMachineDescriptor> CREATOR;
+  }
+
+  public class VirtualMachineException extends java.lang.Exception {
+    ctor public VirtualMachineException(@Nullable String);
+    ctor public VirtualMachineException(@Nullable String, @Nullable Throwable);
+    ctor public VirtualMachineException(@Nullable Throwable);
+  }
+
+  public class VirtualMachineManager {
+    method @NonNull @RequiresPermission(android.system.virtualmachine.VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION) public android.system.virtualmachine.VirtualMachine create(@NonNull String, @NonNull android.system.virtualmachine.VirtualMachineConfig) throws android.system.virtualmachine.VirtualMachineException;
+    method public void delete(@NonNull String) throws android.system.virtualmachine.VirtualMachineException;
+    method @Nullable public android.system.virtualmachine.VirtualMachine get(@NonNull String) throws android.system.virtualmachine.VirtualMachineException;
+    method public int getCapabilities();
+    method @NonNull public android.system.virtualmachine.VirtualMachine getOrCreate(@NonNull String, @NonNull android.system.virtualmachine.VirtualMachineConfig) throws android.system.virtualmachine.VirtualMachineException;
+    method @NonNull public android.system.virtualmachine.VirtualMachine importFromDescriptor(@NonNull String, @NonNull android.system.virtualmachine.VirtualMachineDescriptor) throws android.system.virtualmachine.VirtualMachineException;
+    field public static final int CAPABILITY_NON_PROTECTED_VM = 2; // 0x2
+    field public static final int CAPABILITY_PROTECTED_VM = 1; // 0x1
+  }
+
+}
+
diff --git a/javalib/jarjar-rules.txt b/javalib/jarjar-rules.txt
index dd8ad2d..726f9aa 100644
--- a/javalib/jarjar-rules.txt
+++ b/javalib/jarjar-rules.txt
@@ -1,9 +1,10 @@
 # Rules for the android.system.virtualmachine java_sdk_library.
 
-# This is the root of the API, everything we care about should be
-# reachable from here.
-# (This gets rid of all the android.sysprop classes we don't use.)
+# Keep the API surface, most of it is accessible from VirtualMachineManager
 keep android.system.virtualmachine.VirtualMachineManager
+# VirtualizationModuleFrameworkInitializer is not accessible from
+# VirtualMachineManager, we need to explicitly keep it.
+keep android.system.virtualmachine.VirtualizationFrameworkInitializer
 
 # We statically link PlatformProperties, rename to avoid clashes.
 rule android.sysprop.** com.android.system.virtualmachine.sysprop.@1
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index e2fc33e..ffcdc51 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -16,6 +16,7 @@
 
 package android.system.virtualmachine;
 
+import static android.os.ParcelFileDescriptor.AutoCloseInputStream;
 import static android.os.ParcelFileDescriptor.MODE_READ_ONLY;
 import static android.os.ParcelFileDescriptor.MODE_READ_WRITE;
 import static android.system.virtualmachine.VirtualMachineCallback.ERROR_PAYLOAD_CHANGED;
@@ -47,6 +48,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
+import android.annotation.SystemApi;
 import android.content.Context;
 import android.os.Binder;
 import android.os.IBinder;
@@ -63,6 +65,7 @@
 import android.system.virtualizationservice.VirtualMachineAppConfig;
 import android.system.virtualizationservice.VirtualMachineState;
 import android.util.JsonReader;
+import android.util.Log;
 
 import com.android.internal.annotations.GuardedBy;
 
@@ -75,50 +78,34 @@
 import java.io.InputStreamReader;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
-import java.lang.ref.WeakReference;
+import java.nio.channels.FileChannel;
 import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.WeakHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.zip.ZipFile;
 
 /**
- * A handle to the virtual machine. The virtual machine is local to the app which created the
- * virtual machine.
+ * Represents an VM instance, with its own configuration and state. Instances are persistent and are
+ * created or retrieved via {@link VirtualMachineManager}.
+ *
+ * <p>The {@link #run} method actually starts up the VM and allows the payload code to execute. It
+ * will continue until it exits or {@link #stop} is called. Updates on the state of the VM can be
+ * received using {@link #setCallback}. The app can communicate with the VM using {@link
+ * #connectToVsockServer} or {@link #connectVsock}.
  *
  * @hide
  */
+@SystemApi
 public class VirtualMachine implements AutoCloseable {
-    private static final Map<Context, Map<String, WeakReference<VirtualMachine>>> sInstances =
-            new WeakHashMap<>();
-
-    private static final Object sInstancesLock = new Object();
-
-    /** Name of the directory under the files directory where all VMs created for the app exist. */
-    private static final String VM_DIR = "vm";
-
-    /** Name of the persisted config file for a VM. */
-    private static final String CONFIG_FILE = "config.xml";
-
-    /** Name of the instance image file for a VM. (Not implemented) */
-    private static final String INSTANCE_IMAGE_FILE = "instance.img";
-
-    /** Name of the idsig file for a VM */
-    private static final String IDSIG_FILE = "idsig";
-
-    /** Name of the idsig files for extra APKs. */
-    private static final String EXTRA_IDSIG_FILE_PREFIX = "extra_idsig_";
-
-    /** Name of the virtualization service. */
-    private static final String SERVICE_NAME = "android.system.virtualizationservice";
-
     /** The permission needed to create or run a virtual machine. */
     public static final String MANAGE_VIRTUAL_MACHINE_PERMISSION =
             "android.permission.MANAGE_VIRTUAL_MACHINE";
@@ -129,7 +116,6 @@
     public static final String USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION =
             "android.permission.USE_CUSTOM_VIRTUAL_MACHINE";
 
-
     /**
      * Status of a virtual machine
      *
@@ -156,8 +142,28 @@
      */
     public static final int STATUS_DELETED = 2;
 
-    /** Lock for internal synchronization. */
-    private final Object mLock = new Object();
+    private static final String TAG = "VirtualMachine";
+
+    /** Name of the directory under the files directory where all VMs created for the app exist. */
+    private static final String VM_DIR = "vm";
+
+    /** Name of the persisted config file for a VM. */
+    private static final String CONFIG_FILE = "config.xml";
+
+    /** Name of the instance image file for a VM. (Not implemented) */
+    private static final String INSTANCE_IMAGE_FILE = "instance.img";
+
+    /** Name of the idsig file for a VM */
+    private static final String IDSIG_FILE = "idsig";
+
+    /** Name of the idsig files for extra APKs. */
+    private static final String EXTRA_IDSIG_FILE_PREFIX = "extra_idsig_";
+
+    /** Name of the virtualization service. */
+    private static final String SERVICE_NAME = "android.system.virtualizationservice";
+
+    /** Size of the instance image. 10 MB. */
+    private static final long INSTANCE_FILE_SIZE = 10 * 1024 * 1024;
 
     /** The package which owns this VM. */
     @NonNull private final String mPackageName;
@@ -166,6 +172,11 @@
     @NonNull private final String mName;
 
     /**
+     * Path to the directory containing all the files related to this VM.
+     */
+    @NonNull private final File mVmRootPath;
+
+    /**
      * Path to the config file for this VM. The config file is where the configuration is persisted.
      */
     @NonNull private final File mConfigFilePath;
@@ -176,6 +187,60 @@
     /** Path to the idsig file for this VM. */
     @NonNull private final File mIdsigFilePath;
 
+    /**
+     * Unmodifiable list of extra apks. Apks are specified by the vm config, and corresponding
+     * idsigs are to be generated.
+     */
+    @NonNull private final List<ExtraApkSpec> mExtraApks;
+
+    // A note on lock ordering:
+    // You can take mLock while holding VirtualMachineManager.sCreateLock, but not vice versa.
+    // We never take any other lock while holding mCallbackLock; therefore you can
+    // take mCallbackLock while holding any other lock.
+
+    /** Lock protecting our mutable state (other than callbacks). */
+    private final Object mLock = new Object();
+
+    /** Lock protecting callbacks. */
+    private final Object mCallbackLock = new Object();
+
+
+    /** The configuration that is currently associated with this VM. */
+    @GuardedBy("mLock")
+    @NonNull
+    private VirtualMachineConfig mConfig;
+
+    /** Handle to the "running" VM. */
+    @GuardedBy("mLock")
+    @Nullable
+    private IVirtualMachine mVirtualMachine;
+
+    @GuardedBy("mLock")
+    @Nullable
+    private ParcelFileDescriptor mConsoleReader;
+
+    @GuardedBy("mLock")
+    @Nullable
+    private ParcelFileDescriptor mConsoleWriter;
+
+    @GuardedBy("mLock")
+    @Nullable
+    private ParcelFileDescriptor mLogReader;
+
+    @GuardedBy("mLock")
+    @Nullable
+    private ParcelFileDescriptor mLogWriter;
+
+    /** The registered callback */
+    @GuardedBy("mCallbackLock")
+    @Nullable
+    private VirtualMachineCallback mCallback;
+
+    /** The executor on which the callback will be executed */
+    @GuardedBy("mCallbackLock")
+    @Nullable
+    private Executor mCallbackExecutor;
+
     private static class ExtraApkSpec {
         public final File apk;
         public final File idsig;
@@ -186,39 +251,6 @@
         }
     }
 
-    /**
-     * List of extra apks. Apks are specified by the vm config, and corresponding idsigs are to be
-     * generated.
-     */
-    @NonNull private final List<ExtraApkSpec> mExtraApks;
-
-    /** Size of the instance image. 10 MB. */
-    private static final long INSTANCE_FILE_SIZE = 10 * 1024 * 1024;
-
-    /** The configuration that is currently associated with this VM. */
-    @NonNull private VirtualMachineConfig mConfig;
-
-    /** Handle to the "running" VM. */
-    @Nullable private IVirtualMachine mVirtualMachine;
-
-    /** The registered callback */
-    @GuardedBy("mLock")
-    @Nullable
-    private VirtualMachineCallback mCallback;
-
-    /** The executor on which the callback will be executed */
-    @GuardedBy("mLock")
-    @Nullable
-    private Executor mCallbackExecutor;
-
-    @Nullable private ParcelFileDescriptor mConsoleReader;
-    @Nullable private ParcelFileDescriptor mConsoleWriter;
-
-    @Nullable private ParcelFileDescriptor mLogReader;
-    @Nullable private ParcelFileDescriptor mLogWriter;
-
-    @NonNull private final Context mContext;
-
     static {
         System.loadLibrary("virtualmachine_jni");
     }
@@ -226,123 +258,121 @@
     private VirtualMachine(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
-        mContext = context;
         mPackageName = context.getPackageName();
         mName = requireNonNull(name, "Name must not be null");
         mConfig = requireNonNull(config, "Config must not be null");
-        mConfigFilePath = getConfigFilePath(context, name);
 
-        final File vmRoot = new File(context.getFilesDir(), VM_DIR);
-        final File thisVmDir = new File(vmRoot, mName);
+        File thisVmDir = getVmDir(context, mName);
+        mVmRootPath = thisVmDir;
+        mConfigFilePath = new File(thisVmDir, CONFIG_FILE);
         mInstanceFilePath = new File(thisVmDir, INSTANCE_IMAGE_FILE);
         mIdsigFilePath = new File(thisVmDir, IDSIG_FILE);
         mExtraApks = setupExtraApks(context, config, thisVmDir);
     }
 
     /**
-     * Creates a virtual machine with the given name and config. Once a virtual machine is created
-     * it is persisted until it is deleted by calling {@link #delete()}. The created virtual machine
-     * is in {@link #STATUS_STOPPED} state. To run the VM, call {@link #run()}.
+     * Builds a virtual machine from an {@link VirtualMachineDescriptor} object and associates it
+     * with the given name.
+     *
+     * <p>The new virtual machine will be in the same state as the descriptor indicates.
+     *
+     * <p>Once a virtual machine is imported it is persisted until it is deleted by calling {@link
+     * #delete}. The imported virtual machine is in {@link #STATUS_STOPPED} state. To run the VM,
+     * call {@link #run}.
      */
+    @GuardedBy("VirtualMachineManager.sCreateLock")
+    @NonNull
+    static VirtualMachine fromDescriptor(
+            @NonNull Context context,
+            @NonNull String name,
+            @NonNull VirtualMachineDescriptor vmDescriptor)
+            throws VirtualMachineException {
+        VirtualMachineConfig config = VirtualMachineConfig.from(vmDescriptor.getConfigFd());
+        File vmDir = createVmDir(context, name);
+        try {
+            VirtualMachine vm = new VirtualMachine(context, name, config);
+            config.serialize(vm.mConfigFilePath);
+            try {
+                vm.mInstanceFilePath.createNewFile();
+            } catch (IOException e) {
+                throw new VirtualMachineException("failed to create instance image", e);
+            }
+            vm.importInstanceFrom(vmDescriptor.getInstanceImgFd());
+            return vm;
+        } catch (VirtualMachineException | RuntimeException e) {
+            // If anything goes wrong, delete any files created so far and the VM's directory
+            try {
+                deleteRecursively(vmDir);
+            } catch (IOException innerException) {
+                e.addSuppressed(innerException);
+            }
+            throw e;
+        }
+    }
+
+    /**
+     * Creates a virtual machine with the given name and config. Once a virtual machine is created
+     * it is persisted until it is deleted by calling {@link #delete}. The created virtual machine
+     * is in {@link #STATUS_STOPPED} state. To run the VM, call {@link #run}.
+     */
+    @GuardedBy("VirtualMachineManager.sCreateLock")
     @NonNull
     static VirtualMachine create(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
-        VirtualMachine vm = new VirtualMachine(context, name, config);
+        File vmDir = createVmDir(context, name);
 
         try {
-            final File thisVmDir = vm.mConfigFilePath.getParentFile();
-            Files.createDirectories(thisVmDir.getParentFile().toPath());
-
-            // The checking of the existence of this directory and the creation of it is done
-            // atomically. If the directory already exists (i.e. the VM with the same name was
-            // already created), FileAlreadyExistsException is thrown
-            Files.createDirectory(thisVmDir.toPath());
-
-            try (FileOutputStream output = new FileOutputStream(vm.mConfigFilePath)) {
-                vm.mConfig.serialize(output);
-            }
-        } catch (FileAlreadyExistsException e) {
-            throw new VirtualMachineException("virtual machine already exists", e);
-        } catch (IOException e) {
-            throw new VirtualMachineException(e);
-        }
-
-        try {
-            vm.mInstanceFilePath.createNewFile();
-        } catch (IOException e) {
-            throw new VirtualMachineException("failed to create instance image", e);
-        }
-
-        IVirtualizationService service =
-                IVirtualizationService.Stub.asInterface(
-                        ServiceManager.waitForService(SERVICE_NAME));
-
-        try {
-            service.initializeWritablePartition(
-                    ParcelFileDescriptor.open(vm.mInstanceFilePath, MODE_READ_WRITE),
-                    INSTANCE_FILE_SIZE,
-                    PartitionType.ANDROID_VM_INSTANCE);
-        } catch (FileNotFoundException e) {
-            throw new VirtualMachineException("instance image missing", e);
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        } catch (ServiceSpecificException | IllegalArgumentException e) {
-            throw new VirtualMachineException("failed to create instance partition", e);
-        }
-
-        synchronized (sInstancesLock) {
-            Map<String, WeakReference<VirtualMachine>> instancesMap;
-            if (sInstances.containsKey(context)) {
-                instancesMap = sInstances.get(context);
-            } else {
-                instancesMap = new HashMap<>();
-                sInstances.put(context, instancesMap);
+            VirtualMachine vm = new VirtualMachine(context, name, config);
+            config.serialize(vm.mConfigFilePath);
+            try {
+                vm.mInstanceFilePath.createNewFile();
+            } catch (IOException e) {
+                throw new VirtualMachineException("failed to create instance image", e);
             }
 
-            instancesMap.put(name, new WeakReference<>(vm));
-        }
+            IVirtualizationService service =
+                    IVirtualizationService.Stub.asInterface(
+                            ServiceManager.waitForService(SERVICE_NAME));
 
-        return vm;
+            try {
+                service.initializeWritablePartition(
+                        ParcelFileDescriptor.open(vm.mInstanceFilePath, MODE_READ_WRITE),
+                        INSTANCE_FILE_SIZE,
+                        PartitionType.ANDROID_VM_INSTANCE);
+            } catch (FileNotFoundException e) {
+                throw new VirtualMachineException("instance image missing", e);
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } catch (ServiceSpecificException | IllegalArgumentException e) {
+                throw new VirtualMachineException("failed to create instance partition", e);
+            }
+            return vm;
+        } catch (VirtualMachineException | RuntimeException e) {
+            // If anything goes wrong, delete any files created so far and the VM's directory
+            try {
+                deleteRecursively(vmDir);
+            } catch (IOException innerException) {
+                e.addSuppressed(innerException);
+            }
+            throw e;
+        }
     }
 
     /** Loads a virtual machine that is already created before. */
+    @GuardedBy("VirtualMachineManager.sCreateLock")
     @Nullable
-    static VirtualMachine load(
-            @NonNull Context context, @NonNull String name) throws VirtualMachineException {
-        File configFilePath = getConfigFilePath(context, name);
-        VirtualMachineConfig config;
-        try (FileInputStream input = new FileInputStream(configFilePath)) {
-            config = VirtualMachineConfig.from(input);
-        } catch (FileNotFoundException e) {
+    static VirtualMachine load(@NonNull Context context, @NonNull String name)
+            throws VirtualMachineException {
+        File thisVmDir = getVmDir(context, name);
+        if (!thisVmDir.exists()) {
             // The VM doesn't exist.
             return null;
-        } catch (IOException e) {
-            throw new VirtualMachineException(e);
         }
+        File configFilePath = new File(thisVmDir, CONFIG_FILE);
+        VirtualMachineConfig config = VirtualMachineConfig.from(configFilePath);
+        VirtualMachine vm = new VirtualMachine(context, name, config);
 
-        VirtualMachine vm = null;
-        synchronized (sInstancesLock) {
-            Map<String, WeakReference<VirtualMachine>> instancesMap;
-            if (sInstances.containsKey(context)) {
-                instancesMap = sInstances.get(context);
-            } else {
-                instancesMap = new HashMap<>();
-                sInstances.put(context, instancesMap);
-            }
-
-            if (instancesMap.containsKey(name)) {
-                vm = instancesMap.get(name).get();
-            }
-            if (vm == null) {
-                vm = new VirtualMachine(context, name, config);
-                instancesMap.put(name, new WeakReference<>(vm));
-            }
-        }
-
-        // If config file exists, but the instance image file doesn't, it means that the VM is
-        // corrupted. That's different from the case that the VM doesn't exist. Throw an exception
-        // instead of returning null.
         if (!vm.mInstanceFilePath.exists()) {
             throw new VirtualMachineException("instance image missing");
         }
@@ -350,12 +380,59 @@
         return vm;
     }
 
+    @GuardedBy("VirtualMachineManager.sCreateLock")
+    void delete(Context context, String name) throws VirtualMachineException {
+        synchronized (mLock) {
+            checkStopped();
+            deleteVmDirectory(context, name);
+        }
+    }
+
+    static void deleteVmDirectory(Context context, String name) throws VirtualMachineException {
+        try {
+            deleteRecursively(getVmDir(context, name));
+        } catch (IOException e) {
+            throw new VirtualMachineException(e);
+        }
+    }
+
+    @GuardedBy("VirtualMachineManager.sCreateLock")
+    @NonNull
+    private static File createVmDir(@NonNull Context context, @NonNull String name)
+            throws VirtualMachineException {
+        File vmDir = getVmDir(context, name);
+        try {
+            // We don't need to undo this even if VM creation fails.
+            Files.createDirectories(vmDir.getParentFile().toPath());
+
+            // The checking of the existence of this directory and the creation of it is done
+            // atomically. If the directory already exists (i.e. the VM with the same name was
+            // already created), FileAlreadyExistsException is thrown.
+            Files.createDirectory(vmDir.toPath());
+        } catch (FileAlreadyExistsException e) {
+            throw new VirtualMachineException("virtual machine already exists", e);
+        } catch (IOException e) {
+            throw new VirtualMachineException("failed to create directory for VM", e);
+        }
+        return vmDir;
+    }
+
+    @NonNull
+    private static File getVmDir(@NonNull Context context, @NonNull String name) {
+        if (name.contains(File.separator) || name.equals(".") || name.equals("..")) {
+            throw new IllegalArgumentException("Invalid VM name: " + name);
+        }
+        File vmRoot = new File(context.getDataDir(), VM_DIR);
+        return new File(vmRoot, name);
+    }
+
     /**
      * Returns the name of this virtual machine. The name is unique in the package and can't be
      * changed.
      *
      * @hide
      */
+    @SystemApi
     @NonNull
     public String getName() {
         return mName;
@@ -370,9 +447,12 @@
      *
      * @hide
      */
+    @SystemApi
     @NonNull
     public VirtualMachineConfig getConfig() {
-        return mConfig;
+        synchronized (mLock) {
+            return mConfig;
+        }
     }
 
     /**
@@ -380,29 +460,73 @@
      *
      * @hide
      */
+    @SystemApi
     @Status
     public int getStatus() {
+        IVirtualMachine virtualMachine;
+        synchronized (mLock) {
+            virtualMachine = mVirtualMachine;
+        }
+        if (virtualMachine == null) {
+            return mVmRootPath.exists() ? STATUS_STOPPED : STATUS_DELETED;
+        } else {
+            try {
+                return stateToStatus(virtualMachine.getState());
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            }
+        }
+    }
+
+    private int stateToStatus(@VirtualMachineState int state) {
+        switch (state) {
+            case VirtualMachineState.STARTING:
+            case VirtualMachineState.STARTED:
+            case VirtualMachineState.READY:
+            case VirtualMachineState.FINISHED:
+                return STATUS_RUNNING;
+            case VirtualMachineState.NOT_STARTED:
+            case VirtualMachineState.DEAD:
+            default:
+                return STATUS_STOPPED;
+        }
+    }
+
+    // Throw an appropriate exception if we have a running VM, or the VM has been deleted.
+    @GuardedBy("mLock")
+    private void checkStopped() throws VirtualMachineException {
+        if (!mVmRootPath.exists()) {
+            throw new VirtualMachineException("VM has been deleted");
+        }
+        if (mVirtualMachine == null) {
+            return;
+        }
         try {
-            if (mVirtualMachine != null) {
-                switch (mVirtualMachine.getState()) {
-                    case VirtualMachineState.NOT_STARTED:
-                        return STATUS_STOPPED;
-                    case VirtualMachineState.STARTING:
-                    case VirtualMachineState.STARTED:
-                    case VirtualMachineState.READY:
-                    case VirtualMachineState.FINISHED:
-                        return STATUS_RUNNING;
-                    case VirtualMachineState.DEAD:
-                        return STATUS_STOPPED;
+            if (stateToStatus(mVirtualMachine.getState()) != STATUS_STOPPED) {
+                throw new VirtualMachineException("VM is not in stopped state");
+            }
+        } catch (RemoteException e) {
+            throw e.rethrowAsRuntimeException();
+        }
+    }
+
+    // If we have an IVirtualMachine in the running state return it, otherwise throw.
+    @GuardedBy("mLock")
+    private IVirtualMachine getRunningVm() throws VirtualMachineException {
+        try {
+            if (mVirtualMachine != null
+                    && stateToStatus(mVirtualMachine.getState()) == STATUS_RUNNING) {
+                return mVirtualMachine;
+            } else {
+                if (!mVmRootPath.exists()) {
+                    throw new VirtualMachineException("VM has been deleted");
+                } else {
+                    throw new VirtualMachineException("VM is not in running state");
                 }
             }
         } catch (RemoteException e) {
             throw e.rethrowAsRuntimeException();
         }
-        if (!mConfigFilePath.exists()) {
-            return STATUS_DELETED;
-        }
-        return STATUS_STOPPED;
     }
 
     /**
@@ -411,9 +535,11 @@
      *
      * @hide
      */
-    public void setCallback(@NonNull @CallbackExecutor Executor executor,
+    @SystemApi
+    public void setCallback(
+            @NonNull @CallbackExecutor Executor executor,
             @NonNull VirtualMachineCallback callback) {
-        synchronized (mLock) {
+        synchronized (mCallbackLock) {
             mCallback = callback;
             mCallbackExecutor = executor;
         }
@@ -424,8 +550,9 @@
      *
      * @hide
      */
+    @SystemApi
     public void clearCallback() {
-        synchronized (mLock) {
+        synchronized (mCallbackLock) {
             mCallback = null;
             mCallbackExecutor = null;
         }
@@ -435,7 +562,7 @@
     private void executeCallback(Consumer<VirtualMachineCallback> fn) {
         final VirtualMachineCallback callback;
         final Executor executor;
-        synchronized (mLock) {
+        synchronized (mCallbackLock) {
             callback = mCallback;
             executor = mCallbackExecutor;
         }
@@ -457,122 +584,129 @@
      * calling {@code run()}.
      *
      * @throws VirtualMachineException if the virtual machine is not stopped or could not be
-     *         started.
+     *     started.
      * @hide
      */
+    @SystemApi
     @RequiresPermission(MANAGE_VIRTUAL_MACHINE_PERMISSION)
     public void run() throws VirtualMachineException {
-        if (getStatus() != STATUS_STOPPED) {
-            throw new VirtualMachineException(this + " is not in stopped state");
-        }
+        synchronized (mLock) {
+            checkStopped();
 
-        try {
-            mIdsigFilePath.createNewFile();
-            for (ExtraApkSpec extraApk : mExtraApks) {
-                extraApk.idsig.createNewFile();
-            }
-        } catch (IOException e) {
-            // If the file already exists, exception is not thrown.
-            throw new VirtualMachineException("failed to create idsig file", e);
-        }
-
-        IVirtualizationService service =
-                IVirtualizationService.Stub.asInterface(
-                        ServiceManager.waitForService(SERVICE_NAME));
-
-        try {
-            createVmPipes();
-
-            VirtualMachineAppConfig appConfig = getConfig().toParcel();
-            appConfig.name = mName;
-
-            // Fill the idsig file by hashing the apk
-            service.createOrUpdateIdsigFile(
-                    appConfig.apk, ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_WRITE));
-
-            for (ExtraApkSpec extraApk : mExtraApks) {
-                service.createOrUpdateIdsigFile(
-                        ParcelFileDescriptor.open(extraApk.apk, MODE_READ_ONLY),
-                        ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_WRITE));
-            }
-
-            // Re-open idsig file in read-only mode
-            appConfig.idsig = ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_ONLY);
-            appConfig.instanceImage = ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_WRITE);
-            List<ParcelFileDescriptor> extraIdsigs = new ArrayList<>();
-            for (ExtraApkSpec extraApk : mExtraApks) {
-                extraIdsigs.add(ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_ONLY));
-            }
-            appConfig.extraIdsigs = extraIdsigs;
-
-            android.system.virtualizationservice.VirtualMachineConfig vmConfigParcel =
-                    android.system.virtualizationservice.VirtualMachineConfig.appConfig(appConfig);
-
-            // The VM should only be observed to die once
-            AtomicBoolean onDiedCalled = new AtomicBoolean(false);
-
-            IBinder.DeathRecipient deathRecipient = () -> {
-                if (onDiedCalled.compareAndSet(false, true)) {
-                    executeCallback((cb) -> cb.onStopped(VirtualMachine.this,
-                            VirtualMachineCallback.STOP_REASON_VIRTUALIZATION_SERVICE_DIED));
+            try {
+                mIdsigFilePath.createNewFile();
+                for (ExtraApkSpec extraApk : mExtraApks) {
+                    extraApk.idsig.createNewFile();
                 }
-            };
+            } catch (IOException e) {
+                // If the file already exists, exception is not thrown.
+                throw new VirtualMachineException("failed to create idsig file", e);
+            }
 
-            mVirtualMachine = service.createVm(vmConfigParcel, mConsoleWriter, mLogWriter);
-            mVirtualMachine.registerCallback(
-                    new IVirtualMachineCallback.Stub() {
-                        @Override
-                        public void onPayloadStarted(int cid, ParcelFileDescriptor stream) {
-                            executeCallback(
-                                    (cb) -> cb.onPayloadStarted(VirtualMachine.this, stream));
-                        }
+            IVirtualizationService service =
+                    IVirtualizationService.Stub.asInterface(
+                            ServiceManager.waitForService(SERVICE_NAME));
 
-                        @Override
-                        public void onPayloadReady(int cid) {
-                            executeCallback((cb) -> cb.onPayloadReady(VirtualMachine.this));
-                        }
+            try {
+                createVmPipes();
 
-                        @Override
-                        public void onPayloadFinished(int cid, int exitCode) {
-                            executeCallback(
-                                    (cb) -> cb.onPayloadFinished(VirtualMachine.this, exitCode));
-                        }
+                VirtualMachineAppConfig appConfig = getConfig().toVsConfig();
+                appConfig.name = mName;
 
-                        @Override
-                        public void onError(int cid, int errorCode, String message) {
-                            int translatedError = getTranslatedError(errorCode);
-                            executeCallback(
-                                    (cb) -> cb.onError(VirtualMachine.this, translatedError,
-                                            message));
-                        }
+                // Fill the idsig file by hashing the apk
+                service.createOrUpdateIdsigFile(
+                        appConfig.apk, ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_WRITE));
 
-                        @Override
-                        public void onDied(int cid, int reason) {
-                            service.asBinder().unlinkToDeath(deathRecipient, 0);
-                            int translatedReason = getTranslatedReason(reason);
-                            if (onDiedCalled.compareAndSet(false, true)) {
-                                executeCallback(
-                                        (cb) -> cb.onStopped(VirtualMachine.this,
-                                                translatedReason));
-                            }
-                        }
+                for (ExtraApkSpec extraApk : mExtraApks) {
+                    service.createOrUpdateIdsigFile(
+                            ParcelFileDescriptor.open(extraApk.apk, MODE_READ_ONLY),
+                            ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_WRITE));
+                }
 
-                        @Override
-                        public void onRamdump(int cid, ParcelFileDescriptor ramdump) {
-                            executeCallback(
-                                    (cb) -> cb.onRamdump(VirtualMachine.this, ramdump));
-                        }
+                // Re-open idsig file in read-only mode
+                appConfig.idsig = ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_ONLY);
+                appConfig.instanceImage = ParcelFileDescriptor.open(mInstanceFilePath,
+                        MODE_READ_WRITE);
+                List<ParcelFileDescriptor> extraIdsigs = new ArrayList<>();
+                for (ExtraApkSpec extraApk : mExtraApks) {
+                    extraIdsigs.add(ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_ONLY));
+                }
+                appConfig.extraIdsigs = extraIdsigs;
+
+                android.system.virtualizationservice.VirtualMachineConfig vmConfigParcel =
+                        android.system.virtualizationservice.VirtualMachineConfig.appConfig(
+                                appConfig);
+
+                // The VM should only be observed to die once
+                AtomicBoolean onDiedCalled = new AtomicBoolean(false);
+
+                IBinder.DeathRecipient deathRecipient = () -> {
+                    if (onDiedCalled.compareAndSet(false, true)) {
+                        executeCallback((cb) -> cb.onStopped(VirtualMachine.this,
+                                VirtualMachineCallback.STOP_REASON_VIRTUALIZATION_SERVICE_DIED));
                     }
-            );
-            service.asBinder().linkToDeath(deathRecipient, 0);
-            mVirtualMachine.start();
-        } catch (IOException | IllegalStateException | ServiceSpecificException e) {
-            throw new VirtualMachineException(e);
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
+                };
+
+                mVirtualMachine = service.createVm(vmConfigParcel, mConsoleWriter, mLogWriter);
+                mVirtualMachine.registerCallback(
+                        new IVirtualMachineCallback.Stub() {
+                            @Override
+                            public void onPayloadStarted(int cid) {
+                                executeCallback((cb) -> cb.onPayloadStarted(VirtualMachine.this));
+                            }
+
+                            @Override
+                            public void onPayloadReady(int cid) {
+                                executeCallback((cb) -> cb.onPayloadReady(VirtualMachine.this));
+                            }
+
+                            @Override
+                            public void onPayloadFinished(int cid, int exitCode) {
+                                executeCallback(
+                                        (cb) ->
+                                                cb.onPayloadFinished(
+                                                        VirtualMachine.this, exitCode));
+                            }
+
+                            @Override
+                            public void onError(int cid, int errorCode, String message) {
+                                int translatedError = getTranslatedError(errorCode);
+                                executeCallback(
+                                        (cb) ->
+                                                cb.onError(
+                                                        VirtualMachine.this,
+                                                        translatedError,
+                                                        message));
+                            }
+
+                            @Override
+                            public void onDied(int cid, int reason) {
+                                service.asBinder().unlinkToDeath(deathRecipient, 0);
+                                int translatedReason = getTranslatedReason(reason);
+                                if (onDiedCalled.compareAndSet(false, true)) {
+                                    executeCallback(
+                                            (cb) ->
+                                                    cb.onStopped(
+                                                            VirtualMachine.this, translatedReason));
+                                }
+                            }
+
+                            @Override
+                            public void onRamdump(int cid, ParcelFileDescriptor ramdump) {
+                                executeCallback((cb) -> cb.onRamdump(VirtualMachine.this, ramdump));
+                            }
+                        });
+                service.asBinder().linkToDeath(deathRecipient, 0);
+                mVirtualMachine.start();
+            } catch (IOException | IllegalStateException | ServiceSpecificException e) {
+                throw new VirtualMachineException(e);
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            }
         }
     }
 
+    @GuardedBy("mLock")
     private void createVmPipes() throws VirtualMachineException {
         try {
             if (mConsoleReader == null || mConsoleWriter == null) {
@@ -587,7 +721,235 @@
                 mLogWriter = pipe[1];
             }
         } catch (IOException e) {
-            throw new VirtualMachineException(e);
+            throw new VirtualMachineException("Failed to create stream for VM", e);
+        }
+    }
+
+    /**
+     * Returns the stream object representing the console output from the virtual machine.
+     *
+     * @throws VirtualMachineException if the stream could not be created.
+     * @hide
+     */
+    @SystemApi
+    @NonNull
+    public InputStream getConsoleOutput() throws VirtualMachineException {
+        synchronized (mLock) {
+            createVmPipes();
+            return new FileInputStream(mConsoleReader.getFileDescriptor());
+        }
+    }
+
+    /**
+     * Returns the stream object representing the log output from the virtual machine.
+     *
+     * @throws VirtualMachineException if the stream could not be created.
+     * @hide
+     */
+    @SystemApi
+    @NonNull
+    public InputStream getLogOutput() throws VirtualMachineException {
+        synchronized (mLock) {
+            createVmPipes();
+            return new FileInputStream(mLogReader.getFileDescriptor());
+        }
+    }
+
+    /**
+     * Stops this virtual machine. Stopping a virtual machine is like pulling the plug on a real
+     * computer; the machine halts immediately. Software running on the virtual machine is not
+     * notified of the event. A stopped virtual machine can be re-started by calling {@link #run()}.
+     *
+     * @throws VirtualMachineException if the virtual machine is not running or could not be
+     *     stopped.
+     * @hide
+     */
+    @SystemApi
+    public void stop() throws VirtualMachineException {
+        synchronized (mLock) {
+            if (mVirtualMachine == null) {
+                throw new VirtualMachineException("VM is not running");
+            }
+            try {
+                mVirtualMachine.stop();
+                mVirtualMachine = null;
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } catch (ServiceSpecificException e) {
+                throw new VirtualMachineException(e);
+            }
+        }
+    }
+
+    /**
+     * Stops this virtual machine, if it is running.
+     *
+     * @see #stop()
+     * @hide
+     */
+    @SystemApi
+    @Override
+    public void close() {
+        synchronized (mLock) {
+            if (mVirtualMachine == null) {
+                return;
+            }
+            try {
+                if (stateToStatus(mVirtualMachine.getState()) == STATUS_RUNNING) {
+                    mVirtualMachine.stop();
+                    mVirtualMachine = null;
+                }
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } catch (ServiceSpecificException e) {
+                // Deliberately ignored; this almost certainly means the VM exited just as
+                // we tried to stop it.
+                Log.i(TAG, "Ignoring error on close()", e);
+            }
+        }
+    }
+
+    private static void deleteRecursively(File dir) throws IOException {
+        // Note: This doesn't follow symlinks, which is important. Instead they are just deleted
+        // (and Files.delete deletes the link not the target).
+        Files.walkFileTree(dir.toPath(), new SimpleFileVisitor<>() {
+            @Override
+            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
+                    throws IOException {
+                Files.delete(file);
+                return FileVisitResult.CONTINUE;
+            }
+
+            @Override
+            public FileVisitResult postVisitDirectory(Path dir, IOException e) throws IOException {
+                // Directory is deleted after we've visited (deleted) all its contents, so it
+                // should be empty by now.
+                Files.delete(dir);
+                return FileVisitResult.CONTINUE;
+            }
+        });
+    }
+
+    /**
+     * Returns the CID of this virtual machine, if it is running.
+     *
+     * @throws VirtualMachineException if the virtual machine is not running.
+     * @hide
+     */
+    @SystemApi
+    public int getCid() throws VirtualMachineException {
+        synchronized (mLock) {
+            try {
+                return getRunningVm().getCid();
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            }
+        }
+    }
+
+    /**
+     * Changes the config of this virtual machine to a new one. This can be used to adjust things
+     * like the number of CPU and size of the RAM, depending on the situation (e.g. the size of the
+     * application to run on the virtual machine, etc.)
+     *
+     * <p>The new config must be {@link VirtualMachineConfig#isCompatibleWith compatible with} the
+     * existing config.
+     *
+     * @return the old config
+     * @throws VirtualMachineException if the virtual machine is not stopped, or the new config is
+     *     incompatible.
+     * @hide
+     */
+    @SystemApi
+    @NonNull
+    public VirtualMachineConfig setConfig(@NonNull VirtualMachineConfig newConfig)
+            throws VirtualMachineException {
+        synchronized (mLock) {
+            VirtualMachineConfig oldConfig = mConfig;
+            if (!oldConfig.isCompatibleWith(newConfig)) {
+                throw new VirtualMachineException("incompatible config");
+            }
+            checkStopped();
+
+            // Delete any existing file before recreating; that ensures any VirtualMachineDescriptor
+            // that refers to the old file does not see the new config.
+            mConfigFilePath.delete();
+            newConfig.serialize(mConfigFilePath);
+            mConfig = newConfig;
+            return oldConfig;
+        }
+    }
+
+    @Nullable
+    private static native IBinder nativeConnectToVsockServer(IBinder vmBinder, int port);
+
+    /**
+     * Connect to a VM's binder service via vsock and return the root IBinder object. Guest VMs are
+     * expected to set up vsock servers in their payload. After the host app receives the {@link
+     * VirtualMachineCallback#onPayloadReady(VirtualMachine)}, it can use this method to establish a
+     * connection to the guest VM.
+     *
+     * @throws VirtualMachineException if the virtual machine is not running or the connection
+     *     failed.
+     * @hide
+     */
+    @SystemApi
+    @NonNull
+    public IBinder connectToVsockServer(int port) throws VirtualMachineException {
+        synchronized (mLock) {
+            IBinder iBinder = nativeConnectToVsockServer(getRunningVm().asBinder(), port);
+            if (iBinder == null) {
+                throw new VirtualMachineException("Failed to connect to vsock server");
+            }
+            return iBinder;
+        }
+    }
+
+    /**
+     * Opens a vsock connection to the VM on the given port.
+     *
+     * @throws VirtualMachineException if connecting fails.
+     * @hide
+     */
+    @SystemApi
+    @NonNull
+    public ParcelFileDescriptor connectVsock(int port) throws VirtualMachineException {
+        synchronized (mLock) {
+            try {
+                return getRunningVm().connectVsock(port);
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } catch (ServiceSpecificException e) {
+                throw new VirtualMachineException(e);
+            }
+        }
+    }
+
+    /**
+     * Captures the current state of the VM in a {@link VirtualMachineDescriptor} instance. The VM
+     * needs to be stopped to avoid inconsistency in its state representation.
+     *
+     * <p>The state of the VM is not actually copied until {@link
+     * VirtualMachineManager#importFromDescriptor} is called. It is recommended that the VM not be
+     * started until that operation is complete.
+     *
+     * @return a {@link VirtualMachineDescriptor} instance that represents the VM's state.
+     * @throws VirtualMachineException if the virtual machine is not stopped, or the state could not
+     *     be captured.
+     * @hide
+     */
+    @SystemApi
+    @NonNull
+    public VirtualMachineDescriptor toDescriptor() throws VirtualMachineException {
+        synchronized (mLock) {
+            checkStopped();
+            try {
+                return new VirtualMachineDescriptor(
+                        ParcelFileDescriptor.open(mConfigFilePath, MODE_READ_ONLY),
+                        ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_ONLY));
+            } catch (IOException e) {
+                throw new VirtualMachineException(e);
+            }
         }
     }
 
@@ -645,187 +1007,6 @@
         }
     }
 
-    /**
-     * Returns the stream object representing the console output from the virtual machine.
-     *
-     * @throws VirtualMachineException if the stream could not be created.
-     * @hide
-     */
-    @NonNull
-    public InputStream getConsoleOutputStream() throws VirtualMachineException {
-        createVmPipes();
-        return new FileInputStream(mConsoleReader.getFileDescriptor());
-    }
-
-    /**
-     * Returns the stream object representing the log output from the virtual machine.
-     *
-     * @throws VirtualMachineException if the stream could not be created.
-     * @hide
-     */
-    @NonNull
-    public InputStream getLogOutputStream() throws VirtualMachineException {
-        createVmPipes();
-        return new FileInputStream(mLogReader.getFileDescriptor());
-    }
-
-    /**
-     * Stops this virtual machine. Stopping a virtual machine is like pulling the plug on a real
-     * computer; the machine halts immediately. Software running on the virtual machine is not
-     * notified of the event. A stopped virtual machine can be re-started by calling {@link
-     * #run()}.
-     *
-     * @throws VirtualMachineException if the virtual machine could not be stopped.
-     * @hide
-     */
-    public void stop() throws VirtualMachineException {
-        if (mVirtualMachine == null) return;
-        try {
-            mVirtualMachine.stop();
-            mVirtualMachine = null;
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        } catch (ServiceSpecificException e) {
-            throw new VirtualMachineException(e);
-        }
-    }
-
-    /**
-     * Stops this virtual machine. See {@link #stop()}.
-     *
-     * @throws VirtualMachineException if the virtual machine could not be stopped.
-     * @hide
-     */
-    @Override
-    public void close() throws VirtualMachineException {
-        stop();
-    }
-
-    /**
-     * Deletes this virtual machine. Deleting a virtual machine means deleting any persisted data
-     * associated with it including the per-VM secret. This is an irreversible action. A virtual
-     * machine once deleted can never be restored. A new virtual machine created with the same name
-     * and the same config is different from an already deleted virtual machine.
-     *
-     * @throws VirtualMachineException if the virtual machine is not stopped.
-     * @hide
-     */
-    public void delete() throws VirtualMachineException {
-        if (getStatus() != STATUS_STOPPED) {
-            throw new VirtualMachineException("Virtual machine is not stopped");
-        }
-        final File vmRootDir = mConfigFilePath.getParentFile();
-        for (ExtraApkSpec extraApks : mExtraApks) {
-            extraApks.idsig.delete();
-        }
-        mConfigFilePath.delete();
-        mInstanceFilePath.delete();
-        mIdsigFilePath.delete();
-        vmRootDir.delete();
-
-        synchronized (sInstancesLock) {
-            Map<String, WeakReference<VirtualMachine>> instancesMap = sInstances.get(mContext);
-            if (instancesMap != null) instancesMap.remove(mName);
-        }
-    }
-
-    /**
-     * Returns the CID of this virtual machine, if it is running.
-     *
-     * @throws VirtualMachineException if the virtual machine is not running.
-     * @hide
-     */
-    @NonNull
-    public int getCid() throws VirtualMachineException {
-        if (getStatus() != STATUS_RUNNING) {
-            throw new VirtualMachineException("VM is not running");
-        }
-        try {
-            return mVirtualMachine.getCid();
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        }
-    }
-
-    /**
-     * Changes the config of this virtual machine to a new one. This can be used to adjust things
-     * like the number of CPU and size of the RAM, depending on the situation (e.g. the size of the
-     * application to run on the virtual machine, etc.)
-     *
-     * The new config must be {@link VirtualMachineConfig#isCompatibleWith compatible with} the
-     * existing config.
-     *
-     * @return the old config
-     * @throws VirtualMachineException if the virtual machine is not stopped, or the new config is
-     *         incompatible.
-     * @hide
-     */
-    @NonNull
-    public VirtualMachineConfig setConfig(@NonNull VirtualMachineConfig newConfig)
-            throws VirtualMachineException {
-        final VirtualMachineConfig oldConfig = getConfig();
-        if (!oldConfig.isCompatibleWith(newConfig)) {
-            throw new VirtualMachineException("incompatible config");
-        }
-        if (getStatus() != STATUS_STOPPED) {
-            throw new VirtualMachineException(
-                    "can't change config while virtual machine is not stopped");
-        }
-
-        try {
-            FileOutputStream output = new FileOutputStream(mConfigFilePath);
-            newConfig.serialize(output);
-            output.close();
-        } catch (IOException e) {
-            throw new VirtualMachineException(e);
-        }
-        mConfig = newConfig;
-
-        return oldConfig;
-    }
-
-    @Nullable
-    private static native IBinder nativeConnectToVsockServer(IBinder vmBinder, int port);
-
-    /**
-     * Connect to a VM's binder service via vsock and return the root IBinder object. Guest VMs are
-     * expected to set up vsock servers in their payload. After the host app receives the {@link
-     * VirtualMachineCallback#onPayloadReady(VirtualMachine)}, it can use this method to
-     * establish a connection to the guest VM.
-     *
-     * @throws VirtualMachineException if the virtual machine is not running or the connection
-     *         failed.
-     * @hide
-     */
-    @NonNull
-    public IBinder connectToVsockServer(int port) throws VirtualMachineException {
-        if (getStatus() != STATUS_RUNNING) {
-            throw new VirtualMachineException("VM is not running");
-        }
-        IBinder iBinder = nativeConnectToVsockServer(mVirtualMachine.asBinder(), port);
-        if (iBinder == null) {
-            throw new VirtualMachineException("Failed to connect to vsock server");
-        }
-        return iBinder;
-    }
-
-    /**
-     * Opens a vsock connection to the VM on the given port.
-     *
-     * @throws VirtualMachineException if connecting fails.
-     * @hide
-     */
-    @NonNull
-    public ParcelFileDescriptor connectVsock(int port) throws VirtualMachineException {
-        try {
-            return mVirtualMachine.connectVsock(port);
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        } catch (ServiceSpecificException e) {
-            throw new VirtualMachineException(e);
-        }
-    }
-
     @Override
     public String toString() {
         VirtualMachineConfig config = getConfig();
@@ -916,15 +1097,19 @@
                                 new File(vmDir, EXTRA_IDSIG_FILE_PREFIX + i)));
             }
 
-            return extraApks;
+            return Collections.unmodifiableList(extraApks);
         } catch (IOException e) {
             throw new VirtualMachineException("Couldn't parse extra apks from the vm config", e);
         }
     }
 
-    private static File getConfigFilePath(@NonNull Context context, @NonNull String name) {
-        final File vmRoot = new File(context.getFilesDir(), VM_DIR);
-        final File thisVmDir = new File(vmRoot, name);
-        return new File(thisVmDir, CONFIG_FILE);
+    private void importInstanceFrom(@NonNull ParcelFileDescriptor instanceFd)
+            throws VirtualMachineException {
+        try (FileChannel instance = new FileOutputStream(mInstanceFilePath).getChannel();
+                FileChannel instanceInput = new AutoCloseInputStream(instanceFd).getChannel()) {
+            instance.transferFrom(instanceInput, /*position=*/ 0, instanceInput.size());
+        } catch (IOException e) {
+            throw new VirtualMachineException("failed to transfer instance image", e);
+        }
     }
 }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java b/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
index bb6b2b8..f3c4831 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineCallback.java
@@ -18,8 +18,8 @@
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.annotation.SuppressLint;
+import android.annotation.SystemApi;
 import android.os.ParcelFileDescriptor;
 
 import java.lang.annotation.Retention;
@@ -31,7 +31,8 @@
  *
  * @hide
  */
-@SuppressLint("CallbackInterface")  // Guidance has changed, lint is out of date (b/245552641)
+@SystemApi
+@SuppressLint("CallbackInterface") // Guidance has changed, lint is out of date (b/245552641)
 public interface VirtualMachineCallback {
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
@@ -135,11 +136,8 @@
     /** The VM killed due to hangup */
     int STOP_REASON_HANGUP = 16;
 
-    /**
-     * Called when the payload starts in the VM. The stream, if non-null, provides access
-     * to the stdin/stdout of the VM payload.
-     */
-    void onPayloadStarted(@NonNull VirtualMachine vm, @Nullable ParcelFileDescriptor stream);
+    /** Called when the payload starts in the VM. */
+    void onPayloadStarted(@NonNull VirtualMachine vm);
 
     /**
      * Called when the payload in the VM is ready to serve. See
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index 2dff9bb..b432bde 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -16,14 +16,17 @@
 
 package android.system.virtualmachine;
 
+import static android.os.ParcelFileDescriptor.AutoCloseInputStream;
 import static android.os.ParcelFileDescriptor.MODE_READ_ONLY;
 
 import static java.util.Objects.requireNonNull;
 
 import android.annotation.IntDef;
+import android.annotation.IntRange;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
+import android.annotation.SystemApi;
 import android.content.Context;
 import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
@@ -32,22 +35,26 @@
 import android.system.virtualizationservice.VirtualMachinePayloadConfig;
 
 import java.io.File;
+import java.io.FileInputStream;
 import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.Objects;
+
 /**
  * Represents a configuration of a virtual machine. A configuration consists of hardware
  * configurations like the number of CPUs and the size of RAM, and software configurations like the
- * OS and application to run on the virtual machine.
+ * payload to run on the virtual machine.
  *
  * @hide
  */
+@SystemApi
 public final class VirtualMachineConfig {
-    // These defines the schema of the config file persisted on disk.
+    // These define the schema of the config file persisted on disk.
     private static final int VERSION = 2;
     private static final String KEY_VERSION = "version";
     private static final String KEY_APKPATH = "apkPath";
@@ -58,9 +65,6 @@
     private static final String KEY_MEMORY_MIB = "memoryMib";
     private static final String KEY_NUM_CPUS = "numCpus";
 
-    // Absolute path to the APK file containing the VM payload.
-    @NonNull private final String mApkPath;
-
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
     @IntDef(prefix = "DEBUG_LEVEL_", value = {
@@ -71,12 +75,12 @@
     public @interface DebugLevel {}
 
     /**
-     * Not debuggable at all. No log is exported from the VM. Debugger can't be attached to the
-     * app process running in the VM. This is the default level.
+     * Not debuggable at all. No log is exported from the VM. Debugger can't be attached to the app
+     * process running in the VM. This is the default level.
      *
      * @hide
      */
-    public static final int DEBUG_LEVEL_NONE = 0;
+    @SystemApi public static final int DEBUG_LEVEL_NONE = 0;
 
     /**
      * Only the app is debuggable. Log from the app is exported from the VM. Debugger can be
@@ -84,7 +88,7 @@
      *
      * @hide
      */
-    public static final int DEBUG_LEVEL_APP_ONLY = 1;
+    @SystemApi public static final int DEBUG_LEVEL_APP_ONLY = 1;
 
     /**
      * Fully debuggable. All logs (both logcat and kernel message) are exported. All processes
@@ -92,7 +96,10 @@
      *
      * @hide
      */
-    public static final int DEBUG_LEVEL_FULL = 2;
+    @SystemApi public static final int DEBUG_LEVEL_FULL = 2;
+
+    /** Absolute path to the APK file containing the VM payload. */
+    @NonNull private final String mApkPath;
 
     @DebugLevel private final int mDebugLevel;
 
@@ -133,6 +140,43 @@
         if (!apkPath.startsWith("/")) {
             throw new IllegalArgumentException("APK path must be an absolute path");
         }
+
+        if (memoryMib < 0) {
+            throw new IllegalArgumentException("Memory size cannot be negative");
+        }
+
+        int availableCpus = Runtime.getRuntime().availableProcessors();
+        if (numCpus < 1 || numCpus > availableCpus) {
+            throw new IllegalArgumentException("Number of vCPUs (" + numCpus + ") is out of "
+                    + "range [1, " + availableCpus + "]");
+        }
+
+        if (debugLevel != DEBUG_LEVEL_NONE && debugLevel != DEBUG_LEVEL_APP_ONLY
+                && debugLevel != DEBUG_LEVEL_FULL) {
+            throw new IllegalArgumentException("Invalid debugLevel: " + debugLevel);
+        }
+
+        if (payloadBinaryPath == null) {
+            if (payloadConfigPath == null) {
+                throw new IllegalStateException("setPayloadBinaryPath must be called");
+            }
+        } else {
+            if (payloadConfigPath != null) {
+                throw new IllegalStateException(
+                        "setPayloadBinaryPath and setPayloadConfigPath may not both be called");
+            }
+        }
+
+        if (protectedVm
+                && !HypervisorProperties.hypervisor_protected_vm_supported().orElse(false)) {
+            throw new UnsupportedOperationException(
+                    "Protected VMs are not supported on this device.");
+        }
+        if (!protectedVm && !HypervisorProperties.hypervisor_vm_supported().orElse(false)) {
+            throw new UnsupportedOperationException(
+                    "Unprotected VMs are not supported on this device.");
+        }
+
         mApkPath = apkPath;
         mPayloadConfigPath = payloadConfigPath;
         mPayloadBinaryPath = payloadBinaryPath;
@@ -142,9 +186,30 @@
         mNumCpus = numCpus;
     }
 
+    /** Loads a config from a file. */
+    @NonNull
+    static VirtualMachineConfig from(@NonNull File file) throws VirtualMachineException {
+        try (FileInputStream input = new FileInputStream(file)) {
+            return fromInputStream(input);
+        } catch (IOException e) {
+            throw new VirtualMachineException("Failed to read VM config from file", e);
+        }
+    }
+
+    /** Loads a config from a {@link ParcelFileDescriptor}. */
+    @NonNull
+    static VirtualMachineConfig from(@NonNull ParcelFileDescriptor fd)
+            throws VirtualMachineException {
+        try (AutoCloseInputStream input = new AutoCloseInputStream(fd)) {
+            return fromInputStream(input);
+        } catch (IOException e) {
+            throw new VirtualMachineException("failed to read VM config from file descriptor", e);
+        }
+    }
+
     /** Loads a config from a stream, for example a file. */
     @NonNull
-    static VirtualMachineConfig from(@NonNull InputStream input)
+    private static VirtualMachineConfig fromInputStream(@NonNull InputStream input)
             throws IOException, VirtualMachineException {
         PersistableBundle b = PersistableBundle.readFromStream(input);
         int version = b.getInt(KEY_VERSION);
@@ -176,8 +241,17 @@
                 protectedVm, memoryMib, numCpus);
     }
 
+    /** Persists this config to a file. */
+    void serialize(@NonNull File file) throws VirtualMachineException {
+        try (FileOutputStream output = new FileOutputStream(file)) {
+            serializeOutputStream(output);
+        } catch (IOException e) {
+            throw new VirtualMachineException("failed to write VM config", e);
+        }
+    }
+
     /** Persists this config to a stream, for example a file. */
-    /* package */ void serialize(@NonNull OutputStream output) throws IOException {
+    private void serializeOutputStream(@NonNull OutputStream output) throws IOException {
         PersistableBundle b = new PersistableBundle();
         b.putInt(KEY_VERSION, VERSION);
         b.putString(KEY_APKPATH, mApkPath);
@@ -198,27 +272,31 @@
      *
      * @hide
      */
+    @SystemApi
     @NonNull
     public String getApkPath() {
         return mApkPath;
     }
 
     /**
-     * Returns the path to the payload config within the owning application.
+     * Returns the path within the APK to the payload config file that defines software aspects of
+     * the VM.
      *
      * @hide
      */
+    @SystemApi // TODO(b/243512115): Switch back to @TestApi
     @Nullable
     public String getPayloadConfigPath() {
         return mPayloadConfigPath;
     }
 
     /**
-     * Returns the path within the APK to the payload binary file that will be executed within the
-     * VM.
+     * Returns the path within the {@code lib/<ABI>} directory of the APK to the payload binary file
+     * that will be executed within the VM.
      *
      * @hide
      */
+    @SystemApi
     @Nullable
     public String getPayloadBinaryPath() {
         return mPayloadBinaryPath;
@@ -229,6 +307,7 @@
      *
      * @hide
      */
+    @SystemApi
     @NonNull
     @DebugLevel
     public int getDebugLevel() {
@@ -240,6 +319,7 @@
      *
      * @hide
      */
+    @SystemApi
     public boolean isProtectedVm() {
         return mProtectedVm;
     }
@@ -249,6 +329,8 @@
      *
      * @hide
      */
+    @SystemApi
+    @IntRange(from = 0)
     public int getMemoryMib() {
         return mMemoryMib;
     }
@@ -258,6 +340,8 @@
      *
      * @hide
      */
+    @SystemApi
+    @IntRange(from = 1)
     public int getNumCpus() {
         return mNumCpus;
     }
@@ -270,6 +354,7 @@
      *
      * @hide
      */
+    @SystemApi
     public boolean isCompatibleWith(@NonNull VirtualMachineConfig other) {
         return this.mDebugLevel == other.mDebugLevel
                 && this.mProtectedVm == other.mProtectedVm
@@ -279,41 +364,42 @@
     }
 
     /**
-     * Converts this config object into a parcel. Used when creating a VM via the virtualization
-     * service. Notice that the files are not passed as paths, but as file descriptors because the
-     * service doesn't accept paths as it might not have permission to open app-owned files and that
-     * could be abused to run a VM with software that the calling application doesn't own.
+     * Converts this config object into the parcelable type used when creating a VM via the
+     * virtualization service. Notice that the files are not passed as paths, but as file
+     * descriptors because the service doesn't accept paths as it might not have permission to open
+     * app-owned files and that could be abused to run a VM with software that the calling
+     * application doesn't own.
      */
-    VirtualMachineAppConfig toParcel() throws FileNotFoundException {
-        VirtualMachineAppConfig parcel = new VirtualMachineAppConfig();
-        parcel.apk = ParcelFileDescriptor.open(new File(mApkPath), MODE_READ_ONLY);
+    VirtualMachineAppConfig toVsConfig() throws FileNotFoundException {
+        VirtualMachineAppConfig vsConfig = new VirtualMachineAppConfig();
+        vsConfig.apk = ParcelFileDescriptor.open(new File(mApkPath), MODE_READ_ONLY);
         if (mPayloadBinaryPath != null) {
             VirtualMachinePayloadConfig payloadConfig = new VirtualMachinePayloadConfig();
             payloadConfig.payloadPath = mPayloadBinaryPath;
-            parcel.payload =
+            vsConfig.payload =
                     VirtualMachineAppConfig.Payload.payloadConfig(payloadConfig);
         } else {
-            parcel.payload =
+            vsConfig.payload =
                     VirtualMachineAppConfig.Payload.configPath(mPayloadConfigPath);
         }
         switch (mDebugLevel) {
             case DEBUG_LEVEL_APP_ONLY:
-                parcel.debugLevel = VirtualMachineAppConfig.DebugLevel.APP_ONLY;
+                vsConfig.debugLevel = VirtualMachineAppConfig.DebugLevel.APP_ONLY;
                 break;
             case DEBUG_LEVEL_FULL:
-                parcel.debugLevel = VirtualMachineAppConfig.DebugLevel.FULL;
+                vsConfig.debugLevel = VirtualMachineAppConfig.DebugLevel.FULL;
                 break;
             default:
-                parcel.debugLevel = VirtualMachineAppConfig.DebugLevel.NONE;
+                vsConfig.debugLevel = VirtualMachineAppConfig.DebugLevel.NONE;
                 break;
         }
-        parcel.protectedVm = mProtectedVm;
-        parcel.memoryMib = mMemoryMib;
-        parcel.numCpus = mNumCpus;
+        vsConfig.protectedVm = mProtectedVm;
+        vsConfig.memoryMib = mMemoryMib;
+        vsConfig.numCpus = mNumCpus;
         // Don't allow apps to set task profiles ... at last for now. Also, don't forget to
         // validate the string because these are appended to the cmdline argument.
-        parcel.taskProfiles = new String[0];
-        return parcel;
+        vsConfig.taskProfiles = new String[0];
+        return vsConfig;
     }
 
     /**
@@ -321,6 +407,7 @@
      *
      * @hide
      */
+    @SystemApi
     public static final class Builder {
         private final Context mContext;
         @Nullable private String mApkPath;
@@ -333,12 +420,13 @@
         private int mNumCpus;
 
         /**
-         * Creates a builder for the given context (APK).
+         * Creates a builder for the given context.
          *
          * @hide
          */
+        @SystemApi
         public Builder(@NonNull Context context) {
-            mContext = requireNonNull(context);
+            mContext = requireNonNull(context, "context must not be null");
             mDebugLevel = DEBUG_LEVEL_NONE;
             mNumCpus = 1;
         }
@@ -348,41 +436,15 @@
          *
          * @hide
          */
+        @SystemApi
         @NonNull
         public VirtualMachineConfig build() {
             String apkPath = (mApkPath == null) ? mContext.getPackageCodePath() : mApkPath;
 
-            int availableCpus = Runtime.getRuntime().availableProcessors();
-            if (mNumCpus < 1 || mNumCpus > availableCpus) {
-                throw new IllegalArgumentException("Number of vCPUs (" + mNumCpus + ") is out of "
-                        + "range [1, " + availableCpus + "]");
-            }
-
-            if (mPayloadBinaryPath == null) {
-                if (mPayloadConfigPath == null) {
-                    throw new IllegalStateException("payloadBinaryPath must be set");
-                }
-            } else {
-                if (mPayloadConfigPath != null) {
-                    throw new IllegalStateException(
-                            "payloadBinaryPath and payloadConfigPath may not both be set");
-                }
-            }
-
             if (!mProtectedVmSet) {
                 throw new IllegalStateException("setProtectedVm must be called explicitly");
             }
 
-            if (mProtectedVm
-                    && !HypervisorProperties.hypervisor_protected_vm_supported().orElse(false)) {
-                throw new UnsupportedOperationException(
-                        "Protected VMs are not supported on this device.");
-            }
-            if (!mProtectedVm && !HypervisorProperties.hypervisor_vm_supported().orElse(false)) {
-                throw new UnsupportedOperationException(
-                        "Unprotected VMs are not supported on this device.");
-            }
-
             return new VirtualMachineConfig(
                     apkPath, mPayloadConfigPath, mPayloadBinaryPath, mDebugLevel, mProtectedVm,
                     mMemoryMib, mNumCpus);
@@ -394,6 +456,7 @@
          *
          * @hide
          */
+        @SystemApi
         @NonNull
         public Builder setApkPath(@NonNull String apkPath) {
             mApkPath = requireNonNull(apkPath);
@@ -401,12 +464,14 @@
         }
 
         /**
-         * Sets the path within the APK to the payload config file that defines software aspects
-         * of the VM.
+         * Sets the path within the APK to the payload config file that defines software aspects of
+         * the VM. The file is a JSON file; see
+         * packages/modules/Virtualization/microdroid/payload/config/src/lib.rs for the format.
          *
          * @hide
          */
         @RequiresPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION)
+        @SystemApi // TODO(b/243512115): Switch to @TestApi
         @NonNull
         public Builder setPayloadConfigPath(@NonNull String payloadConfigPath) {
             mPayloadConfigPath = requireNonNull(payloadConfigPath);
@@ -419,6 +484,7 @@
          *
          * @hide
          */
+        @SystemApi
         @NonNull
         public Builder setPayloadBinaryPath(@NonNull String payloadBinaryPath) {
             mPayloadBinaryPath = requireNonNull(payloadBinaryPath);
@@ -426,10 +492,11 @@
         }
 
         /**
-         * Sets the debug level
+         * Sets the debug level. Defaults to {@link #DEBUG_LEVEL_NONE}.
          *
          * @hide
          */
+        @SystemApi
         @NonNull
         public Builder setDebugLevel(@DebugLevel int debugLevel) {
             mDebugLevel = debugLevel;
@@ -437,12 +504,13 @@
         }
 
         /**
-         * Sets whether to protect the VM memory from the host. No default is provided, this
-         * must be set explicitly.
+         * Sets whether to protect the VM memory from the host. No default is provided, this must be
+         * set explicitly.
          *
          * @see VirtualMachineManager#getCapabilities
          * @hide
          */
+        @SystemApi
         @NonNull
         public Builder setProtectedVm(boolean protectedVm) {
             mProtectedVm = protectedVm;
@@ -451,24 +519,27 @@
         }
 
         /**
-         * Sets the amount of RAM to give the VM. If this is zero or negative then the default will
-         * be used.
+         * Sets the amount of RAM to give the VM, in mebibytes. If zero or not explicitly set then a
+         * default size will be used.
          *
          * @hide
          */
+        @SystemApi
         @NonNull
-        public Builder setMemoryMib(int memoryMib) {
+        public Builder setMemoryMib(@IntRange(from = 0) int memoryMib) {
             mMemoryMib = memoryMib;
             return this;
         }
 
         /**
-         * Sets the number of vCPUs in the VM. Defaults to 1.
+         * Sets the number of vCPUs in the VM. Defaults to 1. Cannot be more than the number of real
+         * CPUs (as returned by {@link Runtime#availableProcessors()}).
          *
          * @hide
          */
+        @SystemApi
         @NonNull
-        public Builder setNumCpus(int num) {
+        public Builder setNumCpus(@IntRange(from = 1) int num) {
             mNumCpus = num;
             return this;
         }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineDescriptor.java b/javalib/src/android/system/virtualmachine/VirtualMachineDescriptor.java
new file mode 100644
index 0000000..edaf5b4
--- /dev/null
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineDescriptor.java
@@ -0,0 +1,91 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.system.virtualmachine;
+
+import static java.util.Objects.requireNonNull;
+
+import android.annotation.NonNull;
+import android.annotation.SystemApi;
+import android.os.Parcel;
+import android.os.ParcelFileDescriptor;
+import android.os.Parcelable;
+
+/**
+ * A VM descriptor that captures the state of a Virtual Machine.
+ *
+ * <p>You can capture the current state of VM by creating an instance of this class with {@link
+ * VirtualMachine#toDescriptor()}, optionally pass it to another App, and then build an identical VM
+ * with the descriptor received.
+ *
+ * @hide
+ */
+@SystemApi
+public final class VirtualMachineDescriptor implements Parcelable {
+    @NonNull private final ParcelFileDescriptor mConfigFd;
+    @NonNull private final ParcelFileDescriptor mInstanceImgFd;
+    // TODO(b/243129654): Add trusted storage fd once it is available.
+
+    @Override
+    public int describeContents() {
+        return CONTENTS_FILE_DESCRIPTOR;
+    }
+
+    @Override
+    public void writeToParcel(@NonNull Parcel out, int flags) {
+        mConfigFd.writeToParcel(out, flags);
+        mInstanceImgFd.writeToParcel(out, flags);
+    }
+
+    @NonNull
+    public static final Parcelable.Creator<VirtualMachineDescriptor> CREATOR =
+            new Parcelable.Creator<>() {
+                public VirtualMachineDescriptor createFromParcel(Parcel in) {
+                    return new VirtualMachineDescriptor(in);
+                }
+
+                public VirtualMachineDescriptor[] newArray(int size) {
+                    return new VirtualMachineDescriptor[size];
+                }
+            };
+
+    /**
+     * @return File descriptor of the VM configuration file config.xml.
+     */
+    @NonNull
+    ParcelFileDescriptor getConfigFd() {
+        return mConfigFd;
+    }
+
+    /**
+     * @return File descriptor of the instance.img of the VM.
+     */
+    @NonNull
+    ParcelFileDescriptor getInstanceImgFd() {
+        return mInstanceImgFd;
+    }
+
+    VirtualMachineDescriptor(
+            @NonNull ParcelFileDescriptor configFd, @NonNull ParcelFileDescriptor instanceImgFd) {
+        mConfigFd = configFd;
+        mInstanceImgFd = instanceImgFd;
+    }
+
+    private VirtualMachineDescriptor(Parcel in) {
+        mConfigFd = requireNonNull(in.readFileDescriptor());
+        mInstanceImgFd = requireNonNull(in.readFileDescriptor());
+    }
+}
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineException.java b/javalib/src/android/system/virtualmachine/VirtualMachineException.java
index 88b5ea3..985eb70 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineException.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineException.java
@@ -17,17 +17,15 @@
 package android.system.virtualmachine;
 
 import android.annotation.Nullable;
+import android.annotation.SystemApi;
 
 /**
  * Exception thrown when operations on virtual machines fail.
  *
  * @hide
  */
+@SystemApi
 public class VirtualMachineException extends Exception {
-    public VirtualMachineException() {
-        super();
-    }
-
     public VirtualMachineException(@Nullable String message) {
         super(message);
     }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
index 3f904be..ea0a305 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
@@ -21,31 +21,56 @@
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.RequiresFeature;
 import android.annotation.RequiresPermission;
-import android.annotation.SuppressLint;
+import android.annotation.SystemApi;
 import android.content.Context;
+import android.content.pm.PackageManager;
 import android.sysprop.HypervisorProperties;
+import android.util.ArrayMap;
+
+import com.android.internal.annotations.GuardedBy;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.ref.WeakReference;
 import java.util.Map;
-import java.util.WeakHashMap;
 
 /**
- * Manages {@link VirtualMachine} objects created for an application.
+ * Manages {@link VirtualMachine virtual machine} instances created by an app. Each instance is
+ * created from a {@link VirtualMachineConfig configuration} that defines the shape of the VM (RAM,
+ * CPUs), the code to execute within it, etc.
+ *
+ * <p>Each virtual machine instance is named; the configuration and related state of each is
+ * persisted in the app's private data directory and an instance can be retrieved given the name.
+ * The name must be a valid directory name and must not contain '/'.
+ *
+ * <p>The app can then start, stop and otherwise interact with the VM.
+ *
+ * <p>An instance of {@link VirtualMachineManager} can be obtained by calling {@link
+ * Context#getSystemService(Class)}.
  *
  * @hide
  */
+@SystemApi
+@RequiresFeature(PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK)
 public class VirtualMachineManager {
+    /**
+     * A lock used to synchronize the creation of virtual machines. It protects {@link #mVmsByName},
+     * but is also held throughout VM creation / retrieval / deletion, to prevent these actions
+     * racing with each other.
+     */
+    private static final Object sCreateLock = new Object();
+
     @NonNull private final Context mContext;
 
-    private VirtualMachineManager(@NonNull Context context) {
-        mContext = context;
+    /** @hide */
+    public VirtualMachineManager(@NonNull Context context) {
+        mContext = requireNonNull(context);
     }
 
-    private static final Map<Context, WeakReference<VirtualMachineManager>> sInstances =
-            new WeakHashMap<>();
+    @GuardedBy("sCreateLock")
+    private final Map<String, WeakReference<VirtualMachine>> mVmsByName = new ArrayMap<>();
 
     /**
      * Capabilities of the virtual machine implementation.
@@ -70,35 +95,13 @@
     public static final int CAPABILITY_NON_PROTECTED_VM = 2;
 
     /**
-     * Returns the per-context instance.
-     *
-     * @hide
-     */
-    @NonNull
-    @SuppressLint("ManagerLookup") // Optional API
-    public static VirtualMachineManager getInstance(@NonNull Context context) {
-        requireNonNull(context, "context must not be null");
-        synchronized (sInstances) {
-            VirtualMachineManager vmm =
-                    sInstances.containsKey(context) ? sInstances.get(context).get() : null;
-            if (vmm == null) {
-                vmm = new VirtualMachineManager(context);
-                sInstances.put(context, new WeakReference<>(vmm));
-            }
-            return vmm;
-        }
-    }
-
-    /** A lock used to synchronize the creation of virtual machines */
-    private static final Object sCreateLock = new Object();
-
-    /**
      * Returns a set of flags indicating what this implementation of virtualization is capable of.
      *
      * @see #CAPABILITY_PROTECTED_VM
      * @see #CAPABILITY_NON_PROTECTED_VM
      * @hide
      */
+    @SystemApi
     @Capability
     public int getCapabilities() {
         @Capability int result = 0;
@@ -116,34 +119,81 @@
      * machine with the same name as an existing virtual machine is an error. The existing virtual
      * machine has to be deleted before its name can be reused.
      *
-     * Each successful call to this method creates a new (and different) virtual machine even if the
-     * name and the config are the same as a deleted one. The new virtual machine will initially
+     * <p>Each successful call to this method creates a new (and different) virtual machine even if
+     * the name and the config are the same as a deleted one. The new virtual machine will initially
      * be stopped.
      *
      * @throws VirtualMachineException if the VM cannot be created, or there is an existing VM with
-     *         the given name.
+     *     the given name.
      * @hide
      */
+    @SystemApi
     @NonNull
     @RequiresPermission(VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION)
-    public VirtualMachine create(
-            @NonNull String name, @NonNull VirtualMachineConfig config)
+    public VirtualMachine create(@NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
         synchronized (sCreateLock) {
-            return VirtualMachine.create(mContext, name, config);
+            return createLocked(name, config);
         }
     }
 
+    @NonNull
+    @GuardedBy("sCreateLock")
+    private VirtualMachine createLocked(@NonNull String name, @NonNull VirtualMachineConfig config)
+            throws VirtualMachineException {
+        VirtualMachine vm = VirtualMachine.create(mContext, name, config);
+        mVmsByName.put(name, new WeakReference<>(vm));
+        return vm;
+    }
+
     /**
      * Returns an existing {@link VirtualMachine} with the given name. Returns null if there is no
      * such virtual machine.
      *
-     * @throws VirtualMachineException if the virtual machine could not be successfully retrieved.
+     * @throws VirtualMachineException if the virtual machine exists but could not be successfully
+     *     retrieved.
      * @hide
      */
+    @SystemApi
     @Nullable
     public VirtualMachine get(@NonNull String name) throws VirtualMachineException {
-        return VirtualMachine.load(mContext, name);
+        synchronized (sCreateLock) {
+            return getLocked(name);
+        }
+    }
+
+    @Nullable
+    @GuardedBy("sCreateLock")
+    private VirtualMachine getLocked(@NonNull String name) throws VirtualMachineException {
+        VirtualMachine vm = getVmByName(name);
+        if (vm != null) return vm;
+
+        vm = VirtualMachine.load(mContext, name);
+        if (vm != null) {
+            mVmsByName.put(name, new WeakReference<>(vm));
+        }
+        return vm;
+    }
+
+    /**
+     * Imports a virtual machine from an {@link VirtualMachineDescriptor} object and associates it
+     * with the given name.
+     *
+     * <p>The new virtual machine will be in the same state as the descriptor indicates.
+     *
+     * @throws VirtualMachineException if the VM cannot be imported.
+     * @hide
+     */
+    @NonNull
+    @SystemApi
+    public VirtualMachine importFromDescriptor(
+            @NonNull String name, @NonNull VirtualMachineDescriptor vmDescriptor)
+            throws VirtualMachineException {
+        synchronized (sCreateLock) {
+            VirtualMachine vm = VirtualMachine.fromDescriptor(mContext, name, vmDescriptor);
+            mVmsByName.put(name, new WeakReference<>(vm));
+            return vm;
+        }
     }
 
     /**
@@ -153,17 +203,55 @@
      * @throws VirtualMachineException if the virtual machine could not be created or retrieved.
      * @hide
      */
+    @SystemApi
     @NonNull
-    public VirtualMachine getOrCreate(
-            @NonNull String name, @NonNull VirtualMachineConfig config)
+    public VirtualMachine getOrCreate(@NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
-        VirtualMachine vm;
         synchronized (sCreateLock) {
-            vm = get(name);
-            if (vm == null) {
-                vm = create(name, config);
+            VirtualMachine vm = getLocked(name);
+            if (vm != null) {
+                return vm;
+            } else {
+                return createLocked(name, config);
             }
         }
-        return vm;
+    }
+
+    /**
+     * Deletes an existing {@link VirtualMachine}. Deleting a virtual machine means deleting any
+     * persisted data associated with it including the per-VM secret. This is an irreversible
+     * action. A virtual machine once deleted can never be restored. A new virtual machine created
+     * with the same name is different from an already deleted virtual machine even if it has the
+     * same config.
+     *
+     * @throws VirtualMachineException if the virtual machine does not exist, is not stopped, or
+     *     cannot be deleted.
+     * @hide
+     */
+    @SystemApi
+    public void delete(@NonNull String name) throws VirtualMachineException {
+        synchronized (sCreateLock) {
+            VirtualMachine vm = getVmByName(name);
+            if (vm == null) {
+                VirtualMachine.deleteVmDirectory(mContext, name);
+            } else {
+                vm.delete(mContext, name);
+            }
+            mVmsByName.remove(name);
+        }
+    }
+
+    @Nullable
+    @GuardedBy("sCreateLock")
+    private VirtualMachine getVmByName(@NonNull String name) {
+        requireNonNull(name);
+        WeakReference<VirtualMachine> weakReference = mVmsByName.get(name);
+        if (weakReference != null) {
+            VirtualMachine vm = weakReference.get();
+            if (vm != null && vm.getStatus() != VirtualMachine.STATUS_DELETED) {
+                return vm;
+            }
+        }
+        return null;
     }
 }
diff --git a/javalib/src/android/system/virtualmachine/VirtualizationFrameworkInitializer.java b/javalib/src/android/system/virtualmachine/VirtualizationFrameworkInitializer.java
new file mode 100644
index 0000000..30ac425
--- /dev/null
+++ b/javalib/src/android/system/virtualmachine/VirtualizationFrameworkInitializer.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.system.virtualmachine;
+
+import static android.content.pm.PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK;
+
+import android.annotation.SystemApi;
+import android.app.SystemServiceRegistry;
+import android.content.Context;
+
+/**
+ * Holds initialization code for virtualization module
+ *
+ * @hide
+ */
+@SystemApi(client = SystemApi.Client.MODULE_LIBRARIES)
+public class VirtualizationFrameworkInitializer {
+
+    private VirtualizationFrameworkInitializer() {}
+
+    /**
+     * Called by the static initializer in the {@link SystemServiceRegistry}, and registers {@link
+     * VirtualMachineManager} to the {@link Context}. so that it's accessible from {@link
+     * Context#getSystemService(String)}.
+     */
+    public static void registerServiceWrappers() {
+        // Note: it's important that the getPackageManager().hasSystemFeature() check is executed
+        // in the lambda, and not directly in the registerServiceWrappers method. The
+        // registerServiceWrappers is called during Zygote static initialization, and at that
+        // point the PackageManager is not available yet.
+        //
+        // On the other hand, the lambda is executed after the app calls Context.getSystemService
+        // (VirtualMachineManager.class), at which point the PackageManager is available. The
+        // result of the lambda is cached on per-context basis.
+        SystemServiceRegistry.registerContextAwareService(
+                Context.VIRTUALIZATION_SERVICE,
+                VirtualMachineManager.class,
+                ctx ->
+                        ctx.getPackageManager().hasSystemFeature(FEATURE_VIRTUALIZATION_FRAMEWORK)
+                                ? new VirtualMachineManager(ctx)
+                                : null);
+    }
+}
diff --git a/libs/devicemapper/Android.bp b/libs/devicemapper/Android.bp
index 783fa79..9fa010c 100644
--- a/libs/devicemapper/Android.bp
+++ b/libs/devicemapper/Android.bp
@@ -13,6 +13,7 @@
         "libbitflags",
         "liblibc",
         "libdata_model",
+        "libhex",
         "libnix",
         "libuuid",
     ],
@@ -33,6 +34,7 @@
     defaults: ["libdm_rust.defaults"],
     test_suites: ["general-tests"],
     rustlibs: [
+        "librustutils",
         "libscopeguard",
         "libtempfile",
     ],
diff --git a/libs/devicemapper/src/crypt.rs b/libs/devicemapper/src/crypt.rs
index 9b715a5..b2e677a 100644
--- a/libs/devicemapper/src/crypt.rs
+++ b/libs/devicemapper/src/crypt.rs
@@ -17,10 +17,9 @@
 /// `crypt` module implements the "crypt" target in the device mapper framework. Specifically,
 /// it provides `DmCryptTargetBuilder` struct which is used to construct a `DmCryptTarget` struct
 /// which is then given to `DeviceMapper` to create a mapper device.
-use crate::util::*;
 use crate::DmTargetSpec;
 
-use anyhow::{bail, Context, Result};
+use anyhow::{ensure, Context, Result};
 use data_model::DataInit;
 use std::io::Write;
 use std::mem::size_of;
@@ -32,10 +31,34 @@
 // Documentation/admin-guide/device-mapper/dm-crypt.rst
 
 /// Supported ciphers
+#[derive(Clone, Copy, Debug)]
 pub enum CipherType {
-    // TODO(b/253394457) Include ciphers with authenticated modes as well
+    // AES-256-HCTR2 takes a 32-byte key
+    AES256HCTR2,
+    // XTS requires key of twice the length of the underlying block cipher i.e., 64B for AES256
     AES256XTS,
 }
+impl CipherType {
+    fn get_kernel_crypto_name(&self) -> &str {
+        match *self {
+            // We use "plain64" as the IV/nonce generation algorithm -
+            // which basically is the sector number.
+            CipherType::AES256HCTR2 => "aes-hctr2-plain64",
+            CipherType::AES256XTS => "aes-xts-plain64",
+        }
+    }
+
+    fn get_required_key_size(&self) -> usize {
+        match *self {
+            CipherType::AES256HCTR2 => 32,
+            CipherType::AES256XTS => 64,
+        }
+    }
+
+    fn validata_key_size(&self, key_size: usize) -> bool {
+        key_size == self.get_required_key_size()
+    }
+}
 
 pub struct DmCryptTarget(Box<[u8]>);
 
@@ -59,7 +82,7 @@
 impl<'a> Default for DmCryptTargetBuilder<'a> {
     fn default() -> Self {
         DmCryptTargetBuilder {
-            cipher: CipherType::AES256XTS,
+            cipher: CipherType::AES256HCTR2,
             key: None,
             iv_offset: 0,
             device_path: None,
@@ -112,8 +135,13 @@
             .to_str()
             .context("data device path is not encoded in utf8")?;
 
-        let key =
-            if let Some(key) = self.key { hexstring_from(key) } else { bail!("key is not set") };
+        ensure!(self.key.is_some(), "key is not set");
+        // Unwrap is safe because we already made sure key.is_some()
+        ensure!(
+            self.cipher.validata_key_size(self.key.unwrap().len()),
+            format!("Invalid key size for cipher:{}", self.cipher.get_kernel_crypto_name())
+        );
+        let key = hex::encode(self.key.unwrap());
 
         // Step2: serialize the information according to the spec, which is ...
         // DmTargetSpec{...}
@@ -121,7 +149,7 @@
         // <offset> [<#opt_params> <opt_params>]
         let mut body = String::new();
         use std::fmt::Write;
-        write!(&mut body, "{} ", get_kernel_crypto_name(&self.cipher))?;
+        write!(&mut body, "{} ", self.cipher.get_kernel_crypto_name())?;
         write!(&mut body, "{} ", key)?;
         write!(&mut body, "{} ", self.iv_offset)?;
         write!(&mut body, "{} ", device_path)?;
@@ -145,9 +173,3 @@
         Ok(DmCryptTarget(buf.into_boxed_slice()))
     }
 }
-
-fn get_kernel_crypto_name(cipher: &CipherType) -> &str {
-    match cipher {
-        CipherType::AES256XTS => "aes-xts-plain64",
-    }
-}
diff --git a/libs/devicemapper/src/lib.rs b/libs/devicemapper/src/lib.rs
index b9fb5c3..9069eb2 100644
--- a/libs/devicemapper/src/lib.rs
+++ b/libs/devicemapper/src/lib.rs
@@ -212,7 +212,7 @@
         dm_dev_suspend(self, &mut data).context("failed to activate")?;
 
         // Step 4: wait unti the device is created and return the device path
-        let path = Path::new(MAPPER_DEV_ROOT).join(&name);
+        let path = Path::new(MAPPER_DEV_ROOT).join(name);
         wait_for_path(&path)?;
         Ok(path)
     }
@@ -234,12 +234,28 @@
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crypt::DmCryptTargetBuilder;
+    use crypt::{CipherType, DmCryptTargetBuilder};
+    use rustutils::system_properties;
     use std::fs::{read, File, OpenOptions};
     use std::io::Write;
 
-    const KEY: &[u8; 32] = b"thirtytwobyteslongreallylongword";
-    const DIFFERENT_KEY: &[u8; 32] = b"drowgnolyllaergnolsetybowtytriht";
+    // Just a logical set of keys to make testing easy. This has no real meaning.
+    struct KeySet<'a> {
+        cipher: CipherType,
+        key: &'a [u8],
+        different_key: &'a [u8],
+    }
+
+    const KEY_SET_XTS: KeySet = KeySet {
+        cipher: CipherType::AES256XTS,
+        key: b"sixtyfourbyteslongsentencearerarebutletsgiveitatrycantbethathard",
+        different_key: b"drahtahtebtnacyrtatievigsteltuberareraecnetnesgnolsetybruofytxis",
+    };
+    const KEY_SET_HCTR2: KeySet = KeySet {
+        cipher: CipherType::AES256HCTR2,
+        key: b"thirtytwobyteslongreallylongword",
+        different_key: b"drowgnolyllaergnolsetybowtytriht",
+    };
 
     // Create a file in given temp directory with given size
     fn prepare_tmpfile(test_dir: &Path, filename: &str, sz: u64) -> PathBuf {
@@ -250,18 +266,58 @@
     }
 
     fn write_to_dev(path: &Path, data: &[u8]) {
-        let mut f = OpenOptions::new().read(true).write(true).open(&path).unwrap();
+        let mut f = OpenOptions::new().read(true).write(true).open(path).unwrap();
         f.write_all(data).unwrap();
     }
 
+    // TODO(b/250880499): delete_device() doesn't really delete it even without DM_DEFERRED_REMOVE.
+    // Hence, we have to create a new device with a different name for each test. Retrying
+    // the test on same machine without reboot will also fail.
     fn delete_device(dm: &DeviceMapper, name: &str) -> Result<()> {
         dm.delete_device_deferred(name)?;
-        wait_for_path_disappears(Path::new(MAPPER_DEV_ROOT).join(&name))?;
+        wait_for_path_disappears(Path::new(MAPPER_DEV_ROOT).join(name))?;
         Ok(())
     }
 
+    // TODO(b/260692911): Find a better way to skip a test instead of silently passing it.
+    fn is_hctr2_supported() -> bool {
+        // hctr2 is NOT enabled in kernel 5.10 or lower. We run Microdroid tests on kernel versions
+        // 5.10 or above & therefore,  we don't really care to skip test on other versions.
+        if let Some(version) = system_properties::read("ro.kernel.version")
+            .expect("Unable to read system property ro.kernel.version")
+        {
+            version != "5.10"
+        } else {
+            panic!("Could not read property: kernel.version!!");
+        }
+    }
+
     #[test]
-    fn mapping_again_keeps_data() {
+    fn mapping_again_keeps_data_xts() {
+        mapping_again_keeps_data(&KEY_SET_XTS, "name1");
+    }
+
+    #[test]
+    fn mapping_again_keeps_data_hctr2() {
+        if !is_hctr2_supported() {
+            return;
+        }
+        mapping_again_keeps_data(&KEY_SET_HCTR2, "name2");
+    }
+    #[test]
+    fn data_inaccessible_with_diff_key_xts() {
+        data_inaccessible_with_diff_key(&KEY_SET_XTS, "name3");
+    }
+
+    #[test]
+    fn data_inaccessible_with_diff_key_hctr2() {
+        if !is_hctr2_supported() {
+            return;
+        }
+        data_inaccessible_with_diff_key(&KEY_SET_HCTR2, "name4");
+    }
+
+    fn mapping_again_keeps_data(keyset: &KeySet, device: &str) {
         // This test creates 2 different crypt devices using same key backed by same data_device
         // -> Write data on dev1 -> Check the data is visible & same on dev2
         let dm = DeviceMapper::new().unwrap();
@@ -278,28 +334,33 @@
             /*writable*/ true,
         )
         .unwrap();
+        let device_diff = device.to_owned() + "_diff";
+
         scopeguard::defer! {
             loopdevice::detach(&data_device).unwrap();
-            _ = delete_device(&dm, "crypt1");
-            _ = delete_device(&dm, "crypt2");
+            _ = delete_device(&dm, device);
+            _ = delete_device(&dm, &device_diff);
         }
 
-        let target =
-            DmCryptTargetBuilder::default().data_device(&data_device, sz).key(KEY).build().unwrap();
+        let target = DmCryptTargetBuilder::default()
+            .data_device(&data_device, sz)
+            .cipher(keyset.cipher)
+            .key(keyset.key)
+            .build()
+            .unwrap();
 
-        let mut crypt_device = dm.create_crypt_device("crypt1", &target).unwrap();
+        let mut crypt_device = dm.create_crypt_device(device, &target).unwrap();
         write_to_dev(&crypt_device, inputimg);
 
         // Recreate another device using same target spec & check if the content is the same
-        crypt_device = dm.create_crypt_device("crypt2", &target).unwrap();
+        crypt_device = dm.create_crypt_device(&device_diff, &target).unwrap();
 
         let crypt = read(crypt_device).unwrap();
         assert_eq!(inputimg.len(), crypt.len()); // fail early if the size doesn't match
         assert_eq!(inputimg, crypt.as_slice());
     }
 
-    #[test]
-    fn data_inaccessible_with_diff_key() {
+    fn data_inaccessible_with_diff_key(keyset: &KeySet, device: &str) {
         // This test creates 2 different crypt devices using different keys backed
         // by same data_device -> Write data on dev1 -> Check the data is visible but not the same on dev2
         let dm = DeviceMapper::new().unwrap();
@@ -316,26 +377,32 @@
             /*writable*/ true,
         )
         .unwrap();
+        let device_diff = device.to_owned() + "_diff";
         scopeguard::defer! {
             loopdevice::detach(&data_device).unwrap();
-            _ = delete_device(&dm, "crypt3");
-            _ = delete_device(&dm, "crypt4");
+            _ = delete_device(&dm, device);
+            _ = delete_device(&dm, &device_diff);
         }
 
-        let target =
-            DmCryptTargetBuilder::default().data_device(&data_device, sz).key(KEY).build().unwrap();
+        let target = DmCryptTargetBuilder::default()
+            .data_device(&data_device, sz)
+            .cipher(keyset.cipher)
+            .key(keyset.key)
+            .build()
+            .unwrap();
         let target2 = DmCryptTargetBuilder::default()
             .data_device(&data_device, sz)
-            .key(DIFFERENT_KEY)
+            .cipher(keyset.cipher)
+            .key(keyset.different_key)
             .build()
             .unwrap();
 
-        let mut crypt_device = dm.create_crypt_device("crypt3", &target).unwrap();
+        let mut crypt_device = dm.create_crypt_device(device, &target).unwrap();
 
         write_to_dev(&crypt_device, inputimg);
 
         // Recreate the crypt device again diff key & check if the content is changed
-        crypt_device = dm.create_crypt_device("crypt4", &target2).unwrap();
+        crypt_device = dm.create_crypt_device(&device_diff, &target2).unwrap();
         let crypt = read(crypt_device).unwrap();
         assert_ne!(inputimg, crypt.as_slice());
     }
diff --git a/libs/libfdt/Android.bp b/libs/libfdt/Android.bp
new file mode 100644
index 0000000..72399b0
--- /dev/null
+++ b/libs/libfdt/Android.bp
@@ -0,0 +1,45 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_bindgen {
+    name: "liblibfdt_bindgen",
+    crate_name: "libfdt_bindgen",
+    wrapper_src: "bindgen/fdt.h",
+    source_stem: "bindings",
+    bindgen_flags: [
+        "--size_t-is-usize",
+        "--allowlist-type=fdt_.*",
+        "--allowlist-function=fdt_.*",
+        "--allowlist-var=FDT_.*",
+        "--use-core",
+        "--raw-line=#![no_std]",
+        "--ctypes-prefix=core::ffi",
+    ],
+    static_libs: [
+        "libfdt",
+    ],
+    apex_available: ["com.android.virt"],
+}
+
+rust_library_rlib {
+    name: "liblibfdt",
+    crate_name: "libfdt",
+    srcs: [
+        "src/lib.rs",
+        ":liblibfdt_bindgen",
+    ],
+    edition: "2021",
+    no_stdlibs: true,
+    prefer_rlib: true,
+    stdlibs: [
+        "libcore.rust_sysroot",
+    ],
+    rustlibs: [
+        "liblibfdt_bindgen",
+    ],
+    whole_static_libs: [
+        "libfdt",
+    ],
+    apex_available: ["com.android.virt"],
+}
diff --git a/tests/helper/src/java/com/android/virt/VirtualizationTestHelper.java b/libs/libfdt/bindgen/fdt.h
similarity index 66%
rename from tests/helper/src/java/com/android/virt/VirtualizationTestHelper.java
rename to libs/libfdt/bindgen/fdt.h
index 4c27915..16f2784 100644
--- a/tests/helper/src/java/com/android/virt/VirtualizationTestHelper.java
+++ b/libs/libfdt/bindgen/fdt.h
@@ -5,7 +5,7 @@
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *      http://www.apache.org/licenses/LICENSE-2.0
+ *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
@@ -13,10 +13,5 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.virt;
 
-public abstract class VirtualizationTestHelper {
-    public static boolean isCuttlefish(String vendorDeviceName) {
-        return vendorDeviceName != null && vendorDeviceName.startsWith("vsoc_");
-    }
-}
+#include <libfdt.h>
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
new file mode 100644
index 0000000..ff1db63
--- /dev/null
+++ b/libs/libfdt/src/lib.rs
@@ -0,0 +1,663 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Wrapper around libfdt library. Provides parsing/generating functionality
+//! to a bare-metal environment.
+
+#![no_std]
+#![feature(let_else)] // Stabilized in 1.65.0
+
+use core::ffi::{c_int, c_void, CStr};
+use core::fmt;
+use core::mem;
+use core::ops::Range;
+use core::result;
+use core::slice;
+
+/// Error type corresponding to libfdt error codes.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum FdtError {
+    /// FDT_ERR_NOTFOUND
+    NotFound,
+    /// FDT_ERR_EXISTS
+    Exists,
+    /// FDT_ERR_NOSPACE
+    NoSpace,
+    /// FDT_ERR_BADOFFSET
+    BadOffset,
+    /// FDT_ERR_BADPATH
+    BadPath,
+    /// FDT_ERR_BADPHANDLE
+    BadPhandle,
+    /// FDT_ERR_BADSTATE
+    BadState,
+    /// FDT_ERR_TRUNCATED
+    Truncated,
+    /// FDT_ERR_BADMAGIC
+    BadMagic,
+    /// FDT_ERR_BADVERSION
+    BadVersion,
+    /// FDT_ERR_BADSTRUCTURE
+    BadStructure,
+    /// FDT_ERR_BADLAYOUT
+    BadLayout,
+    /// FDT_ERR_INTERNAL
+    Internal,
+    /// FDT_ERR_BADNCELLS
+    BadNCells,
+    /// FDT_ERR_BADVALUE
+    BadValue,
+    /// FDT_ERR_BADOVERLAY
+    BadOverlay,
+    /// FDT_ERR_NOPHANDLES
+    NoPhandles,
+    /// FDT_ERR_BADFLAGS
+    BadFlags,
+    /// FDT_ERR_ALIGNMENT
+    Alignment,
+    /// Unexpected error code
+    Unknown(i32),
+}
+
+impl fmt::Display for FdtError {
+    /// Prints error messages from libfdt.h documentation.
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::NotFound => write!(f, "The requested node or property does not exist"),
+            Self::Exists => write!(f, "Attempted to create an existing node or property"),
+            Self::NoSpace => write!(f, "Insufficient buffer space to contain the expanded tree"),
+            Self::BadOffset => write!(f, "Structure block offset is out-of-bounds or invalid"),
+            Self::BadPath => write!(f, "Badly formatted path"),
+            Self::BadPhandle => write!(f, "Invalid phandle length or value"),
+            Self::BadState => write!(f, "Received incomplete device tree"),
+            Self::Truncated => write!(f, "Device tree or sub-block is improperly terminated"),
+            Self::BadMagic => write!(f, "Device tree header missing its magic number"),
+            Self::BadVersion => write!(f, "Device tree has a version which can't be handled"),
+            Self::BadStructure => write!(f, "Device tree has a corrupt structure block"),
+            Self::BadLayout => write!(f, "Device tree sub-blocks in unsupported order"),
+            Self::Internal => write!(f, "libfdt has failed an internal assertion"),
+            Self::BadNCells => write!(f, "Bad format or value of #address-cells or #size-cells"),
+            Self::BadValue => write!(f, "Unexpected property value"),
+            Self::BadOverlay => write!(f, "Overlay cannot be applied"),
+            Self::NoPhandles => write!(f, "Device tree doesn't have any phandle available anymore"),
+            Self::BadFlags => write!(f, "Invalid flag or invalid combination of flags"),
+            Self::Alignment => write!(f, "Device tree base address is not 8-byte aligned"),
+            Self::Unknown(e) => write!(f, "Unknown libfdt error '{e}'"),
+        }
+    }
+}
+
+/// Result type with FdtError enum.
+pub type Result<T> = result::Result<T, FdtError>;
+
+fn fdt_err(val: c_int) -> Result<c_int> {
+    if val >= 0 {
+        Ok(val)
+    } else {
+        Err(match -val as _ {
+            libfdt_bindgen::FDT_ERR_NOTFOUND => FdtError::NotFound,
+            libfdt_bindgen::FDT_ERR_EXISTS => FdtError::Exists,
+            libfdt_bindgen::FDT_ERR_NOSPACE => FdtError::NoSpace,
+            libfdt_bindgen::FDT_ERR_BADOFFSET => FdtError::BadOffset,
+            libfdt_bindgen::FDT_ERR_BADPATH => FdtError::BadPath,
+            libfdt_bindgen::FDT_ERR_BADPHANDLE => FdtError::BadPhandle,
+            libfdt_bindgen::FDT_ERR_BADSTATE => FdtError::BadState,
+            libfdt_bindgen::FDT_ERR_TRUNCATED => FdtError::Truncated,
+            libfdt_bindgen::FDT_ERR_BADMAGIC => FdtError::BadMagic,
+            libfdt_bindgen::FDT_ERR_BADVERSION => FdtError::BadVersion,
+            libfdt_bindgen::FDT_ERR_BADSTRUCTURE => FdtError::BadStructure,
+            libfdt_bindgen::FDT_ERR_BADLAYOUT => FdtError::BadLayout,
+            libfdt_bindgen::FDT_ERR_INTERNAL => FdtError::Internal,
+            libfdt_bindgen::FDT_ERR_BADNCELLS => FdtError::BadNCells,
+            libfdt_bindgen::FDT_ERR_BADVALUE => FdtError::BadValue,
+            libfdt_bindgen::FDT_ERR_BADOVERLAY => FdtError::BadOverlay,
+            libfdt_bindgen::FDT_ERR_NOPHANDLES => FdtError::NoPhandles,
+            libfdt_bindgen::FDT_ERR_BADFLAGS => FdtError::BadFlags,
+            libfdt_bindgen::FDT_ERR_ALIGNMENT => FdtError::Alignment,
+            _ => FdtError::Unknown(val),
+        })
+    }
+}
+
+fn fdt_err_expect_zero(val: c_int) -> Result<()> {
+    match fdt_err(val)? {
+        0 => Ok(()),
+        _ => Err(FdtError::Unknown(val)),
+    }
+}
+
+fn fdt_err_or_option(val: c_int) -> Result<Option<c_int>> {
+    match fdt_err(val) {
+        Ok(val) => Ok(Some(val)),
+        Err(FdtError::NotFound) => Ok(None),
+        Err(e) => Err(e),
+    }
+}
+
+/// Value of a #address-cells property.
+#[derive(Copy, Clone, Debug)]
+enum AddrCells {
+    Single = 1,
+    Double = 2,
+}
+
+impl TryFrom<c_int> for AddrCells {
+    type Error = FdtError;
+
+    fn try_from(res: c_int) -> Result<Self> {
+        match fdt_err(res)? {
+            x if x == Self::Single as c_int => Ok(Self::Single),
+            x if x == Self::Double as c_int => Ok(Self::Double),
+            _ => Err(FdtError::BadNCells),
+        }
+    }
+}
+
+/// Value of a #size-cells property.
+#[derive(Copy, Clone, Debug)]
+enum SizeCells {
+    None = 0,
+    Single = 1,
+    Double = 2,
+}
+
+impl TryFrom<c_int> for SizeCells {
+    type Error = FdtError;
+
+    fn try_from(res: c_int) -> Result<Self> {
+        match fdt_err(res)? {
+            x if x == Self::None as c_int => Ok(Self::None),
+            x if x == Self::Single as c_int => Ok(Self::Single),
+            x if x == Self::Double as c_int => Ok(Self::Double),
+            _ => Err(FdtError::BadNCells),
+        }
+    }
+}
+
+/// Iterator over cells of a DT property.
+#[derive(Debug)]
+pub struct CellIterator<'a> {
+    chunks: slice::ChunksExact<'a, u8>,
+}
+
+impl<'a> CellIterator<'a> {
+    fn new(bytes: &'a [u8]) -> Self {
+        const CHUNK_SIZE: usize = mem::size_of::<<CellIterator as Iterator>::Item>();
+
+        Self { chunks: bytes.chunks_exact(CHUNK_SIZE) }
+    }
+}
+
+impl<'a> Iterator for CellIterator<'a> {
+    type Item = u32;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        Some(Self::Item::from_be_bytes(self.chunks.next()?.try_into().ok()?))
+    }
+}
+
+/// Iterator over a 'reg' property of a DT node.
+#[derive(Debug)]
+pub struct RegIterator<'a> {
+    cells: CellIterator<'a>,
+    addr_cells: AddrCells,
+    size_cells: SizeCells,
+}
+
+/// Represents a contiguous region within the address space defined by the parent bus.
+/// Commonly means the offsets and lengths of MMIO blocks, but may have a different meaning on some
+/// bus types. Addresses in the address space defined by the root node are CPU real addresses.
+#[derive(Copy, Clone, Debug)]
+pub struct Reg<T> {
+    /// Base address of the region.
+    pub addr: T,
+    /// Size of the region (optional).
+    pub size: Option<T>,
+}
+
+impl<'a> RegIterator<'a> {
+    fn new(cells: CellIterator<'a>, addr_cells: AddrCells, size_cells: SizeCells) -> Self {
+        Self { cells, addr_cells, size_cells }
+    }
+}
+
+impl<'a> Iterator for RegIterator<'a> {
+    type Item = Reg<u64>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let make_double = |a, b| (u64::from(a) << 32) | u64::from(b);
+
+        let addr = match self.addr_cells {
+            AddrCells::Single => self.cells.next()?.into(),
+            AddrCells::Double => make_double(self.cells.next()?, self.cells.next()?),
+        };
+        // If the parent node specifies a value of 0 for #size-cells, 'size' shall be omitted.
+        let size = match self.size_cells {
+            SizeCells::None => None,
+            SizeCells::Single => Some(self.cells.next()?.into()),
+            SizeCells::Double => Some(make_double(self.cells.next()?, self.cells.next()?)),
+        };
+
+        Some(Self::Item { addr, size })
+    }
+}
+
+/// Iterator over the address ranges defined by the /memory/ node.
+#[derive(Debug)]
+pub struct MemRegIterator<'a> {
+    reg: RegIterator<'a>,
+}
+
+impl<'a> MemRegIterator<'a> {
+    fn new(reg: RegIterator<'a>) -> Self {
+        Self { reg }
+    }
+}
+
+impl<'a> Iterator for MemRegIterator<'a> {
+    type Item = Range<usize>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let next = self.reg.next()?;
+        let addr = usize::try_from(next.addr).ok()?;
+        let size = usize::try_from(next.size?).ok()?;
+
+        Some(addr..addr.checked_add(size)?)
+    }
+}
+
+/// DT node.
+#[derive(Clone, Copy)]
+pub struct FdtNode<'a> {
+    fdt: &'a Fdt,
+    offset: c_int,
+}
+
+impl<'a> FdtNode<'a> {
+    /// Find parent node.
+    pub fn parent(&self) -> Result<Self> {
+        // SAFETY - Accesses (read-only) are constrained to the DT totalsize.
+        let ret = unsafe { libfdt_bindgen::fdt_parent_offset(self.fdt.as_ptr(), self.offset) };
+
+        Ok(Self { fdt: self.fdt, offset: fdt_err(ret)? })
+    }
+
+    /// Retrieve the standard (deprecated) device_type <string> property.
+    pub fn device_type(&self) -> Result<Option<&CStr>> {
+        self.getprop_str(CStr::from_bytes_with_nul(b"device_type\0").unwrap())
+    }
+
+    /// Retrieve the standard reg <prop-encoded-array> property.
+    pub fn reg(&self) -> Result<Option<RegIterator<'a>>> {
+        let reg = CStr::from_bytes_with_nul(b"reg\0").unwrap();
+
+        if let Some(cells) = self.getprop_cells(reg)? {
+            let parent = self.parent()?;
+
+            let addr_cells = parent.address_cells()?;
+            let size_cells = parent.size_cells()?;
+
+            Ok(Some(RegIterator::new(cells, addr_cells, size_cells)))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Retrieve the value of a given <string> property.
+    pub fn getprop_str(&self, name: &CStr) -> Result<Option<&CStr>> {
+        let value = if let Some(bytes) = self.getprop(name)? {
+            Some(CStr::from_bytes_with_nul(bytes).map_err(|_| FdtError::BadValue)?)
+        } else {
+            None
+        };
+        Ok(value)
+    }
+
+    /// Retrieve the value of a given property as an array of cells.
+    pub fn getprop_cells(&self, name: &CStr) -> Result<Option<CellIterator<'a>>> {
+        if let Some(cells) = self.getprop(name)? {
+            Ok(Some(CellIterator::new(cells)))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Retrieve the value of a given <u32> property.
+    pub fn getprop_u32(&self, name: &CStr) -> Result<Option<u32>> {
+        let value = if let Some(bytes) = self.getprop(name)? {
+            Some(u32::from_be_bytes(bytes.try_into().map_err(|_| FdtError::BadValue)?))
+        } else {
+            None
+        };
+        Ok(value)
+    }
+
+    /// Retrieve the value of a given <u64> property.
+    pub fn getprop_u64(&self, name: &CStr) -> Result<Option<u64>> {
+        let value = if let Some(bytes) = self.getprop(name)? {
+            Some(u64::from_be_bytes(bytes.try_into().map_err(|_| FdtError::BadValue)?))
+        } else {
+            None
+        };
+        Ok(value)
+    }
+
+    /// Retrieve the value of a given property.
+    pub fn getprop(&self, name: &CStr) -> Result<Option<&'a [u8]>> {
+        let mut len: i32 = 0;
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor) and the
+        // function respects the passed number of characters.
+        let prop = unsafe {
+            libfdt_bindgen::fdt_getprop_namelen(
+                self.fdt.as_ptr(),
+                self.offset,
+                name.as_ptr(),
+                // *_namelen functions don't include the trailing nul terminator in 'len'.
+                name.to_bytes().len().try_into().map_err(|_| FdtError::BadPath)?,
+                &mut len as *mut i32,
+            )
+        } as *const u8;
+
+        let Some(len) = fdt_err_or_option(len)? else {
+            return Ok(None); // Property was not found.
+        };
+        let len = usize::try_from(len).map_err(|_| FdtError::Internal)?;
+
+        if prop.is_null() {
+            // We expected an error code in len but still received a valid value?!
+            return Err(FdtError::Internal);
+        }
+
+        let offset =
+            (prop as usize).checked_sub(self.fdt.as_ptr() as usize).ok_or(FdtError::Internal)?;
+
+        Ok(Some(self.fdt.buffer.get(offset..(offset + len)).ok_or(FdtError::Internal)?))
+    }
+
+    /// Get reference to the containing device tree.
+    pub fn fdt(&self) -> &Fdt {
+        self.fdt
+    }
+
+    fn next_compatible(self, compatible: &CStr) -> Result<Option<Self>> {
+        // SAFETY - Accesses (read-only) are constrained to the DT totalsize.
+        let ret = unsafe {
+            libfdt_bindgen::fdt_node_offset_by_compatible(
+                self.fdt.as_ptr(),
+                self.offset,
+                compatible.as_ptr(),
+            )
+        };
+
+        Ok(fdt_err_or_option(ret)?.map(|offset| Self { fdt: self.fdt, offset }))
+    }
+
+    fn address_cells(&self) -> Result<AddrCells> {
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor).
+        unsafe { libfdt_bindgen::fdt_address_cells(self.fdt.as_ptr(), self.offset) }
+            .try_into()
+            .map_err(|_| FdtError::Internal)
+    }
+
+    fn size_cells(&self) -> Result<SizeCells> {
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor).
+        unsafe { libfdt_bindgen::fdt_size_cells(self.fdt.as_ptr(), self.offset) }
+            .try_into()
+            .map_err(|_| FdtError::Internal)
+    }
+}
+
+/// Mutable FDT node.
+pub struct FdtNodeMut<'a> {
+    fdt: &'a mut Fdt,
+    offset: c_int,
+}
+
+impl<'a> FdtNodeMut<'a> {
+    /// Append a property name-value (possibly empty) pair to the given node.
+    pub fn appendprop<T: AsRef<[u8]>>(&mut self, name: &CStr, value: &T) -> Result<()> {
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor).
+        let ret = unsafe {
+            libfdt_bindgen::fdt_appendprop(
+                self.fdt.as_mut_ptr(),
+                self.offset,
+                name.as_ptr(),
+                value.as_ref().as_ptr().cast::<c_void>(),
+                value.as_ref().len().try_into().map_err(|_| FdtError::BadValue)?,
+            )
+        };
+
+        fdt_err_expect_zero(ret)
+    }
+
+    /// Append a (address, size) pair property to the given node.
+    pub fn appendprop_addrrange(&mut self, name: &CStr, addr: u64, size: u64) -> Result<()> {
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor).
+        let ret = unsafe {
+            libfdt_bindgen::fdt_appendprop_addrrange(
+                self.fdt.as_mut_ptr(),
+                self.parent()?.offset,
+                self.offset,
+                name.as_ptr(),
+                addr,
+                size,
+            )
+        };
+
+        fdt_err_expect_zero(ret)
+    }
+
+    /// Get reference to the containing device tree.
+    pub fn fdt(&mut self) -> &mut Fdt {
+        self.fdt
+    }
+
+    /// Add a new subnode to the given node and return it as a FdtNodeMut on success.
+    pub fn add_subnode(&'a mut self, name: &CStr) -> Result<Self> {
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor).
+        let ret = unsafe {
+            libfdt_bindgen::fdt_add_subnode(self.fdt.as_mut_ptr(), self.offset, name.as_ptr())
+        };
+
+        Ok(Self { fdt: self.fdt, offset: fdt_err(ret)? })
+    }
+
+    fn parent(&'a self) -> Result<FdtNode<'a>> {
+        // SAFETY - Accesses (read-only) are constrained to the DT totalsize.
+        let ret = unsafe { libfdt_bindgen::fdt_parent_offset(self.fdt.as_ptr(), self.offset) };
+
+        Ok(FdtNode { fdt: &*self.fdt, offset: fdt_err(ret)? })
+    }
+}
+
+/// Iterator over nodes sharing a same compatible string.
+pub struct CompatibleIterator<'a> {
+    node: FdtNode<'a>,
+    compatible: &'a CStr,
+}
+
+impl<'a> CompatibleIterator<'a> {
+    fn new(fdt: &'a Fdt, compatible: &'a CStr) -> Result<Self> {
+        let node = fdt.root()?;
+        Ok(Self { node, compatible })
+    }
+}
+
+impl<'a> Iterator for CompatibleIterator<'a> {
+    type Item = FdtNode<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let next = self.node.next_compatible(self.compatible).ok()?;
+
+        if let Some(node) = next {
+            self.node = node;
+        }
+
+        next
+    }
+}
+
+/// Wrapper around low-level libfdt functions.
+#[repr(transparent)]
+pub struct Fdt {
+    buffer: [u8],
+}
+
+impl Fdt {
+    /// Wraps a slice containing a Flattened Device Tree.
+    ///
+    /// Fails if the FDT does not pass validation.
+    pub fn from_slice(fdt: &[u8]) -> Result<&Self> {
+        // SAFETY - The FDT will be validated before it is returned.
+        let fdt = unsafe { Self::unchecked_from_slice(fdt) };
+        fdt.check_full()?;
+        Ok(fdt)
+    }
+
+    /// Wraps a mutable slice containing a Flattened Device Tree.
+    ///
+    /// Fails if the FDT does not pass validation.
+    pub fn from_mut_slice(fdt: &mut [u8]) -> Result<&mut Self> {
+        // SAFETY - The FDT will be validated before it is returned.
+        let fdt = unsafe { Self::unchecked_from_mut_slice(fdt) };
+        fdt.check_full()?;
+        Ok(fdt)
+    }
+
+    /// Wraps a slice containing a Flattened Device Tree.
+    ///
+    /// # Safety
+    ///
+    /// The returned FDT might be invalid, only use on slices containing a valid DT.
+    pub unsafe fn unchecked_from_slice(fdt: &[u8]) -> &Self {
+        mem::transmute::<&[u8], &Self>(fdt)
+    }
+
+    /// Wraps a mutable slice containing a Flattened Device Tree.
+    ///
+    /// # Safety
+    ///
+    /// The returned FDT might be invalid, only use on slices containing a valid DT.
+    pub unsafe fn unchecked_from_mut_slice(fdt: &mut [u8]) -> &mut Self {
+        mem::transmute::<&mut [u8], &mut Self>(fdt)
+    }
+
+    /// Make the whole slice containing the DT available to libfdt.
+    pub fn unpack(&mut self) -> Result<()> {
+        // SAFETY - "Opens" the DT in-place (supported use-case) by updating its header and
+        // internal structures to make use of the whole self.fdt slice but performs no accesses
+        // outside of it and leaves the DT in a state that will be detected by other functions.
+        let ret = unsafe {
+            libfdt_bindgen::fdt_open_into(
+                self.as_ptr(),
+                self.as_mut_ptr(),
+                self.capacity().try_into().map_err(|_| FdtError::Internal)?,
+            )
+        };
+        fdt_err_expect_zero(ret)
+    }
+
+    /// Pack the DT to take a minimum amount of memory.
+    ///
+    /// Doesn't shrink the underlying memory slice.
+    pub fn pack(&mut self) -> Result<()> {
+        // SAFETY - "Closes" the DT in-place by updating its header and relocating its structs.
+        let ret = unsafe { libfdt_bindgen::fdt_pack(self.as_mut_ptr()) };
+        fdt_err_expect_zero(ret)
+    }
+
+    /// Return an iterator of memory banks specified the "/memory" node.
+    ///
+    /// NOTE: This does not support individual "/memory@XXXX" banks.
+    pub fn memory(&self) -> Result<Option<MemRegIterator>> {
+        let memory = CStr::from_bytes_with_nul(b"/memory\0").unwrap();
+        let device_type = CStr::from_bytes_with_nul(b"memory\0").unwrap();
+
+        if let Some(node) = self.node(memory)? {
+            if node.device_type()? != Some(device_type) {
+                return Err(FdtError::BadValue);
+            }
+            let reg = node.reg()?.ok_or(FdtError::BadValue)?;
+
+            Ok(Some(MemRegIterator::new(reg)))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Retrieve the standard /chosen node.
+    pub fn chosen(&self) -> Result<Option<FdtNode>> {
+        self.node(CStr::from_bytes_with_nul(b"/chosen\0").unwrap())
+    }
+
+    /// Get the root node of the tree.
+    pub fn root(&self) -> Result<FdtNode> {
+        self.node(CStr::from_bytes_with_nul(b"/\0").unwrap())?.ok_or(FdtError::Internal)
+    }
+
+    /// Find a tree node by its full path.
+    pub fn node(&self, path: &CStr) -> Result<Option<FdtNode>> {
+        Ok(self.path_offset(path)?.map(|offset| FdtNode { fdt: self, offset }))
+    }
+
+    /// Iterate over nodes with a given compatible string.
+    pub fn compatible_nodes<'a>(&'a self, compatible: &'a CStr) -> Result<CompatibleIterator<'a>> {
+        CompatibleIterator::new(self, compatible)
+    }
+
+    /// Get the mutable root node of the tree.
+    pub fn root_mut(&mut self) -> Result<FdtNodeMut> {
+        self.node_mut(CStr::from_bytes_with_nul(b"/\0").unwrap())?.ok_or(FdtError::Internal)
+    }
+
+    /// Find a mutable tree node by its full path.
+    pub fn node_mut(&mut self, path: &CStr) -> Result<Option<FdtNodeMut>> {
+        Ok(self.path_offset(path)?.map(|offset| FdtNodeMut { fdt: self, offset }))
+    }
+
+    fn path_offset(&self, path: &CStr) -> Result<Option<c_int>> {
+        let len = path.to_bytes().len().try_into().map_err(|_| FdtError::BadPath)?;
+        // SAFETY - Accesses are constrained to the DT totalsize (validated by ctor) and the
+        // function respects the passed number of characters.
+        let ret = unsafe {
+            // *_namelen functions don't include the trailing nul terminator in 'len'.
+            libfdt_bindgen::fdt_path_offset_namelen(self.as_ptr(), path.as_ptr(), len)
+        };
+
+        fdt_err_or_option(ret)
+    }
+
+    fn check_full(&self) -> Result<()> {
+        let len = self.buffer.len();
+        // SAFETY - Only performs read accesses within the limits of the slice. If successful, this
+        // call guarantees to other unsafe calls that the header contains a valid totalsize (w.r.t.
+        // 'len' i.e. the self.fdt slice) that those C functions can use to perform bounds
+        // checking. The library doesn't maintain an internal state (such as pointers) between
+        // calls as it expects the client code to keep track of the objects (DT, nodes, ...).
+        let ret = unsafe { libfdt_bindgen::fdt_check_full(self.as_ptr(), len) };
+        fdt_err_expect_zero(ret)
+    }
+
+    fn as_ptr(&self) -> *const c_void {
+        self as *const _ as *const c_void
+    }
+
+    fn as_mut_ptr(&mut self) -> *mut c_void {
+        self as *mut _ as *mut c_void
+    }
+
+    fn capacity(&self) -> usize {
+        self.buffer.len()
+    }
+}
diff --git a/libs/vbmeta/Android.bp b/libs/vbmeta/Android.bp
index c5078c2..a487097 100644
--- a/libs/vbmeta/Android.bp
+++ b/libs/vbmeta/Android.bp
@@ -27,7 +27,11 @@
         "libanyhow",
         "libtempfile",
     ],
-    data: ["tests/data/*"],
+    data: [
+        ":avb_testkey_rsa2048",
+        ":avb_testkey_rsa4096",
+        ":avb_testkey_rsa8192",
+    ],
     required: ["avbtool"],
     test_suites: ["general-tests"],
     test_options: {
diff --git a/libs/vbmeta/src/lib.rs b/libs/vbmeta/src/lib.rs
index 887844c..8e81ea4 100644
--- a/libs/vbmeta/src/lib.rs
+++ b/libs/vbmeta/src/lib.rs
@@ -238,7 +238,7 @@
         Ok(())
     }
 
-    fn test_signed_image(algorithm: &str, key: &str) -> Result<()> {
+    fn signed_image_has_valid_vbmeta(algorithm: &str, key: &str) -> Result<()> {
         let test_dir = TempDir::new().unwrap();
         let test_file = test_dir.path().join("test.img");
         let mut cmd = Command::new("./avbtool");
@@ -289,16 +289,16 @@
 
     #[test]
     fn test_rsa2048_signed_image() -> Result<()> {
-        test_signed_image("SHA256_RSA2048", "tests/data/testkey_rsa2048.pem")
+        signed_image_has_valid_vbmeta("SHA256_RSA2048", "data/testkey_rsa2048.pem")
     }
 
     #[test]
     fn test_rsa4096_signed_image() -> Result<()> {
-        test_signed_image("SHA256_RSA4096", "tests/data/testkey_rsa4096.pem")
+        signed_image_has_valid_vbmeta("SHA256_RSA4096", "data/testkey_rsa4096.pem")
     }
 
     #[test]
     fn test_rsa8192_signed_image() -> Result<()> {
-        test_signed_image("SHA256_RSA8192", "tests/data/testkey_rsa8192.pem")
+        signed_image_has_valid_vbmeta("SHA256_RSA8192", "data/testkey_rsa8192.pem")
     }
 }
diff --git a/libs/vbmeta/tests/data/testkey_rsa2048.pem b/libs/vbmeta/tests/data/testkey_rsa2048.pem
deleted file mode 100644
index 867dcff..0000000
--- a/libs/vbmeta/tests/data/testkey_rsa2048.pem
+++ /dev/null
@@ -1,27 +0,0 @@
------BEGIN RSA PRIVATE KEY-----
-MIIEowIBAAKCAQEAxlVR3TIkouAOvH79vaJTgFhpfvVKQIeVkFRZPVXK/zY0Gvrh
-4JAqGjJoW/PfrQv5sdD36qtHH3a+G5hLZ6Ni+t/mtfjucxZfuLGC3kmJ1T3XqEKZ
-gXXI2IR7vVSoImREvDQGEDyJwtHzLANlkbGg0cghVhWZSCAndO8BenalC2v94/rt
-DfkPekH6dgU3Sf40T0sBSeSY94mOzTaqOR2pfV1rWlLRdWmo33zeHBv52Rlbt0dM
-uXAureXWiHztkm5GCBC1dgM+CaxNtizNEgC91KcD0xuRCCM2WxH+r1lpszyIJDct
-YbrFmVEYl/kjQpafhy7Nsk1fqSTyRdriZSYmTQIDAQABAoIBAQC+kJgaCuX8wYAn
-SXWQ0fmdZlXnMNRpcF0a0pD0SAzGb1RdYBXMaXiqtyhiwc53PPxsCDdNecjayIMd
-jJVXPTwLhTruOgMS/bp3gcgWwV34UHV4LJXGOGAE+jbS0hbDBMiudOYmj6RmVshp
-z9G1zZCSQNMXHaWsEYkX59XpzzoB384nRul2QgEtwzUNR9XlpzgtJBLk3SACkvsN
-mQ/DW8IWHXLg8vLn1LzVJ2e3B16H4MoE2TCHxqfMgr03IDRRJogkenQuQsFhevYT
-o/mJyHSWavVgzMHG9I5m+eepF4Wyhj1Y4WyKAuMI+9dHAX/h7Lt8XFCQCh5DbkVG
-zGr34sWBAoGBAOs7n7YZqNaaguovfIdRRsxxZr1yJAyDsr6w3yGImDZYju4c4WY9
-5esO2kP3FA4p0c7FhQF5oOb1rBuHEPp36cpL4aGeK87caqTfq63WZAujoTZpr9Lp
-BRbkL7w/xG7jpQ/clpA8sHzHGQs/nelxoOtC7E118FiRgvD/jdhlMyL9AoGBANfX
-vyoN1pplfT2xR8QOjSZ+Q35S/+SAtMuBnHx3l0qH2bbBjcvM1MNDWjnRDyaYhiRu
-i+KA7tqfib09+XpB3g5D6Ov7ls/Ldx0S/VcmVWtia2HK8y8iLGtokoBZKQ5AaFX2
-iQU8+tC4h69GnJYQKqNwgCUzh8+gHX5Y46oDiTmRAoGAYpOx8lX+czB8/Da6MNrW
-mIZNT8atZLEsDs2ANEVRxDSIcTCZJId7+m1W+nRoaycLTWNowZ1+2ErLvR10+AGY
-b7Ys79Wg9idYaY9yGn9lnZsMzAiuLeyIvXcSqgjvAKlVWrhOQFOughvNWvFl85Yy
-oWSCMlPiTLtt7CCsCKsgKuECgYBgdIp6GZsIfkgclKe0hqgvRoeU4TR3gcjJlM9A
-lBTo+pKhaBectplx9RxR8AnsPobbqwcaHnIfAuKDzjk5mEvKZjClnFXF4HAHbyAF
-nRzZEy9XkWFhc80T5rRpZO7C7qdxmu2aiKixM3V3L3/0U58qULEDbubHMw9bEhAT
-PudI8QKBgHEEiMm/hr9T41hbQi/LYanWnlFw1ue+osKuF8bXQuxnnHNuFT/c+9/A
-vWhgqG6bOEHu+p/IPrYm4tBMYlwsyh4nXCyGgDJLbLIfzKwKAWCtH9LwnyDVhOow
-GH9shdR+sW3Ew97xef02KAH4VlNANEmBV4sQNqWWvsYrcFm2rOdL
------END RSA PRIVATE KEY-----
diff --git a/libs/vbmeta/tests/data/testkey_rsa4096.pem b/libs/vbmeta/tests/data/testkey_rsa4096.pem
deleted file mode 100644
index 26db5c3..0000000
--- a/libs/vbmeta/tests/data/testkey_rsa4096.pem
+++ /dev/null
@@ -1,51 +0,0 @@
------BEGIN RSA PRIVATE KEY-----
-MIIJKQIBAAKCAgEA2ASv49OEbH4NiT3CjNMSVeliyfEPXswWcqtEfCxlSpS1FisA
-uwbvEwdTTPlkuSh6G4SYiNhnpCP5p0vcSg/3OhiuVKgV/rCtrDXaO60nvK/o0y83
-NNZRK2xaJ9eWBq9ruIDK+jC0sYWzTaqqwxY0Grjnx/r5CXerl5PrRK7PILzwgBHb
-IwxHcblt1ntgR4cWVpO3wiqasEwBDDDYk4fw7W6LvjBb9qav3YB8RV6PkZNeRP64
-ggfuecq/MXNiWOPNxLzCER2hSr/+J32h9jWjXsrcVy8+8Mldhmr4r2an7c247aFf
-upuFGtUJrpROO8/LXMl5gPfMpkqoatjTMRH59gJjKhot0RpmGxZBvb33TcBK5SdJ
-X39Y4yct5clmDlI4Fjj7FutTP+b96aJeJVnYeUX/A0wmogBajsJRoRX5e/RcgZsY
-RzXYLQXprQ81dBWjjovMJ9p8XeT6BNMFC7o6sklFL0fHDUE/l4BNP8G1u3Bfpzev
-SCISRS71D4eS4oQB+RIPFBUkzomZ7rnEF3BwFeq+xmwfYrP0LRaH+1YeRauuMuRe
-ke1TZl697a3mEjkNg8noa2wtpe7EWmaujJfXDWxJx/XEkjGLCe4z2qk3tkkY+A5g
-Rcgzke8gVxC+eC2DJtbKYfkv4L8FMFJaEhwAp13MfC7FlYujO/BDLl7dANsCAwEA
-AQKCAgAWoL8P/WsktjuSwb5sY/vKtgzcHH1Ar942GsysuTXPDy686LpF3R8T/jNy
-n7k2UBAia8xSoWCR6BbRuHeV5oA+PLGeOpE7QaSfonB+yc+cy0x3Or3ssfqEsu/q
-toGHp75/8DXS6WE0K04x94u1rdC9b9sPrrGBlWCLGzqM0kbuJfyHXdd3n2SofAUO
-b5QRSgxD+2tHUpEroHqHnWJCaf4J0QegX45yktlfOYNK/PHLDQXV8ly/ejc32M4Y
-Tv7hUtOOJTuq8VCg9OWZm2Zo1QuM9XEJTPCp5l3+o5vzO6yhk2gotDvD32CdA+3k
-tLJRP54M1Sn+IXb1gGKN9rKAtGJbenWIPlNObhQgkbwG89Qd+5rfMXsiPv1Hl1tK
-+tqwjD82/H3/ElaaMnwHCpeoGSp95OblAoBjzjMP2KsbvKSdL8O/rf1c3uOw9+DF
-cth0SA8y3ZzI11gJtb2QMGUrCny5n4sPGGbc3x38NdLhwbkPKZy60OiT4g2kNpdY
-dIitmAML2otttiF4AJM6AraPk8YVzkPLTksoL3azPBya5lIoDI2H3QvTtSvpXkXP
-yKchsDSWYbdqfplqC/X0Djp2/Zd8jpN5I6+1aSmpTmbwx/JTllY1N89FRZLIdxoh
-2k81LPiXhE6uRbjioJUlbnEWIpY2y2N2Clmxpjh0/IcXd1XImQKCAQEA7Zai+yjj
-8xit24aO9Tf3mZBXBjSaDodjC2KS1yCcAIXp6S7aH0wZipyZpQjys3zaBQyMRYFG
-bQqIfVAa6inWyDoofbAJHMu5BVcHFBPZvSS5YhDjc8XZ5dqSCxzIz9opIqAbm+b4
-aEV/3A3Jki5Dy8y/5j21GAK4Y4mqQOYzne7bDGi3Hyu041MGM4qfIcIkS5N1eHW4
-sDZJh6+K5tuxN5TX3nDZSpm9luNH8mLGgKAZ15b1LqXAtM5ycoBY9Hv082suPPom
-O+r0ybdRX6nDSH8+11y2KiP2kdVIUHCGkwlqgrux5YZyjCZPwOvEPhzSoOS+vBiF
-UVXA8idnxNLk1QKCAQEA6MIihDSXx+350fWqhQ/3Qc6gA/t2C15JwJ9+uFWA+gjd
-c/hn5HcmnmBJN4R04nLG/aU9SQur87a4mnC/Mp9JIARjHlZ/WNT4U0sJyPEVRg5U
-Z9VajAucWwi0JyJYCO1EMMy68Jp8qlTriK/L7nbD86JJ5ASxjojiN/0psK/Pk60F
-Rr+shKPi3jRQ1BDjDtAxOfo4ctf/nFbUM4bY0FNPQMP7WesoSKU0NBCRR6d0d2tq
-YflMjIQHx+N74P5jEdSCHTVGQm+dj47pUt3lLPLWc0bX1G/GekwXP4NUsR/70Hsi
-bwxkNnK2TSGzkt2rcOnutP125rJu6WpV7SNrq9rm7wKCAQAfMROcnbWviKHqnDPQ
-hdR/2K9UJTvEhInASOS2UZWpi+s1rez9BuSjigOx4wbaAZ4t44PW7C3uyt84dHfU
-HkIQb3I5bg8ENMrJpK9NN33ykwuzkDwMSwFcZ+Gci97hSubzoMl/IkeiiN1MapL4
-GhLUgsD+3UMVL+Y9SymK8637IgyoCGdiND6/SXsa8SwLJo3VTjqx4eKpX7cvlSBL
-RrRxc50TmwUsAhsd4CDl9YnSATLjVvJBeYlfM2tbFPaYwl1aR8v+PWkfnK0efm60
-fHki33HEnGteBPKuGq4vwVYpn6bYGwQz+f6335/A2DMfZHFSpjVURHPcRcHbCMla
-0cUxAoIBAQC25eYNkO478mo+bBbEXJlkoqLmvjAyGrNFo48F9lpVH6Y0vNuWkXJN
-PUgLUhAu6RYotjGENqG17rz8zt/PPY9Ok2P3sOx8t00y1mIn/hlDZXs55FM0fOMu
-PZaiscAPs7HDzvyOmDah+fzi+ZD8H2M3DS2W+YE0iaeJa2vZJS2t02W0BGXiDI33
-IZDqMyLYvwwPjOnShJydEzXID4xLl0tNjzLxo3GSNA7jYqlmbtV8CXIc7rMSL6WV
-ktIDKKJcnmpn3TcKeX6MEjaSIT82pNOS3fY3PmXuL+CMzfw8+u77Eecq78fHaTiL
-P5JGM93F6mzi19EY0tmInUBMCWtQLcENAoIBAQCg0KaOkb8T36qzPrtgbfou0E2D
-ufdpL1ugmD4edOFKQB5fDFQhLnSEVSJq3KUg4kWsXapQdsBd6kLdxS+K6MQrLBzr
-4tf0c7UCF1AzWk6wXMExZ8mRb2RkGZYQB2DdyhFB3TPmnq9CW8JCq+6kxg/wkU4s
-vM4JXzgcqVoSf42QJl+B9waeWhg0BTWx01lal4ds88HvEKmE0ik5GwiDbr7EvDDw
-E6UbZtQcIoSTIIZDgYqVFfR2DAho3wXJRsOXh433lEJ8X7cCDzrngFbQnlKrpwML
-Xgm0SIUc+Nf5poMM3rfLFK77t/ob4w+5PwRKcoSniyAxrHd6bwykYA8Vuydv
------END RSA PRIVATE KEY-----
diff --git a/libs/vbmeta/tests/data/testkey_rsa8192.pem b/libs/vbmeta/tests/data/testkey_rsa8192.pem
deleted file mode 100644
index a383428..0000000
--- a/libs/vbmeta/tests/data/testkey_rsa8192.pem
+++ /dev/null
@@ -1,99 +0,0 @@
------BEGIN RSA PRIVATE KEY-----
-MIISKgIBAAKCBAEA0D3T+dISsmCHm797wsX0vVfqUWDJ/3mvDYozlCabDhnGLlSE
-pAQbf1Z8Ts+OM4pVRHOJUJL0WebNdmPPGjsyWQz6zZE96lQZL3avCEXqYVQR66V5
-3wdK/ohaMSRnGyEMBrqkVVbF3gCr+/irxD3YK+VowO2WKs/6GrMdqTA8Y5CTF/Je
-ptwsSg5MMjr6UaK4qDcrej3hkgBVGvRV3cj1snK6Br8HuYdFnpGGTS0d7UJlHFgl
-trGHU/CBO923hkHgJaWEjC0giSGjhKKtLzrVcpDV2y/lWQP9T/T4djEAIaHqQ++P
-SdOSR6psIGR6hVgSigt7HCnE7nW711/rfV5Ur9EiVpB040mDImKZcy8//TMnXydN
-1KYTVd/34fdpzMpSw5iblErbwOLXVTUmOztYnpl41feHSv/jPesHstPlfklIF2vo
-GZEohf9scQvcuM7wEBfC/aTA9K39zMmkBbcvSZjLyhmcSZWMPPOZyIcl3zY53QhW
-QC/abmIcBfI1S4+r7mC4i2Jn++oEvuGNVGr2SY2Z0ZZxXGL1HI/08D/3+Tcumrcn
-4YjPK/DMFi0F+e+1x41lipuf+cx/2qRNQX/m02STrLYdM6e0g33KvlnFdi2b752y
-/OIaMwxDaJvunMh6EMDWKM1AHbY/ioAoK7eS26HeJLEDllqO4+SWP37c8lMvSEWy
-1GiErR0HcsOj/QwWGPFseoVroMiA2sUQ0Ic/tgVjCTlXg+12XpUnouIweCi8KcL/
-ad2zJkju9hBhJLBQ/2GnivJi3lFgF4Gd//TSJ6rgWuXFfMKt/9z2Sz35ohEX4yA0
-flqlCeLInFEoevbz+XT9aRfDe65MZ79yw3TfP9CrV74hf1RRzveD4zpi3F+hcY2i
-JWsH7gROZeCm6fAX5Trecd3hOxJOfA4N4rvSSCq6BwCvebT8FY25Z/VF7cQrHYDS
-ij5w6lqhMzXHeUEY90Ga9AK4XzaWwGgezq+R7Zs00YSKqFv9qYNKdR7tz3cjijWf
-9q/3R1uh6EQKTMZKo4SEClJiGyjOBvmPK09jMFZTJv00hDxagDPZBl7XpLDJ5/Ln
-1uppvLCNWWY1zeJfaElMyq3/PqKZLidF9rVoA1SIwk2lpdUvPote2oFiwCZoXlwZ
-J2ncjmXgQNs76/8unDJA0rj4JPqccw4M5GxQ7okbgm3F4rmzriCuv8BeMSCkr2ry
-0mY3UhpohX4wCMq0G4x5sEUAz9FVVPZKjxnYBmLDzrJAR+4+G7gZsct01XDJYgDd
-JVYInFP22/cIre8VrFWYtHbgOFdNqUiVq58de6PdZG/E+uaWmEThSlRrgEjTxupi
-OXfgdKW/20j1qAtjOlqFwsY094Q5rqULQ6wPxQIDAQABAoIEAQChmkmlhrRBv42d
-fYUiyxK52b8ath0saJdDz6tlXmxYDgJxM9/XlORt9oTzeDknoEO5olu+rrx4BBgQ
-tzYiaiwRVXRREVTWQ7tjzRvaNL/GFkLt93XTccpuKwyrNE/bitLVagRbwcI+HZFa
-MknCOihHMHoRto8h3FKAY94xzSAgODMek1WG8jhgpCXXmVNnBPt+d4oDDIDAGAfz
-qgf03J5nhIb+80KgZOzPOKnbvJaL6EmlLHbgB3c42dzAw7hHtVmofYGWcvLb2MIY
-DVKO435/sQx1U/8NDH6JjVdACZjLgObXH9K3/Tt46DWPEcrPLmD8xhoc6gFM+Qr0
-AhkzKoBYDNk0CljbhdIBXjktXU6wRQFZ45uP2e4JZ4zrzGBLr/t4lTavZ0SQtLld
-A6kOsGh+dCWFDtnshxYnl/xad/yR+3a5zmDJbo/fJTBXrlf1B4rfQkFtK20etOPQ
-B++FC/rjh3Mm/Kb/p9Gz/2upZdArH97ZvD2LBFfj77lFmAhqAi3wCRlN+ekuYxaZ
-t1pBV9yXig8Dyldg1d7X8pOn2kyrF3rQUDDf4pa7x9vpnbkUlEUifoV9gnYsmdni
-qDzYBtTv2g6MKqwQySXaIUW0YOBPbOellWEwxJqGYQ7y4IfVHfM0iyHnehk2tZcr
-+XazLnwGe+Bz4vcguFhJXLyIu//lAOhZtbk6r1QJEUuxaOOQX3wzyceE6nkDsgmr
-P5dj3Zpd7fS2VV2vyGHIFnBJ88LRxreVvgr6Q28UT27SB82zMb7mRZTVE2zeuubT
-5D2D1XbZ0wBo6WiK6eRRrDQ2Haeetkj/uoRy6PWXwnAaTmmIrrXwLqaoJh/U1e+D
-tfsDLWd6IxLjfXvGglrHsrtAz0oprpixUTeVhgTrGk9IQRd5rvxuGUYhFujVaYI6
-+QUf+33AFdtncb8y9C9jZmgx8AKbJk+e73SLhB5JVos+WteU7b8d/Mim5mALjnO6
-Z1n/uimsT79sSDqy3XSymtKWXo/22UlrvGCpoEuELPMb6dSFWR7vwrsvhFngY4/K
-UnitnvxboEflQnaIQ4IfRLRzZsX+sC5Esqw9U5tHt4oI+91Dv3KbdbcERgV73K6B
-ZQgC4lkAQquFXiZ5AICkxjiMyZwTtU9KJ7xv17Xu6oywF/3AtbVGETW1D+3maHsD
-y3DASWojyqZdLj+WGzKQRa+swgCDAYKeek2fIAXFSdF63zxJ2RxOJ4GijSaoh+mr
-4HVvcpDaTj+A8T1+QdByM4s98gu4GD7kVtVQGBZdWjutyHvh0hWv1gtVmbhQ/413
-gDMFFDzHIjLTYGYes4hHL22169jVR9sZ1eQxwvTIg3N4pD5cFm0rRuZZTS+oJToF
-G27aBFihAoICAQDyVB62ZDnbxQthk+zITKIzRUrJbLoXrUcANcSHfaN7inF87Ova
-ze7ejT9DNSEhbtfZFJ1G6diOYoSw+2MzFXv0gEkLKY0dETydKgHEu6nVq5eivMgv
-D4hc9YkJMHDSlmv2FDkpL3AXCAmnW9rKp+ddttBZECnmlPEpHLoj6xgBw3pNa1Xs
-IcLVfdugH86Hexj6o0oKgYfcqrX8UUHtUI2/XQqgFrIj8ksjf1fFVWJRJFWmBXqp
-nMEsYarzATeM1kQ/kDeT1ZUpoGPQt02/XqXT4B5A3ATiEtpM2u+l48xtogWWg2Ry
-G9l938StAmhUiW1m7GnKE6EIFvQY85WvbzxOR0JYVUSr7MrasF6nnQlhYxFuIJoJ
-2h/KJQao5GCTvG4+GtbJJm4c2nyZgwyhizMsdgsdcls79aXiMkrZZkamLVUZWOtE
-3pA/oBuz2qnO9HwjbH1HGOccq0TXfmpFScEV3CQGYJdno6Fy7cbmupaL4U9agQ4e
-w+ygL18nq5HV++LStFnVrgs5YijjskfRdE9GUMVDh5pCsd9Y23Fymaad4O/2SRCC
-YkSsyH5OvyDOLpoyUJ6g6Q+45Hqm/3lG4YjNpzFUiMcnp7+3xU35qC0LK8xEfeei
-Ms1mTVEiHNIp6xH/TqRdX73WD7+YuKZSLIfRG7dgrirU6w+mhhvxD51uHQKCAgEA
-2/1mBCR5qm3/0Lt++RQbeyE3tiw40UeyQqucG/+VvY77sSLkI/Lx8iwRlywXcLBn
-+A4TvgukmAdWzCs8ndgKNxPA+gfohvBsMOGN9KOB1Ug5vvg2J2kiI64vwYCwzhdZ
-NTUUmL+GMFHUqSsWYg6i7iBFcZmznr4W2T3bBxyTMZki7JStB86e35KXrzc2/W/b
-+/p5U2HCSazDHI5mMyuClHc6GmUSVJ7f7LHjL94jviNqobp0Vj603tScHISmNrZw
-TBavkvZGYXsoWKvqavk7jBB9QzaBL+unaFRslg5jTaiKnISj44Us1fjFKu84xifL
-nJaEzjDPt7PBxko7LPgEY7wF39nM9VpoetI7bwR6NwDLSX8UU97MGd+HY+MO1Wi1
-pd2Lapwrx/EK7Oxz335VRK4Je0aZna4j2TyQdMJac9fsGPXv4ZsLfDLj/wD6l1j+
-lLLbBv3ImdSj32LBbhsgF4iCGeXO8HpPO+Q/h9XVsnY52Um2XdNMn03PCGm6ZvtM
-7DXiS+lPF90HjolJVHZTBNtdVRrLr53zLuWEfqT4FeKrDaxdtiXkxLjrB+5/VYu7
-ntyk01ZQ63VNfEwS1irmKl9+qZkTHk3HHV9jNV5RzWViwmJI7Wpr1YzBwmcKCB1O
-oGUADDs8QpnkCz0xkMVtYwHj9qKZlqfbHzrFDUUcF8kCggIAdYvUcgjf//ju8mA8
-5VQ3AcPE6TvycPW+kR2DvW12VcDsF/sc1UA7dHzziPhGn98SmNxlBjb8suSbFPZ8
-QhVT0WBBDkcTilwIGPx9ax7U3S6lGW2VdS6FqQH5fRmgQKZyrCVXLOEz8BgYBrSJ
-xu/3TQAWxH0QtibdbGHg8Pdi58gYlWFRhn9B8Slh1aRYHGPb1AhNLBd0/ddY+5G2
-9xSyDXdmZg1cUA+B3zAwNSqbzFxhp2zU+V1uXsbpk4KtnYV6CZM9QlrCRjTk9iNU
-dVXF/qaiRjfzrm4SsmEpCkEbsrp7F22Y1bkooORglMOsNAWNqfVXw4wN+syXj1ro
-6vZ8PERYrFyAOR1dsQMIhymnmTPjCpaJ4emKrhWTy20sY71thHakZWJc22YoNpbZ
-E6tgIVsJPTlxg/4+fyCCKj5wWr92nhsB1KBZPGO/zFhvMlJpvQ0tH8W2pbN2a0mI
-5x9FqALm/qjwCHfZItSwPM+ZozSht3cOkGHdcD5KXAXfcfsDJc4SHZKVIzq4NusN
-504R/jvD1GP8sglyG7omp75ckgzAmakLdxOP2HhQvIX9tcXpSirNJ6Sl2bwKuuMF
-wxo3r/o/9Y97e4LlfpEYp9eqMdcG+NpR993IwK0UhAWS9H5wdnWBSUHd5e4xtDUt
-iILNRuO46g7R/AIhz1cSSraWWQkCggIBAMhhPP5C9yt9PIm1b0eTwCBctnFSQIKo
-KsA9rll2ab+bMLk9jc8M6MLszy0CtWso09sHf4YY9tifvrkEHRethEh8zscwUuYu
-sm2n1fTixk0ul6LSVgl54uXbMJayENn4PIKRkew8cA8tSma43497w37hmD+MgCb1
-ALzqcco9hfmkgkI6fo1g8Ce3UEECKy2YKSmREdgYcK9JFQO61W6AkFWJcDxAmfzI
-JjFkKwsb7TSw79zWiEdSoM9jm7sCPKATd6Bm/ZAAkUUTuEFkfobn9Ax1rJN/Xxb2
-MKuAUtQv0NYY0gEVdG62jItuKLId6nncH8PG+rsRjPLIYpWqYdJpKx5pUnR+4AkQ
-S6CsRASwcF4PdBvDDBIFG6XpjFo4pPdQhDzL2sTF8b8SWSBLlJQbb7G6UNqgCSau
-SusCFpazvU5NfDmUMuctob2EYVaSXq9jGaj6bTUmDwXHwWilfIk9XfLxnYfXYrJ6
-xhdIpXGmHhuLQtAgK2O1JtLoPc9s9qP8/SkfP7xjjG6xHsP/WvL7QE1pPs9ZM/UI
-C01JNHFi9LKCn8o5mbZjN8jUowi7ffK+76wZUG1L7zM5ytWQOYwo0TQBfc8fpmFw
-+RBRJX2kJyDO27ExczoGOKjwqEDaODIB9+9zcCK0BgSoRibSm4ZBvoxzWWD65Kls
-xdPhZUHcFGW5AoICAQC8iG27aD8aRUt94Oek66gFOJx84QVZehWPqtZjWyVenDuc
-T8dink8oejGjcK2UJuQDa83azv90ocVqE0n0ronYyszt9Ib1jlYC+CK1Ar9TYGFg
-WU5OWEDyCzCpqW/w/aG68U8qhKm0MvkLJR+G6evan9TwEhFEVAm3iWllNXs9x29s
-BucwyMMC23zsimxYlS7dA4DtyvVA+zL1omLpSWHbU/qtuI3HV1NeJzsy+gC4mwPh
-j52tdl669fyWLzHzBRLeq6dVOedjnCo+jlU3dL20DEk9SaW08D1CPuZekV1jVPMw
-JoaDcIRh4KLtQ0BYZ7UJeFUTsx1CS/+UqzqYSPOi57a5kvr0Y8YwRnSB8dHVFttX
-JTv83wTQXHPFSBgfnHNe7lsRTfIQfuIkr2bpiU7h85UQ7LsqcI6YHaC07URcsGFF
-FrLWGh91qzAd1diSHla2RnY3n8PPuMnCkguNhLUrYdmyMol7FfWFa9lwplsuTzBq
-B6yj8iaiE3LL+Q/eulJ7S6QPfAI2bU0UJO23Y4koeoIibEEDMSCQ6KYZ2NClRRRT
-ga5fS1YfkDFEcHUQ1/KIkdYHGBKBjoKGExzi8+CgiSySVSYDZl6wIOhLjH2OZ3ol
-ldPN7iNAHirrxg9v8QO6OQlpLUk5Lhp/1dSlZ6sy3UjFqvax3tw6ZjrL88YP5g==
------END RSA PRIVATE KEY-----
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index c3e2692..af6031a 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -72,7 +72,6 @@
         "debuggerd",
         "linker",
         "linkerconfig",
-        "servicemanager.microdroid",
         "tombstoned.microdroid",
         "tombstone_transmit.microdroid",
         "cgroups.json",
@@ -84,7 +83,7 @@
         "microdroid_manifest",
         "microdroid_plat_sepolicy_and_mapping.sha256",
         "microdroid_property_contexts",
-        "microdroid_service_contexts",
+        "mke2fs",
 
         // TODO(b/195425111) these should be added automatically
         "libcrypto", // used by many (init_second_stage, microdroid_manager, toybox, etc)
@@ -108,6 +107,7 @@
                 "apkdmverity",
                 "authfs",
                 "authfs_service",
+                "encryptedstore",
                 "microdroid_crashdump_kernel",
                 "microdroid_kexec",
                 "microdroid_manager",
diff --git a/microdroid/README.md b/microdroid/README.md
index 2519416..41278a5 100644
--- a/microdroid/README.md
+++ b/microdroid/README.md
@@ -141,7 +141,7 @@
 PATH_TO_YOUR_APP \
 $TEST_ROOT/MyApp.apk.idsig \
 $TEST_ROOT/instance.img \
-assets/VM_CONFIG_FILE
+--config-path assets/VM_CONFIG_FILE
 ```
 
 The last command lets you know the CID assigned to the VM. The console output
diff --git a/microdroid/init.rc b/microdroid/init.rc
index 9c62782..7d04557 100644
--- a/microdroid/init.rc
+++ b/microdroid/init.rc
@@ -26,20 +26,6 @@
     setprop ro.log.file_logger.path /dev/null
 
 on init
-    # Mount binderfs
-    mkdir /dev/binderfs
-    mount binder binder /dev/binderfs stats=global
-    chmod 0755 /dev/binderfs
-
-    symlink /dev/binderfs/binder /dev/binder
-    symlink /dev/binderfs/vndbinder /dev/vndbinder
-
-    chmod 0666 /dev/binderfs/binder
-    chmod 0666 /dev/binderfs/vndbinder
-
-    start servicemanager
-
-on init
     mkdir /mnt/apk 0755 system system
     mkdir /mnt/extra-apk 0755 root root
     # Microdroid_manager starts apkdmverity/zipfuse/apexd
@@ -88,6 +74,12 @@
     # some services can be started.
     trigger late-fs
 
+    # Wait for microdroid_manager to finish setting up sysprops from the payload config.
+    # Some further actions in the boot sequence might depend on the sysprops from the payloag,
+    # e.g. microdroid.config.enable_authfs configures whether to run authfs_service after
+    # /data is mounted.
+    wait_for_prop microdroid_manager.config_done 1
+
     trigger post-fs-data
 
     # Load persist properties and override properties (if enabled) from /data.
@@ -146,6 +138,13 @@
     mkdir /data/local 0751 root root
     mkdir /data/local/tmp 0771 shell shell
 
+on post-fs-data && property:microdroid_manager.authfs.enabled=1
+    start authfs_service
+
+on boot
+    # Mark boot completed. This will notify microdroid_manager to run payload.
+    setprop dev.bootcomplete 1
+
 service tombstone_transmit /system/bin/tombstone_transmit.microdroid -cid 2 -port 2000 -remove_tombstones_after_transmitting
     user system
     group system
diff --git a/microdroid/kernel/README.md b/microdroid/kernel/README.md
index 1f6b2ee..44f828c 100644
--- a/microdroid/kernel/README.md
+++ b/microdroid/kernel/README.md
@@ -24,6 +24,7 @@
 
 For x86\_64,
 ```bash
+tools/bazel clean
 tools/bazel run --config=fast --lto=thin //common-modules/virtual-device:microdroid_x86_64_dist -- --dist_dir=out/dist
 ```
 
@@ -36,7 +37,15 @@
 
 ### Change the kernel configs
 
+For ARM64
+```bash
+tools/bazel run //common-modules/virtual-device:microdroid_aarch64_config -- menuconfig
+```
 
+For x86\_64
+```bash
+tools/bazel run //common-modules/virtual-device:microdroid_x86_64_config -- menuconfig
+```
 
 ## How to update Microdroid kernel prebuilts
 
diff --git a/microdroid/vm_payload/Android.bp b/microdroid/vm_payload/Android.bp
deleted file mode 100644
index 8d78444..0000000
--- a/microdroid/vm_payload/Android.bp
+++ /dev/null
@@ -1,45 +0,0 @@
-package {
-    default_applicable_licenses: ["Android-Apache-2.0"],
-}
-
-rust_ffi_shared {
-    name: "libvm_payload",
-    crate_name: "vm_payload",
-    srcs: ["src/*.rs"],
-    include_dirs: ["include"],
-    prefer_rlib: true,
-    rustlibs: [
-        "android.system.virtualization.payload-rust",
-        "libandroid_logger",
-        "libanyhow",
-        "libbinder_rs",
-        "liblazy_static",
-        "liblog_rust",
-        "librpcbinder_rs",
-    ],
-    apex_available: [
-        "com.android.compos",
-    ],
-    // The sanitize section below fixes the fuzzer build in b/256166339.
-    // TODO(b/250854486): Remove the sanitize section once the bug is fixed.
-    sanitize: {
-        address: false,
-    },
-}
-
-rust_bindgen {
-    name: "libvm_payload_bindgen",
-    wrapper_src: "include/vm_payload.h",
-    crate_name: "vm_payload_bindgen",
-    source_stem: "bindings",
-    apex_available: ["com.android.compos"],
-    visibility: ["//packages/modules/Virtualization/compos"],
-    shared_libs: [
-        "libvm_payload",
-    ],
-}
-
-cc_library_headers {
-    name: "vm_payload_headers",
-    export_include_dirs: ["include"],
-}
diff --git a/microdroid/vm_payload/include/vm_payload.h b/microdroid/vm_payload/include/vm_payload.h
deleted file mode 100644
index 2aeb44e..0000000
--- a/microdroid/vm_payload/include/vm_payload.h
+++ /dev/null
@@ -1,112 +0,0 @@
-/*
- * Copyright (C) 2022 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#pragma once
-
-#include <stdbool.h>
-#include <stddef.h>
-
-#include "vm_main.h"
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-struct AIBinder;
-typedef struct AIBinder AIBinder;
-
-/**
- * Notifies the host that the payload is ready.
- *
- * \return true if the notification succeeds else false.
- */
-bool AVmPayload_notifyPayloadReady(void);
-
-/**
- * Runs a binder RPC server, serving the supplied binder service implementation on the given vsock
- * port.
- *
- * If and when the server is ready for connections (it is listening on the port), `on_ready` is
- * called to allow appropriate action to be taken - e.g. to notify clients that they may now
- * attempt to connect with `AVmPayload_notifyPayloadReady`.
- *
- * The current thread is joined to the binder thread pool to handle incoming messages.
- *
- * \param service the service to bind to the given port.
- * \param port vsock port.
- * \param on_ready the callback to execute once the server is ready for connections. The callback
- *                 will be called at most once.
- * \param param param for the `on_ready` callback.
- *
- * \return true if the server has shutdown normally, false if it failed in some way.
- */
-bool AVmPayload_runVsockRpcServer(AIBinder *service, unsigned int port,
-                                  void (*on_ready)(void *param), void *param);
-
-/**
- * Get a secret that is uniquely bound to this VM instance. The secrets are 32-byte values and the
- * value associated with an identifier will not change over the lifetime of the VM instance.
- *
- * \param identifier identifier of the secret to return.
- * \param identifier_size size of the secret identifier.
- * \param secret pointer to size bytes where the secret is written.
- * \param size number of bytes of the secret to get, up to the secret size.
- *
- * \return true on success and false on failure.
- */
-bool AVmPayload_getVmInstanceSecret(const void *identifier, size_t identifier_size, void *secret,
-                                    size_t size);
-
-/**
- * Get the VM's DICE attestation chain.
- *
- * This function will fail if the use of restricted APIs is not permitted.
- *
- * \param data pointer to size bytes where the chain is written.
- * \param size number of bytes that can be written to data.
- * \param total outputs the total size of the chain if the function succeeds
- *
- * \return true on success and false on failure.
- */
-bool AVmPayload_getDiceAttestationChain(void *data, size_t size, size_t *total);
-
-/**
- * Get the VM's DICE attestation CDI.
- *
- * This function will fail if the use of restricted APIs is not permitted.
- *
- * \param data pointer to size bytes where the CDI is written.
- * \param size number of bytes that can be written to data.
- * \param total outputs the total size of the CDI if the function succeeds
- *
- * \return true on success and false on failure.
- */
-bool AVmPayload_getDiceAttestationCdi(void *data, size_t size, size_t *total);
-
-/**
- * Gets the path to the APK contents. It is a directory, under which are
- * the unzipped contents of the APK containing the payload, all read-only
- * but accessible to the payload.
- *
- * \return the path to the APK contents. The returned string should not be
- * deleted or freed by the application. The string remains valid for the
- * lifetime of the VM.
- */
-const char *AVmPayload_getApkContentsPath(void);
-
-#ifdef __cplusplus
-} // extern "C"
-#endif
diff --git a/microdroid/vm_payload/src/vm_payload_service.rs b/microdroid/vm_payload/src/vm_payload_service.rs
deleted file mode 100644
index b0dd891..0000000
--- a/microdroid/vm_payload/src/vm_payload_service.rs
+++ /dev/null
@@ -1,208 +0,0 @@
-// Copyright 2022, The Android Open Source Project
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//! This module handles the interaction with virtual machine payload service.
-
-use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
-    IVmPayloadService, VM_PAYLOAD_SERVICE_NAME, VM_APK_CONTENTS_PATH};
-use anyhow::{Context, Result};
-use binder::{wait_for_interface, Strong, unstable_api::{AIBinder, new_spibinder}};
-use lazy_static::lazy_static;
-use log::{error, info, Level};
-use rpcbinder::run_vsock_rpc_server;
-use std::ffi::CString;
-use std::os::raw::{c_char, c_void};
-
-lazy_static! {
-    static ref VM_APK_CONTENTS_PATH_C: CString =
-        CString::new(VM_APK_CONTENTS_PATH).expect("CString::new failed");
-}
-
-/// Notifies the host that the payload is ready.
-/// Returns true if the notification succeeds else false.
-#[no_mangle]
-pub extern "C" fn AVmPayload_notifyPayloadReady() -> bool {
-    android_logger::init_once(
-        android_logger::Config::default().with_tag("vm_payload").with_min_level(Level::Debug),
-    );
-    if let Err(e) = try_notify_payload_ready() {
-        error!("{:?}", e);
-        false
-    } else {
-        info!("Notified host payload ready successfully");
-        true
-    }
-}
-
-/// Notifies the host that the payload is ready.
-/// Returns a `Result` containing error information if failed.
-fn try_notify_payload_ready() -> Result<()> {
-    get_vm_payload_service()?.notifyPayloadReady().context("Cannot notify payload ready")
-}
-
-/// Runs a binder RPC server, serving the supplied binder service implementation on the given vsock
-/// port.
-///
-/// If and when the server is ready for connections (it is listening on the port), `on_ready` is
-/// called to allow appropriate action to be taken - e.g. to notify clients that they may now
-/// attempt to connect.
-///
-/// The current thread is joined to the binder thread pool to handle incoming messages.
-///
-/// Returns true if the server has shutdown normally, false if it failed in some way.
-///
-/// # Safety
-///
-/// The `on_ready` callback is only called inside `run_vsock_rpc_server`, within the lifetime of
-/// `ReadyNotifier` (the last parameter of `run_vsock_rpc_server`). If `on_ready` is called with
-/// wrong param, the callback execution could go wrong.
-#[no_mangle]
-pub unsafe extern "C" fn AVmPayload_runVsockRpcServer(
-    service: *mut AIBinder,
-    port: u32,
-    on_ready: Option<unsafe extern "C" fn(param: *mut c_void)>,
-    param: *mut c_void,
-) -> bool {
-    // SAFETY: AIBinder returned has correct reference count, and the ownership can
-    // safely be taken by new_spibinder.
-    let service = new_spibinder(service);
-    if let Some(service) = service {
-        run_vsock_rpc_server(service, port, || {
-            if let Some(on_ready) = on_ready {
-                on_ready(param);
-            }
-        })
-    } else {
-        error!("Failed to convert the given service from AIBinder to SpIBinder.");
-        false
-    }
-}
-
-/// Get a secret that is uniquely bound to this VM instance.
-///
-/// # Safety
-///
-/// Behavior is undefined if any of the following conditions are violated:
-///
-/// * `identifier` must be [valid] for reads of `identifier_size` bytes.
-/// * `secret` must be [valid] for writes of `size` bytes.
-///
-/// [valid]: std::ptr#safety
-#[no_mangle]
-pub unsafe extern "C" fn AVmPayload_getVmInstanceSecret(
-    identifier: *const u8,
-    identifier_size: usize,
-    secret: *mut u8,
-    size: usize,
-) -> bool {
-    let identifier = std::slice::from_raw_parts(identifier, identifier_size);
-    match try_get_vm_instance_secret(identifier, size) {
-        Err(e) => {
-            error!("{:?}", e);
-            false
-        }
-        Ok(vm_secret) => {
-            if vm_secret.len() != size {
-                return false;
-            }
-            std::ptr::copy_nonoverlapping(vm_secret.as_ptr(), secret, size);
-            true
-        }
-    }
-}
-
-fn try_get_vm_instance_secret(identifier: &[u8], size: usize) -> Result<Vec<u8>> {
-    get_vm_payload_service()?
-        .getVmInstanceSecret(identifier, i32::try_from(size)?)
-        .context("Cannot get VM instance secret")
-}
-
-/// Get the VM's attestation chain.
-/// Returns true on success, else false.
-///
-/// # Safety
-///
-/// Behavior is undefined if any of the following conditions are violated:
-///
-/// * `data` must be [valid] for writes of `size` bytes.
-/// * `total` must be [valid] for writes.
-///
-/// [valid]: std::ptr#safety
-#[no_mangle]
-pub unsafe extern "C" fn AVmPayload_getDiceAttestationChain(
-    data: *mut u8,
-    size: usize,
-    total: *mut usize,
-) -> bool {
-    match try_get_dice_attestation_chain() {
-        Err(e) => {
-            error!("{:?}", e);
-            false
-        }
-        Ok(chain) => {
-            total.write(chain.len());
-            std::ptr::copy_nonoverlapping(chain.as_ptr(), data, std::cmp::min(chain.len(), size));
-            true
-        }
-    }
-}
-
-fn try_get_dice_attestation_chain() -> Result<Vec<u8>> {
-    get_vm_payload_service()?.getDiceAttestationChain().context("Cannot get attestation chain")
-}
-
-/// Get the VM's attestation CDI.
-/// Returns true on success, else false.
-///
-/// # Safety
-///
-/// Behavior is undefined if any of the following conditions are violated:
-///
-/// * `data` must be [valid] for writes of `size` bytes.
-/// * `total` must be [valid] for writes.
-///
-/// [valid]: std::ptr#safety
-#[no_mangle]
-pub unsafe extern "C" fn AVmPayload_getDiceAttestationCdi(
-    data: *mut u8,
-    size: usize,
-    total: *mut usize,
-) -> bool {
-    match try_get_dice_attestation_cdi() {
-        Err(e) => {
-            error!("{:?}", e);
-            false
-        }
-        Ok(cdi) => {
-            total.write(cdi.len());
-            std::ptr::copy_nonoverlapping(cdi.as_ptr(), data, std::cmp::min(cdi.len(), size));
-            true
-        }
-    }
-}
-
-/// Gets the path to the APK contents.
-#[no_mangle]
-pub extern "C" fn AVmPayload_getApkContentsPath() -> *const c_char {
-    (*VM_APK_CONTENTS_PATH_C).as_ptr()
-}
-
-fn try_get_dice_attestation_cdi() -> Result<Vec<u8>> {
-    get_vm_payload_service()?.getDiceAttestationCdi().context("Cannot get attestation CDI")
-}
-
-fn get_vm_payload_service() -> Result<Strong<dyn IVmPayloadService>> {
-    wait_for_interface(VM_PAYLOAD_SERVICE_NAME)
-        .context(format!("Failed to connect to service: {}", VM_PAYLOAD_SERVICE_NAME))
-}
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index 4b06b3e..44b4c01 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -23,6 +23,7 @@
         "libdiced_sample_inputs",
         "libdiced_utils",
         "libglob",
+        "libhex",
         "libitertools",
         "libkernlog",
         "libkeystore2_crypto_rust",
diff --git a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
index 4823bb8..f8e7d34 100644
--- a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
+++ b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
@@ -21,8 +21,8 @@
  * Microdroid Manager for execution.
  */
 interface IVmPayloadService {
-    /** Name of the service IVmPayloadService. */
-    const String VM_PAYLOAD_SERVICE_NAME = "virtual_machine_payload_service";
+    /** Socket name of the service IVmPayloadService. */
+    const String VM_PAYLOAD_SERVICE_SOCKET_NAME = "vm_payload_service";
 
     /** Path to the APK contents path. */
     const String VM_APK_CONTENTS_PATH = "/mnt/apk";
diff --git a/microdroid_manager/microdroid_manager.rc b/microdroid_manager/microdroid_manager.rc
index 74a219d..c41ee38 100644
--- a/microdroid_manager/microdroid_manager.rc
+++ b/microdroid_manager/microdroid_manager.rc
@@ -1,8 +1,12 @@
 service microdroid_manager /system/bin/microdroid_manager
     disabled
+    # print android log to kmsg
     file /dev/kmsg w
+    # redirect stdout/stderr to kmsg_debug
+    stdio_to_kmsg
     setenv RUST_LOG info
     # TODO(jooyung) remove this when microdroid_manager becomes a daemon
     oneshot
     # SYS_BOOT is required to exec kexecload from microdroid_manager
     capabilities AUDIT_CONTROL SYS_ADMIN SYS_BOOT
+    socket vm_payload_service stream 0666 system system
diff --git a/microdroid_manager/src/dice.rs b/microdroid_manager/src/dice.rs
index 3881db3..499835f 100644
--- a/microdroid_manager/src/dice.rs
+++ b/microdroid_manager/src/dice.rs
@@ -38,6 +38,16 @@
     pub bcc: Vec<u8>,
 }
 
+impl DiceContext {
+    pub fn get_sealing_key(&self, salt: &[u8], identifier: &[u8], keysize: u32) -> Result<ZVec> {
+        // Deterministically derive a key to use for sealing data based on salt. Use different salt
+        // for different keys.
+        let mut key = ZVec::new(keysize as usize)?;
+        hkdf(&mut key, Md::sha256(), &self.cdi_seal, salt, identifier)?;
+        Ok(key)
+    }
+}
+
 /// Artifacts that are mapped into the process address space from the driver.
 pub enum DiceDriver<'a> {
     Real {
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index b8e85e7..4f94bb4 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -18,7 +18,6 @@
 mod instance;
 mod ioutil;
 mod payload;
-mod procutil;
 mod swap;
 mod vm_payload_service;
 
@@ -26,45 +25,42 @@
 use crate::instance::{ApexData, ApkData, InstanceDisk, MicrodroidData, RootHash};
 use crate::vm_payload_service::register_vm_payload_service;
 use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::ErrorCode::ErrorCode;
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
-    IVirtualMachineService::{
-        IVirtualMachineService, VM_BINDER_SERVICE_PORT, VM_STREAM_SERVICE_PORT,
-    },
-    VirtualMachineCpuStatus::VirtualMachineCpuStatus,
-    VirtualMachineMemStatus::VirtualMachineMemStatus,
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
+use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
+    VM_APK_CONTENTS_PATH,
+    VM_PAYLOAD_SERVICE_SOCKET_NAME,
 };
-use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::VM_APK_CONTENTS_PATH;
 use anyhow::{anyhow, bail, ensure, Context, Error, Result};
 use apkverify::{get_public_key_der, verify, V4Signature};
-use binder::{ProcessState, Strong};
+use binder::Strong;
 use diced_utils::cbor::{encode_header, encode_number};
 use glob::glob;
 use itertools::sorted;
 use libc::VMADDR_CID_HOST;
-use log::{error, info};
+use log::{error, info, warn};
 use microdroid_metadata::{write_metadata, Metadata, PayloadMetadata};
 use microdroid_payload_config::{OsConfig, Task, TaskType, VmPayloadConfig};
+use nix::fcntl::{fcntl, F_SETFD, FdFlag};
+use nix::sys::signal::Signal;
 use openssl::sha::Sha512;
 use payload::{get_apex_data_from_payload, load_metadata, to_metadata};
-use procutil::{get_cpu_time, get_mem_info};
 use rand::Fill;
 use rpcbinder::get_vsock_rpc_interface;
+use rustutils::sockets::android_get_control_socket;
 use rustutils::system_properties;
 use rustutils::system_properties::PropertyWatcher;
 use std::borrow::Cow::{Borrowed, Owned};
 use std::convert::TryInto;
-use std::fs::{self, create_dir, File, OpenOptions};
+use std::env;
+use std::fs::{self, create_dir, OpenOptions};
 use std::io::Write;
-use std::os::unix::io::{FromRawFd, IntoRawFd};
+use std::os::unix::process::ExitStatusExt;
 use std::path::Path;
 use std::process::{Child, Command, Stdio};
 use std::str;
-use std::thread;
 use std::time::{Duration, SystemTime};
-use vsock::VsockStream;
 
 const WAIT_TIMEOUT: Duration = Duration::from_secs(10);
-const SENDING_VM_STATUS_CYCLE_PERIOD: Duration = Duration::from_secs(60);
 const MAIN_APK_PATH: &str = "/dev/block/by-name/microdroid-apk";
 const MAIN_APK_IDSIG_PATH: &str = "/dev/block/by-name/microdroid-apk-idsig";
 const MAIN_APK_DEVICE_NAME: &str = "microdroid-apk";
@@ -85,6 +81,13 @@
 // SYNC WITH virtualizationservice/src/crosvm.rs
 const FAILURE_SERIAL_DEVICE: &str = "/dev/ttyS1";
 
+/// Identifier for the key used for encrypted store.
+const ENCRYPTEDSTORE_BACKING_DEVICE: &str = "/dev/block/by-name/encryptedstore";
+const ENCRYPTEDSTORE_BIN: &str = "/system/bin/encryptedstore";
+const ENCRYPTEDSTORE_KEY_IDENTIFIER: &str = "encryptedstore_key";
+const ENCRYPTEDSTORE_KEYSIZE: u32 = 32;
+const ENCRYPTEDSTORE_MOUNTPOINT: &str = "/mnt/encryptedstore";
+
 #[derive(thiserror::Error, Debug)]
 enum MicrodroidError {
     #[error("Cannot connect to virtualization service: {0}")]
@@ -97,42 +100,6 @@
     InvalidConfig(String),
 }
 
-fn send_vm_status(service: &Strong<dyn IVirtualMachineService>) -> Result<()> {
-    // Collect VM CPU time information and creating VmCpuStatus atom for metrics.
-    let cpu_time = get_cpu_time()?;
-    let vm_cpu_status = VirtualMachineCpuStatus {
-        cpu_time_user: cpu_time.user,
-        cpu_time_nice: cpu_time.nice,
-        cpu_time_sys: cpu_time.sys,
-        cpu_time_idle: cpu_time.idle,
-    };
-    service.notifyCpuStatus(&vm_cpu_status).expect("Can't send information about VM CPU status");
-
-    // Collect VM memory information and creating VmMemStatus atom for metrics.
-    let mem_info = get_mem_info()?;
-    let vm_mem_status = VirtualMachineMemStatus {
-        mem_total: mem_info.total,
-        mem_free: mem_info.free,
-        mem_available: mem_info.available,
-        mem_buffer: mem_info.buffer,
-        mem_cached: mem_info.cached,
-    };
-    service.notifyMemStatus(&vm_mem_status).expect("Can't send information about VM memory status");
-
-    Ok(())
-}
-
-fn send_vm_status_periodically() -> Result<()> {
-    let service = get_vms_rpc_binder()
-        .context("cannot connect to VirtualMachineService")
-        .map_err(|e| MicrodroidError::FailedToConnectToVirtualizationService(e.to_string()))?;
-
-    loop {
-        send_vm_status(&service)?;
-        thread::sleep(SENDING_VM_STATUS_CYCLE_PERIOD);
-    }
-}
-
 fn translate_error(err: &Error) -> (ErrorCode, String) {
     if let Some(e) = err.downcast_ref::<MicrodroidError>() {
         match e {
@@ -191,11 +158,19 @@
 }
 
 fn get_vms_rpc_binder() -> Result<Strong<dyn IVirtualMachineService>> {
-    get_vsock_rpc_interface(VMADDR_CID_HOST, VM_BINDER_SERVICE_PORT as u32)
-        .context("Cannot connect to RPC service")
+    // The host is running a VirtualMachineService for this VM on a port equal
+    // to the CID of this VM.
+    let port = vsock::get_local_cid().context("Could not determine local CID")?;
+    get_vsock_rpc_interface(VMADDR_CID_HOST, port)
+        .context("Could not connect to IVirtualMachineService")
 }
 
 fn main() -> Result<()> {
+    // If debuggable, print full backtrace to console log with stdio_to_kmsg
+    if system_properties::read_bool(APP_DEBUGGABLE_PROP, true)? {
+        env::set_var("RUST_BACKTRACE", "full");
+    }
+
     scopeguard::defer! {
         info!("Shutting down...");
         if let Err(e) = system_properties::write("sys.powerctl", "shutdown") {
@@ -212,10 +187,22 @@
     })
 }
 
+fn set_cloexec_on_vm_payload_service_socket() -> Result<()> {
+    let fd = android_get_control_socket(VM_PAYLOAD_SERVICE_SOCKET_NAME)?;
+
+    fcntl(fd, F_SETFD(FdFlag::FD_CLOEXEC))?;
+
+    Ok(())
+}
+
 fn try_main() -> Result<()> {
     let _ = kernlog::init();
     info!("started.");
 
+    if let Err(e) = set_cloexec_on_vm_payload_service_socket() {
+        warn!("Failed to set cloexec on vm payload socket: {:?}", e);
+    }
+
     load_crashkernel_if_supported().context("Failed to load crashkernel")?;
 
     swap::init_swap().context("Failed to initialise swap")?;
@@ -225,12 +212,6 @@
         .context("cannot connect to VirtualMachineService")
         .map_err(|e| MicrodroidError::FailedToConnectToVirtualizationService(e.to_string()))?;
 
-    thread::spawn(move || {
-        if let Err(e) = send_vm_status_periodically() {
-            error!("failed to get virtual machine status: {:?}", e);
-        }
-    });
-
     match try_run_payload(&service) {
         Ok(code) => {
             info!("notifying payload finished");
@@ -392,7 +373,15 @@
 
     // To minimize the exposure to untrusted data, derive dice profile as soon as possible.
     info!("DICE derivation for payload");
-    let dice = dice_derivation(dice, &verified_data, &payload_metadata)?;
+    let dice_context = dice_derivation(dice, &verified_data, &payload_metadata)?;
+
+    // Run encryptedstore binary to prepare the storage
+    let encryptedstore_child = if Path::new(ENCRYPTEDSTORE_BACKING_DEVICE).exists() {
+        info!("Preparing encryptedstore ...");
+        Some(prepare_encryptedstore(&dice_context).context("encryptedstore run")?)
+    } else {
+        None
+    };
 
     // Before reading a file from the APK, start zipfuse
     run_zipfuse(
@@ -428,9 +417,10 @@
     mount_extra_apks(&config)?;
 
     // Wait until apex config is done. (e.g. linker configuration for apexes)
-    // TODO(jooyung): wait until sys.boot_completed?
     wait_for_apex_config_done()?;
 
+    setup_config_sysprops(&config)?;
+
     // Start tombstone_transmit if enabled
     if config.export_tombstones {
         control_service("start", "tombstone_transmit")?;
@@ -438,20 +428,18 @@
         control_service("stop", "tombstoned")?;
     }
 
-    // Start authfs if enabled
-    if config.enable_authfs {
-        control_service("start", "authfs_service")?;
-    }
-
     // Wait until zipfuse has mounted the APK so we can access the payload
     wait_for_property_true(APK_MOUNT_DONE_PROP).context("Failed waiting for APK mount done")?;
 
-    register_vm_payload_service(allow_restricted_apis, service.clone(), dice)?;
-    ProcessState::start_thread_pool();
+    register_vm_payload_service(allow_restricted_apis, service.clone(), dice_context)?;
 
-    system_properties::write("dev.bootcomplete", "1").context("set dev.bootcomplete")?;
-    send_vm_status(service)?;
+    if let Some(mut child) = encryptedstore_child {
+        let exitcode = child.wait().context("Wait for encryptedstore child")?;
+        ensure!(exitcode.success(), "Unable to prepare encrypted storage. Exitcode={}", exitcode);
+    }
 
+    wait_for_property_true("dev.bootcomplete").context("failed waiting for dev.bootcomplete")?;
+    info!("boot completed, time to run payload");
     exec_task(task, service).context("Failed to run payload")
 }
 
@@ -470,8 +458,6 @@
 fn run_apkdmverity(args: &[ApkDmverityArgument]) -> Result<Child> {
     let mut cmd = Command::new(APKDMVERITY_BIN);
 
-    cmd.stdin(Stdio::null()).stdout(Stdio::null()).stderr(Stdio::null());
-
     for argument in args {
         cmd.arg("--apk").arg(argument.apk).arg(argument.idsig).arg(argument.name);
         if let Some(root_hash) = argument.saved_root_hash {
@@ -503,15 +489,7 @@
     if let Some(property_name) = ready_prop {
         cmd.args(["-p", property_name]);
     }
-    cmd.arg("-o")
-        .arg(option)
-        .arg(zip_path)
-        .arg(mount_dir)
-        .stdin(Stdio::null())
-        .stdout(Stdio::null())
-        .stderr(Stdio::null())
-        .spawn()
-        .context("Spawn zipfuse")
+    cmd.arg("-o").arg(option).arg(zip_path).arg(mount_dir).spawn().context("Spawn zipfuse")
 }
 
 fn write_apex_payload_data(
@@ -701,6 +679,16 @@
     Ok(())
 }
 
+fn setup_config_sysprops(config: &VmPayloadConfig) -> Result<()> {
+    if config.enable_authfs {
+        system_properties::write("microdroid_manager.authfs.enabled", "1")
+            .context("failed to write microdroid_manager.authfs.enabled")?;
+    }
+    system_properties::write("microdroid_manager.config_done", "1")
+        .context("failed to write microdroid_manager.config_done")?;
+    Ok(())
+}
+
 // Waits until linker config is generated
 fn wait_for_apex_config_done() -> Result<()> {
     wait_for_property_true(APEX_CONFIG_DONE_PROP).context("Failed waiting for apex config done")
@@ -773,21 +761,9 @@
     Ok(())
 }
 
-/// Executes the given task. Stdout of the task is piped into the vsock stream to the
-/// virtualizationservice in the host side.
+/// Executes the given task.
 fn exec_task(task: &Task, service: &Strong<dyn IVirtualMachineService>) -> Result<i32> {
     info!("executing main task {:?}...", task);
-    let mut command = build_command(task)?;
-
-    info!("notifying payload started");
-    service.notifyPayloadStarted()?;
-
-    let exit_status = command.spawn()?.wait()?;
-    send_vm_status(service)?;
-    exit_status.code().ok_or_else(|| anyhow!("Failed to get exit_code from the paylaod."))
-}
-
-fn build_command(task: &Task) -> Result<Command> {
     let mut command = match task.type_ {
         TaskType::Executable => Command::new(&task.command),
         TaskType::MicrodroidLauncher => {
@@ -796,29 +772,23 @@
             command
         }
     };
+    command.stdin(Stdio::null()).stdout(Stdio::null()).stderr(Stdio::null());
 
-    match VsockStream::connect_with_cid_port(VMADDR_CID_HOST, VM_STREAM_SERVICE_PORT as u32) {
-        Ok(stream) => {
-            // SAFETY: the ownership of the underlying file descriptor is transferred from stream
-            // to the file object, and then into the Command object. When the command is finished,
-            // the file descriptor is closed.
-            let file = unsafe { File::from_raw_fd(stream.into_raw_fd()) };
-            command
-                .stdin(Stdio::from(file.try_clone()?))
-                .stdout(Stdio::from(file.try_clone()?))
-                .stderr(Stdio::from(file));
-        }
-        Err(e) => {
-            error!("failed to connect to virtualization service: {}", e);
-            // Don't fail hard here. Even if we failed to connect to the virtualizationservice,
-            // we keep executing the task. This can happen if the owner of the VM doesn't register
-            // callback to accept the stream. Use /dev/null as the stream so that the task can
-            // make progress without waiting for someone to consume the output.
-            command.stdin(Stdio::null()).stdout(Stdio::null()).stderr(Stdio::null());
-        }
+    info!("notifying payload started");
+    service.notifyPayloadStarted()?;
+
+    let exit_status = command.spawn()?.wait()?;
+    match exit_status.code() {
+        Some(exit_code) => Ok(exit_code),
+        None => Err(match exit_status.signal() {
+            Some(signal) => anyhow!(
+                "Payload exited due to signal: {} ({})",
+                signal,
+                Signal::try_from(signal).map_or("unknown", |s| s.as_str())
+            ),
+            None => anyhow!("Payload has neither exit code nor signal"),
+        }),
     }
-
-    Ok(command)
 }
 
 fn find_library_path(name: &str) -> Result<String> {
@@ -838,3 +808,28 @@
 fn to_hex_string(buf: &[u8]) -> String {
     buf.iter().map(|b| format!("{:02X}", b)).collect()
 }
+
+fn prepare_encryptedstore(dice: &DiceContext) -> Result<Child> {
+    // Use a fixed salt to scope the derivation to this API.
+    // Generated using hexdump -vn32 -e'14/1 "0x%02X, " 1 "\n"' /dev/urandom
+    // TODO(b/241541860) : Move this (& other salts) to a salt container, i.e. a global enum
+    let salt = [
+        0xFC, 0x1D, 0x35, 0x7B, 0x96, 0xF3, 0xEF, 0x17, 0x78, 0x7D, 0x70, 0xED, 0xEA, 0xFE, 0x1D,
+        0x6F, 0xB3, 0xF9, 0x40, 0xCE, 0xDD, 0x99, 0x40, 0xAA, 0xA7, 0x0E, 0x92, 0x73, 0x90, 0x86,
+        0x4A, 0x75,
+    ];
+    let key = dice.get_sealing_key(
+        &salt,
+        ENCRYPTEDSTORE_KEY_IDENTIFIER.as_bytes(),
+        ENCRYPTEDSTORE_KEYSIZE,
+    )?;
+
+    let mut cmd = Command::new(ENCRYPTEDSTORE_BIN);
+    cmd.arg("--blkdevice")
+        .arg(ENCRYPTEDSTORE_BACKING_DEVICE)
+        .arg("--key")
+        .arg(hex::encode(&*key))
+        .args(["--mountpoint", ENCRYPTEDSTORE_MOUNTPOINT])
+        .spawn()
+        .context("encryptedstore failed")
+}
diff --git a/microdroid_manager/src/procutil.rs b/microdroid_manager/src/procutil.rs
deleted file mode 100644
index f9479c4..0000000
--- a/microdroid_manager/src/procutil.rs
+++ /dev/null
@@ -1,127 +0,0 @@
-// Copyright 2022, The Android Open Source Project
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use anyhow::{bail, Result};
-use libc::{sysconf, _SC_CLK_TCK};
-use std::fs::File;
-use std::io::{BufRead, BufReader};
-
-const MILLIS_PER_SEC: i64 = 1000;
-
-pub struct CpuTime {
-    pub user: i64,
-    pub nice: i64,
-    pub sys: i64,
-    pub idle: i64,
-}
-
-pub struct MemInfo {
-    pub total: i64,
-    pub free: i64,
-    pub available: i64,
-    pub buffer: i64,
-    pub cached: i64,
-}
-
-// Get CPU time information from /proc/stat
-//
-// /proc/stat example(omitted):
-//   cpu  24790952 21104390 10771070 10480973587 1700955 0 410931 0 316532 0
-//   cpu0 169636 141307 61153 81785791 9605 0 183524 0 1345 0
-//   cpu1 182431 198327 68273 81431817 10445 0 32392 0 2616 0
-//   cpu2 183209 174917 68591 81933935 12239 0 10042 0 2415 0
-//   cpu3 183413 177758 69908 81927474 13354 0 5853 0 2491 0
-//   intr 7913477443 39 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
-//   ctxt 10326710014
-//   btime 1664123605
-//   processes 9225712
-//   procs_running 1
-//   procs_blocked 0
-//   softirq 2683914305 14595298 304837101 1581 327291100 16397051 0 208857783 1024640365 787932 786506094
-//
-// expected output:
-//   user: 24790952
-//   nice: 21104390
-//   sys: 10771070
-//   idle: 10480973587
-pub fn get_cpu_time() -> Result<CpuTime> {
-    let mut proc_stat = BufReader::new(File::open("/proc/stat")?);
-    let mut line = String::new();
-    proc_stat.read_line(&mut line)?;
-    let data_list: Vec<_> = line.split_whitespace().filter_map(|s| s.parse::<i64>().ok()).collect();
-    if data_list.len() < 4 {
-        bail!("Failed to extract numeric values in /proc/stat :\n{}", line);
-    }
-
-    let ticks_per_sec = unsafe { sysconf(_SC_CLK_TCK) } as i64;
-    let cpu_time = CpuTime {
-        user: data_list[0] * MILLIS_PER_SEC / ticks_per_sec,
-        nice: data_list[1] * MILLIS_PER_SEC / ticks_per_sec,
-        sys: data_list[2] * MILLIS_PER_SEC / ticks_per_sec,
-        idle: data_list[3] * MILLIS_PER_SEC / ticks_per_sec,
-    };
-    Ok(cpu_time)
-}
-
-// Get memory information from /proc/meminfo
-//
-// /proc/meminfo example(omitted):
-//   MemTotal:       263742736 kB
-//   MemFree:        37144204 kB
-//   MemAvailable:   249168700 kB
-//   Buffers:        10231296 kB
-//   Cached:         189502836 kB
-//   SwapCached:       113848 kB
-//   Active:         132266424 kB
-//   Inactive:       73587504 kB
-//   Active(anon):    1455240 kB
-//   Inactive(anon):  6993584 kB
-//   Active(file):   130811184 kB
-//   Inactive(file): 66593920 kB
-//   Unevictable:       56436 kB
-//   Mlocked:           56436 kB
-//   SwapTotal:      255123452 kB
-//   SwapFree:       254499068 kB
-//   Dirty:               596 kB
-//   Writeback:             0 kB
-//   AnonPages:       5295864 kB
-//   Mapped:          3512608 kB
-//
-// expected output:
-//   total: 263742736
-//   free: 37144204
-//   available: 249168700
-//   buffer: 10231296
-//   cached: 189502836
-pub fn get_mem_info() -> Result<MemInfo> {
-    let mut proc_stat = BufReader::new(File::open("/proc/meminfo")?);
-    let mut lines = String::new();
-    for _ in 0..5 {
-        proc_stat.read_line(&mut lines)?;
-    }
-    let data_list: Vec<_> =
-        lines.split_whitespace().filter_map(|s| s.parse::<i64>().ok()).collect();
-    if data_list.len() != 5 {
-        bail!("Failed to extract numeric values in /proc/meminfo :\n{}", lines);
-    }
-
-    let mem_info = MemInfo {
-        total: data_list[0],
-        free: data_list[1],
-        available: data_list[2],
-        buffer: data_list[3],
-        cached: data_list[4],
-    };
-    Ok(mem_info)
-}
diff --git a/microdroid_manager/src/swap.rs b/microdroid_manager/src/swap.rs
index 0fd1468..d7916db 100644
--- a/microdroid_manager/src/swap.rs
+++ b/microdroid_manager/src/swap.rs
@@ -43,20 +43,21 @@
     let sysfs_size = format!("/sys/{}/size", dev);
     let len = read_to_string(&sysfs_size)?
         .trim()
-        .parse::<u32>()
-        .context(format!("No u32 in {}", &sysfs_size))?
-        * 512;
+        .parse::<u64>()
+        .context(format!("No u64 in {}", &sysfs_size))?
+        .checked_mul(512)
+        .ok_or_else(|| anyhow!("sysfs_size too large"))?;
 
-    let pagesize: libc::c_uint;
     // safe because we give a constant and known-valid sysconf parameter
-    unsafe {
-        pagesize = libc::sysconf(libc::_SC_PAGE_SIZE) as libc::c_uint;
-    }
+    let pagesize = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) as u64 };
 
     let mut f = OpenOptions::new().read(false).write(true).open(format!("/dev/{}", dev))?;
 
+    let last_page = len / pagesize - 1;
+
     // Write the info fields: [ version, last_page ]
-    let info: [u32; 2] = [1, (len / pagesize) - 1];
+    let info: [u32; 2] = [1, last_page.try_into().context("Number of pages out of range")?];
+
     f.seek(SeekFrom::Start(1024))?;
     f.write_all(&info.iter().flat_map(|v| v.to_ne_bytes()).collect::<Vec<u8>>())?;
 
diff --git a/microdroid_manager/src/vm_payload_service.rs b/microdroid_manager/src/vm_payload_service.rs
index 159bf67..98b9f2b 100644
--- a/microdroid_manager/src/vm_payload_service.rs
+++ b/microdroid_manager/src/vm_payload_service.rs
@@ -16,13 +16,14 @@
 
 use crate::dice::DiceContext;
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
-    BnVmPayloadService, IVmPayloadService, VM_PAYLOAD_SERVICE_NAME};
+    BnVmPayloadService, IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME};
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
-use anyhow::{Context, Result};
-use binder::{Interface, BinderFeatures, ExceptionCode, Status, Strong, add_service};
-use log::error;
+use anyhow::Result;
+use binder::{Interface, BinderFeatures, ExceptionCode, Status, Strong};
+use log::{error, info};
 use openssl::hkdf::hkdf;
 use openssl::md::Md;
+use rpcbinder::RpcServer;
 
 /// Implementation of `IVmPayloadService`.
 struct VmPayloadService {
@@ -97,8 +98,16 @@
         VmPayloadService::new(allow_restricted_apis, vm_service, dice),
         BinderFeatures::default(),
     );
-    add_service(VM_PAYLOAD_SERVICE_NAME, vm_payload_binder.as_binder())
-        .with_context(|| format!("Failed to register service {}", VM_PAYLOAD_SERVICE_NAME))?;
-    log::info!("{} is running", VM_PAYLOAD_SERVICE_NAME);
+
+    let server = RpcServer::new_init_unix_domain(
+        vm_payload_binder.as_binder(),
+        VM_PAYLOAD_SERVICE_SOCKET_NAME,
+    )?;
+    info!("The RPC server '{}' is running.", VM_PAYLOAD_SERVICE_SOCKET_NAME);
+
+    // Move server reference into a background thread and run it forever.
+    std::thread::spawn(move || {
+        server.join();
+    });
     Ok(())
 }
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index b644905..b78e077 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -8,11 +8,21 @@
     defaults: ["vmbase_ffi_defaults"],
     srcs: ["src/main.rs"],
     edition: "2021",
+    features: [
+        "legacy",
+    ],
     rustlibs: [
+        "libaarch64_paging",
+        "libbuddy_system_allocator",
+        "liblibfdt",
         "liblog_rust_nostd",
         "libpvmfw_embedded_key",
+        "libtinyvec_nostd",
         "libvmbase",
     ],
+    static_libs: [
+        "libarm-optimized-routines-mem",
+    ],
     apex_available: ["com.android.virt"],
 }
 
diff --git a/pvmfw/idmap.S b/pvmfw/idmap.S
index ec3ceaf..2ef0d42 100644
--- a/pvmfw/idmap.S
+++ b/pvmfw/idmap.S
@@ -40,13 +40,9 @@
 	/* level 1 */
 	.quad		.L_BLOCK_DEV | 0x0		// 1 GB of device mappings
 	.quad		.L_TT_TYPE_TABLE + 0f		// Unmapped device memory, and pVM firmware
-	.quad		.L_TT_TYPE_TABLE + 1f		// up to 1 GB of DRAM
-	.fill		509, 8, 0x0			// 509 GB of remaining VA space
+	.fill		510, 8, 0x0			// 510 GB of remaining VA space
 
 	/* level 2 */
 0:	.fill		510, 8, 0x0
 	.quad		.L_BLOCK_MEM_XIP | 0x7fc00000	// pVM firmware image
 	.quad		.L_BLOCK_MEM	 | 0x7fe00000	// Writable memory for stack, heap &c.
-1:	.quad		.L_BLOCK_RO	 | 0x80000000	// DT provided by VMM
-	.quad		.L_BLOCK_RO	 | 0x80200000	// 2 MB of DRAM containing payload image
-	.fill		510, 8, 0x0
diff --git a/pvmfw/src/config.rs b/pvmfw/src/config.rs
new file mode 100644
index 0000000..0f2a39c
--- /dev/null
+++ b/pvmfw/src/config.rs
@@ -0,0 +1,200 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Support for the pvmfw configuration data format.
+
+use crate::helpers;
+use core::fmt;
+use core::mem;
+use core::num::NonZeroUsize;
+use core::ops;
+use core::result;
+
+#[repr(C, packed)]
+#[derive(Clone, Copy, Debug)]
+struct Header {
+    magic: u32,
+    version: u32,
+    total_size: u32,
+    flags: u32,
+    entries: [HeaderEntry; Entry::COUNT],
+}
+
+#[derive(Debug)]
+pub enum Error {
+    /// Reserved region can't fit configuration header.
+    BufferTooSmall,
+    /// Header doesn't contain the expect magic value.
+    InvalidMagic,
+    /// Version of the header isn't supported.
+    UnsupportedVersion(u16, u16),
+    /// Header sets flags incorrectly or uses reserved flags.
+    InvalidFlags(u32),
+    /// Header describes configuration data that doesn't fit in the expected buffer.
+    InvalidSize(usize),
+    /// Header entry is invalid.
+    InvalidEntry(Entry),
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::BufferTooSmall => write!(f, "Reserved region is smaller than config header"),
+            Self::InvalidMagic => write!(f, "Wrong magic number"),
+            Self::UnsupportedVersion(x, y) => write!(f, "Version {x}.{y} not supported"),
+            Self::InvalidFlags(v) => write!(f, "Flags value {v:#x} is incorrect or reserved"),
+            Self::InvalidSize(sz) => write!(f, "Total size ({sz:#x}) overflows reserved region"),
+            Self::InvalidEntry(e) => write!(f, "Entry {e:?} is invalid"),
+        }
+    }
+}
+
+pub type Result<T> = result::Result<T, Error>;
+
+impl Header {
+    const MAGIC: u32 = u32::from_ne_bytes(*b"pvmf");
+    const PADDED_SIZE: usize =
+        helpers::unchecked_align_up(mem::size_of::<Self>(), mem::size_of::<u64>());
+
+    pub const fn version(major: u16, minor: u16) -> u32 {
+        ((major as u32) << 16) | (minor as u32)
+    }
+
+    pub const fn version_tuple(&self) -> (u16, u16) {
+        ((self.version >> 16) as u16, self.version as u16)
+    }
+
+    pub fn total_size(&self) -> usize {
+        self.total_size as usize
+    }
+
+    pub fn body_size(&self) -> usize {
+        self.total_size() - Self::PADDED_SIZE
+    }
+
+    fn get(&self, entry: Entry) -> HeaderEntry {
+        self.entries[entry as usize]
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub enum Entry {
+    Bcc = 0,
+    DebugPolicy = 1,
+}
+
+impl Entry {
+    const COUNT: usize = 2;
+}
+
+#[repr(packed)]
+#[derive(Clone, Copy, Debug)]
+struct HeaderEntry {
+    offset: u32,
+    size: u32,
+}
+
+impl HeaderEntry {
+    pub fn is_empty(&self) -> bool {
+        self.offset() == 0 && self.size() == 0
+    }
+
+    pub fn fits_in(&self, max_size: usize) -> bool {
+        (Header::PADDED_SIZE..max_size).contains(&self.offset())
+            && NonZeroUsize::new(self.size())
+                .and_then(|s| s.checked_add(self.offset()))
+                .filter(|&x| x.get() <= max_size)
+                .is_some()
+    }
+
+    pub fn as_body_range(&self) -> ops::Range<usize> {
+        let start = self.offset() - Header::PADDED_SIZE;
+
+        start..(start + self.size())
+    }
+
+    pub fn offset(&self) -> usize {
+        self.offset as usize
+    }
+
+    pub fn size(&self) -> usize {
+        self.size as usize
+    }
+}
+
+#[derive(Debug)]
+pub struct Config<'a> {
+    header: &'a Header,
+    body: &'a mut [u8],
+}
+
+impl<'a> Config<'a> {
+    /// Take ownership of a pvmfw configuration consisting of its header and following entries.
+    ///
+    /// SAFETY - 'data' should respect the alignment of Header.
+    pub unsafe fn new(data: &'a mut [u8]) -> Result<Self> {
+        let header = data.get(..Header::PADDED_SIZE).ok_or(Error::BufferTooSmall)?;
+
+        let header = &*(header.as_ptr() as *const Header);
+
+        if header.magic != Header::MAGIC {
+            return Err(Error::InvalidMagic);
+        }
+
+        if header.version != Header::version(1, 0) {
+            let (major, minor) = header.version_tuple();
+            return Err(Error::UnsupportedVersion(major, minor));
+        }
+
+        if header.flags != 0 {
+            return Err(Error::InvalidFlags(header.flags));
+        }
+
+        let total_size = header.total_size();
+
+        // BCC is a mandatory entry of the configuration data.
+        if !header.get(Entry::Bcc).fits_in(total_size) {
+            return Err(Error::InvalidEntry(Entry::Bcc));
+        }
+
+        // Debug policy is optional.
+        let dp = header.get(Entry::DebugPolicy);
+        if !dp.is_empty() && !dp.fits_in(total_size) {
+            return Err(Error::InvalidEntry(Entry::DebugPolicy));
+        }
+
+        let body = data
+            .get_mut(Header::PADDED_SIZE..)
+            .ok_or(Error::BufferTooSmall)?
+            .get_mut(..header.body_size())
+            .ok_or(Error::InvalidSize(total_size))?;
+
+        Ok(Self { header, body })
+    }
+
+    /// Get slice containing the platform BCC.
+    pub fn get_bcc_mut(&mut self) -> &mut [u8] {
+        &mut self.body[self.header.get(Entry::Bcc).as_body_range()]
+    }
+
+    /// Get slice containing the platform debug policy.
+    pub fn get_debug_policy(&mut self) -> Option<&mut [u8]> {
+        let entry = self.header.get(Entry::DebugPolicy);
+        if entry.is_empty() {
+            None
+        } else {
+            Some(&mut self.body[entry.as_body_range()])
+        }
+    }
+}
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 15a8c03..7859ff3 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -14,17 +14,37 @@
 
 //! Low-level entry and exit points of pvmfw.
 
+use crate::config;
+use crate::fdt;
+use crate::heap;
 use crate::helpers;
+use crate::memory::MemoryTracker;
 use crate::mmio_guard;
+use crate::mmu;
 use core::arch::asm;
+use core::num::NonZeroUsize;
 use core::slice;
-use log::{debug, error, LevelFilter};
-use vmbase::{console, logger, main, power::reboot};
+use log::debug;
+use log::error;
+use log::info;
+use log::warn;
+use log::LevelFilter;
+use vmbase::{console, layout, logger, main, power::reboot};
 
 #[derive(Debug, Clone)]
 enum RebootReason {
+    /// A malformed BCC was received.
+    InvalidBcc,
+    /// An invalid configuration was appended to pvmfw.
+    InvalidConfig,
     /// An unexpected internal error happened.
     InternalError,
+    /// The provided FDT was invalid.
+    InvalidFdt,
+    /// The provided payload was invalid.
+    InvalidPayload,
+    /// The provided ramdisk was invalid.
+    InvalidRamdisk,
 }
 
 main!(start);
@@ -43,6 +63,98 @@
     // if we reach this point and return, vmbase::entry::rust_entry() will call power::shutdown().
 }
 
+struct MemorySlices<'a> {
+    fdt: &'a mut libfdt::Fdt,
+    kernel: &'a [u8],
+    ramdisk: Option<&'a [u8]>,
+}
+
+impl<'a> MemorySlices<'a> {
+    fn new(
+        fdt: usize,
+        payload: usize,
+        payload_size: usize,
+        memory: &mut MemoryTracker,
+    ) -> Result<Self, RebootReason> {
+        // SAFETY - SIZE_2MB is non-zero.
+        const FDT_SIZE: NonZeroUsize = unsafe { NonZeroUsize::new_unchecked(helpers::SIZE_2MB) };
+        // TODO - Only map the FDT as read-only, until we modify it right before jump_to_payload()
+        // e.g. by generating a DTBO for a template DT in main() and, on return, re-map DT as RW,
+        // overwrite with the template DT and apply the DTBO.
+        let range = memory.alloc_mut(fdt, FDT_SIZE).map_err(|e| {
+            error!("Failed to allocate the FDT range: {e}");
+            RebootReason::InternalError
+        })?;
+
+        // SAFETY - The tracker validated the range to be in main memory, mapped, and not overlap.
+        let fdt = unsafe { slice::from_raw_parts_mut(range.start as *mut u8, range.len()) };
+        let fdt = libfdt::Fdt::from_mut_slice(fdt).map_err(|e| {
+            error!("Failed to spawn the FDT wrapper: {e}");
+            RebootReason::InvalidFdt
+        })?;
+
+        debug!("Fdt passed validation!");
+
+        let memory_range = fdt
+            .memory()
+            .map_err(|e| {
+                error!("Failed to get /memory from the DT: {e}");
+                RebootReason::InvalidFdt
+            })?
+            .ok_or_else(|| {
+                error!("Node /memory was found empty");
+                RebootReason::InvalidFdt
+            })?
+            .next()
+            .ok_or_else(|| {
+                error!("Failed to read the memory size from the FDT");
+                RebootReason::InternalError
+            })?;
+
+        debug!("Resizing MemoryTracker to range {memory_range:#x?}");
+
+        memory.shrink(&memory_range).map_err(|_| {
+            error!("Failed to use memory range value from DT: {memory_range:#x?}");
+            RebootReason::InvalidFdt
+        })?;
+
+        let payload_size = NonZeroUsize::new(payload_size).ok_or_else(|| {
+            error!("Invalid payload size: {payload_size:#x}");
+            RebootReason::InvalidPayload
+        })?;
+
+        let payload_range = memory.alloc(payload, payload_size).map_err(|e| {
+            error!("Failed to obtain the payload range: {e}");
+            RebootReason::InternalError
+        })?;
+        // SAFETY - The tracker validated the range to be in main memory, mapped, and not overlap.
+        let kernel =
+            unsafe { slice::from_raw_parts(payload_range.start as *const u8, payload_range.len()) };
+
+        let ramdisk_range = fdt::initrd_range(fdt).map_err(|e| {
+            error!("An error occurred while locating the ramdisk in the device tree: {e}");
+            RebootReason::InternalError
+        })?;
+
+        let ramdisk = if let Some(r) = ramdisk_range {
+            debug!("Located ramdisk at {r:?}");
+            let r = memory.alloc_range(&r).map_err(|e| {
+                error!("Failed to obtain the initrd range: {e}");
+                RebootReason::InvalidRamdisk
+            })?;
+
+            // SAFETY - The region was validated by memory to be in main memory, mapped, and
+            // not overlap.
+            Some(unsafe { slice::from_raw_parts(r.start as *const u8, r.len()) })
+        } else {
+            info!("Couldn't locate the ramdisk from the device tree");
+            None
+        };
+
+        Ok(Self { fdt, kernel, ramdisk })
+    }
+}
+
 /// Sets up the environment for main() and wraps its result for start().
 ///
 /// Provide the abstractions necessary for start() to abort the pVM boot and for main() to run with
@@ -52,15 +164,11 @@
     // - only access MMIO once (and while) it has been mapped and configured
     // - only perform logging once the logger has been initialized
     // - only access non-pvmfw memory once (and while) it has been mapped
-    logger::init(LevelFilter::Info).map_err(|_| RebootReason::InternalError)?;
 
-    const FDT_MAX_SIZE: usize = helpers::SIZE_2MB;
-    // TODO: Check that the FDT is fully contained in RAM.
-    // SAFETY - We trust the VMM, for now.
-    let fdt = unsafe { slice::from_raw_parts_mut(fdt as *mut u8, FDT_MAX_SIZE) };
-    // TODO: Check that the payload is fully contained in RAM and doesn't overlap with the FDT.
-    // SAFETY - We trust the VMM, for now.
-    let payload = unsafe { slice::from_raw_parts(payload as *const u8, payload_size) };
+    // SAFETY - This function should and will only be called once, here.
+    unsafe { heap::init() };
+
+    logger::init(LevelFilter::Info).map_err(|_| RebootReason::InternalError)?;
 
     // Use debug!() to avoid printing to the UART if we failed to configure it as only local
     // builds that have tweaked the logger::init() call will actually attempt to log the message.
@@ -75,8 +183,49 @@
         RebootReason::InternalError
     })?;
 
+    // SAFETY - We only get the appended payload from here, once. It is mapped and the linker
+    // script prevents it from overlapping with other objects.
+    let appended_data = unsafe { get_appended_data_slice() };
+
+    // Up to this point, we were using the built-in static (from .rodata) page tables.
+
+    let mut page_table = mmu::PageTable::from_static_layout().map_err(|e| {
+        error!("Failed to set up the dynamic page tables: {e}");
+        RebootReason::InternalError
+    })?;
+
+    const CONSOLE_LEN: usize = 1; // vmbase::uart::Uart only uses one u8 register.
+    let uart_range = console::BASE_ADDRESS..(console::BASE_ADDRESS + CONSOLE_LEN);
+    page_table.map_device(&uart_range).map_err(|e| {
+        error!("Failed to remap the UART as a dynamic page table entry: {e}");
+        RebootReason::InternalError
+    })?;
+
+    // SAFETY - We only get the appended payload from here, once. It is statically mapped and the
+    // linker script prevents it from overlapping with other objects.
+    let mut appended = unsafe { AppendedPayload::new(appended_data) }.ok_or_else(|| {
+        error!("No valid configuration found");
+        RebootReason::InvalidConfig
+    })?;
+
+    let bcc = appended.get_bcc_mut().ok_or_else(|| {
+        error!("Invalid BCC");
+        RebootReason::InvalidBcc
+    })?;
+
+    debug!("Activating dynamic page table...");
+    // SAFETY - page_table duplicates the static mappings for everything that the Rust code is
+    // aware of so activating it shouldn't have any visible effect.
+    unsafe { page_table.activate() };
+    debug!("... Success!");
+
+    let mut memory = MemoryTracker::new(page_table);
+    let slices = MemorySlices::new(fdt, payload, payload_size, &mut memory)?;
+
     // This wrapper allows main() to be blissfully ignorant of platform details.
-    crate::main(fdt, payload);
+    crate::main(slices.fdt, slices.kernel, slices.ramdisk, bcc);
+
+    // TODO: Overwrite BCC before jumping to payload to avoid leaking our sealing key.
 
     mmio_guard::unmap(console::BASE_ADDRESS).map_err(|e| {
         error!("Failed to unshare the UART: {e}");
@@ -145,3 +294,59 @@
         );
     };
 }
+
+unsafe fn get_appended_data_slice() -> &'static mut [u8] {
+    let base = helpers::align_up(layout::binary_end(), helpers::SIZE_4KB).unwrap();
+    // pvmfw is contained in a 2MiB region so the payload can't be larger than the 2MiB alignment.
+    let size = helpers::align_up(base, helpers::SIZE_2MB).unwrap() - base;
+
+    slice::from_raw_parts_mut(base as *mut u8, size)
+}
+
+enum AppendedPayload<'a> {
+    /// Configuration data.
+    Config(config::Config<'a>),
+    /// Deprecated raw BCC, as used in Android T.
+    LegacyBcc(&'a mut [u8]),
+}
+
+impl<'a> AppendedPayload<'a> {
+    /// SAFETY - 'data' should respect the alignment of config::Header.
+    unsafe fn new(data: &'a mut [u8]) -> Option<Self> {
+        if Self::is_valid_config(data) {
+            Some(Self::Config(config::Config::new(data).unwrap()))
+        } else if cfg!(feature = "legacy") {
+            const BCC_SIZE: usize = helpers::SIZE_4KB;
+            warn!("Assuming the appended data at {:?} to be a raw BCC", data.as_ptr());
+            Some(Self::LegacyBcc(&mut data[..BCC_SIZE]))
+        } else {
+            None
+        }
+    }
+
+    unsafe fn is_valid_config(data: &mut [u8]) -> bool {
+        // This function is necessary to prevent the borrow checker from getting confused
+        // about the ownership of data in new(); see https://users.rust-lang.org/t/78467.
+        let addr = data.as_ptr();
+        config::Config::new(data)
+            .map_err(|e| warn!("Invalid configuration data at {addr:?}: {e}"))
+            .is_ok()
+    }
+
+    #[allow(dead_code)] // TODO(b/232900974)
+    fn get_debug_policy(&mut self) -> Option<&mut [u8]> {
+        match self {
+            Self::Config(ref mut cfg) => cfg.get_debug_policy(),
+            Self::LegacyBcc(_) => None,
+        }
+    }
+
+    fn get_bcc_mut(&mut self) -> Option<&mut [u8]> {
+        let bcc = match self {
+            Self::LegacyBcc(ref mut bcc) => bcc,
+            Self::Config(ref mut cfg) => cfg.get_bcc_mut(),
+        };
+        // TODO(b/256148034): return None if BccHandoverParse(bcc) != kDiceResultOk.
+        Some(bcc)
+    }
+}
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
new file mode 100644
index 0000000..5b9efd2
--- /dev/null
+++ b/pvmfw/src/fdt.rs
@@ -0,0 +1,32 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! High-level FDT functions.
+
+use core::ffi::CStr;
+use core::ops::Range;
+
+/// Extract from /chosen the address range containing the pre-loaded ramdisk.
+pub fn initrd_range(fdt: &libfdt::Fdt) -> libfdt::Result<Option<Range<usize>>> {
+    let start = CStr::from_bytes_with_nul(b"linux,initrd-start\0").unwrap();
+    let end = CStr::from_bytes_with_nul(b"linux,initrd-end\0").unwrap();
+
+    if let Some(chosen) = fdt.chosen()? {
+        if let (Some(start), Some(end)) = (chosen.getprop_u32(start)?, chosen.getprop_u32(end)?) {
+            return Ok(Some((start as usize)..(end as usize)));
+        }
+    }
+
+    Ok(None)
+}
diff --git a/microdroid/vm_payload/src/lib.rs b/pvmfw/src/heap.rs
similarity index 67%
copy from microdroid/vm_payload/src/lib.rs
copy to pvmfw/src/heap.rs
index be6cf93..bfa8320 100644
--- a/microdroid/vm_payload/src/lib.rs
+++ b/pvmfw/src/heap.rs
@@ -12,11 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! Library for payload to communicate with the Microdroid Manager.
+//! Heap implementation.
 
-mod vm_payload_service;
+use buddy_system_allocator::LockedHeap;
 
-pub use vm_payload_service::{
-    AVmPayload_getDiceAttestationCdi, AVmPayload_getDiceAttestationChain,
-    AVmPayload_getVmInstanceSecret, AVmPayload_notifyPayloadReady,
-};
+#[global_allocator]
+static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
+
+static mut HEAP: [u8; 65536] = [0; 65536];
+
+pub unsafe fn init() {
+    HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
+}
diff --git a/pvmfw/src/helpers.rs b/pvmfw/src/helpers.rs
index adfc189..f1ff36d 100644
--- a/pvmfw/src/helpers.rs
+++ b/pvmfw/src/helpers.rs
@@ -14,15 +14,64 @@
 
 //! Miscellaneous helper functions.
 
+use core::arch::asm;
+
 pub const SIZE_4KB: usize = 4 << 10;
 pub const SIZE_2MB: usize = 2 << 20;
 
-/// Computes the address of the page containing a given address.
-pub const fn page_of(addr: usize, page_size: usize) -> usize {
-    addr & !(page_size - 1)
+/// Computes the largest multiple of the provided alignment smaller or equal to the address.
+///
+/// Note: the result is undefined if alignment isn't a power of two.
+pub const fn unchecked_align_down(addr: usize, alignment: usize) -> usize {
+    addr & !(alignment - 1)
+}
+
+/// Computes the smallest multiple of the provided alignment larger or equal to the address.
+///
+/// Note: the result is undefined if alignment isn't a power of two and may wrap to 0.
+pub const fn unchecked_align_up(addr: usize, alignment: usize) -> usize {
+    unchecked_align_down(addr + alignment - 1, alignment)
+}
+
+/// Safe wrapper around unchecked_align_up() that validates its assumptions and doesn't wrap.
+pub const fn align_up(addr: usize, alignment: usize) -> Option<usize> {
+    if !alignment.is_power_of_two() {
+        None
+    } else if let Some(s) = addr.checked_add(alignment - 1) {
+        Some(unchecked_align_down(s, alignment))
+    } else {
+        None
+    }
 }
 
 /// Computes the address of the 4KiB page containing a given address.
 pub const fn page_4kb_of(addr: usize) -> usize {
-    page_of(addr, SIZE_4KB)
+    unchecked_align_down(addr, SIZE_4KB)
+}
+
+#[inline]
+fn min_dcache_line_size() -> usize {
+    const DMINLINE_SHIFT: usize = 16;
+    const DMINLINE_MASK: usize = 0xf;
+    let ctr_el0: usize;
+
+    unsafe { asm!("mrs {x}, ctr_el0", x = out(reg) ctr_el0) }
+
+    // DminLine: log2 of the number of words in the smallest cache line of all the data caches.
+    let dminline = (ctr_el0 >> DMINLINE_SHIFT) & DMINLINE_MASK;
+
+    1 << dminline
+}
+
+/// Flush `size` bytes of data cache by virtual address.
+#[inline]
+pub fn flush_region(start: usize, size: usize) {
+    let line_size = min_dcache_line_size();
+    let end = start + size;
+    let start = unchecked_align_down(start, line_size);
+
+    for line in (start..end).step_by(line_size) {
+        // SAFETY - Clearing cache lines shouldn't have Rust-visible side effects.
+        unsafe { asm!("dc cvau, {x}", x = in(reg) line) }
+    }
 }
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index 870172d..6810fda 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -16,25 +16,34 @@
 
 #![no_main]
 #![no_std]
+#![feature(default_alloc_error_handler)]
+#![feature(ptr_const_cast)] // Stabilized in 1.65.0
 
 mod avb;
+mod config;
 mod entry;
 mod exceptions;
+mod fdt;
+mod heap;
 mod helpers;
+mod memory;
 mod mmio_guard;
+mod mmu;
 mod smccc;
 
 use avb::PUBLIC_KEY;
 use log::{debug, info};
 
-fn main(fdt: &mut [u8], payload: &[u8]) {
+fn main(fdt: &libfdt::Fdt, signed_kernel: &[u8], ramdisk: Option<&[u8]>, bcc: &[u8]) {
     info!("pVM firmware");
-    debug!(
-        "fdt_address={:#018x}, payload_start={:#018x}, payload_size={:#018x}",
-        fdt.as_ptr() as usize,
-        payload.as_ptr() as usize,
-        payload.len(),
-    );
+    debug!("FDT: {:?}", fdt as *const libfdt::Fdt);
+    debug!("Signed kernel: {:?} ({:#x} bytes)", signed_kernel.as_ptr(), signed_kernel.len());
+    if let Some(rd) = ramdisk {
+        debug!("Ramdisk: {:?} ({:#x} bytes)", rd.as_ptr(), rd.len());
+    } else {
+        debug!("Ramdisk: None");
+    }
+    debug!("BCC: {:?} ({:#x} bytes)", bcc.as_ptr(), bcc.len());
     debug!("AVB public key: addr={:?}, size={:#x} ({1})", PUBLIC_KEY.as_ptr(), PUBLIC_KEY.len());
     info!("Starting payload...");
 }
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
new file mode 100644
index 0000000..ca0b886
--- /dev/null
+++ b/pvmfw/src/memory.rs
@@ -0,0 +1,190 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Low-level allocation and tracking of main memory.
+
+use crate::helpers;
+use crate::mmu;
+use core::cmp::max;
+use core::cmp::min;
+use core::fmt;
+use core::num::NonZeroUsize;
+use core::ops::Range;
+use core::result;
+use log::error;
+use tinyvec::ArrayVec;
+
+type MemoryRange = Range<usize>;
+
+#[derive(Clone, Copy, Debug, Default)]
+enum MemoryType {
+    #[default]
+    ReadOnly,
+    ReadWrite,
+}
+
+#[derive(Clone, Debug, Default)]
+struct MemoryRegion {
+    range: MemoryRange,
+    mem_type: MemoryType,
+}
+
+impl MemoryRegion {
+    /// True if the instance overlaps with the passed range.
+    pub fn overlaps(&self, range: &MemoryRange) -> bool {
+        let our: &MemoryRange = self.as_ref();
+        max(our.start, range.start) < min(our.end, range.end)
+    }
+
+    /// True if the instance is fully contained within the passed range.
+    pub fn is_within(&self, range: &MemoryRange) -> bool {
+        let our: &MemoryRange = self.as_ref();
+        self.as_ref() == &(max(our.start, range.start)..min(our.end, range.end))
+    }
+}
+
+impl AsRef<MemoryRange> for MemoryRegion {
+    fn as_ref(&self) -> &MemoryRange {
+        &self.range
+    }
+}
+
+/// Tracks non-overlapping slices of main memory.
+pub struct MemoryTracker {
+    regions: ArrayVec<[MemoryRegion; MemoryTracker::CAPACITY]>,
+    total: MemoryRange,
+    page_table: mmu::PageTable,
+}
+
+/// Errors for MemoryTracker operations.
+#[derive(Debug, Clone)]
+pub enum MemoryTrackerError {
+    /// Tried to modify the memory base address.
+    DifferentBaseAddress,
+    /// Tried to shrink to a larger memory size.
+    SizeTooLarge,
+    /// Tracked regions would not fit in memory size.
+    SizeTooSmall,
+    /// Reached limit number of tracked regions.
+    Full,
+    /// Region is out of the tracked memory address space.
+    OutOfRange,
+    /// New region overlaps with tracked regions.
+    Overlaps,
+    /// Region couldn't be mapped.
+    FailedToMap,
+}
+
+impl fmt::Display for MemoryTrackerError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::DifferentBaseAddress => write!(f, "Received different base address"),
+            Self::SizeTooLarge => write!(f, "Tried to shrink to a larger memory size"),
+            Self::SizeTooSmall => write!(f, "Tracked regions would not fit in memory size"),
+            Self::Full => write!(f, "Reached limit number of tracked regions"),
+            Self::OutOfRange => write!(f, "Region is out of the tracked memory address space"),
+            Self::Overlaps => write!(f, "New region overlaps with tracked regions"),
+            Self::FailedToMap => write!(f, "Failed to map the new region"),
+        }
+    }
+}
+
+type Result<T> = result::Result<T, MemoryTrackerError>;
+
+impl MemoryTracker {
+    const CAPACITY: usize = 5;
+    /// Base of the system's contiguous "main" memory.
+    const BASE: usize = 0x8000_0000;
+    /// First address that can't be translated by a level 1 TTBR0_EL1.
+    const MAX_ADDR: usize = 1 << 39;
+
+    /// Create a new instance from an active page table, covering the maximum RAM size.
+    pub fn new(page_table: mmu::PageTable) -> Self {
+        Self { total: Self::BASE..Self::MAX_ADDR, page_table, regions: ArrayVec::new() }
+    }
+
+    /// Resize the total RAM size.
+    ///
+    /// This function fails if it contains regions that are not included within the new size.
+    pub fn shrink(&mut self, range: &MemoryRange) -> Result<()> {
+        if range.start != self.total.start {
+            return Err(MemoryTrackerError::DifferentBaseAddress);
+        }
+        if self.total.end < range.end {
+            return Err(MemoryTrackerError::SizeTooLarge);
+        }
+        if !self.regions.iter().all(|r| r.is_within(range)) {
+            return Err(MemoryTrackerError::SizeTooSmall);
+        }
+
+        self.total = range.clone();
+        Ok(())
+    }
+
+    /// Allocate the address range for a const slice; returns None if failed.
+    pub fn alloc_range(&mut self, range: &MemoryRange) -> Result<MemoryRange> {
+        self.page_table.map_rodata(range).map_err(|e| {
+            error!("Error during range allocation: {e}");
+            MemoryTrackerError::FailedToMap
+        })?;
+        self.add(MemoryRegion { range: range.clone(), mem_type: MemoryType::ReadOnly })
+    }
+
+    /// Allocate the address range for a mutable slice; returns None if failed.
+    pub fn alloc_range_mut(&mut self, range: &MemoryRange) -> Result<MemoryRange> {
+        self.page_table.map_data(range).map_err(|e| {
+            error!("Error during mutable range allocation: {e}");
+            MemoryTrackerError::FailedToMap
+        })?;
+        self.add(MemoryRegion { range: range.clone(), mem_type: MemoryType::ReadWrite })
+    }
+
+    /// Allocate the address range for a const slice; returns None if failed.
+    pub fn alloc(&mut self, base: usize, size: NonZeroUsize) -> Result<MemoryRange> {
+        self.alloc_range(&(base..(base + size.get())))
+    }
+
+    /// Allocate the address range for a mutable slice; returns None if failed.
+    pub fn alloc_mut(&mut self, base: usize, size: NonZeroUsize) -> Result<MemoryRange> {
+        self.alloc_range_mut(&(base..(base + size.get())))
+    }
+
+    fn add(&mut self, region: MemoryRegion) -> Result<MemoryRange> {
+        if !region.is_within(&self.total) {
+            return Err(MemoryTrackerError::OutOfRange);
+        }
+        if self.regions.iter().any(|r| r.overlaps(region.as_ref())) {
+            return Err(MemoryTrackerError::Overlaps);
+        }
+        if self.regions.try_push(region).is_some() {
+            return Err(MemoryTrackerError::Full);
+        }
+
+        Ok(self.regions.last().unwrap().as_ref().clone())
+    }
+}
+
+impl Drop for MemoryTracker {
+    fn drop(&mut self) {
+        for region in self.regions.iter() {
+            match region.mem_type {
+                MemoryType::ReadWrite => {
+                    // TODO: Use page table's dirty bit to only flush pages that were touched.
+                    helpers::flush_region(region.range.start, region.range.len())
+                }
+                MemoryType::ReadOnly => {}
+            }
+        }
+    }
+}
diff --git a/pvmfw/src/mmu.rs b/pvmfw/src/mmu.rs
new file mode 100644
index 0000000..fa94e85
--- /dev/null
+++ b/pvmfw/src/mmu.rs
@@ -0,0 +1,86 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Memory management.
+
+use crate::helpers;
+use aarch64_paging::idmap::IdMap;
+use aarch64_paging::paging::Attributes;
+use aarch64_paging::paging::MemoryRegion;
+use aarch64_paging::MapError;
+use core::ops::Range;
+use vmbase::layout;
+
+// We assume that:
+// - MAIR_EL1.Attr0 = "Device-nGnRE memory" (0b0000_0100)
+// - MAIR_EL1.Attr1 = "Normal memory, Outer & Inner WB Non-transient, R/W-Allocate" (0b1111_1111)
+const MEMORY: Attributes = Attributes::NORMAL.union(Attributes::NON_GLOBAL);
+const DEVICE: Attributes = Attributes::DEVICE_NGNRE.union(Attributes::EXECUTE_NEVER);
+const CODE: Attributes = MEMORY.union(Attributes::READ_ONLY);
+const DATA: Attributes = MEMORY.union(Attributes::EXECUTE_NEVER);
+const RODATA: Attributes = DATA.union(Attributes::READ_ONLY);
+
+/// High-level API for managing MMU mappings.
+pub struct PageTable {
+    idmap: IdMap,
+}
+
+fn appended_payload_range() -> Range<usize> {
+    let start = helpers::align_up(layout::binary_end(), helpers::SIZE_4KB).unwrap();
+    // pvmfw is contained in a 2MiB region so the payload can't be larger than the 2MiB alignment.
+    let end = helpers::align_up(start, helpers::SIZE_2MB).unwrap();
+
+    start..end
+}
+
+impl PageTable {
+    const ASID: usize = 1;
+    const ROOT_LEVEL: usize = 1;
+
+    /// Creates an instance pre-populated with pvmfw's binary layout.
+    pub fn from_static_layout() -> Result<Self, MapError> {
+        let mut page_table = Self { idmap: IdMap::new(Self::ASID, Self::ROOT_LEVEL) };
+
+        page_table.map_code(&layout::text_range())?;
+        page_table.map_data(&layout::writable_region())?;
+        page_table.map_rodata(&layout::rodata_range())?;
+        page_table.map_data(&appended_payload_range())?;
+
+        Ok(page_table)
+    }
+
+    pub unsafe fn activate(&mut self) {
+        self.idmap.activate()
+    }
+
+    pub fn map_device(&mut self, range: &Range<usize>) -> Result<(), MapError> {
+        self.map_range(range, DEVICE)
+    }
+
+    pub fn map_data(&mut self, range: &Range<usize>) -> Result<(), MapError> {
+        self.map_range(range, DATA)
+    }
+
+    pub fn map_code(&mut self, range: &Range<usize>) -> Result<(), MapError> {
+        self.map_range(range, CODE)
+    }
+
+    pub fn map_rodata(&mut self, range: &Range<usize>) -> Result<(), MapError> {
+        self.map_range(range, RODATA)
+    }
+
+    fn map_range(&mut self, range: &Range<usize>, attr: Attributes) -> Result<(), MapError> {
+        self.idmap.map_range(&MemoryRegion::new(range.start, range.end), attr)
+    }
+}
diff --git a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
index 260f804..16e4893 100644
--- a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
@@ -25,7 +25,7 @@
      *
      * @return The read rate in MB/s.
      */
-    double measureReadRate(String filename, long fileSizeBytes, boolean isRand);
+    double measureReadRate(String filename, boolean isRand);
 
     /** Returns an entry from /proc/meminfo. */
     long getMemInfoEntry(String name);
diff --git a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
index e8c435f..eda4f75 100644
--- a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
@@ -36,4 +36,7 @@
 
     /* get the APK contents path. */
     String getApkContentsPath();
+
+    /* get the encrypted storage path. */
+    String getEncryptedStoragePath();
 }
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
index e6f39f8..9d2b6c7 100644
--- a/tests/benchmark/Android.bp
+++ b/tests/benchmark/Android.bp
@@ -16,46 +16,28 @@
         "com.android.microdroid.testservice-java",
         "truth-prebuilt",
     ],
-    libs: ["android.system.virtualmachine"],
     jni_libs: [
         "MicrodroidBenchmarkNativeLib",
         "MicrodroidIdleNativeLib",
         "libiovsock_host_jni",
     ],
-    platform_apis: true,
+    jni_uses_platform_apis: true,
+    sdk_version: "test_current",
     use_embedded_native_libs: true,
     compile_multilib: "64",
 }
 
 cc_library_shared {
-    name: "MicrodroidIdleNativeLib",
-    srcs: ["src/native/idlebinary.cpp"],
-    header_libs: ["vm_payload_headers"],
-    shared_libs: [
-        "libbase",
-    ],
-}
-
-cc_library_shared {
     name: "MicrodroidBenchmarkNativeLib",
-    srcs: ["src/native/benchmarkbinary.cpp"],
+    srcs: ["src/native/*.cpp"],
+    local_include_dirs: ["src/native/include"],
     static_libs: [
         "com.android.microdroid.testservice-ndk",
-        "libiobenchmark",
     ],
     shared_libs: [
         "libbase",
         "libbinder_ndk",
         "liblog",
-        "libvm_payload",
-    ],
-}
-
-cc_library {
-    name: "libiobenchmark",
-    srcs: ["src/native/io_vsock.cpp"],
-    export_include_dirs: ["src/native/include"],
-    shared_libs: [
-        "libbase",
+        "libvm_payload#current",
     ],
 }
diff --git a/tests/benchmark/AndroidManifest.xml b/tests/benchmark/AndroidManifest.xml
index c39b91c..8a7366a 100644
--- a/tests/benchmark/AndroidManifest.xml
+++ b/tests/benchmark/AndroidManifest.xml
@@ -18,7 +18,6 @@
     <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
     <uses-permission android:name="android.permission.USE_CUSTOM_VIRTUAL_MACHINE" />
     <application>
-        <uses-library android:name="android.system.virtualmachine" android:required="false" />
     </application>
     <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
         android:targetPackage="com.android.microdroid.benchmark"
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index db74358..14a0e39 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -88,16 +88,17 @@
 
     private boolean canBootMicrodroidWithMemory(int mem)
             throws VirtualMachineException, InterruptedException, IOException {
-        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidIdleNativeLib.so")
-                .setDebugLevel(DEBUG_LEVEL_NONE)
-                .setMemoryMib(mem)
-                .build();
+        VirtualMachineConfig normalConfig =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidIdleNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .setMemoryMib(mem)
+                        .build();
 
         // returns true if succeeded at least once.
         final int trialCount = 5;
         for (int i = 0; i < trialCount; i++) {
-            mInner.forceCreateNewVirtualMachine("test_vm_minimum_memory", normalConfig);
+            forceCreateNewVirtualMachine("test_vm_minimum_memory", normalConfig);
 
             if (tryBootVm(TAG, "test_vm_minimum_memory").payloadStarted) return true;
         }
@@ -147,12 +148,13 @@
         for (int i = 0; i < trialCount; i++) {
 
             // To grab boot events from log, set debug mode to FULL
-            VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                    .setPayloadBinaryPath("MicrodroidIdleNativeLib.so")
-                    .setDebugLevel(DEBUG_LEVEL_FULL)
-                    .setMemoryMib(256)
-                    .build();
-            mInner.forceCreateNewVirtualMachine("test_vm_boot_time", normalConfig);
+            VirtualMachineConfig normalConfig =
+                    newVmConfigBuilder()
+                            .setPayloadBinaryPath("MicrodroidIdleNativeLib.so")
+                            .setDebugLevel(DEBUG_LEVEL_FULL)
+                            .setMemoryMib(256)
+                            .build();
+            forceCreateNewVirtualMachine("test_vm_boot_time", normalConfig);
 
             BootResult result = tryBootVm(TAG, "test_vm_boot_time");
             assertThat(result.payloadStarted).isTrue();
@@ -195,17 +197,17 @@
 
     @Test
     public void testVsockTransferFromHostToVM() throws Exception {
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config_io.json")
-                .setDebugLevel(DEBUG_LEVEL_FULL)
-                .build();
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config_io.json")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
         List<Double> transferRates = new ArrayList<>(IO_TEST_TRIAL_COUNT);
 
         for (int i = 0; i < IO_TEST_TRIAL_COUNT; ++i) {
             int port = (mProtectedVm ? 5666 : 6666) + i;
             String vmName = "test_vm_io_" + i;
-            mInner.forceCreateNewVirtualMachine(vmName, config);
-            VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
+            VirtualMachine vm = forceCreateNewVirtualMachine(vmName, config);
             BenchmarkVmListener.create(new VsockListener(transferRates, port)).runToFinish(TAG, vm);
         }
         reportMetrics(transferRates, "vsock/transfer_host_to_vm", "mb_per_sec");
@@ -222,10 +224,11 @@
     }
 
     private void testVirtioBlkReadRate(boolean isRand) throws Exception {
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config_io.json")
-                .setDebugLevel(DEBUG_LEVEL_FULL)
-                .build();
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config_io.json")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
         List<Double> readRates = new ArrayList<>(IO_TEST_TRIAL_COUNT);
 
         for (int i = 0; i < IO_TEST_TRIAL_COUNT + 1; ++i) {
@@ -236,8 +239,7 @@
                 readRates.clear();
             }
             String vmName = "test_vm_io_" + i;
-            mInner.forceCreateNewVirtualMachine(vmName, config);
-            VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
+            VirtualMachine vm = forceCreateNewVirtualMachine(vmName, config);
             BenchmarkVmListener.create(new VirtioBlkListener(readRates, isRand))
                     .runToFinish(TAG, vm);
         }
@@ -257,18 +259,10 @@
     private static class VirtioBlkListener implements BenchmarkVmListener.InnerListener {
         private static final String FILENAME = APEX_ETC_FS + "microdroid_super.img";
 
-        private final long mFileSizeBytes;
         private final List<Double> mReadRates;
         private final boolean mIsRand;
 
         VirtioBlkListener(List<Double> readRates, boolean isRand) {
-            File file = new File(FILENAME);
-            try {
-                mFileSizeBytes = Files.size(file.toPath());
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
-            assertThat(mFileSizeBytes).isGreaterThan((long) SIZE_MB);
             mReadRates = readRates;
             mIsRand = isRand;
         }
@@ -276,7 +270,7 @@
         @Override
         public void onPayloadReady(VirtualMachine vm, IBenchmarkService benchmarkService)
                 throws RemoteException {
-            double readRate = benchmarkService.measureReadRate(FILENAME, mFileSizeBytes, mIsRand);
+            double readRate = benchmarkService.measureReadRate(FILENAME, mIsRand);
             mReadRates.add(readRate);
         }
     }
@@ -288,13 +282,13 @@
     @Test
     public void testMemoryUsage() throws Exception {
         final String vmName = "test_vm_mem_usage";
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config_io.json")
-                .setDebugLevel(DEBUG_LEVEL_NONE)
-                .setMemoryMib(256)
-                .build();
-        mInner.forceCreateNewVirtualMachine(vmName, config);
-        VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config_io.json")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .setMemoryMib(256)
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine(vmName, config);
         MemoryUsageListener listener = new MemoryUsageListener(this::executeCommand);
         BenchmarkVmListener.create(listener).runToFinish(TAG, vm);
 
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 18293f3..70ec7db 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -23,6 +23,8 @@
 #include <fcntl.h>
 #include <linux/vm_sockets.h>
 #include <stdio.h>
+#include <sys/stat.h>
+#include <sys/types.h>
 #include <time.h>
 #include <unistd.h>
 #include <vm_main.h>
@@ -56,9 +58,9 @@
 
 class IOBenchmarkService : public aidl::com::android::microdroid::testservice::BnBenchmarkService {
 public:
-    ndk::ScopedAStatus measureReadRate(const std::string& filename, int64_t fileSizeBytes,
-                                       bool isRand, double* out) override {
-        auto res = measure_read_rate(filename, fileSizeBytes, isRand);
+    ndk::ScopedAStatus measureReadRate(const std::string& filename, bool isRand,
+                                       double* out) override {
+        auto res = measure_read_rate(filename, isRand);
         if (res.ok()) {
             *out = res.value();
         }
@@ -90,15 +92,23 @@
     }
 
 private:
-    /** Measures the read rate for reading the given file. */
-    Result<double> measure_read_rate(const std::string& filename, int64_t fileSizeBytes,
-                                     bool is_rand) {
-        const int64_t block_count = fileSizeBytes / kBlockSizeBytes;
-        std::vector<uint64_t> offsets;
+    /**
+     * Measures the read rate for reading the given file.
+     * @return The read rate in MB/s.
+     */
+    Result<double> measure_read_rate(const std::string& filename, bool is_rand) {
+        struct stat file_stats;
+        if (stat(filename.c_str(), &file_stats) == -1) {
+            return Error() << "failed to get file stats";
+        }
+        const int64_t file_size_bytes = file_stats.st_size;
+        const int64_t block_count = file_size_bytes / kBlockSizeBytes;
+        std::vector<uint64_t> offsets(block_count);
+        for (auto i = 0; i < block_count; ++i) {
+            offsets[i] = i * kBlockSizeBytes;
+        }
         if (is_rand) {
             std::mt19937 rd{std::random_device{}()};
-            offsets.reserve(block_count);
-            for (auto i = 0; i < block_count; ++i) offsets.push_back(i * kBlockSizeBytes);
             std::shuffle(offsets.begin(), offsets.end(), rd);
         }
         char buf[kBlockSizeBytes];
@@ -109,12 +119,7 @@
             return ErrnoError() << "Read: opening " << filename << " failed";
         }
         for (auto i = 0; i < block_count; ++i) {
-            if (is_rand) {
-                if (lseek(fd.get(), offsets[i], SEEK_SET) == -1) {
-                    return ErrnoError() << "failed to lseek";
-                }
-            }
-            auto bytes = read(fd.get(), buf, kBlockSizeBytes);
+            auto bytes = pread(fd, buf, kBlockSizeBytes, offsets[i]);
             if (bytes == 0) {
                 return Error() << "unexpected end of file";
             } else if (bytes == -1) {
@@ -122,8 +127,8 @@
             }
         }
         double elapsed_seconds = ((double)clock() - start) / CLOCKS_PER_SEC;
-        double read_rate = (double)fileSizeBytes / kNumBytesPerMB / elapsed_seconds;
-        return {read_rate};
+        double file_size_mb = (double)file_size_bytes / kNumBytesPerMB;
+        return {file_size_mb / elapsed_seconds};
     }
 
     Result<size_t> read_meminfo_entry(const std::string& stat) {
@@ -155,16 +160,9 @@
 
 Result<void> run_io_benchmark_tests() {
     auto test_service = ndk::SharedRefBase::make<IOBenchmarkService>();
-    auto callback = []([[maybe_unused]] void* param) {
-        if (!AVmPayload_notifyPayloadReady()) {
-            LOG(ERROR) << "failed to notify payload ready to virtualizationservice";
-            abort();
-        }
-    };
-    if (!AVmPayload_runVsockRpcServer(test_service->asBinder().get(), test_service->SERVICE_PORT,
-                                      callback, nullptr)) {
-        return Error() << "RPC Server failed to run";
-    }
+    auto callback = []([[maybe_unused]] void* param) { AVmPayload_notifyPayloadReady(); };
+    AVmPayload_runVsockRpcServer(test_service->asBinder().get(), test_service->SERVICE_PORT,
+                                 callback, nullptr);
     return {};
 }
 } // Anonymous namespace
diff --git a/tests/benchmark_hostside/Android.bp b/tests/benchmark_hostside/Android.bp
index c5c7571..e053406 100644
--- a/tests/benchmark_hostside/Android.bp
+++ b/tests/benchmark_hostside/Android.bp
@@ -17,4 +17,7 @@
     test_suites: [
         "general-tests",
     ],
+    data: [
+        ":MicrodroidTestApp",
+    ],
 }
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index 5e57b32..c47e915 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -16,6 +16,7 @@
 
 package android.avf.test;
 
+import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 
 import static com.google.common.truth.Truth.assertWithMessage;
@@ -31,6 +32,8 @@
 import com.android.microdroid.test.host.CommandRunner;
 import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.util.CommandResult;
@@ -45,7 +48,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -85,6 +87,9 @@
     @Before
     public void setUp() throws Exception {
         testIfDeviceIsCapable(getDevice());
+
+        getDevice().installPackage(findTestFile(APK_NAME), /* reinstall */ false);
+
         mMetricsProcessor = new MetricsProcessor(getMetricPrefix() + "hostside/");
     }
 
@@ -144,7 +149,9 @@
     }
 
     private void appStartupHelper(String launchIntentPackage) throws Exception {
-        assumeTrue("Skip on non-protected VMs", isProtectedVmSupported());
+        assumeTrue(
+                "Skip on non-protected VMs",
+                ((TestDevice) getDevice()).supportsMicrodroid(/*protectedVm=*/ true));
 
         StartupTimeMetricCollection mCollection =
                 new StartupTimeMetricCollection(getPackageName(launchIntentPackage), ROUND_COUNT);
@@ -196,10 +203,6 @@
         return null;
     }
 
-    private void microdroidWaitForBootComplete() {
-        runOnMicrodroidForResult("watch -e \"getprop dev.bootcomplete | grep '^0$'\"");
-    }
-
     private AmStartupTimeCmdParser getColdRunStartupTimes(CommandRunner android, String pkgName)
             throws DeviceNotAvailableException, InterruptedException {
         unlockScreen(android);
@@ -216,9 +219,7 @@
     // and the time measured after running the VM.
     private void getAppStartupTime(String pkgName, StartupTimeMetricCollection metricColector)
             throws Exception {
-        final String configPath = "assets/vm_config.json";
-        final String cid;
-        final int vm_mem_mb;
+        TestDevice device = (TestDevice) getDevice();
 
         // 1. Reboot the device to run the test without stage2 fragmentation
         getDevice().rebootUntilOnline();
@@ -240,30 +241,30 @@
         android.tryRun("rm", "-rf", MicrodroidHostTestCaseBase.TEST_ROOT);
 
         // Donate 80% of the available device memory to the VM
-        vm_mem_mb = getFreeMemoryInfoMb(android) * 80 / 100;
-        cid = startMicrodroid(
-                            getDevice(),
-                            getBuild(),
-                            APK_NAME,
-                            PACKAGE_NAME,
-                            configPath,
-                            true,
-                            vm_mem_mb,
-                            Optional.of(NUM_VCPUS));
-        adbConnectToMicrodroid(getDevice(), cid);
-        microdroidWaitForBootComplete();
+        final String configPath = "assets/vm_config.json";
+        final int vm_mem_mb = getFreeMemoryInfoMb(android) * 80 / 100;
+        ITestDevice microdroidDevice =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(vm_mem_mb)
+                        .numCpus(NUM_VCPUS)
+                        .build(device);
+        microdroidDevice.waitForBootComplete(30000);
+        microdroidDevice.enableAdbRoot();
 
-        rootMicrodroid();
+        CommandRunner microdroid = new CommandRunner(microdroidDevice);
 
-        runOnMicrodroid("mkdir -p /mnt/ramdisk && chmod 777 /mnt/ramdisk");
-        runOnMicrodroid("mount -t tmpfs -o size=32G tmpfs /mnt/ramdisk");
+        microdroid.run("mkdir -p /mnt/ramdisk && chmod 777 /mnt/ramdisk");
+        microdroid.run("mount -t tmpfs -o size=32G tmpfs /mnt/ramdisk");
 
         // Allocate memory for the VM until it fails and make sure that we touch
         // the allocated memory in the guest to be able to create stage2 fragmentation.
         try {
-            runOnMicrodroidForResult(String.format("cd /mnt/ramdisk && truncate -s %dM sprayMemory"
-                    + " && dd if=/dev/zero of=sprayMemory bs=1MB count=%d",
-                    vm_mem_mb , vm_mem_mb));
+            microdroid.tryRun(
+                    String.format(
+                            "cd /mnt/ramdisk && truncate -s %dM sprayMemory"
+                                    + " && dd if=/dev/zero of=sprayMemory bs=1MB count=%d",
+                            vm_mem_mb, vm_mem_mb));
         } catch (Exception ex) {
         }
 
@@ -273,7 +274,7 @@
             metricColector.addStartupTimeMetricDuringVmRun(duringVmStartApp);
         }
 
-        shutdownMicrodroid(getDevice(), cid);
+        device.shutdownMicrodroid(microdroidDevice);
 
         // Run the app after the VM run and collect cold startup time.
         for (int i = 0; i < ROUND_COUNT; i++) {
diff --git a/tests/helper/Android.bp b/tests/helper/Android.bp
index 60d4be1..61c5dcd 100644
--- a/tests/helper/Android.bp
+++ b/tests/helper/Android.bp
@@ -3,15 +3,10 @@
 }
 
 java_library_static {
-    name: "VirtualizationTestHelper",
-    srcs: ["src/java/com/android/virt/**/*.java"],
-    host_supported: true,
-}
-
-java_library_static {
     name: "MicrodroidTestHelper",
     srcs: ["src/java/com/android/microdroid/test/common/*.java"],
     host_supported: true,
+    sdk_version: "system_current",
 }
 
 java_library_static {
@@ -21,8 +16,7 @@
         "androidx.test.runner",
         "androidx.test.ext.junit",
         "MicrodroidTestHelper",
-        "VirtualizationTestHelper",
         "truth-prebuilt",
     ],
-    libs: ["android.system.virtualmachine"],
+    sdk_version: "system_current",
 }
diff --git a/tests/helper/src/java/com/android/microdroid/test/common/DeviceProperties.java b/tests/helper/src/java/com/android/microdroid/test/common/DeviceProperties.java
new file mode 100644
index 0000000..94f7e99
--- /dev/null
+++ b/tests/helper/src/java/com/android/microdroid/test/common/DeviceProperties.java
@@ -0,0 +1,63 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.microdroid.test.common;
+
+import static java.util.Objects.requireNonNull;
+
+/** This class can be used in both host tests and device tests to get the device properties. */
+public final class DeviceProperties {
+    /** PropertyGetter is used to get the property associated to a given key. */
+    public interface PropertyGetter {
+        String getProperty(String key) throws Exception;
+    }
+
+    private static final String KEY_VENDOR_DEVICE = "ro.product.vendor.device";
+    private static final String KEY_METRICS_TAG = "debug.hypervisor.metrics_tag";
+
+    private static final String CUTTLEFISH_DEVICE_PREFIX = "vsoc_";
+
+    private final PropertyGetter mPropertyGetter;
+
+    private DeviceProperties(PropertyGetter propertyGetter) {
+        mPropertyGetter = requireNonNull(propertyGetter);
+    }
+
+    /** Creates a new instance of {@link DeviceProperties}. */
+    public static DeviceProperties create(PropertyGetter propertyGetter) {
+        return new DeviceProperties(propertyGetter);
+    }
+
+    /**
+     * @return whether the device is a cuttlefish device.
+     */
+    public boolean isCuttlefish() {
+        String vendorDeviceName = getProperty(KEY_VENDOR_DEVICE);
+        return vendorDeviceName != null && vendorDeviceName.startsWith(CUTTLEFISH_DEVICE_PREFIX);
+    }
+
+    public String getMetricsTag() {
+        return getProperty(KEY_METRICS_TAG);
+    }
+
+    private String getProperty(String key) {
+        try {
+            return mPropertyGetter.getProperty(key);
+        } catch (Exception e) {
+            throw new IllegalArgumentException("Cannot get property for the key: " + key, e);
+        }
+    }
+}
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index 66cd211..bd5b180 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -15,9 +15,9 @@
  */
 package com.android.microdroid.test.device;
 
-import static com.google.common.truth.TruthJUnit.assume;
+import static android.content.pm.PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK;
 
-import static org.junit.Assume.assumeNoException;
+import static com.google.common.truth.TruthJUnit.assume;
 
 import android.app.Instrumentation;
 import android.app.UiAutomation;
@@ -35,8 +35,8 @@
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import com.android.microdroid.test.common.DeviceProperties;
 import com.android.microdroid.test.common.MetricsProcessor;
-import com.android.virt.VirtualizationTestHelper;
 
 import java.io.BufferedReader;
 import java.io.ByteArrayOutputStream;
@@ -51,13 +51,12 @@
 
 public abstract class MicrodroidDeviceTestBase {
     public static boolean isCuttlefish() {
-        return VirtualizationTestHelper.isCuttlefish(
-                SystemProperties.get("ro.product.vendor.device"));
+        return DeviceProperties.create(SystemProperties::get).isCuttlefish();
     }
 
     public static String getMetricPrefix() {
         return MetricsProcessor.getMetricPrefix(
-                SystemProperties.get("debug.hypervisor.metrics_tag"));
+                DeviceProperties.create(SystemProperties::get).getMetricsTag());
     }
 
     protected final void grantPermission(String permission) {
@@ -74,64 +73,48 @@
                 permission);
     }
 
-    // TODO(b/220920264): remove Inner class; this is a hack to hide virt APEX types
-    protected static class Inner {
-        private final boolean mProtectedVm;
-        private final Context mContext;
-        private final VirtualMachineManager mVmm;
-
-        public Inner(Context context, boolean protectedVm, VirtualMachineManager vmm) {
-            mProtectedVm = protectedVm;
-            mVmm = vmm;
-            mContext = context;
-        }
-
-        public VirtualMachineManager getVirtualMachineManager() {
-            return mVmm;
-        }
-
-        public Context getContext() {
-            return mContext;
-        }
-
-        public VirtualMachineConfig.Builder newVmConfigBuilder() {
-            return new VirtualMachineConfig.Builder(mContext).setProtectedVm(mProtectedVm);
-        }
-
-        /**
-         * Creates a new virtual machine, potentially removing an existing virtual machine with
-         * given name.
-         */
-        public VirtualMachine forceCreateNewVirtualMachine(String name, VirtualMachineConfig config)
-                throws VirtualMachineException {
-            VirtualMachine existingVm = mVmm.get(name);
-            if (existingVm != null) {
-                existingVm.delete();
-            }
-            return mVmm.create(name, config);
-        }
-    }
-
-    protected Inner mInner;
+    private Context mCtx;
+    private boolean mProtectedVm;
 
     protected Context getContext() {
-        return mInner.getContext();
+        return mCtx;
+    }
+
+    public VirtualMachineManager getVirtualMachineManager() {
+        return mCtx.getSystemService(VirtualMachineManager.class);
+    }
+
+    public VirtualMachineConfig.Builder newVmConfigBuilder() {
+        return new VirtualMachineConfig.Builder(mCtx).setProtectedVm(mProtectedVm);
+    }
+
+    protected final boolean isProtectedVm() {
+        return mProtectedVm;
+    }
+
+    /**
+     * Creates a new virtual machine, potentially removing an existing virtual machine with given
+     * name.
+     */
+    public VirtualMachine forceCreateNewVirtualMachine(String name, VirtualMachineConfig config)
+            throws VirtualMachineException {
+        final VirtualMachineManager vmm = getVirtualMachineManager();
+        VirtualMachine existingVm = vmm.get(name);
+        if (existingVm != null) {
+            vmm.delete(name);
+        }
+        return vmm.create(name, config);
     }
 
     public void prepareTestSetup(boolean protectedVm) {
-        // In case when the virt APEX doesn't exist on the device, classes in the
-        // android.system.virtualmachine package can't be loaded. Therefore, before using the
-        // classes, check the existence of a class in the package and skip this test if not exist.
-        try {
-            Class.forName("android.system.virtualmachine.VirtualMachineManager");
-        } catch (ClassNotFoundException e) {
-            assumeNoException(e);
-            return;
-        }
-        Context context = ApplicationProvider.getApplicationContext();
-        mInner = new Inner(context, protectedVm, VirtualMachineManager.getInstance(context));
+        mCtx = ApplicationProvider.getApplicationContext();
+        assume().withMessage("Device doesn't support AVF")
+                .that(mCtx.getPackageManager().hasSystemFeature(FEATURE_VIRTUALIZATION_FRAMEWORK))
+                .isTrue();
 
-        int capabilities = mInner.getVirtualMachineManager().getCapabilities();
+        mProtectedVm = protectedVm;
+
+        int capabilities = getVirtualMachineManager().getCapabilities();
         if (protectedVm) {
             assume().withMessage("Skip where protected VMs aren't supported")
                     .that(capabilities & VirtualMachineManager.CAPABILITY_PROTECTED_VM)
@@ -202,8 +185,8 @@
                 throws VirtualMachineException, InterruptedException {
             vm.setCallback(mExecutorService, this);
             vm.run();
-            logVmOutputAndMonitorBootEvents(logTag, vm.getConsoleOutputStream(), "Console");
-            logVmOutput(logTag, vm.getLogOutputStream(), "Log");
+            logVmOutputAndMonitorBootEvents(logTag, vm.getConsoleOutput(), "Console");
+            logVmOutput(logTag, vm.getLogOutput(), "Log");
             mExecutorService.awaitTermination(300, TimeUnit.SECONDS);
         }
 
@@ -232,7 +215,7 @@
         }
 
         @Override
-        public void onPayloadStarted(VirtualMachine vm, ParcelFileDescriptor stream) {}
+        public void onPayloadStarted(VirtualMachine vm) {}
 
         @Override
         public void onPayloadReady(VirtualMachine vm) {}
@@ -320,14 +303,14 @@
 
     public BootResult tryBootVm(String logTag, String vmName)
             throws VirtualMachineException, InterruptedException {
-        VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
+        VirtualMachine vm = getVirtualMachineManager().get(vmName);
         final CompletableFuture<Boolean> payloadStarted = new CompletableFuture<>();
         final CompletableFuture<Integer> deathReason = new CompletableFuture<>();
         final CompletableFuture<Long> endTime = new CompletableFuture<>();
         VmEventListener listener =
                 new VmEventListener() {
                     @Override
-                    public void onPayloadStarted(VirtualMachine vm, ParcelFileDescriptor stream) {
+                    public void onPayloadStarted(VirtualMachine vm) {
                         endTime.complete(System.nanoTime());
                         payloadStarted.complete(true);
                         forceStop(vm);
diff --git a/tests/hostside/helper/Android.bp b/tests/hostside/helper/Android.bp
index b2333ab..6196ec5 100644
--- a/tests/hostside/helper/Android.bp
+++ b/tests/hostside/helper/Android.bp
@@ -12,6 +12,5 @@
     ],
     static_libs: [
         "MicrodroidTestHelper",
-        "VirtualizationTestHelper",
     ],
 }
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
index 1a1dbeb..8816dbd 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
@@ -26,38 +26,31 @@
 import static org.junit.Assume.assumeTrue;
 
 import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
+import com.android.microdroid.test.common.DeviceProperties;
 import com.android.microdroid.test.common.MetricsProcessor;
 import com.android.tradefed.build.IBuildInfo;
+import com.android.tradefed.config.Option;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.junit4.BaseHostJUnit4Test;
 import com.android.tradefed.util.CommandResult;
 import com.android.tradefed.util.RunUtil;
-import com.android.virt.VirtualizationTestHelper;
 
 import java.io.File;
 import java.io.FileNotFoundException;
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Optional;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 
 public abstract class MicrodroidHostTestCaseBase extends BaseHostJUnit4Test {
+
     protected static final String TEST_ROOT = "/data/local/tmp/virt/";
-    protected static final String VIRT_APEX = "/apex/com.android.virt/";
     protected static final String LOG_PATH = TEST_ROOT + "log.txt";
     protected static final String CONSOLE_PATH = TEST_ROOT + "console.txt";
     private static final int TEST_VM_ADB_PORT = 8000;
     private static final String MICRODROID_SERIAL = "localhost:" + TEST_VM_ADB_PORT;
     private static final String INSTANCE_IMG = "instance.img";
-
-    // This is really slow on GCE (2m 40s) but fast on localhost or actual Android phones (< 10s).
-    // Then there is time to run the actual task. Set the maximum timeout value big enough.
-    private static final long MICRODROID_MAX_LIFETIME_MINUTES = 20;
+    private static final String PVMFW_IMG_PATH = TEST_ROOT + "pvmfw.img";
+    private static final String PVMFW_IMG_PATH_PROP = "hypervisor.pvmfw.path";
 
     private static final long MICRODROID_ADB_CONNECT_TIMEOUT_MINUTES = 5;
     protected static final long MICRODROID_COMMAND_TIMEOUT_MILLIS = 30000;
@@ -66,6 +59,19 @@
             (int) (MICRODROID_ADB_CONNECT_TIMEOUT_MINUTES * 60 * 1000
                 / MICRODROID_COMMAND_RETRY_INTERVAL_MILLIS);
 
+    @Option(
+            name = "pvmfw",
+            description =
+                    "Custom pvmfw.img path on host device."
+                            + " If present, it will be pushed to "
+                            + PVMFW_IMG_PATH,
+            mandatory = false)
+    private static String sCustomPvmfwPathOnHost = "";
+
+    private static boolean isEmptyText(String str) {
+        return str == null || str.length() == 0;
+    }
+
     public static void prepareVirtualizationTestSetup(ITestDevice androidDevice)
             throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(androidDevice);
@@ -78,6 +84,13 @@
 
         // remove any leftover files under test root
         android.tryRun("rm", "-rf", TEST_ROOT + "*");
+
+        // prepare custom pvmfw.img if necessary
+        if (!isEmptyText(sCustomPvmfwPathOnHost)) {
+            runOnHost("adb", "root");
+            runOnHost("adb", "push", sCustomPvmfwPathOnHost, PVMFW_IMG_PATH);
+            runOnHost("adb", "shell", "setprop", PVMFW_IMG_PATH_PROP, PVMFW_IMG_PATH);
+        }
     }
 
     public static void cleanUpVirtualizationTestSetup(ITestDevice androidDevice)
@@ -91,21 +104,27 @@
         android.tryRun("killall", "crosvm");
         android.tryRun("stop", "virtualizationservice");
         android.tryRun("rm", "-rf", "/data/misc/virtualizationservice/*");
+
+        if (!isEmptyText(sCustomPvmfwPathOnHost)) {
+            runOnHost("adb", "shell", "setprop", PVMFW_IMG_PATH_PROP, "\"\"");
+        }
     }
 
-    protected boolean isCuttlefish() throws Exception {
-        return VirtualizationTestHelper.isCuttlefish(
-            getDevice().getProperty("ro.product.vendor.device"));
+    protected boolean isCuttlefish() {
+        return DeviceProperties.create(getDevice()::getProperty).isCuttlefish();
     }
 
-    protected String getMetricPrefix() throws Exception {
+    protected String getMetricPrefix() {
         return MetricsProcessor.getMetricPrefix(
-                getDevice().getProperty("debug.hypervisor.metrics_tag"));
+                DeviceProperties.create(getDevice()::getProperty).getMetricsTag());
     }
 
     public static void testIfDeviceIsCapable(ITestDevice androidDevice) throws Exception {
         assumeTrue("Need an actual TestDevice", androidDevice instanceof TestDevice);
         TestDevice testDevice = (TestDevice) androidDevice;
+        assumeTrue(
+                "Requires VM support",
+                testDevice.hasFeature("android.software.virtualization_framework"));
         assumeTrue("Requires VM support", testDevice.supportsMicrodroid());
     }
 
@@ -160,53 +179,12 @@
         return result.getStdout().trim();
     }
 
-    // Same as runOnMicrodroid, but keeps retrying on error for maximum attempts times
-    // Each attempt with timeoutMs
-    public static String runOnMicrodroidRetryingOnFailure(
-            long timeoutMs, int attempts, String... cmd) {
-        CommandResult result = RunUtil.getDefault()
-                .runTimedCmdRetry(timeoutMs, MICRODROID_COMMAND_RETRY_INTERVAL_MILLIS, attempts,
-                        "adb", "-s", MICRODROID_SERIAL, "shell", join(cmd));
-        assertWithMessage("Command `" + Arrays.toString(cmd) + "` has failed")
-                .about(command_results())
-                .that(result)
-                .isSuccess();
-        return result.getStdout().trim();
-    }
-
     public static CommandResult runOnMicrodroidForResult(String... cmd) {
         final long timeoutMs = 30000; // 30 sec. Microdroid is extremely slow on GCE-on-CF.
         return RunUtil.getDefault()
                 .runTimedCmd(timeoutMs, "adb", "-s", MICRODROID_SERIAL, "shell", join(cmd));
     }
 
-    public static void pullMicrodroidFile(String path, File target) {
-        final long timeoutMs = 30000; // 30 sec. Microdroid is extremely slow on GCE-on-CF.
-        CommandResult result =
-                RunUtil.getDefault()
-                        .runTimedCmd(
-                                timeoutMs,
-                                "adb",
-                                "-s",
-                                MICRODROID_SERIAL,
-                                "pull",
-                                path,
-                                target.getPath());
-        assertWithMessage("pulling " + path + " from microdroid")
-                .about(command_results())
-                .that(result)
-                .isSuccess();
-    }
-
-    // Asserts the command will fail on Microdroid.
-    public static void assertFailedOnMicrodroid(String... cmd) {
-        CommandResult result = runOnMicrodroidForResult(cmd);
-        assertWithMessage("Microdroid command `" + join(cmd) + "` did not fail expectedly")
-                .about(command_results())
-                .that(result)
-                .isFailed();
-    }
-
     private static String join(String... strs) {
         return String.join(" ", Arrays.asList(strs));
     }
@@ -240,153 +218,6 @@
         return pathLine.substring("package:".length());
     }
 
-    public static String startMicrodroid(
-            ITestDevice androidDevice,
-            IBuildInfo buildInfo,
-            String apkName,
-            String packageName,
-            String configPath,
-            boolean debug,
-            int memoryMib,
-            Optional<Integer> numCpus)
-            throws DeviceNotAvailableException {
-        return startMicrodroid(androidDevice, buildInfo, apkName, packageName, null, configPath,
-                debug, memoryMib, numCpus);
-    }
-
-    public static String startMicrodroid(
-            ITestDevice androidDevice,
-            IBuildInfo buildInfo,
-            String apkName,
-            String packageName,
-            String[] extraIdsigPaths,
-            String configPath,
-            boolean debug,
-            int memoryMib,
-            Optional<Integer> numCpus)
-            throws DeviceNotAvailableException {
-        return startMicrodroid(androidDevice, buildInfo, apkName, null, packageName,
-                extraIdsigPaths, configPath, debug,
-                memoryMib, numCpus);
-    }
-
-    private static void forwardFileToLog(CommandRunner android, String path, String tag)
-            throws DeviceNotAvailableException {
-        android.runWithTimeout(
-                MICRODROID_MAX_LIFETIME_MINUTES * 60 * 1000,
-                "logwrapper",
-                "sh",
-                "-c",
-                "\"$'tail -f -n +0 " + path
-                        + " | sed \\'s/^/" + tag + ": /g\\''\""); // add tags in front of lines
-    }
-
-    public static String startMicrodroid(
-            ITestDevice androidDevice,
-            IBuildInfo buildInfo,
-            String apkName,
-            String apkPath,
-            String packageName,
-            String[] extraIdsigPaths,
-            String configPath,
-            boolean debug,
-            int memoryMib,
-            Optional<Integer> numCpus)
-            throws DeviceNotAvailableException {
-        CommandRunner android = new CommandRunner(androidDevice);
-
-        // Install APK if necessary
-        if (apkName != null) {
-            File apkFile = findTestFile(buildInfo, apkName);
-            androidDevice.installPackage(apkFile, /* reinstall */ true);
-        }
-
-        if (apkPath == null) {
-            apkPath = getPathForPackage(androidDevice, packageName);
-        }
-
-        android.run("mkdir", "-p", TEST_ROOT);
-
-        // This file is not what we provide. It will be created by the vm tool.
-        final String outApkIdsigPath = TEST_ROOT + apkName + ".idsig";
-
-        final String instanceImg = TEST_ROOT + INSTANCE_IMG;
-        final String logPath = LOG_PATH;
-        final String consolePath = CONSOLE_PATH;
-        final String debugFlag = debug ? "--debug full" : "";
-
-        // Run the VM
-        ArrayList<String> args = new ArrayList<>(Arrays.asList(
-                VIRT_APEX + "bin/vm",
-                "run-app",
-                "--daemonize",
-                "--log " + logPath,
-                "--console " + consolePath,
-                "--mem " + memoryMib,
-                numCpus.isPresent() ? "--cpus " + numCpus.get() : "",
-                debugFlag,
-                apkPath,
-                outApkIdsigPath,
-                instanceImg,
-                configPath));
-        if (extraIdsigPaths != null) {
-            for (String path : extraIdsigPaths) {
-                args.add("--extra-idsig");
-                args.add(path);
-            }
-        }
-        String ret = android.run(args.toArray(new String[0]));
-
-        // Redirect log.txt and console.txt to logd using logwrapper
-        // Keep redirecting as long as the expecting maximum test time. When an adb
-        // command times out, it may trigger the device recovery process, which
-        // disconnect adb, which terminates any live adb commands. See an example at
-        // b/194974010#comment25.
-        ExecutorService executor = Executors.newFixedThreadPool(2);
-        executor.execute(
-                () -> {
-                    try {
-                        forwardFileToLog(android, logPath, "MicrodroidLog");
-                    } catch (Exception e) {
-                        // Consume
-                    }
-                });
-
-        executor.execute(
-                () -> {
-                    try {
-                        forwardFileToLog(android, consolePath, "MicrodroidConsole");
-                    } catch (Exception e) {
-                        // Consume
-                    }
-                });
-
-        // Retrieve the CID from the vm tool output
-        Pattern pattern = Pattern.compile("with CID (\\d+)");
-        Matcher matcher = pattern.matcher(ret);
-        assertWithMessage("Failed to find CID").that(matcher.find()).isTrue();
-        return matcher.group(1);
-    }
-
-    public static void shutdownMicrodroid(ITestDevice androidDevice, String cid)
-            throws DeviceNotAvailableException {
-        CommandRunner android = new CommandRunner(androidDevice);
-
-        // Shutdown the VM
-        android.run(VIRT_APEX + "bin/vm", "stop", cid);
-    }
-
-    public static void rootMicrodroid() throws InterruptedException {
-        runOnHostRetryingOnFailure(MICRODROID_COMMAND_TIMEOUT_MILLIS,
-                MICRODROID_ADB_CONNECT_MAX_ATTEMPTS, "adb", "-s", MICRODROID_SERIAL, "root");
-        // adbd reboots after root. Some commands (including wait-for-device) following this fails
-        // with error: closed. Hence, we disconnect and re-connect to the device before returning.
-        runOnHostRetryingOnFailure(MICRODROID_COMMAND_TIMEOUT_MILLIS,
-                MICRODROID_ADB_CONNECT_MAX_ATTEMPTS, "adb", "disconnect", MICRODROID_SERIAL);
-        runOnHostRetryingOnFailure(MICRODROID_COMMAND_TIMEOUT_MILLIS,
-                MICRODROID_ADB_CONNECT_MAX_ATTEMPTS, "adb", "connect", MICRODROID_SERIAL);
-    }
-
     // Establish an adb connection to microdroid by letting Android forward the connection to
     // microdroid. Wait until the connection is established and microdroid is booted.
     public static void adbConnectToMicrodroid(ITestDevice androidDevice, String cid) {
@@ -434,8 +265,4 @@
                 .stdoutTrimmed()
                 .isEqualTo("microdroid");
     }
-
-    public boolean isProtectedVmSupported() throws DeviceNotAvailableException {
-        return getDevice().getBooleanProperty("ro.boot.hypervisor.protected_vm.supported", false);
-    }
 }
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 33788ed..11b3e84 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -17,12 +17,14 @@
 package com.android.microdroid.test;
 
 import static com.android.microdroid.test.host.CommandResultSubject.command_results;
+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 com.google.common.truth.Truth.assertWithMessage;
 
 import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assume.assumeTrue;
 
@@ -38,6 +40,9 @@
 import com.android.os.AtomsProto;
 import com.android.os.StatsLog;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
+import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 import com.android.tradefed.util.CommandResult;
@@ -46,7 +51,6 @@
 import com.android.tradefed.util.xml.AbstractXmlParser;
 
 import org.json.JSONArray;
-import org.json.JSONException;
 import org.json.JSONObject;
 import org.junit.After;
 import org.junit.Before;
@@ -58,14 +62,21 @@
 import org.xml.sax.Attributes;
 import org.xml.sax.helpers.DefaultHandler;
 
+import java.io.BufferedReader;
 import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.PipedInputStream;
+import java.io.PipedOutputStream;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -74,6 +85,7 @@
     private static final String APK_NAME = "MicrodroidTestApp.apk";
     private static final String PACKAGE_NAME = "com.android.microdroid.test";
     private static final String SHELL_PACKAGE_NAME = "com.android.shell";
+    private static final String VIRT_APEX = "/apex/com.android.virt/";
 
     private static final int MIN_MEM_ARM64 = 145;
     private static final int MIN_MEM_X86_64 = 196;
@@ -81,12 +93,26 @@
     // Number of vCPUs for testing purpose
     private static final int NUM_VCPUS = 3;
 
+    private static final int BOOT_COMPLETE_TIMEOUT = 30000; // 30 seconds
+
+    private static class VmInfo {
+        final Process mProcess;
+        final String mCid;
+
+        VmInfo(Process process, String cid) {
+            mProcess = process;
+            mCid = cid;
+        }
+    }
+
     @Rule public TestLogData mTestLogs = new TestLogData();
     @Rule public TestName mTestName = new TestName();
     @Rule public TestMetrics mMetrics = new TestMetrics();
 
     private String mMetricPrefix;
 
+    private ITestDevice mMicrodroidDevice;
+
     private int minMemorySize() throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(getDevice());
         String abi = android.run("getprop", "ro.product.cpu.abi");
@@ -99,10 +125,6 @@
         throw new AssertionError("Unsupported ABI: " + abi);
     }
 
-    private void waitForBootComplete() {
-        runOnMicrodroidForResult("watch -e \"getprop dev.bootcomplete | grep '^0$'\"");
-    }
-
     private static JSONObject newPartition(String label, String path) {
         return new JSONObject(Map.of("label", label, "path", path));
     }
@@ -124,23 +146,16 @@
                                 .collect(toList())));
         FileUtil.writeToFile(config.toString(), configFile);
 
-        File mkPayload = findTestFile("mk_payload");
         RunUtil runUtil = new RunUtil();
-        // Set the parent dir on the PATH (e.g. <workdir>/bin)
-        String separator = System.getProperty("path.separator");
-        String path = mkPayload.getParentFile().getPath() + separator + System.getenv("PATH");
-        runUtil.setEnvVariable("PATH", path);
-
-        List<String> command = new ArrayList<>();
-        command.add("mk_payload");
-        command.add("--metadata-only");
-        command.add(configFile.toString());
-        command.add(payloadMetadata.toString());
-
-        CommandResult result =
-                runUtil.runTimedCmd(
-                        // mk_payload should run fast enough
-                        5 * 1000, "/bin/bash", "-c", String.join(" ", command));
+        String command =
+                String.join(
+                        " ",
+                        findTestFile("mk_payload").getAbsolutePath(),
+                        "--metadata-only",
+                        configFile.getAbsolutePath(),
+                        payloadMetadata.getAbsolutePath());
+        // mk_payload should run fast enough
+        CommandResult result = runUtil.runTimedCmd(5000, "/bin/bash", "-c", command);
         String out = result.getStdout();
         String err = result.getStderr();
         assertWithMessage(
@@ -164,12 +179,10 @@
         runUtil.setEnvVariable("PATH", path);
 
         List<String> command = new ArrayList<>();
-        command.add("sign_virt_apex");
-        for (Map.Entry<String, File> entry : keyOverrides.entrySet()) {
-            String filename = entry.getKey();
-            File overridingKey = entry.getValue();
-            command.add("--key_override " + filename + "=" + overridingKey.getPath());
-        }
+        command.add(signVirtApex.getAbsolutePath());
+        keyOverrides.forEach(
+                (filename, keyFile) ->
+                        command.add("--key_override " + filename + "=" + keyFile.getPath()));
         command.add(signingKey.getPath());
         command.add(virtApexDir.getPath());
 
@@ -191,18 +204,11 @@
             long timeoutMillis, Callable<T> callable, org.hamcrest.Matcher<T> matcher)
             throws Exception {
         long start = System.currentTimeMillis();
-        while (true) {
-            try {
-                assertThat(callable.call(), matcher);
-                return;
-            } catch (Throwable e) {
-                if (System.currentTimeMillis() - start < timeoutMillis) {
-                    Thread.sleep(500);
-                } else {
-                    throw e;
-                }
-            }
+        while ((System.currentTimeMillis() - start < timeoutMillis)
+                && !matcher.matches(callable.call())) {
+            Thread.sleep(500);
         }
+        assertThat(callable.call(), matcher);
     }
 
     static class ActiveApexInfo {
@@ -225,12 +231,10 @@
         }
 
         ActiveApexInfo get(String apexName) {
-            for (ActiveApexInfo info : mList) {
-                if (info.name.equals(apexName)) {
-                    return info;
-                }
-            }
-            return null;
+            return mList.stream()
+                    .filter(info -> apexName.equals(info.name))
+                    .findFirst()
+                    .orElse(null);
         }
 
         List<ActiveApexInfo> getSharedLibApexes() {
@@ -262,13 +266,8 @@
         return new ActiveApexInfoList(list);
     }
 
-    private String runMicrodroidWithResignedImages(
-            File key,
-            Map<String, File> keyOverrides,
-            boolean isProtected,
-            boolean daemonize,
-            String consolePath)
-            throws Exception {
+    private VmInfo runMicrodroidWithResignedImages(
+            File key, Map<String, File> keyOverrides, boolean isProtected) throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
 
         File virtApexDir = FileUtil.createTempDir("virt_apex");
@@ -371,19 +370,40 @@
         final String configPath = TEST_ROOT + "raw_config.json";
         getDevice().pushString(config.toString(), configPath);
 
-        final String logPath = LOG_PATH;
-        final String ret =
-                android.runWithTimeout(
-                        60 * 1000,
+        List<String> args =
+                Arrays.asList(
+                        "adb",
+                        "-s",
+                        getDevice().getSerialNumber(),
+                        "shell",
                         VIRT_APEX + "bin/vm run",
-                        daemonize ? "--daemonize" : "",
-                        (consolePath != null) ? "--console " + consolePath : "",
-                        "--log " + logPath,
+                        "--console " + CONSOLE_PATH,
+                        "--log " + LOG_PATH,
                         configPath);
+
+        PipedInputStream pis = new PipedInputStream();
+        Process process = RunUtil.getDefault().runCmdInBackground(args, new PipedOutputStream(pis));
+        return new VmInfo(process, extractCidFrom(pis));
+    }
+
+    private static String extractCidFrom(InputStream input) throws IOException {
+        String cid = null;
         Pattern pattern = Pattern.compile("with CID (\\d+)");
-        Matcher matcher = pattern.matcher(ret);
-        assertWithMessage("Failed to find CID").that(matcher.find()).isTrue();
-        return matcher.group(1);
+        String line;
+        try (BufferedReader out = new BufferedReader(new InputStreamReader(input))) {
+            while ((line = out.readLine()) != null) {
+                CLog.i("VM output: " + line);
+                Matcher matcher = pattern.matcher(line);
+                if (matcher.find()) {
+                    cid = matcher.group(1);
+                    break;
+                }
+            }
+        }
+        assertWithMessage("The output does not contain the expected pattern for CID.")
+                .that(cid)
+                .isNotNull();
+        return cid;
     }
 
     @Test
@@ -391,15 +411,16 @@
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
     public void testBootFailsWhenProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
-        assumeTrue("Protected VMs are not supported", isProtectedVmSupported());
+        boolean protectedVm = true;
+        assumeTrue(
+                "Skip if protected VMs are not supported",
+                getAndroidDevice().supportsMicrodroid(protectedVm));
 
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
-        boolean isProtected = true;
-        boolean daemonize = false; // VM should shut down due to boot failure.
-        String consolePath = TEST_ROOT + "console";
-        runMicrodroidWithResignedImages(key, keyOverrides, isProtected, daemonize, consolePath);
-        assertThat(getDevice().pullFileContents(consolePath), containsString("pvmfw boot failed"));
+        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, protectedVm);
+        vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
+        assertThat(getDevice().pullFileContents(CONSOLE_PATH), containsString("pvmfw boot failed"));
     }
 
     // TODO(b/245277660): Resigning the system/vendor image changes the vbmeta hash.
@@ -411,15 +432,10 @@
             throws Exception {
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
-        boolean isProtected = false;
-        boolean daemonize = true;
-        String consolePath = TEST_ROOT + "console";
-        String cid =
-                runMicrodroidWithResignedImages(
-                        key, keyOverrides, isProtected, daemonize, consolePath);
-        // Adb connection to the microdroid means that boot succeeded.
-        adbConnectToMicrodroid(getDevice(), cid);
-        shutdownMicrodroid(getDevice(), cid);
+        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, /*isProtected=*/ false);
+        // Device online means that boot must have succeeded.
+        adbConnectToMicrodroid(getDevice(), vmInfo.mCid);
+        vmInfo.mProcess.destroy();
     }
 
     @Test
@@ -429,37 +445,36 @@
         File key = findTestFile("test.com.android.virt.pem");
         File key2 = findTestFile("test2.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of("microdroid_vbmeta.img", key2);
-        boolean isProtected = false; // Not interested in pvwfw
-        boolean daemonize = true; // Bootloader fails and enters prompts.
         // To be able to stop it, it should be a daemon.
-        String consolePath = TEST_ROOT + "console";
-        String cid =
-                runMicrodroidWithResignedImages(
-                        key, keyOverrides, isProtected, daemonize, consolePath);
+        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, /*isProtected=*/ false);
         // Wait so that init can print errors to console (time in cuttlefish >> in real device)
         assertThatEventually(
                 100000,
-                () -> getDevice().pullFileContents(consolePath),
+                () -> getDevice().pullFileContents(CONSOLE_PATH),
                 containsString("init: [libfs_avb]Failed to verify vbmeta digest"));
-        shutdownMicrodroid(getDevice(), cid);
+        vmInfo.mProcess.destroy();
     }
 
     private boolean isTombstoneGeneratedWithConfig(String configPath) throws Exception {
         // Note this test relies on logcat values being printed by tombstone_transmit on
         // and the reeceiver on host (virtualization_service)
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
+        mMicrodroidDevice =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(minMemorySize())
+                        .numCpus(NUM_VCPUS)
+                        .build(getAndroidDevice());
+        mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
+        mMicrodroidDevice.enableAdbRoot();
+
+        CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
+        microdroid.run("kill", "-SIGSEGV", "$(pidof microdroid_launcher)");
+
         // check until microdroid is shut down
         CommandRunner android = new CommandRunner(getDevice());
-        android.runWithTimeout(15000, "logcat", "-m", "1", "-e", "'crosvm has exited normally'");
+        // TODO: improve crosvm exit check. b/258848245
+        android.runWithTimeout(15000, "logcat", "-m", "1", "-e",
+                              "'virtualizationservice::crosvm.*exited with status exit status: 0'");
         // Check that tombstone is received (from host logcat)
         String result =
                 runOnHost(
@@ -485,7 +500,7 @@
     }
 
     @Test
-    public void testTelemetryPushedAtomsOfEventMetrics() throws Exception {
+    public void testTelemetryPushedAtoms() throws Exception {
         // Reset statsd config and report before the test
         ConfigUtils.removeConfig(getDevice());
         ReportUtils.clearReports(getDevice());
@@ -499,22 +514,28 @@
         ConfigUtils.uploadConfigForPushedAtoms(getDevice(), PACKAGE_NAME, atomIds);
 
         // Create VM with microdroid
+        TestDevice device = getAndroidDevice();
         final String configPath = "assets/vm_config_apex.json"; // path inside the APK
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
+        ITestDevice microdroid =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(minMemorySize())
+                        .numCpus(NUM_VCPUS)
+                        .build(device);
+        microdroid.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
+        device.shutdownMicrodroid(microdroid);
 
-        // Check VmCreationRequested atom and clear the statsd report
-        List<StatsLog.EventMetricData> data;
-        data = ReportUtils.getEventMetricDataList(getDevice());
-        assertThat(data).hasSize(1);
+        List<StatsLog.EventMetricData> data = new ArrayList<>();
+        assertThatEventually(
+                10000,
+                () -> {
+                    data.addAll(ReportUtils.getEventMetricDataList(getDevice()));
+                    return data.size();
+                },
+                is(3)
+        );
+
+        // Check VmCreationRequested atom
         assertThat(data.get(0).getAtom().getPushedCase().getNumber()).isEqualTo(
                 AtomsProto.Atom.VM_CREATION_REQUESTED_FIELD_NUMBER);
         AtomsProto.VmCreationRequested atomVmCreationRequested =
@@ -532,29 +553,16 @@
         assertThat(atomVmCreationRequested.getApexes())
                 .isEqualTo("com.android.art:com.android.compos:com.android.sdkext");
 
-        // Boot VM with microdroid
-        adbConnectToMicrodroid(getDevice(), cid);
-        waitForBootComplete();
-
-        // Check VmBooted atom and clear the statsd report
-        data = ReportUtils.getEventMetricDataList(getDevice());
-        assertThat(data).hasSize(1);
-        assertThat(data.get(0).getAtom().getPushedCase().getNumber())
+        // Check VmBooted atom
+        assertThat(data.get(1).getAtom().getPushedCase().getNumber())
                 .isEqualTo(AtomsProto.Atom.VM_BOOTED_FIELD_NUMBER);
-        AtomsProto.VmBooted atomVmBooted = data.get(0).getAtom().getVmBooted();
+        AtomsProto.VmBooted atomVmBooted = data.get(1).getAtom().getVmBooted();
         assertThat(atomVmBooted.getVmIdentifier()).isEqualTo("VmRunApp");
 
-        // Shutdown VM with microdroid
-        shutdownMicrodroid(getDevice(), cid);
-        // TODO: make sure the VM is completely shut down while 'vm stop' command running.
-        Thread.sleep(1000);
-
-        // Check VmExited atom and clear the statsd report
-        data = ReportUtils.getEventMetricDataList(getDevice());
-        assertThat(data).hasSize(1);
-        assertThat(data.get(0).getAtom().getPushedCase().getNumber())
+        // Check VmExited atom
+        assertThat(data.get(2).getAtom().getPushedCase().getNumber())
                 .isEqualTo(AtomsProto.Atom.VM_EXITED_FIELD_NUMBER);
-        AtomsProto.VmExited atomVmExited = data.get(0).getAtom().getVmExited();
+        AtomsProto.VmExited atomVmExited = data.get(2).getAtom().getVmExited();
         assertThat(atomVmExited.getVmIdentifier()).isEqualTo("VmRunApp");
         assertThat(atomVmExited.getDeathReason()).isEqualTo(AtomsProto.VmExited.DeathReason.KILLED);
 
@@ -566,100 +574,53 @@
     }
 
     @Test
-    public void testTelemetryPushedAtomsOfValueMetrics() throws Exception {
-        // Reset statsd config and report before the test
-        ConfigUtils.removeConfig(getDevice());
-        ReportUtils.clearReports(getDevice());
-
-        // Setup statsd config
-        int[] atomIds = {
-            AtomsProto.Atom.VM_CPU_STATUS_REPORTED_FIELD_NUMBER,
-            AtomsProto.Atom.VM_MEM_STATUS_REPORTED_FIELD_NUMBER,
-        };
-        ConfigUtils.uploadConfigForPushedAtoms(getDevice(), PACKAGE_NAME, atomIds);
-
-        // Create VM with microdroid
-        final String configPath = "assets/vm_config_apex.json"; // path inside the APK
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
-
-        // Boot VM with microdroid
-        adbConnectToMicrodroid(getDevice(), cid);
-        waitForBootComplete();
-
-        // Check VmCpuStatusReported and VmMemStatusReported atoms and clear the statsd report
-        List<StatsLog.EventMetricData> data;
-        data = ReportUtils.getEventMetricDataList(getDevice());
-        assertThat(data.size() >= 2).isTrue();
-        assertThat(data.get(0).getAtom().getPushedCase().getNumber())
-                .isEqualTo(AtomsProto.Atom.VM_CPU_STATUS_REPORTED_FIELD_NUMBER);
-        assertThat(data.get(1).getAtom().getPushedCase().getNumber())
-                .isEqualTo(AtomsProto.Atom.VM_MEM_STATUS_REPORTED_FIELD_NUMBER);
-
-        // Shutdown VM with microdroid
-        shutdownMicrodroid(getDevice(), cid);
-    }
-
-    @Test
     @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
     public void testMicrodroidBoots() throws Exception {
         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));
-        adbConnectToMicrodroid(getDevice(), cid);
-        waitForBootComplete();
+        mMicrodroidDevice =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(minMemorySize())
+                        .numCpus(NUM_VCPUS)
+                        .build(getAndroidDevice());
+        mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
+        CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
+
         // Test writing to /data partition
-        runOnMicrodroid("echo MicrodroidTest > /data/local/tmp/test.txt");
-        assertThat(runOnMicrodroid("cat /data/local/tmp/test.txt")).isEqualTo("MicrodroidTest");
+        microdroid.run("echo MicrodroidTest > /data/local/tmp/test.txt");
+        assertThat(microdroid.run("cat /data/local/tmp/test.txt")).isEqualTo("MicrodroidTest");
 
         // Check if the APK & its idsig partitions exist
         final String apkPartition = "/dev/block/by-name/microdroid-apk";
-        assertThat(runOnMicrodroid("ls", apkPartition)).isEqualTo(apkPartition);
+        assertThat(microdroid.run("ls", apkPartition)).isEqualTo(apkPartition);
         final String apkIdsigPartition = "/dev/block/by-name/microdroid-apk-idsig";
-        assertThat(runOnMicrodroid("ls", apkIdsigPartition)).isEqualTo(apkIdsigPartition);
+        assertThat(microdroid.run("ls", apkIdsigPartition)).isEqualTo(apkIdsigPartition);
         // Check the vm-instance partition as well
         final String vmInstancePartition = "/dev/block/by-name/vm-instance";
-        assertThat(runOnMicrodroid("ls", vmInstancePartition)).isEqualTo(vmInstancePartition);
+        assertThat(microdroid.run("ls", vmInstancePartition)).isEqualTo(vmInstancePartition);
 
         // Check if the native library in the APK is has correct filesystem info
-        final String[] abis = runOnMicrodroid("getprop", "ro.product.cpu.abilist").split(",");
+        final String[] abis = microdroid.run("getprop", "ro.product.cpu.abilist").split(",");
         assertThat(abis).hasLength(1);
         final String testLib = "/mnt/apk/lib/" + abis[0] + "/MicrodroidTestNativeLib.so";
         final String label = "u:object_r:system_file:s0";
-        assertThat(runOnMicrodroid("ls", "-Z", testLib)).isEqualTo(label + " " + testLib);
+        assertThat(microdroid.run("ls", "-Z", testLib)).isEqualTo(label + " " + testLib);
 
         // Check that no denials have happened so far
         CommandRunner android = new CommandRunner(getDevice());
         assertThat(android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", LOG_PATH)).isNull();
+        assertThat(android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", CONSOLE_PATH)).isNull();
 
-        assertThat(runOnMicrodroid("cat /proc/cpuinfo | grep processor | wc -l"))
+        assertThat(microdroid.run("cat /proc/cpuinfo | grep processor | wc -l"))
                 .isEqualTo(Integer.toString(NUM_VCPUS));
 
         // Check that selinux is enabled
-        assertThat(runOnMicrodroid("getenforce")).isEqualTo("Enforcing");
+        assertThat(microdroid.run("getenforce")).isEqualTo("Enforcing");
 
         // TODO(b/176805428): adb is broken for nested VM
         if (!isCuttlefish()) {
             // Check neverallow rules on microdroid
-            File policyFile = FileUtil.createTempFile("microdroid_sepolicy", "");
-            pullMicrodroidFile("/sys/fs/selinux/policy", policyFile);
-
+            File policyFile = mMicrodroidDevice.pullFile("/sys/fs/selinux/policy");
             File generalPolicyConfFile = findTestFile("microdroid_general_sepolicy.conf");
             File sepolicyAnalyzeBin = findTestFile("sepolicy-analyze");
 
@@ -678,38 +639,41 @@
                     .that(result)
                     .isSuccess();
         }
-
-        shutdownMicrodroid(getDevice(), cid);
     }
 
     @Test
     public void testMicrodroidRamUsage() throws Exception {
         final String configPath = "assets/vm_config.json";
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
-        adbConnectToMicrodroid(getDevice(), cid);
-        waitForBootComplete();
-        rootMicrodroid();
+        mMicrodroidDevice =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(minMemorySize())
+                        .numCpus(NUM_VCPUS)
+                        .build(getAndroidDevice());
+        mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
+        mMicrodroidDevice.enableAdbRoot();
+
+        CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
+        Function<String, String> microdroidExec =
+                (cmd) -> {
+                    try {
+                        return microdroid.run(cmd);
+                    } catch (Exception ex) {
+                        throw new IllegalStateException(ex);
+                    }
+                };
 
         for (Map.Entry<String, Long> stat :
-                ProcessUtil.getProcessMemoryMap(cmd -> runOnMicrodroid(cmd)).entrySet()) {
+                ProcessUtil.getProcessMemoryMap(microdroidExec).entrySet()) {
             mMetrics.addTestMetric(
                     mMetricPrefix + "meminfo/" + stat.getKey().toLowerCase(),
                     stat.getValue().toString());
         }
 
         for (Map.Entry<Integer, String> proc :
-                ProcessUtil.getProcessMap(cmd -> runOnMicrodroid(cmd)).entrySet()) {
+                ProcessUtil.getProcessMap(microdroidExec).entrySet()) {
             for (Map.Entry<String, Long> stat :
-                    ProcessUtil.getProcessSmapsRollup(proc.getKey(), cmd -> runOnMicrodroid(cmd))
+                    ProcessUtil.getProcessSmapsRollup(proc.getKey(), microdroidExec)
                             .entrySet()) {
                 String name = stat.getKey().toLowerCase();
                 mMetrics.addTestMetric(
@@ -717,14 +681,13 @@
                         stat.getValue().toString());
             }
         }
-
-        shutdownMicrodroid(getDevice(), cid);
     }
 
     @Test
-    public void testCustomVirtualMachinePermission()
-            throws DeviceNotAvailableException, IOException, JSONException {
-        assumeTrue("Protected VMs are not supported", isProtectedVmSupported());
+    public void testCustomVirtualMachinePermission() throws Exception {
+        assumeTrue(
+                "Protected VMs are not supported",
+                getAndroidDevice().supportsMicrodroid(/*protectedVm=*/ true));
         CommandRunner android = new CommandRunner(getDevice());
 
         // Pull etc/microdroid.json
@@ -760,6 +723,7 @@
     public void setUp() throws Exception {
         testIfDeviceIsCapable(getDevice());
         mMetricPrefix = getMetricPrefix() + "microdroid/";
+        mMicrodroidDevice = null;
 
         prepareVirtualizationTestSetup(getDevice());
 
@@ -771,6 +735,10 @@
 
     @After
     public void shutdown() throws Exception {
+        if (mMicrodroidDevice != null) {
+            getAndroidDevice().shutdownMicrodroid(mMicrodroidDevice);
+        }
+
         cleanUpVirtualizationTestSetup(getDevice());
 
         archiveLogThenDelete(
@@ -786,4 +754,10 @@
                         SHELL_PACKAGE_NAME,
                         "android.permission.USE_CUSTOM_VIRTUAL_MACHINE");
     }
+
+    private TestDevice getAndroidDevice() {
+        TestDevice androidDevice = (TestDevice) getDevice();
+        assertThat(androidDevice).isNotNull();
+        return androidDevice;
+    }
 }
diff --git a/tests/no_avf/Android.bp b/tests/no_avf/Android.bp
new file mode 100644
index 0000000..fd0d5e2
--- /dev/null
+++ b/tests/no_avf/Android.bp
@@ -0,0 +1,21 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_test {
+    name: "CtsMicrodroidDisabledTestCases",
+    test_suites: [
+        "cts",
+        "general-tests",
+    ],
+    srcs: ["src/**/*.java"],
+    static_libs: [
+        "androidx.test.runner",
+        "androidx.test.ext.junit",
+        "compatibility-common-util-devicesidelib",
+        "truth-prebuilt",
+    ],
+    sdk_version: "test_current",
+    compile_multilib: "both",
+    min_sdk_version: "UpsideDownCake",
+}
diff --git a/tests/no_avf/AndroidManifest.xml b/tests/no_avf/AndroidManifest.xml
new file mode 100644
index 0000000..4a1304e
--- /dev/null
+++ b/tests/no_avf/AndroidManifest.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2022 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+      package="com.android.nomicrodroid.test">
+    <application />
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.nomicrodroid.test"
+        android:label="CTS test for devices without Microdroid support" />
+</manifest>
diff --git a/tests/no_avf/AndroidTest.xml b/tests/no_avf/AndroidTest.xml
new file mode 100644
index 0000000..1e93887
--- /dev/null
+++ b/tests/no_avf/AndroidTest.xml
@@ -0,0 +1,29 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<configuration description="Runs Microdroid device-side tests.">
+    <option name="test-suite-tag" value="cts" />
+    <option name="config-descriptor:metadata" key="component" value="security" />
+    <option name="config-descriptor:metadata" key="parameter" value="not_instant_app" />
+    <option name="config-descriptor:metadata" key="parameter" value="not_multi_abi" />
+    <option name="config-descriptor:metadata" key="parameter" value="secondary_user" />
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="test-file-name" value="CtsMicrodroidDisabledTestCases.apk" />
+    </target_preparer>
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.nomicrodroid.test" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+    </test>
+</configuration>
diff --git a/tests/no_avf/README.md b/tests/no_avf/README.md
new file mode 100644
index 0000000..b96dc97
--- /dev/null
+++ b/tests/no_avf/README.md
@@ -0,0 +1 @@
+CTS tests for devices that don't support AVF.
\ No newline at end of file
diff --git a/tests/no_avf/src/com/android/nomicrodroid/test/NoMicrodroidTest.java b/tests/no_avf/src/com/android/nomicrodroid/test/NoMicrodroidTest.java
new file mode 100644
index 0000000..0982e35
--- /dev/null
+++ b/tests/no_avf/src/com/android/nomicrodroid/test/NoMicrodroidTest.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.nomicrodroid.test;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.TruthJUnit.assume;
+
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.system.virtualmachine.VirtualMachineManager;
+
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import com.android.compatibility.common.util.CddTest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests to validate that devices without support for AVF (Android Virtualization Framework) are set
+ * up correctly.
+ */
+@RunWith(JUnit4.class)
+public class NoMicrodroidTest {
+
+    @Before
+    public void setUp() {
+        final PackageManager pm =
+                InstrumentationRegistry.getInstrumentation().getTargetContext().getPackageManager();
+        assume().withMessage("Device supports AVF")
+                .that(pm.hasSystemFeature(PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK))
+                .isFalse();
+    }
+
+    @CddTest(requirements = {"9.17/C-1-1"})
+    @Test
+    public void testVirtualMachineManagerLookupReturnsNull() {
+        final Context ctx = InstrumentationRegistry.getInstrumentation().getTargetContext();
+        assertThat(ctx.getSystemService(VirtualMachineManager.class)).isNull();
+    }
+}
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index 42abbbf..4dc9489 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -19,12 +19,12 @@
         "truth-prebuilt",
         "compatibility-common-util-devicesidelib",
     ],
-    libs: ["android.system.virtualmachine"],
+    sdk_version: "test_current",
     jni_libs: [
         "MicrodroidTestNativeLib",
-        "MicrodroidTestNativeCrashLib",
+        "MicrodroidIdleNativeLib",
     ],
-    platform_apis: true,
+    jni_uses_platform_apis: true,
     use_embedded_native_libs: true,
     // We only support 64-bit ABI, but CTS demands all APKs to be multi-ABI.
     compile_multilib: "both",
@@ -35,10 +35,11 @@
     name: "MicrodroidTestNativeLib",
     srcs: ["src/native/testbinary.cpp"],
     stl: "libc++_static",
+    header_libs: ["vm_payload_restricted_headers"],
     shared_libs: [
         "libbinder_ndk",
         "MicrodroidTestNativeLibSub",
-        "libvm_payload",
+        "libvm_payload#current",
     ],
     static_libs: [
         "com.android.microdroid.testservice-ndk",
@@ -50,14 +51,14 @@
 }
 
 cc_library_shared {
-    name: "MicrodroidTestNativeCrashLib",
-    header_libs: ["vm_payload_headers"],
-    srcs: ["src/native/crashbinary.cpp"],
+    name: "MicrodroidTestNativeLibSub",
+    srcs: ["src/native/testlib.cpp"],
     stl: "libc++_static",
 }
 
 cc_library_shared {
-    name: "MicrodroidTestNativeLibSub",
-    srcs: ["src/native/testlib.cpp"],
+    name: "MicrodroidIdleNativeLib",
+    srcs: ["src/native/idlebinary.cpp"],
+    header_libs: ["vm_payload_headers"],
     stl: "libc++_static",
 }
diff --git a/tests/testapk/AndroidManifest.xml b/tests/testapk/AndroidManifest.xml
index ab22546..fefd20a 100644
--- a/tests/testapk/AndroidManifest.xml
+++ b/tests/testapk/AndroidManifest.xml
@@ -18,8 +18,8 @@
     <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
     <uses-permission android:name="android.permission.USE_CUSTOM_VIRTUAL_MACHINE" />
     <uses-sdk android:minSdkVersion="33" android:targetSdkVersion="33" />
+    <uses-feature android:name="android.software.virtualization_framework" android:required="false" />
     <application>
-        <uses-library android:name="android.system.virtualmachine" android:required="false" />
     </application>
     <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
         android:targetPackage="com.android.microdroid.test"
diff --git a/tests/testapk/assets/vm_config_crash.json b/tests/testapk/assets/vm_config_crash.json
index 2951fdf..3ec34a3 100644
--- a/tests/testapk/assets/vm_config_crash.json
+++ b/tests/testapk/assets/vm_config_crash.json
@@ -4,7 +4,7 @@
     },
     "task": {
       "type": "microdroid_launcher",
-      "command": "MicrodroidTestNativeCrashLib.so"
+      "command": "MicrodroidIdleNativeLib.so"
     },
     "export_tombstones": true
   }
diff --git a/tests/testapk/assets/vm_config_crash_no_tombstone.json b/tests/testapk/assets/vm_config_crash_no_tombstone.json
index 583f87b..9678e38 100644
--- a/tests/testapk/assets/vm_config_crash_no_tombstone.json
+++ b/tests/testapk/assets/vm_config_crash_no_tombstone.json
@@ -4,7 +4,7 @@
     },
     "task": {
       "type": "microdroid_launcher",
-      "command": "MicrodroidTestNativeCrashLib.so"
+      "command": "MicrodroidIdleNativeLib.so"
     },
     "export_tombstones": false
   }
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 e5052bf..679f6e7 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -15,6 +15,10 @@
  */
 package com.android.microdroid.test;
 
+import static android.system.virtualmachine.VirtualMachine.STATUS_DELETED;
+import static android.system.virtualmachine.VirtualMachine.STATUS_RUNNING;
+import static android.system.virtualmachine.VirtualMachine.STATUS_STOPPED;
+import static android.system.virtualmachine.VirtualMachineConfig.DEBUG_LEVEL_APP_ONLY;
 import static android.system.virtualmachine.VirtualMachineConfig.DEBUG_LEVEL_FULL;
 import static android.system.virtualmachine.VirtualMachineConfig.DEBUG_LEVEL_NONE;
 
@@ -25,14 +29,19 @@
 
 import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 
+import android.content.Context;
 import android.os.Build;
-import android.os.ParcelFileDescriptor;
 import android.os.SystemProperties;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineCallback;
 import android.system.virtualmachine.VirtualMachineConfig;
+import android.system.virtualmachine.VirtualMachineDescriptor;
+import android.system.virtualmachine.VirtualMachineException;
+import android.system.virtualmachine.VirtualMachineManager;
 import android.util.Log;
 
+import androidx.test.core.app.ApplicationProvider;
+
 import com.android.compatibility.common.util.CddTest;
 import com.android.microdroid.test.device.MicrodroidDeviceTestBase;
 import com.android.microdroid.testservice.ITestService;
@@ -52,6 +61,9 @@
 import java.io.IOException;
 import java.io.RandomAccessFile;
 import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
 import java.util.List;
 import java.util.OptionalLong;
 import java.util.UUID;
@@ -93,18 +105,16 @@
     private static final int MIN_MEM_X86_64 = 196;
 
     @Test
-    @CddTest(requirements = {
-            "9.17/C-1-1",
-            "9.17/C-2-1"
-    })
-    public void connectToVmService() throws Exception {
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-1"})
+    public void createAndConnectToVm() throws Exception {
         assumeSupportedKernel();
 
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                .setMemoryMib(minMemoryRequired())
-                .build();
-        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm_extra_apk", config);
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setMemoryMib(minMemoryRequired())
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
 
         TestResults testResults = runVmTestService(vm);
         assertThat(testResults.mException).isNull();
@@ -112,31 +122,128 @@
         assertThat(testResults.mAppRunProp).isEqualTo("true");
         assertThat(testResults.mSublibRunProp).isEqualTo("true");
         assertThat(testResults.mApkContentsPath).isEqualTo("/mnt/apk");
+        assertThat(testResults.mEncryptedStoragePath).isEqualTo("");
     }
 
     @Test
-    @CddTest(requirements = {
-            "9.17/C-1-1",
-            "9.17/C-1-2",
-            "9.17/C-1-4",
-    })
-    public void createVmRequiresPermission() throws Exception {
+    @CddTest(
+            requirements = {
+                "9.17/C-1-1",
+                "9.17/C-1-2",
+                "9.17/C-1-4",
+            })
+    public void createVmRequiresPermission() {
         assumeSupportedKernel();
 
         revokePermission(VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION);
 
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                .setMemoryMib(minMemoryRequired())
-                .build();
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setMemoryMib(minMemoryRequired())
+                        .build();
 
-        SecurityException e = assertThrows(SecurityException.class,
-                () -> mInner.forceCreateNewVirtualMachine("test_vm_requires_permission", config));
+        SecurityException e =
+                assertThrows(
+                        SecurityException.class,
+                        () -> forceCreateNewVirtualMachine("test_vm_requires_permission", config));
         assertThat(e).hasMessageThat()
                 .contains("android.permission.MANAGE_VIRTUAL_MACHINE permission");
     }
 
     @Test
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-1"})
+    public void autoCloseVm() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setMemoryMib(minMemoryRequired())
+                        .build();
+
+        try (VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config)) {
+            assertThat(vm.getStatus()).isEqualTo(STATUS_STOPPED);
+            // close() implicitly called on stopped VM.
+        }
+
+        try (VirtualMachine vm = getVirtualMachineManager().get("test_vm")) {
+            vm.run();
+            assertThat(vm.getStatus()).isEqualTo(STATUS_RUNNING);
+            // close() implicitly called on running VM.
+        }
+
+        try (VirtualMachine vm = getVirtualMachineManager().get("test_vm")) {
+            assertThat(vm.getStatus()).isEqualTo(STATUS_STOPPED);
+            getVirtualMachineManager().delete("test_vm");
+            assertThat(vm.getStatus()).isEqualTo(STATUS_DELETED);
+            // close() implicitly called on deleted VM.
+        }
+    }
+
+    @Test
+    @CddTest(requirements = {"9.17/C-1-1"})
+    public void vmConfigUnitTests() {
+        VirtualMachineConfig minimal =
+                newVmConfigBuilder().setPayloadBinaryPath("binary/path").build();
+
+        assertThat(minimal.getApkPath()).isEqualTo(getContext().getPackageCodePath());
+        assertThat(minimal.getDebugLevel()).isEqualTo(DEBUG_LEVEL_NONE);
+        assertThat(minimal.getMemoryMib()).isEqualTo(0);
+        assertThat(minimal.getNumCpus()).isEqualTo(1);
+        assertThat(minimal.getPayloadBinaryPath()).isEqualTo("binary/path");
+        assertThat(minimal.getPayloadConfigPath()).isNull();
+        assertThat(minimal.isProtectedVm()).isEqualTo(isProtectedVm());
+
+        int maxCpus = Runtime.getRuntime().availableProcessors();
+        VirtualMachineConfig.Builder maximalBuilder =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("config/path")
+                        .setApkPath("/apk/path")
+                        .setNumCpus(maxCpus)
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .setMemoryMib(42);
+        VirtualMachineConfig maximal = maximalBuilder.build();
+
+        assertThat(maximal.getApkPath()).isEqualTo("/apk/path");
+        assertThat(maximal.getDebugLevel()).isEqualTo(DEBUG_LEVEL_FULL);
+        assertThat(maximal.getMemoryMib()).isEqualTo(42);
+        assertThat(maximal.getNumCpus()).isEqualTo(maxCpus);
+        assertThat(maximal.getPayloadBinaryPath()).isNull();
+        assertThat(maximal.getPayloadConfigPath()).isEqualTo("config/path");
+        assertThat(maximal.isProtectedVm()).isEqualTo(isProtectedVm());
+
+        assertThat(minimal.isCompatibleWith(maximal)).isFalse();
+        assertThat(minimal.isCompatibleWith(minimal)).isTrue();
+        assertThat(maximal.isCompatibleWith(maximal)).isTrue();
+
+        VirtualMachineConfig compatible = maximalBuilder.setNumCpus(1).setMemoryMib(99).build();
+        assertThat(compatible.isCompatibleWith(maximal)).isTrue();
+    }
+
+    @Test
+    @CddTest(requirements = {"9.17/C-1-1"})
+    public void vmUnitTests() throws Exception {
+        VirtualMachineConfig.Builder builder =
+                newVmConfigBuilder().setPayloadBinaryPath("binary/path");
+        VirtualMachineConfig config = builder.build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("vm_name", config);
+
+        assertThat(vm.getName()).isEqualTo("vm_name");
+        assertThat(vm.getConfig().getPayloadBinaryPath()).isEqualTo("binary/path");
+        assertThat(vm.getConfig().getMemoryMib()).isEqualTo(0);
+
+        VirtualMachineConfig compatibleConfig = builder.setMemoryMib(42).build();
+        vm.setConfig(compatibleConfig);
+
+        assertThat(vm.getName()).isEqualTo("vm_name");
+        assertThat(vm.getConfig().getPayloadBinaryPath()).isEqualTo("binary/path");
+        assertThat(vm.getConfig().getMemoryMib()).isEqualTo(42);
+
+        assertThat(getVirtualMachineManager().get("vm_name")).isSameInstanceAs(vm);
+    }
+
+    @Test
     @CddTest(requirements = {
             "9.17/C-1-1",
             "9.17/C-1-2",
@@ -145,13 +252,14 @@
     public void createVmWithConfigRequiresPermission() throws Exception {
         assumeSupportedKernel();
 
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config.json")
-                .setMemoryMib(minMemoryRequired())
-                .build();
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config.json")
+                        .setMemoryMib(minMemoryRequired())
+                        .build();
 
-        VirtualMachine vm = mInner.forceCreateNewVirtualMachine(
-                "test_vm_config_requires_permission", config);
+        VirtualMachine vm =
+                forceCreateNewVirtualMachine("test_vm_config_requires_permission", config);
 
         SecurityException e = assertThrows(SecurityException.class, () -> runVmTestService(vm));
         assertThat(e).hasMessageThat()
@@ -162,17 +270,44 @@
     @CddTest(requirements = {
             "9.17/C-1-1",
     })
+    public void deleteVm() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setMemoryMib(minMemoryRequired())
+                        .build();
+
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_delete", config);
+        VirtualMachineManager vmm = getVirtualMachineManager();
+        vmm.delete("test_vm_delete");
+
+        // VM should no longer exist
+        assertThat(vmm.get("test_vm_delete")).isNull();
+
+        // Can't start the VM even with an existing reference
+        assertThrows(VirtualMachineException.class, vm::run);
+
+        // Can't delete the VM since it no longer exists
+        assertThrows(VirtualMachineException.class, () -> vmm.delete("test_vm_delete"));
+    }
+
+    @Test
+    @CddTest(requirements = {
+            "9.17/C-1-1",
+    })
     public void validApkPathIsAccepted() throws Exception {
         assumeSupportedKernel();
 
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                .setApkPath(getContext().getPackageCodePath())
-                .setMemoryMib(minMemoryRequired())
-                .build();
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setApkPath(getContext().getPackageCodePath())
+                        .setMemoryMib(minMemoryRequired())
+                        .build();
 
-        VirtualMachine vm = mInner.forceCreateNewVirtualMachine(
-                "test_vm_explicit_apk_path", config);
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_explicit_apk_path", config);
 
         TestResults testResults = runVmTestService(vm);
         assertThat(testResults.mException).isNull();
@@ -183,16 +318,23 @@
             "9.17/C-1-1",
     })
     public void invalidApkPathIsRejected() {
-        assumeSupportedKernel();
-
-        VirtualMachineConfig.Builder builder = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                .setApkPath("relative/path/to.apk")
-                .setMemoryMib(minMemoryRequired());
+        VirtualMachineConfig.Builder builder =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setApkPath("relative/path/to.apk")
+                        .setMemoryMib(minMemoryRequired());
         assertThrows(IllegalArgumentException.class, () -> builder.build());
     }
 
     @Test
+    @CddTest(requirements = {"9.17/C-1-1"})
+    public void invalidVmNameIsRejected() {
+        VirtualMachineManager vmm = getVirtualMachineManager();
+        assertThrows(IllegalArgumentException.class, () -> vmm.get("../foo"));
+        assertThrows(IllegalArgumentException.class, () -> vmm.get(".."));
+    }
+
+    @Test
     @CddTest(requirements = {
             "9.17/C-1-1",
             "9.17/C-2-1"
@@ -201,11 +343,12 @@
         assumeSupportedKernel();
 
         grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config_extra_apk.json")
-                .setMemoryMib(minMemoryRequired())
-                .build();
-        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm_extra_apk", config);
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config_extra_apk.json")
+                        .setMemoryMib(minMemoryRequired())
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_extra_apk", config);
 
         TestResults testResults = runVmTestService(vm);
         assertThat(testResults.mExtraApkTestProp).isEqualTo("PASS");
@@ -214,12 +357,13 @@
     @Test
     public void bootFailsWhenLowMem() throws Exception {
         for (int memMib : new int[]{ 10, 20, 40 }) {
-            VirtualMachineConfig lowMemConfig = mInner.newVmConfigBuilder()
-                    .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                    .setMemoryMib(memMib)
-                    .setDebugLevel(DEBUG_LEVEL_NONE)
-                    .build();
-            VirtualMachine vm = mInner.forceCreateNewVirtualMachine("low_mem", lowMemConfig);
+            VirtualMachineConfig lowMemConfig =
+                    newVmConfigBuilder()
+                            .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                            .setMemoryMib(memMib)
+                            .setDebugLevel(DEBUG_LEVEL_NONE)
+                            .build();
+            VirtualMachine vm = forceCreateNewVirtualMachine("low_mem", lowMemConfig);
             final CompletableFuture<Boolean> onPayloadReadyExecuted = new CompletableFuture<>();
             final CompletableFuture<Boolean> onStoppedExecuted = new CompletableFuture<>();
             VmEventListener listener =
@@ -243,36 +387,47 @@
     }
 
     @Test
-    @CddTest(requirements = {
-            "9.17/C-1-1",
-            "9.17/C-2-7"
-    })
-    public void changingDebugLevelInvalidatesVmIdentity() throws Exception {
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-7"})
+    public void changingNonDebuggableVmDebuggableInvalidatesVmIdentity() throws Exception {
+        changeDebugLevel(DEBUG_LEVEL_NONE, DEBUG_LEVEL_FULL);
+        changeDebugLevel(DEBUG_LEVEL_NONE, DEBUG_LEVEL_APP_ONLY);
+    }
+
+    @Test
+    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-2-7"})
+    @Ignore("b/260067026")
+    public void changingAppDebuggableVmFullyDebuggableInvalidatesVmIdentity() throws Exception {
+        assume().withMessage("Skip for non-protected VM. b/239158757").that(mProtectedVm).isTrue();
+        changeDebugLevel(DEBUG_LEVEL_APP_ONLY, DEBUG_LEVEL_FULL);
+    }
+
+    private void changeDebugLevel(int fromLevel, int toLevel) throws Exception {
         assumeSupportedKernel();
 
-        VirtualMachineConfig.Builder builder = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                .setDebugLevel(DEBUG_LEVEL_NONE);
+        VirtualMachineConfig.Builder builder =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(fromLevel);
         VirtualMachineConfig normalConfig = builder.build();
-        mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        forceCreateNewVirtualMachine("test_vm", normalConfig);
         assertThat(tryBootVm(TAG, "test_vm").payloadStarted).isTrue();
 
         // Try to run the VM again with the previous instance.img
         // We need to make sure that no changes on config don't invalidate the identity, to compare
         // the result with the below "different debug level" test.
-        File vmRoot = new File(getContext().getFilesDir(), "vm");
-        File vmInstance = new File(new File(vmRoot, "test_vm"), "instance.img");
+        File vmInstance = getVmFile("test_vm", "instance.img");
         File vmInstanceBackup = File.createTempFile("instance", ".img");
         Files.copy(vmInstance.toPath(), vmInstanceBackup.toPath(), REPLACE_EXISTING);
-        mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        forceCreateNewVirtualMachine("test_vm", normalConfig);
         Files.copy(vmInstanceBackup.toPath(), vmInstance.toPath(), REPLACE_EXISTING);
         assertThat(tryBootVm(TAG, "test_vm").payloadStarted).isTrue();
 
-        // Launch the same VM with different debug level. The Java API prohibits this (thankfully).
+        // Launch the same VM with a different debug level. The Java API prohibits this
+        // (thankfully).
         // For testing, we do that by creating a new VM with debug level, and copy the old instance
         // image to the new VM instance image.
-        VirtualMachineConfig debugConfig = builder.setDebugLevel(DEBUG_LEVEL_FULL).build();
-        mInner.forceCreateNewVirtualMachine("test_vm", debugConfig);
+        VirtualMachineConfig debugConfig = builder.setDebugLevel(toLevel).build();
+        forceCreateNewVirtualMachine("test_vm", debugConfig);
         Files.copy(vmInstanceBackup.toPath(), vmInstance.toPath(), REPLACE_EXISTING);
         assertThat(tryBootVm(TAG, "test_vm").payloadStarted).isFalse();
     }
@@ -283,7 +438,7 @@
     }
 
     private VmCdis launchVmAndGetCdis(String instanceName) throws Exception {
-        VirtualMachine vm = mInner.getVirtualMachineManager().get(instanceName);
+        VirtualMachine vm = getVirtualMachineManager().get(instanceName);
         final VmCdis vmCdis = new VmCdis();
         final CompletableFuture<Exception> exception = new CompletableFuture<>();
         VmEventListener listener =
@@ -303,7 +458,10 @@
                     }
                 };
         listener.runToFinish(TAG, vm);
-        assertThat(exception.getNow(null)).isNull();
+        Exception e = exception.getNow(null);
+        if (e != null) {
+            throw new RuntimeException(e);
+        }
         return vmCdis;
     }
 
@@ -316,12 +474,13 @@
         assumeSupportedKernel();
 
         grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
-        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config.json")
-                .setDebugLevel(DEBUG_LEVEL_FULL)
-                .build();
-        mInner.forceCreateNewVirtualMachine("test_vm_a", normalConfig);
-        mInner.forceCreateNewVirtualMachine("test_vm_b", normalConfig);
+        VirtualMachineConfig normalConfig =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config.json")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        forceCreateNewVirtualMachine("test_vm_a", normalConfig);
+        forceCreateNewVirtualMachine("test_vm_b", normalConfig);
         VmCdis vm_a_cdis = launchVmAndGetCdis("test_vm_a");
         VmCdis vm_b_cdis = launchVmAndGetCdis("test_vm_b");
         assertThat(vm_a_cdis.cdiAttest).isNotNull();
@@ -339,13 +498,15 @@
     })
     public void sameInstanceKeepsSameCdis() throws Exception {
         assumeSupportedKernel();
+        assume().withMessage("Skip on CF. Too Slow. b/257270529").that(isCuttlefish()).isFalse();
 
         grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
-        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config.json")
-                .setDebugLevel(DEBUG_LEVEL_FULL)
-                .build();
-        mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        VirtualMachineConfig normalConfig =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config.json")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        forceCreateNewVirtualMachine("test_vm", normalConfig);
 
         VmCdis first_boot_cdis = launchVmAndGetCdis("test_vm");
         VmCdis second_boot_cdis = launchVmAndGetCdis("test_vm");
@@ -364,11 +525,12 @@
         assumeSupportedKernel();
 
         grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
-        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config.json")
-                .setDebugLevel(DEBUG_LEVEL_FULL)
-                .build();
-        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("bcc_vm", normalConfig);
+        VirtualMachineConfig normalConfig =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config.json")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("bcc_vm", normalConfig);
         final CompletableFuture<byte[]> bcc = new CompletableFuture<>();
         final CompletableFuture<Exception> exception = new CompletableFuture<>();
         VmEventListener listener =
@@ -404,6 +566,25 @@
         }
     }
 
+    @Test
+    @CddTest(requirements = {
+            "9.17/C-1-1",
+            "9.17/C-1-2"
+    })
+    public void accessToCdisIsRestricted() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        forceCreateNewVirtualMachine("test_vm", config);
+
+        assertThrows(Exception.class, () -> launchVmAndGetCdis("test_vm"));
+    }
+
+
     private static final UUID MICRODROID_PARTITION_UUID =
             UUID.fromString("cf9afe9a-0662-11ec-a329-c32663a09d75");
     private static final UUID U_BOOT_AVB_PARTITION_UUID =
@@ -440,17 +621,15 @@
     }
 
     private RandomAccessFile prepareInstanceImage(String vmName) throws Exception {
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                .setDebugLevel(DEBUG_LEVEL_FULL)
-                .build();
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
 
-        mInner.forceCreateNewVirtualMachine(vmName, config);
+        forceCreateNewVirtualMachine(vmName, config);
         assertThat(tryBootVm(TAG, vmName).payloadStarted).isTrue();
-
-        File vmRoot = new File(getContext().getFilesDir(), "vm");
-        File vmDir = new File(vmRoot, vmName);
-        File instanceImgPath = new File(vmDir, "instance.img");
+        File instanceImgPath = getVmFile(vmName, "instance.img");
         return new RandomAccessFile(instanceImgPath, "rw");
     }
 
@@ -503,11 +682,12 @@
     @Test
     public void bootFailsWhenConfigIsInvalid() throws Exception {
         grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
-        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder()
-                .setPayloadConfigPath("assets/vm_config_no_task.json")
-                .setDebugLevel(DEBUG_LEVEL_FULL)
-                .build();
-        mInner.forceCreateNewVirtualMachine("test_vm_invalid_config", normalConfig);
+        VirtualMachineConfig normalConfig =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config_no_task.json")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        forceCreateNewVirtualMachine("test_vm_invalid_config", normalConfig);
 
         BootResult bootResult = tryBootVm(TAG, "test_vm_invalid_config");
         assertThat(bootResult.payloadStarted).isFalse();
@@ -517,10 +697,10 @@
 
     @Test
     public void bootFailsWhenBinaryPathIsInvalid() throws Exception {
-        VirtualMachineConfig.Builder builder = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("DoesNotExist.so");
+        VirtualMachineConfig.Builder builder =
+                newVmConfigBuilder().setPayloadBinaryPath("DoesNotExist.so");
         VirtualMachineConfig normalConfig = builder.setDebugLevel(DEBUG_LEVEL_FULL).build();
-        mInner.forceCreateNewVirtualMachine("test_vm_invalid_binary_path", normalConfig);
+        forceCreateNewVirtualMachine("test_vm_invalid_binary_path", normalConfig);
 
         BootResult bootResult = tryBootVm(TAG, "test_vm_invalid_binary_path");
         assertThat(bootResult.payloadStarted).isFalse();
@@ -530,22 +710,104 @@
 
     @Test
     public void sameInstancesShareTheSameVmObject() throws Exception {
-        VirtualMachineConfig config = mInner.newVmConfigBuilder()
-                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
-                .setDebugLevel(DEBUG_LEVEL_NONE)
-                .build();
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .build();
 
-        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm", config);
-        VirtualMachine vm2 = mInner.getVirtualMachineManager().get("test_vm");
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
+        VirtualMachine vm2 = getVirtualMachineManager().get("test_vm");
         assertThat(vm).isEqualTo(vm2);
 
-        VirtualMachine newVm = mInner.forceCreateNewVirtualMachine("test_vm", config);
-        VirtualMachine newVm2 = mInner.getVirtualMachineManager().get("test_vm");
+        VirtualMachine newVm = forceCreateNewVirtualMachine("test_vm", config);
+        VirtualMachine newVm2 = getVirtualMachineManager().get("test_vm");
         assertThat(newVm).isEqualTo(newVm2);
 
         assertThat(vm).isNotEqualTo(newVm);
     }
 
+    @Test
+    public void importedVmAndOriginalVmHaveTheSameCdi() throws Exception {
+        assumeSupportedKernel();
+        // Arrange
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config.json")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        String vmNameOrig = "test_vm_orig";
+        String vmNameImport = "test_vm_import";
+        VirtualMachine vmOrig = forceCreateNewVirtualMachine(vmNameOrig, config);
+        VmCdis origCdis = launchVmAndGetCdis(vmNameOrig);
+        assertThat(origCdis.instanceSecret).isNotNull();
+        VirtualMachineDescriptor descriptor = vmOrig.toDescriptor();
+        VirtualMachineManager vmm = getVirtualMachineManager();
+        if (vmm.get(vmNameImport) != null) {
+            vmm.delete(vmNameImport);
+        }
+
+        // Action
+        // The imported VM will be fetched by name later.
+        VirtualMachine unusedVmImport = vmm.importFromDescriptor(vmNameImport, descriptor);
+
+        // Asserts
+        VmCdis importCdis = launchVmAndGetCdis(vmNameImport);
+        assertThat(origCdis.instanceSecret).isEqualTo(importCdis.instanceSecret);
+    }
+
+    @Test
+    public void importedVmIsEqualToTheOriginalVm() throws Exception {
+        // Arrange
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .build();
+        String vmNameOrig = "test_vm_orig";
+        String vmNameImport = "test_vm_import";
+        VirtualMachine vmOrig = forceCreateNewVirtualMachine(vmNameOrig, config);
+        // Run something to make the instance.img different with the initialized one.
+        TestResults origTestResults = runVmTestService(vmOrig);
+        assertThat(origTestResults.mException).isNull();
+        assertThat(origTestResults.mAddInteger).isEqualTo(123 + 456);
+        VirtualMachineDescriptor descriptor = vmOrig.toDescriptor();
+        VirtualMachineManager vmm = getVirtualMachineManager();
+        if (vmm.get(vmNameImport) != null) {
+            vmm.delete(vmNameImport);
+        }
+
+        // Action
+        VirtualMachine vmImport = vmm.importFromDescriptor(vmNameImport, descriptor);
+
+        // Asserts
+        assertFileContentsAreEqualInTwoVms("config.xml", vmNameOrig, vmNameImport);
+        assertFileContentsAreEqualInTwoVms("instance.img", vmNameOrig, vmNameImport);
+        assertThat(vmImport).isNotEqualTo(vmOrig);
+        vmm.delete(vmNameOrig);
+        assertThat(vmImport).isEqualTo(vmm.get(vmNameImport));
+        TestResults testResults = runVmTestService(vmImport);
+        assertThat(testResults.mException).isNull();
+        assertThat(testResults.mAddInteger).isEqualTo(123 + 456);
+    }
+
+    private void assertFileContentsAreEqualInTwoVms(String fileName, String vmName1, String vmName2)
+            throws IOException {
+        File file1 = getVmFile(vmName1, fileName);
+        File file2 = getVmFile(vmName2, fileName);
+        try (FileInputStream input1 = new FileInputStream(file1);
+                FileInputStream input2 = new FileInputStream(file2)) {
+            assertThat(Arrays.equals(input1.readAllBytes(), input2.readAllBytes())).isTrue();
+        }
+    }
+
+    private File getVmFile(String vmName, String fileName) {
+        Context context = ApplicationProvider.getApplicationContext();
+        Path filePath = Paths.get(context.getDataDir().getPath(), "vm", vmName, fileName);
+        return filePath.toFile();
+    }
+
     private int minMemoryRequired() {
         if (Build.SUPPORTED_ABIS.length > 0) {
             String primaryAbi = Build.SUPPORTED_ABIS[0];
@@ -573,6 +835,7 @@
         String mSublibRunProp;
         String mExtraApkTestProp;
         String mApkContentsPath;
+        String mEncryptedStoragePath;
     }
 
     private TestResults runVmTestService(VirtualMachine vm) throws Exception {
@@ -583,8 +846,9 @@
                 new VmEventListener() {
                     private void testVMService(VirtualMachine vm) {
                         try {
-                            ITestService testService = ITestService.Stub.asInterface(
-                                    vm.connectToVsockServer(ITestService.SERVICE_PORT));
+                            ITestService testService =
+                                    ITestService.Stub.asInterface(
+                                            vm.connectToVsockServer(ITestService.SERVICE_PORT));
                             testResults.mAddInteger = testService.addInteger(123, 456);
                             testResults.mAppRunProp =
                                     testService.readProperty("debug.microdroid.app.run");
@@ -593,6 +857,8 @@
                             testResults.mExtraApkTestProp =
                                     testService.readProperty("debug.microdroid.test.extra_apk");
                             testResults.mApkContentsPath = testService.getApkContentsPath();
+                            testResults.mEncryptedStoragePath =
+                                    testService.getEncryptedStoragePath();
                         } catch (Exception e) {
                             testResults.mException = e;
                         }
@@ -607,11 +873,9 @@
                     }
 
                     @Override
-                    public void onPayloadStarted(VirtualMachine vm, ParcelFileDescriptor stream) {
+                    public void onPayloadStarted(VirtualMachine vm) {
                         Log.i(TAG, "onPayloadStarted");
                         payloadStarted.complete(true);
-                        logVmOutput(TAG, new FileInputStream(stream.getFileDescriptor()),
-                                "Payload");
                     }
                 };
         listener.runToFinish(TAG, vm);
diff --git a/tests/testapk/src/native/crashbinary.cpp b/tests/testapk/src/native/crashbinary.cpp
deleted file mode 100644
index a0edc40..0000000
--- a/tests/testapk/src/native/crashbinary.cpp
+++ /dev/null
@@ -1,26 +0,0 @@
-/*
- * Copyright 2022 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <stdio.h>
-#include <stdlib.h>
-
-#include "vm_main.h"
-
-// A VM payload that crashes as soon as it starts, to allow us to exercise that error path.
-extern "C" int AVmPayload_main() {
-    printf("test crash!!!!\n");
-    abort();
-}
diff --git a/tests/benchmark/src/native/idlebinary.cpp b/tests/testapk/src/native/idlebinary.cpp
similarity index 100%
rename from tests/benchmark/src/native/idlebinary.cpp
rename to tests/testapk/src/native/idlebinary.cpp
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index 1a3e940..c0a8c0e 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -28,7 +28,7 @@
 #include <sys/system_properties.h>
 #include <unistd.h>
 #include <vm_main.h>
-#include <vm_payload.h>
+#include <vm_payload_restricted.h>
 
 #include <string>
 
@@ -76,40 +76,22 @@
         ndk::ScopedAStatus insecurelyExposeVmInstanceSecret(std::vector<uint8_t>* out) override {
             const uint8_t identifier[] = {1, 2, 3, 4};
             out->resize(32);
-            if (!AVmPayload_getVmInstanceSecret(identifier, sizeof(identifier), out->data(),
-                                                out->size())) {
-                return ndk::ScopedAStatus::
-                        fromServiceSpecificErrorWithMessage(0, "Failed to VM instance secret");
-            }
+            AVmPayload_getVmInstanceSecret(identifier, sizeof(identifier), out->data(),
+                                           out->size());
             return ndk::ScopedAStatus::ok();
         }
 
         ndk::ScopedAStatus insecurelyExposeAttestationCdi(std::vector<uint8_t>* out) override {
-            size_t cdi_size;
-            if (!AVmPayload_getDiceAttestationCdi(nullptr, 0, &cdi_size)) {
-                return ndk::ScopedAStatus::
-                        fromServiceSpecificErrorWithMessage(0, "Failed to measure attestation cdi");
-            }
+            size_t cdi_size = AVmPayload_getDiceAttestationCdi(nullptr, 0);
             out->resize(cdi_size);
-            if (!AVmPayload_getDiceAttestationCdi(out->data(), out->size(), &cdi_size)) {
-                return ndk::ScopedAStatus::
-                        fromServiceSpecificErrorWithMessage(0, "Failed to get attestation cdi");
-            }
+            AVmPayload_getDiceAttestationCdi(out->data(), out->size());
             return ndk::ScopedAStatus::ok();
         }
 
         ndk::ScopedAStatus getBcc(std::vector<uint8_t>* out) override {
-            size_t bcc_size;
-            if (!AVmPayload_getDiceAttestationChain(nullptr, 0, &bcc_size)) {
-                return ndk::ScopedAStatus::
-                        fromServiceSpecificErrorWithMessage(0,
-                                                            "Failed to measure attestation chain");
-            }
+            size_t bcc_size = AVmPayload_getDiceAttestationChain(nullptr, 0);
             out->resize(bcc_size);
-            if (!AVmPayload_getDiceAttestationChain(out->data(), out->size(), &bcc_size)) {
-                return ndk::ScopedAStatus::
-                        fromServiceSpecificErrorWithMessage(0, "Failed to get attestation chain");
-            }
+            AVmPayload_getDiceAttestationChain(out->data(), out->size());
             return ndk::ScopedAStatus::ok();
         }
 
@@ -123,19 +105,22 @@
             *out = path;
             return ndk::ScopedAStatus::ok();
         }
+
+        ndk::ScopedAStatus getEncryptedStoragePath(std::string* out) override {
+            const char* path_c = AVmPayload_getEncryptedStoragePath();
+            if (path_c == nullptr) {
+                out->clear();
+            } else {
+                *out = path_c;
+            }
+            return ndk::ScopedAStatus::ok();
+        }
     };
     auto testService = ndk::SharedRefBase::make<TestService>();
 
-    auto callback = []([[maybe_unused]] void* param) {
-        if (!AVmPayload_notifyPayloadReady()) {
-            std::cerr << "failed to notify payload ready to virtualizationservice" << std::endl;
-            abort();
-        }
-    };
-    if (!AVmPayload_runVsockRpcServer(testService->asBinder().get(), testService->SERVICE_PORT,
-                                      callback, nullptr)) {
-        return Error() << "RPC Server failed to run";
-    }
+    auto callback = []([[maybe_unused]] void* param) { AVmPayload_notifyPayloadReady(); };
+    AVmPayload_runVsockRpcServer(testService->asBinder().get(), testService->SERVICE_PORT, callback,
+                                 nullptr);
 
     return {};
 }
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 26d41c9..b767013 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -22,6 +22,7 @@
     rustlibs: [
         "android.system.virtualizationcommon-rust",
         "android.system.virtualizationservice-rust",
+        "android.system.virtualizationservice_internal-rust",
         "android.system.virtualmachineservice-rust",
         "android.os.permissions_aidl-rust",
         "libandroid_logger",
@@ -71,5 +72,8 @@
 rust_test {
     name: "virtualizationservice_device_test",
     defaults: ["virtualizationservice_defaults"],
+    rustlibs: [
+        "libtempfile",
+    ],
     test_suites: ["general-tests"],
 }
diff --git a/virtualizationservice/aidl/Android.bp b/virtualizationservice/aidl/Android.bp
index 4d5326a..a0bbc00 100644
--- a/virtualizationservice/aidl/Android.bp
+++ b/virtualizationservice/aidl/Android.bp
@@ -34,11 +34,31 @@
 }
 
 aidl_interface {
+    name: "android.system.virtualizationservice_internal",
+    srcs: ["android/system/virtualizationservice_internal/**/*.aidl"],
+    unstable: true,
+    backend: {
+        java: {
+            sdk_version: "module_current",
+        },
+        rust: {
+            enabled: true,
+            apex_available: [
+                "com.android.virt",
+            ],
+        },
+    },
+}
+
+aidl_interface {
     name: "android.system.virtualmachineservice",
     srcs: ["android/system/virtualmachineservice/**/*.aidl"],
     imports: ["android.system.virtualizationcommon"],
     unstable: true,
     backend: {
+        java: {
+            sdk_version: "module_current",
+        },
         rust: {
             enabled: true,
             apex_available: [
@@ -55,6 +75,7 @@
     unstable: true,
     backend: {
         java: {
+            sdk_version: "module_current",
             apex_available: ["com.android.virt"],
         },
         ndk: {
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
index 8d6ed08..a329fa6 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
@@ -24,13 +24,9 @@
  */
 oneway interface IVirtualMachineCallback {
     /**
-     * Called when the payload starts in the VM. `stream` is the input/output port of the payload.
-     *
-     * <p>Note: when the virtual machine object is shared to multiple processes and they register
-     * this callback to the same virtual machine object, the processes will compete to access the
-     * same payload stream. Keep only one process to access the stream.
+     * Called when the payload starts in the VM.
      */
-    void onPayloadStarted(int cid, in @nullable ParcelFileDescriptor stream);
+    void onPayloadStarted(int cid);
 
     /**
      * Called when the payload in the VM is ready to serve.
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/PartitionType.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/PartitionType.aidl
index f25e674..774681a 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/PartitionType.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/PartitionType.aidl
@@ -29,4 +29,8 @@
      * The partition is initialized as an instance image which is formatted to hold per-VM secrets
      */
     ANDROID_VM_INSTANCE = 1,
+    /**
+     * The partition is initialized to back encryptedstore disk image formatted to indicate intent
+     */
+    ENCRYPTEDSTORE = 2,
 }
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineCpuStatus.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IGlobalVmContext.aidl
similarity index 77%
rename from virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineCpuStatus.aidl
rename to virtualizationservice/aidl/android/system/virtualizationservice_internal/IGlobalVmContext.aidl
index 307c3f9..1a7aa4a 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineCpuStatus.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IGlobalVmContext.aidl
@@ -13,11 +13,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package android.system.virtualmachineservice;
+package android.system.virtualizationservice_internal;
 
-parcelable VirtualMachineCpuStatus {
-    long cpu_time_user;
-    long cpu_time_nice;
-    long cpu_time_sys;
-    long cpu_time_idle;
+interface IGlobalVmContext {
+    /** Get the CID allocated to the VM. */
+    int getCid();
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
new file mode 100644
index 0000000..851ddf4
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package android.system.virtualizationservice_internal;
+
+import android.system.virtualizationservice_internal.IGlobalVmContext;
+
+interface IVirtualizationServiceInternal {
+    /**
+     * Allocates global context for a new VM.
+     *
+     * This allocates VM's globally unique resources such as the CID.
+     * The resources will not be recycled as long as there is a strong reference
+     * to the returned object.
+     */
+    IGlobalVmContext allocateGlobalVmContext();
+}
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
index 4fa5fa0..3fdb48a 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -16,25 +16,11 @@
 package android.system.virtualmachineservice;
 
 import android.system.virtualizationcommon.ErrorCode;
-import android.system.virtualmachineservice.VirtualMachineCpuStatus;
-import android.system.virtualmachineservice.VirtualMachineMemStatus;
 
 /** {@hide} */
 interface IVirtualMachineService {
     /**
      * Port number that VirtualMachineService listens on connections from the guest VMs for the
-     * payload input and output.
-     */
-    const int VM_STREAM_SERVICE_PORT = 3000;
-
-    /**
-     * Port number that VirtualMachineService listens on connections from the guest VMs for the
-     * VirtualMachineService binder service.
-     */
-    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;
@@ -55,17 +41,7 @@
     void notifyPayloadFinished(int exitCode);
 
     /**
-     * Notifies that an error has occurred inside the VM..
+     * Notifies that an error has occurred inside the VM.
      */
     void notifyError(ErrorCode errorCode, in String message);
-
-    /**
-     * Notifies the current CPU status of the VM.
-     */
-    void notifyCpuStatus(in VirtualMachineCpuStatus cpuStatus);
-
-    /**
-     * Notifies the current memory status of the VM.
-     */
-    void notifyMemStatus(in VirtualMachineMemStatus memStatus);
 }
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineMemStatus.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineMemStatus.aidl
deleted file mode 100644
index 3de57c6..0000000
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineMemStatus.aidl
+++ /dev/null
@@ -1,24 +0,0 @@
-/*
- * Copyright 2022 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package android.system.virtualmachineservice;
-
-parcelable VirtualMachineMemStatus {
-    long mem_total;
-    long mem_free;
-    long mem_available;
-    long mem_buffer;
-    long mem_cached;
-}
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index b4ce9d2..040c0d8 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -14,12 +14,9 @@
 
 //! Implementation of the AIDL interface of the VirtualizationService.
 
-use crate::atom::{
-    write_vm_booted_stats, write_vm_cpu_status_stats, write_vm_creation_stats,
-    write_vm_mem_status_stats,
-};
+use crate::atom::{write_vm_booted_stats, write_vm_creation_stats};
 use crate::composite::make_composite_image;
-use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmInstance, VmState};
+use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmContext, VmInstance, VmState};
 use crate::payload::{add_microdroid_payload_images, add_microdroid_system_images};
 use crate::selinux::{getfilecon, SeContext};
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
@@ -39,27 +36,27 @@
     VirtualMachineRawConfig::VirtualMachineRawConfig,
     VirtualMachineState::VirtualMachineState,
 };
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
-    IVirtualMachineService::{
-        BnVirtualMachineService, IVirtualMachineService, VM_BINDER_SERVICE_PORT,
-        VM_STREAM_SERVICE_PORT, VM_TOMBSTONES_SERVICE_PORT,
-    },
-    VirtualMachineCpuStatus::VirtualMachineCpuStatus,
-    VirtualMachineMemStatus::VirtualMachineMemStatus,
+use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::{
+    IGlobalVmContext::{BnGlobalVmContext, IGlobalVmContext},
+    IVirtualizationServiceInternal::{BnVirtualizationServiceInternal, IVirtualizationServiceInternal},
+};
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
+        BnVirtualMachineService, IVirtualMachineService, VM_TOMBSTONES_SERVICE_PORT,
 };
 use anyhow::{anyhow, bail, Context, Result};
 use apkverify::{HashAlgorithm, V4Signature};
 use binder::{
     self, BinderFeatures, ExceptionCode, Interface, LazyServiceGuard, ParcelFileDescriptor,
-    SpIBinder, Status, StatusCode, Strong, ThreadState,
+    Status, StatusCode, Strong, ThreadState,
 };
 use disk::QcowFile;
 use libc::VMADDR_CID_HOST;
 use log::{debug, error, info, warn};
 use microdroid_payload_config::{OsConfig, Task, TaskType, VmPayloadConfig};
-use rpcbinder::run_vsock_rpc_server_with_factory;
+use rpcbinder::RpcServer;
 use rustutils::system_properties;
 use semver::VersionReq;
+use std::collections::HashMap;
 use std::convert::TryInto;
 use std::ffi::CStr;
 use std::fs::{create_dir, File, OpenOptions};
@@ -83,7 +80,8 @@
 
 /// The first CID to assign to a guest VM managed by the VirtualizationService. CIDs lower than this
 /// are reserved for the host or other usage.
-const FIRST_GUEST_CID: Cid = 10;
+const GUEST_CID_MIN: Cid = 2048;
+const GUEST_CID_MAX: Cid = 65535;
 
 const SYSPROP_LAST_CID: &str = "virtualizationservice.state.last_cid";
 
@@ -101,10 +99,149 @@
 
 const MICRODROID_OS_NAME: &str = "microdroid";
 
-/// Implementation of `IVirtualizationService`, the entry point of the AIDL service.
+const UNFORMATTED_STORAGE_MAGIC: &str = "UNFORMATTED-STORAGE";
+
+fn is_valid_guest_cid(cid: Cid) -> bool {
+    (GUEST_CID_MIN..=GUEST_CID_MAX).contains(&cid)
+}
+
+fn next_guest_cid(cid: Cid) -> Cid {
+    assert!(is_valid_guest_cid(cid));
+    if cid == GUEST_CID_MAX {
+        GUEST_CID_MIN
+    } else {
+        cid + 1
+    }
+}
+
+/// Singleton service for allocating globally-unique VM resources, such as the CID, and running
+/// singleton servers, like tombstone receiver.
 #[derive(Debug, Default)]
+pub struct VirtualizationServiceInternal {
+    state: Arc<Mutex<GlobalState>>,
+}
+
+impl VirtualizationServiceInternal {
+    pub fn init() -> VirtualizationServiceInternal {
+        let service = VirtualizationServiceInternal::default();
+
+        std::thread::spawn(|| {
+            if let Err(e) = handle_stream_connection_tombstoned() {
+                warn!("Error receiving tombstone from guest or writing them. Error: {:?}", e);
+            }
+        });
+
+        service
+    }
+}
+
+impl Interface for VirtualizationServiceInternal {}
+
+impl IVirtualizationServiceInternal for VirtualizationServiceInternal {
+    fn allocateGlobalVmContext(&self) -> binder::Result<Strong<dyn IGlobalVmContext>> {
+        let state = &mut *self.state.lock().unwrap();
+        let cid = state.allocate_cid().map_err(|e| {
+            Status::new_exception_str(ExceptionCode::ILLEGAL_STATE, Some(e.to_string()))
+        })?;
+        Ok(GlobalVmContext::create(cid))
+    }
+}
+
+/// The mutable state of the VirtualizationServiceInternal. There should only be one instance
+/// of this struct.
+#[derive(Debug, Default)]
+struct GlobalState {
+    /// CIDs currently allocated to running VMs. A CID is never recycled as long
+    /// as there is a strong reference held by a GlobalVmContext.
+    held_cids: HashMap<Cid, Weak<Cid>>,
+}
+
+impl GlobalState {
+    /// Get the next available CID, or an error if we have run out. The last CID used is stored in
+    /// a system property so that restart of virtualizationservice doesn't reuse CID while the host
+    /// Android is up.
+    fn allocate_cid(&mut self) -> Result<Arc<Cid>> {
+        // Garbage collect unused CIDs.
+        self.held_cids.retain(|_, cid| cid.strong_count() > 0);
+
+        // Start trying to find a CID from the last used CID + 1. This ensures
+        // that we do not eagerly recycle CIDs. It makes debugging easier but
+        // also means that retrying to allocate a CID, eg. because it is
+        // erroneously occupied by a process, will not recycle the same CID.
+        let last_cid_prop =
+            system_properties::read(SYSPROP_LAST_CID)?.and_then(|val| match val.parse::<Cid>() {
+                Ok(num) => {
+                    if is_valid_guest_cid(num) {
+                        Some(num)
+                    } else {
+                        error!("Invalid value '{}' of property '{}'", num, SYSPROP_LAST_CID);
+                        None
+                    }
+                }
+                Err(_) => {
+                    error!("Invalid value '{}' of property '{}'", val, SYSPROP_LAST_CID);
+                    None
+                }
+            });
+
+        let first_cid = if let Some(last_cid) = last_cid_prop {
+            next_guest_cid(last_cid)
+        } else {
+            GUEST_CID_MIN
+        };
+
+        let cid = self
+            .find_available_cid(first_cid..=GUEST_CID_MAX)
+            .or_else(|| self.find_available_cid(GUEST_CID_MIN..first_cid));
+
+        if let Some(cid) = cid {
+            let cid_arc = Arc::new(cid);
+            self.held_cids.insert(cid, Arc::downgrade(&cid_arc));
+            system_properties::write(SYSPROP_LAST_CID, &format!("{}", cid))?;
+            Ok(cid_arc)
+        } else {
+            Err(anyhow!("Could not find an available CID."))
+        }
+    }
+
+    fn find_available_cid<I>(&self, mut range: I) -> Option<Cid>
+    where
+        I: Iterator<Item = Cid>,
+    {
+        range.find(|cid| !self.held_cids.contains_key(cid))
+    }
+}
+
+/// Implementation of the AIDL `IGlobalVmContext` interface.
+#[derive(Debug, Default)]
+struct GlobalVmContext {
+    /// The unique CID assigned to the VM for vsock communication.
+    cid: Arc<Cid>,
+    /// Keeps our service process running as long as this VM context exists.
+    #[allow(dead_code)]
+    lazy_service_guard: LazyServiceGuard,
+}
+
+impl GlobalVmContext {
+    fn create(cid: Arc<Cid>) -> Strong<dyn IGlobalVmContext> {
+        let binder = GlobalVmContext { cid, ..Default::default() };
+        BnGlobalVmContext::new_binder(binder, BinderFeatures::default())
+    }
+}
+
+impl Interface for GlobalVmContext {}
+
+impl IGlobalVmContext for GlobalVmContext {
+    fn getCid(&self) -> binder::Result<i32> {
+        Ok(*self.cid as i32)
+    }
+}
+
+/// Implementation of `IVirtualizationService`, the entry point of the AIDL service.
+#[derive(Debug)]
 pub struct VirtualizationService {
     state: Arc<Mutex<State>>,
+    global_service: Strong<dyn IVirtualizationServiceInternal>,
 }
 
 impl Interface for VirtualizationService {
@@ -180,6 +317,7 @@
         match partition_type {
             PartitionType::RAW => Ok(()),
             PartitionType::ANDROID_VM_INSTANCE => format_as_android_vm_instance(&mut part),
+            PartitionType::ENCRYPTEDSTORE => format_as_encryptedstore(&mut part),
             _ => Err(Error::new(
                 ErrorKind::Unsupported,
                 format!("Unsupported partition type {:?}", partition_type),
@@ -257,8 +395,10 @@
 }
 
 fn handle_stream_connection_tombstoned() -> Result<()> {
+    // Should not listen for tombstones on a guest VM's port.
+    assert!(!is_valid_guest_cid(VM_TOMBSTONES_SERVICE_PORT as Cid));
     let listener =
-        VsockListener::bind_with_cid_port(VMADDR_CID_HOST, VM_TOMBSTONES_SERVICE_PORT as u32)?;
+        VsockListener::bind_with_cid_port(VMADDR_CID_HOST, VM_TOMBSTONES_SERVICE_PORT as Cid)?;
     for incoming_stream in listener.incoming() {
         let mut incoming_stream = match incoming_stream {
             Err(e) => {
@@ -306,34 +446,35 @@
 
 impl VirtualizationService {
     pub fn init() -> VirtualizationService {
-        let service = VirtualizationService::default();
+        let global_service = VirtualizationServiceInternal::init();
+        let global_service =
+            BnVirtualizationServiceInternal::new_binder(global_service, BinderFeatures::default());
 
-        // server for payload output
-        let state = service.state.clone(); // reference to state (not the state itself) is copied
-        std::thread::spawn(move || {
-            handle_stream_connection_from_vm(state).unwrap();
-        });
+        VirtualizationService { global_service, state: Default::default() }
+    }
 
-        std::thread::spawn(|| {
-            if let Err(e) = handle_stream_connection_tombstoned() {
-                warn!("Error receiving tombstone from guest or writing them. Error: {:?}", e);
+    fn create_vm_context(&self) -> Result<(VmContext, Cid)> {
+        const NUM_ATTEMPTS: usize = 5;
+
+        for _ in 0..NUM_ATTEMPTS {
+            let global_context = self.global_service.allocateGlobalVmContext()?;
+            let cid = global_context.getCid()? as Cid;
+            let service = VirtualMachineService::new_binder(self.state.clone(), cid).as_binder();
+
+            // Start VM service listening for connections from the new CID on port=CID.
+            // TODO(b/245727626): Only accept connections from the new VM.
+            let port = cid;
+            match RpcServer::new_vsock(service, port) {
+                Ok(vm_server) => {
+                    vm_server.start();
+                    return Ok((VmContext::new(global_context, vm_server), cid));
+                }
+                Err(err) => {
+                    warn!("Could not start RpcServer on port {}: {}", port, err);
+                }
             }
-        });
-
-        // binder server for vm
-        // reference to state (not the state itself) is copied
-        let state = service.state.clone();
-        std::thread::spawn(move || {
-            debug!("VirtualMachineService is starting as an RPC service.");
-            if run_vsock_rpc_server_with_factory(VM_BINDER_SERVICE_PORT as u32, |cid| {
-                VirtualMachineService::factory(cid, &state)
-            }) {
-                debug!("RPC server has shut down gracefully");
-            } else {
-                panic!("Premature termination of RPC server");
-            }
-        });
-        service
+        }
+        bail!("Too many attempts to create VM context failed.");
     }
 
     fn create_vm_internal(
@@ -359,12 +500,19 @@
             check_use_custom_virtual_machine()?;
         }
 
+        let (vm_context, cid) = self.create_vm_context().map_err(|e| {
+            error!("Failed to create VmContext: {:?}", e);
+            Status::new_service_specific_error_str(
+                -1,
+                Some(format!("Failed to create VmContext: {:?}", e)),
+            )
+        })?;
+
         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_debug_pid = ThreadState::get_calling_pid();
-        let cid = state.next_cid().or(Err(ExceptionCode::ILLEGAL_STATE))?;
 
         // Counter to generate unique IDs for temporary image files.
         let mut next_temporary_image_id = 0;
@@ -481,47 +629,26 @@
             detect_hangup: is_app_config,
         };
         let instance = Arc::new(
-            VmInstance::new(crosvm_config, temporary_directory, requester_uid, requester_debug_pid)
-                .map_err(|e| {
-                    error!("Failed to create VM with config {:?}: {:?}", config, e);
-                    Status::new_service_specific_error_str(
-                        -1,
-                        Some(format!("Failed to create VM: {:?}", e)),
-                    )
-                })?,
+            VmInstance::new(
+                crosvm_config,
+                temporary_directory,
+                requester_uid,
+                requester_debug_pid,
+                vm_context,
+            )
+            .map_err(|e| {
+                error!("Failed to create VM with config {:?}: {:?}", config, e);
+                Status::new_service_specific_error_str(
+                    -1,
+                    Some(format!("Failed to create VM: {:?}", e)),
+                )
+            })?,
         );
         state.add_vm(Arc::downgrade(&instance));
         Ok(VirtualMachine::create(instance))
     }
 }
 
-/// Waits for incoming connections from VM. If a new connection is made, stores the stream in the
-/// corresponding `VmInstance`.
-fn handle_stream_connection_from_vm(state: Arc<Mutex<State>>) -> Result<()> {
-    let listener =
-        VsockListener::bind_with_cid_port(VMADDR_CID_HOST, VM_STREAM_SERVICE_PORT as u32)?;
-    for stream in listener.incoming() {
-        let stream = match stream {
-            Err(e) => {
-                warn!("invalid incoming connection: {:?}", e);
-                continue;
-            }
-            Ok(s) => s,
-        };
-        if let Ok(addr) = stream.peer_addr() {
-            let cid = addr.cid();
-            let port = addr.port();
-            info!("payload stream connected from cid={}, port={}", cid, port);
-            if let Some(vm) = state.lock().unwrap().get_vm(cid) {
-                *vm.stream.lock().unwrap() = Some(stream);
-            } else {
-                error!("connection from cid={} is not from a guest VM", cid);
-            }
-        }
-    }
-    Ok(())
-}
-
 fn write_zero_filler(zero_filler_path: &Path) -> Result<()> {
     let file = OpenOptions::new()
         .create_new(true)
@@ -539,9 +666,13 @@
     part.flush()
 }
 
+fn format_as_encryptedstore(part: &mut dyn Write) -> std::io::Result<()> {
+    part.write_all(UNFORMATTED_STORAGE_MAGIC.as_bytes())?;
+    part.flush()
+}
+
 fn prepare_ramdump_file(ramdump_path: &Path) -> Result<File> {
-    File::create(&ramdump_path)
-        .context(format!("Failed to create ramdump file {:?}", &ramdump_path))
+    File::create(ramdump_path).context(format!("Failed to create ramdump file {:?}", &ramdump_path))
 }
 
 /// Given the configuration for a disk image, assembles the `DiskFile` to pass to crosvm.
@@ -861,11 +992,10 @@
 
 impl VirtualMachineCallbacks {
     /// Call all registered callbacks to notify that the payload has started.
-    pub fn notify_payload_started(&self, cid: Cid, stream: Option<VsockStream>) {
+    pub fn notify_payload_started(&self, cid: Cid) {
         let callbacks = &*self.0.lock().unwrap();
-        let pfd = stream.map(vsock_stream_to_pfd);
         for callback in callbacks {
-            if let Err(e) = callback.onPayloadStarted(cid as i32, pfd.as_ref()) {
+            if let Err(e) = callback.onPayloadStarted(cid as i32) {
                 error!("Error notifying payload start event from VM CID {}: {:?}", cid, e);
             }
         }
@@ -975,27 +1105,6 @@
         let vm = self.debug_held_vms.swap_remove(pos);
         Some(vm)
     }
-
-    /// Get the next available CID, or an error if we have run out. The last CID used is stored in
-    /// a system property so that restart of virtualizationservice doesn't reuse CID while the host
-    /// Android is up.
-    fn next_cid(&mut self) -> Result<Cid> {
-        let next = if let Some(val) = system_properties::read(SYSPROP_LAST_CID)? {
-            if let Ok(num) = val.parse::<u32>() {
-                num.checked_add(1).ok_or_else(|| anyhow!("run out of CID"))?
-            } else {
-                error!("Invalid last CID {}. Using {}", &val, FIRST_GUEST_CID);
-                FIRST_GUEST_CID
-            }
-        } else {
-            // First VM since the boot
-            FIRST_GUEST_CID
-        };
-        // Persist the last value for next use
-        let str_val = format!("{}", next);
-        system_properties::write(SYSPROP_LAST_CID, &str_val)?;
-        Ok(next)
-    }
 }
 
 /// Gets the `VirtualMachineState` of the given `VmInstance`.
@@ -1079,11 +1188,10 @@
             vm.update_payload_state(PayloadState::Started).map_err(|e| {
                 Status::new_exception_str(ExceptionCode::ILLEGAL_STATE, Some(e.to_string()))
             })?;
-            let stream = vm.stream.lock().unwrap().take();
-            vm.callbacks.notify_payload_started(cid, stream);
+            vm.callbacks.notify_payload_started(cid);
 
-            let vm_start_timestamp = vm.vm_start_timestamp.lock().unwrap();
-            write_vm_booted_stats(vm.requester_uid as i32, &vm.name, *vm_start_timestamp);
+            let vm_start_timestamp = vm.vm_metric.lock().unwrap().start_timestamp;
+            write_vm_booted_stats(vm.requester_uid as i32, &vm.name, vm_start_timestamp);
             Ok(())
         } else {
             error!("notifyPayloadStarted is called from an unknown CID {}", cid);
@@ -1147,50 +1255,9 @@
             ))
         }
     }
-
-    fn notifyCpuStatus(&self, status: &VirtualMachineCpuStatus) -> binder::Result<()> {
-        let cid = self.cid;
-        if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
-            info!("VM with CID {} reported its CPU status", cid);
-            write_vm_cpu_status_stats(vm.requester_uid as i32, &vm.name, status);
-            Ok(())
-        } else {
-            error!("notifyCurrentStatus is called from an unknown CID {}", cid);
-            Err(Status::new_service_specific_error_str(
-                -1,
-                Some(format!("cannot find a VM with CID {}", cid)),
-            ))
-        }
-    }
-
-    fn notifyMemStatus(&self, status: &VirtualMachineMemStatus) -> binder::Result<()> {
-        let cid = self.cid;
-        if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
-            info!("VM with CID {} reported its memory status", cid);
-            write_vm_mem_status_stats(vm.requester_uid as i32, &vm.name, status);
-            Ok(())
-        } else {
-            error!("notifyCurrentStatus is called from an unknown CID {}", cid);
-            Err(Status::new_service_specific_error_str(
-                -1,
-                Some(format!("cannot find a VM with CID {}", cid)),
-            ))
-        }
-    }
 }
 
 impl VirtualMachineService {
-    fn factory(cid: Cid, state: &Arc<Mutex<State>>) -> Option<SpIBinder> {
-        if let Some(vm) = state.lock().unwrap().get_vm(cid) {
-            let mut vm_service = vm.vm_service.lock().unwrap();
-            let service = vm_service.get_or_insert_with(|| Self::new_binder(state.clone(), cid));
-            Some(service.as_binder())
-        } else {
-            error!("connection from cid={} is not from a guest VM", cid);
-            None
-        }
-    }
-
     fn new_binder(state: Arc<Mutex<State>>, cid: Cid) -> Strong<dyn IVirtualMachineService> {
         BnVirtualMachineService::new_binder(
             VirtualMachineService { state, cid },
diff --git a/virtualizationservice/src/atom.rs b/virtualizationservice/src/atom.rs
index 8c46ac5..9c74d1e 100644
--- a/virtualizationservice/src/atom.rs
+++ b/virtualizationservice/src/atom.rs
@@ -15,6 +15,7 @@
 //! Functions for creating and collecting atoms.
 
 use crate::aidl::clone_file;
+use crate::crosvm::VmMetric;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
     DeathReason::DeathReason,
     IVirtualMachine::IVirtualMachine,
@@ -22,17 +23,13 @@
     VirtualMachineConfig::VirtualMachineConfig,
 };
 use android_system_virtualizationservice::binder::{Status, Strong};
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
-    VirtualMachineCpuStatus::VirtualMachineCpuStatus,
-    VirtualMachineMemStatus::VirtualMachineMemStatus,
-};
 use anyhow::{anyhow, Result};
 use binder::{ParcelFileDescriptor, ThreadState};
 use log::{trace, warn};
 use microdroid_payload_config::VmPayloadConfig;
-use statslog_virtualization_rust::{
-    vm_booted, vm_cpu_status_reported, vm_creation_requested, vm_exited, vm_mem_status_reported,
-};
+use rustutils::system_properties;
+use statslog_virtualization_rust::{vm_booted, vm_creation_requested, vm_exited};
+use std::thread;
 use std::time::{Duration, SystemTime};
 use zip::ZipArchive;
 
@@ -92,17 +89,16 @@
             binder_exception_code = e.exception_code() as i32;
         }
     }
-
     let (vm_identifier, config_type, num_cpus, memory_mib, apexes) = match config {
         VirtualMachineConfig::AppConfig(config) => (
-            &config.name,
+            config.name.clone(),
             vm_creation_requested::ConfigType::VirtualMachineAppConfig,
             config.numCpus,
             config.memoryMib,
             get_apex_list(config),
         ),
         VirtualMachineConfig::RawConfig(config) => (
-            &config.name,
+            config.name.clone(),
             vm_creation_requested::ConfigType::VirtualMachineRawConfig,
             config.numCpus,
             config.memoryMib,
@@ -110,150 +106,145 @@
         ),
     };
 
-    let vm_creation_requested = vm_creation_requested::VmCreationRequested {
-        uid: ThreadState::get_calling_uid() as i32,
-        vm_identifier,
-        hypervisor: vm_creation_requested::Hypervisor::Pkvm,
-        is_protected,
-        creation_succeeded,
-        binder_exception_code,
-        config_type,
-        num_cpus,
-        cpu_affinity: "", // deprecated
-        memory_mib,
-        apexes: &apexes,
-        // TODO(seungjaeyoo) Fill information about task_profile
-        // TODO(seungjaeyoo) Fill information about disk_image for raw config
-    };
+    let uid = ThreadState::get_calling_uid() as i32;
+    thread::spawn(move || {
+        let vm_creation_requested = vm_creation_requested::VmCreationRequested {
+            uid,
+            vm_identifier: &vm_identifier,
+            hypervisor: vm_creation_requested::Hypervisor::Pkvm,
+            is_protected,
+            creation_succeeded,
+            binder_exception_code,
+            config_type,
+            num_cpus,
+            cpu_affinity: "", // deprecated
+            memory_mib,
+            apexes: &apexes,
+            // TODO(seungjaeyoo) Fill information about task_profile
+            // TODO(seungjaeyoo) Fill information about disk_image for raw config
+        };
 
-    match vm_creation_requested.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
+        wait_for_statsd().unwrap_or_else(|e| warn!("failed to wait for statsd with error: {}", e));
+        match vm_creation_requested.stats_write() {
+            Err(e) => {
+                warn!("statslog_rust failed with error: {}", e);
+            }
+            Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
         }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
+    });
 }
 
 /// Write the stats of VM boot to statsd
+/// The function creates a separate thread which waits fro statsd to start to push atom
 pub fn write_vm_booted_stats(
     uid: i32,
-    vm_identifier: &String,
+    vm_identifier: &str,
     vm_start_timestamp: Option<SystemTime>,
 ) {
+    let vm_identifier = vm_identifier.to_owned();
     let duration = get_duration(vm_start_timestamp);
-    let vm_booted = vm_booted::VmBooted {
-        uid,
-        vm_identifier,
-        elapsed_time_millis: duration.as_millis() as i64,
-    };
-    match vm_booted.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
+    thread::spawn(move || {
+        let vm_booted = vm_booted::VmBooted {
+            uid,
+            vm_identifier: &vm_identifier,
+            elapsed_time_millis: duration.as_millis() as i64,
+        };
+        wait_for_statsd().unwrap_or_else(|e| warn!("failed to wait for statsd with error: {}", e));
+        match vm_booted.stats_write() {
+            Err(e) => {
+                warn!("statslog_rust failed with error: {}", e);
+            }
+            Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
         }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
+    });
 }
 
 /// Write the stats of VM exit to statsd
+/// The function creates a separate thread which waits fro statsd to start to push atom
 pub fn write_vm_exited_stats(
     uid: i32,
-    vm_identifier: &String,
+    vm_identifier: &str,
     reason: DeathReason,
-    vm_start_timestamp: Option<SystemTime>,
+    vm_metric: &VmMetric,
 ) {
-    let duration = get_duration(vm_start_timestamp);
-    let vm_exited = vm_exited::VmExited {
-        uid,
-        vm_identifier,
-        elapsed_time_millis: duration.as_millis() as i64,
-        death_reason: match reason {
-            DeathReason::INFRASTRUCTURE_ERROR => vm_exited::DeathReason::InfrastructureError,
-            DeathReason::KILLED => vm_exited::DeathReason::Killed,
-            DeathReason::UNKNOWN => vm_exited::DeathReason::Unknown,
-            DeathReason::SHUTDOWN => vm_exited::DeathReason::Shutdown,
-            DeathReason::ERROR => vm_exited::DeathReason::Error,
-            DeathReason::REBOOT => vm_exited::DeathReason::Reboot,
-            DeathReason::CRASH => vm_exited::DeathReason::Crash,
-            DeathReason::PVM_FIRMWARE_PUBLIC_KEY_MISMATCH => {
-                vm_exited::DeathReason::PvmFirmwarePublicKeyMismatch
+    let vm_identifier = vm_identifier.to_owned();
+    let elapsed_time_millis = get_duration(vm_metric.start_timestamp).as_millis() as i64;
+    let guest_time_millis = vm_metric.cpu_guest_time.unwrap_or_default();
+    let rss = vm_metric.rss.unwrap_or_default();
+
+    thread::spawn(move || {
+        let vm_exited = vm_exited::VmExited {
+            uid,
+            vm_identifier: &vm_identifier,
+            elapsed_time_millis,
+            death_reason: match reason {
+                DeathReason::INFRASTRUCTURE_ERROR => vm_exited::DeathReason::InfrastructureError,
+                DeathReason::KILLED => vm_exited::DeathReason::Killed,
+                DeathReason::UNKNOWN => vm_exited::DeathReason::Unknown,
+                DeathReason::SHUTDOWN => vm_exited::DeathReason::Shutdown,
+                DeathReason::ERROR => vm_exited::DeathReason::Error,
+                DeathReason::REBOOT => vm_exited::DeathReason::Reboot,
+                DeathReason::CRASH => vm_exited::DeathReason::Crash,
+                DeathReason::PVM_FIRMWARE_PUBLIC_KEY_MISMATCH => {
+                    vm_exited::DeathReason::PvmFirmwarePublicKeyMismatch
+                }
+                DeathReason::PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED => {
+                    vm_exited::DeathReason::PvmFirmwareInstanceImageChanged
+                }
+                DeathReason::BOOTLOADER_PUBLIC_KEY_MISMATCH => {
+                    vm_exited::DeathReason::BootloaderPublicKeyMismatch
+                }
+                DeathReason::BOOTLOADER_INSTANCE_IMAGE_CHANGED => {
+                    vm_exited::DeathReason::BootloaderInstanceImageChanged
+                }
+                DeathReason::MICRODROID_FAILED_TO_CONNECT_TO_VIRTUALIZATION_SERVICE => {
+                    vm_exited::DeathReason::MicrodroidFailedToConnectToVirtualizationService
+                }
+                DeathReason::MICRODROID_PAYLOAD_HAS_CHANGED => {
+                    vm_exited::DeathReason::MicrodroidPayloadHasChanged
+                }
+                DeathReason::MICRODROID_PAYLOAD_VERIFICATION_FAILED => {
+                    vm_exited::DeathReason::MicrodroidPayloadVerificationFailed
+                }
+                DeathReason::MICRODROID_INVALID_PAYLOAD_CONFIG => {
+                    vm_exited::DeathReason::MicrodroidInvalidPayloadConfig
+                }
+                DeathReason::MICRODROID_UNKNOWN_RUNTIME_ERROR => {
+                    vm_exited::DeathReason::MicrodroidUnknownRuntimeError
+                }
+                DeathReason::HANGUP => vm_exited::DeathReason::Hangup,
+                _ => vm_exited::DeathReason::Unknown,
+            },
+            guest_time_millis,
+            rss_vm_kb: rss.vm,
+            rss_crosvm_kb: rss.crosvm,
+        };
+        wait_for_statsd().unwrap_or_else(|e| warn!("failed to wait for statsd with error: {}", e));
+        match vm_exited.stats_write() {
+            Err(e) => {
+                warn!("statslog_rust failed with error: {}", e);
             }
-            DeathReason::PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED => {
-                vm_exited::DeathReason::PvmFirmwareInstanceImageChanged
-            }
-            DeathReason::BOOTLOADER_PUBLIC_KEY_MISMATCH => {
-                vm_exited::DeathReason::BootloaderPublicKeyMismatch
-            }
-            DeathReason::BOOTLOADER_INSTANCE_IMAGE_CHANGED => {
-                vm_exited::DeathReason::BootloaderInstanceImageChanged
-            }
-            DeathReason::MICRODROID_FAILED_TO_CONNECT_TO_VIRTUALIZATION_SERVICE => {
-                vm_exited::DeathReason::MicrodroidFailedToConnectToVirtualizationService
-            }
-            DeathReason::MICRODROID_PAYLOAD_HAS_CHANGED => {
-                vm_exited::DeathReason::MicrodroidPayloadHasChanged
-            }
-            DeathReason::MICRODROID_PAYLOAD_VERIFICATION_FAILED => {
-                vm_exited::DeathReason::MicrodroidPayloadVerificationFailed
-            }
-            DeathReason::MICRODROID_INVALID_PAYLOAD_CONFIG => {
-                vm_exited::DeathReason::MicrodroidInvalidPayloadConfig
-            }
-            DeathReason::MICRODROID_UNKNOWN_RUNTIME_ERROR => {
-                vm_exited::DeathReason::MicrodroidUnknownRuntimeError
-            }
-            DeathReason::HANGUP => vm_exited::DeathReason::Hangup,
-            _ => vm_exited::DeathReason::Unknown,
-        },
-    };
-    match vm_exited.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
+            Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
         }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
+    });
 }
 
-/// Write the stats of VM cpu status to statsd
-pub fn write_vm_cpu_status_stats(
-    uid: i32,
-    vm_identifier: &String,
-    cpu_status: &VirtualMachineCpuStatus,
-) {
-    let vm_cpu_status_reported = vm_cpu_status_reported::VmCpuStatusReported {
-        uid,
-        vm_identifier,
-        cpu_time_user_millis: cpu_status.cpu_time_user,
-        cpu_time_nice_millis: cpu_status.cpu_time_nice,
-        cpu_time_sys_millis: cpu_status.cpu_time_sys,
-        cpu_time_idle_millis: cpu_status.cpu_time_idle,
-    };
-    match vm_cpu_status_reported.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
+fn wait_for_statsd() -> Result<()> {
+    let mut prop = system_properties::PropertyWatcher::new("init.svc.statsd")?;
+    loop {
+        prop.wait()?;
+        match system_properties::read("init.svc.statsd")? {
+            Some(s) => {
+                if s == "running" {
+                    break;
+                }
+            }
+            None => {
+                // This case never really happens because
+                // prop.wait() waits for property to be non-null.
+                break;
+            }
         }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
     }
-}
-
-/// Write the stats of VM memory status to statsd
-pub fn write_vm_mem_status_stats(
-    uid: i32,
-    vm_identifier: &String,
-    mem_status: &VirtualMachineMemStatus,
-) {
-    let vm_mem_status_reported = vm_mem_status_reported::VmMemStatusReported {
-        uid,
-        vm_identifier,
-        mem_total_kb: mem_status.mem_total,
-        mem_free_kb: mem_status.mem_free,
-        mem_available_kb: mem_status.mem_available,
-        mem_buffer_kb: mem_status.mem_buffer,
-        mem_cached_kb: mem_status.mem_cached,
-    };
-    match vm_mem_status_reported.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
-        }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
+    Ok(())
 }
diff --git a/virtualizationservice/src/composite.rs b/virtualizationservice/src/composite.rs
index c9a68ac..fe17ff4 100644
--- a/virtualizationservice/src/composite.rs
+++ b/virtualizationservice/src/composite.rs
@@ -51,7 +51,7 @@
         OpenOptions::new().create_new(true).read(true).write(true).open(footer_path).with_context(
             || format!("Failed to create composite image header {:?}", footer_path),
         )?;
-    let zero_filler_file = File::open(&zero_filler_path).with_context(|| {
+    let zero_filler_file = File::open(zero_filler_path).with_context(|| {
         format!("Failed to open composite image zero filler {:?}", zero_filler_path)
     })?;
 
@@ -66,7 +66,7 @@
     )?;
 
     // Re-open the composite image as read-only.
-    let composite_image = File::open(&output_path)
+    let composite_image = File::open(output_path)
         .with_context(|| format!("Failed to open composite image {:?}", output_path))?;
 
     files.push(header_file);
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 6f646b7..13e5c70 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -16,16 +16,19 @@
 
 use crate::aidl::{Cid, VirtualMachineCallbacks};
 use crate::atom::write_vm_exited_stats;
-use anyhow::{anyhow, bail, Context, Error};
+use anyhow::{anyhow, bail, Context, Error, Result};
 use command_fds::CommandFdExt;
 use lazy_static::lazy_static;
+use libc::{sysconf, _SC_CLK_TCK};
 use log::{debug, error, info};
 use semver::{Version, VersionReq};
 use nix::{fcntl::OFlag, unistd::pipe2};
 use regex::{Captures, Regex};
+use rustutils::system_properties;
 use shared_child::SharedChild;
 use std::borrow::Cow;
-use std::fs::{remove_dir_all, File};
+use std::cmp::max;
+use std::fs::{read_to_string, remove_dir_all, File};
 use std::io::{self, Read};
 use std::mem;
 use std::num::NonZeroU32;
@@ -35,11 +38,12 @@
 use std::sync::{Arc, Condvar, Mutex};
 use std::time::{Duration, SystemTime};
 use std::thread;
-use vsock::VsockStream;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::DeathReason::DeathReason;
+use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IGlobalVmContext::IGlobalVmContext;
 use binder::Strong;
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
 use tombstoned_client::{TombstonedConnection, DebuggerdDumpType};
+use rpcbinder::RpcServer;
 
 /// external/crosvm
 use base::UnixSeqpacketListener;
@@ -61,6 +65,10 @@
 /// The exit status which crosvm returns when vcpu is stalled.
 const CROSVM_WATCHDOG_REBOOT_STATUS: i32 = 36;
 
+const MILLIS_PER_SEC: i64 = 1000;
+
+const SYSPROP_CUSTOM_PVMFW_PATH: &str = "hypervisor.pvmfw.path";
+
 lazy_static! {
     /// If the VM doesn't move to the Started state within this amount time, a hang-up error is
     /// triggered.
@@ -133,6 +141,24 @@
     Failed,
 }
 
+/// RSS values of VM and CrosVM process itself.
+#[derive(Copy, Clone, Debug, Default)]
+pub struct Rss {
+    pub vm: i64,
+    pub crosvm: i64,
+}
+
+/// Metrics regarding the VM.
+#[derive(Debug, Default)]
+pub struct VmMetric {
+    /// Recorded timestamp when the VM is started.
+    pub start_timestamp: Option<SystemTime>,
+    /// Update most recent guest_time periodically from /proc/[crosvm pid]/stat while VM is running.
+    pub cpu_guest_time: Option<i64>,
+    /// Update maximum RSS values periodically from /proc/[crosvm pid]/smaps while VM is running.
+    pub rss: Option<Rss>,
+}
+
 impl VmState {
     /// Tries to start the VM, if it is in the `NotStarted` state.
     ///
@@ -147,6 +173,12 @@
             let child =
                 Arc::new(run_vm(config, &instance.temporary_directory, failure_pipe_write)?);
 
+            let instance_monitor_status = instance.clone();
+            let child_monitor_status = child.clone();
+            thread::spawn(move || {
+                instance_monitor_status.clone().monitor_vm_status(child_monitor_status);
+            });
+
             let child_clone = child.clone();
             let instance_clone = instance.clone();
             thread::spawn(move || {
@@ -170,11 +202,30 @@
     }
 }
 
+/// Internal struct that holds the handles to globally unique resources of a VM.
+#[derive(Debug)]
+pub struct VmContext {
+    #[allow(dead_code)] // Keeps the global context alive
+    global_context: Strong<dyn IGlobalVmContext>,
+    #[allow(dead_code)] // Keeps the server alive
+    vm_server: RpcServer,
+}
+
+impl VmContext {
+    /// Construct new VmContext.
+    pub fn new(global_context: Strong<dyn IGlobalVmContext>, vm_server: RpcServer) -> VmContext {
+        VmContext { global_context, vm_server }
+    }
+}
+
 /// Information about a particular instance of a VM which may be running.
 #[derive(Debug)]
 pub struct VmInstance {
     /// The current state of the VM.
     pub vm_state: Mutex<VmState>,
+    /// Global resources allocated for this VM.
+    #[allow(dead_code)] // Keeps the context alive
+    vm_context: VmContext,
     /// The CID assigned to the VM for vsock communication.
     pub cid: Cid,
     /// The name of the VM.
@@ -190,12 +241,10 @@
     pub requester_debug_pid: i32,
     /// Callbacks to clients of the VM.
     pub callbacks: VirtualMachineCallbacks,
-    /// Input/output stream of the payload run in the VM.
-    pub stream: Mutex<Option<VsockStream>>,
     /// VirtualMachineService binder object for the VM.
     pub vm_service: Mutex<Option<Strong<dyn IVirtualMachineService>>>,
-    /// Recorded timestamp when the VM is started.
-    pub vm_start_timestamp: Mutex<Option<SystemTime>>,
+    /// Recorded metrics of VM such as timestamp or cpu / memory usage.
+    pub vm_metric: Mutex<VmMetric>,
     /// The latest lifecycle state which the payload reported itself to be in.
     payload_state: Mutex<PayloadState>,
     /// Represents the condition that payload_state was updated
@@ -209,6 +258,7 @@
         temporary_directory: PathBuf,
         requester_uid: u32,
         requester_debug_pid: i32,
+        vm_context: VmContext,
     ) -> Result<VmInstance, Error> {
         validate_config(&config)?;
         let cid = config.cid;
@@ -216,6 +266,7 @@
         let protected = config.protected;
         Ok(VmInstance {
             vm_state: Mutex::new(VmState::NotStarted { config }),
+            vm_context,
             cid,
             name,
             protected,
@@ -223,9 +274,8 @@
             requester_uid,
             requester_debug_pid,
             callbacks: Default::default(),
-            stream: Mutex::new(None),
             vm_service: Mutex::new(None),
-            vm_start_timestamp: Mutex::new(None),
+            vm_metric: Mutex::new(Default::default()),
             payload_state: Mutex::new(PayloadState::Starting),
             payload_state_updated: Condvar::new(),
         })
@@ -234,7 +284,8 @@
     /// Starts an instance of `crosvm` to manage the VM. The `crosvm` instance will be killed when
     /// the `VmInstance` is dropped.
     pub fn start(self: &Arc<Self>) -> Result<(), Error> {
-        *self.vm_start_timestamp.lock().unwrap() = Some(SystemTime::now());
+        let mut vm_metric = self.vm_metric.lock().unwrap();
+        vm_metric.start_timestamp = Some(SystemTime::now());
         self.vm_state.lock().unwrap().start(self.clone())
     }
 
@@ -283,13 +334,8 @@
         let death_reason = death_reason(&result, &failure_reason);
         self.callbacks.callback_on_died(self.cid, death_reason);
 
-        let vm_start_timestamp = self.vm_start_timestamp.lock().unwrap();
-        write_vm_exited_stats(
-            self.requester_uid as i32,
-            &self.name,
-            death_reason,
-            *vm_start_timestamp,
-        );
+        let vm_metric = self.vm_metric.lock().unwrap();
+        write_vm_exited_stats(self.requester_uid as i32, &self.name, death_reason, &*vm_metric);
 
         // Delete temporary files.
         if let Err(e) = remove_dir_all(&self.temporary_directory) {
@@ -321,6 +367,42 @@
         }
     }
 
+    fn monitor_vm_status(&self, child: Arc<SharedChild>) {
+        let pid = child.id();
+
+        loop {
+            {
+                // Check VM state
+                let vm_state = &*self.vm_state.lock().unwrap();
+                if let VmState::Dead = vm_state {
+                    break;
+                }
+
+                let mut vm_metric = self.vm_metric.lock().unwrap();
+
+                // Get CPU Information
+                // TODO: Collect it once right before VM dies using SIGCHLD
+                if let Ok(guest_time) = get_guest_time(pid) {
+                    vm_metric.cpu_guest_time = Some(guest_time);
+                } else {
+                    error!("Failed to parse /proc/[pid]/stat");
+                }
+
+                // Get Memory Information
+                if let Ok(rss) = get_rss(pid) {
+                    vm_metric.rss = match &vm_metric.rss {
+                        Some(x) => Some(Rss::extract_max(x, &rss)),
+                        None => Some(rss),
+                    }
+                } else {
+                    error!("Failed to parse /proc/[pid]/smaps");
+                }
+            }
+
+            thread::sleep(Duration::from_secs(1));
+        }
+    }
+
     /// Returns the last reported state of the VM payload.
     pub fn payload_state(&self) -> PayloadState {
         *self.payload_state.lock().unwrap()
@@ -391,6 +473,60 @@
     }
 }
 
+impl Rss {
+    fn extract_max(x: &Rss, y: &Rss) -> Rss {
+        Rss { vm: max(x.vm, y.vm), crosvm: max(x.crosvm, y.crosvm) }
+    }
+}
+
+// Get guest time from /proc/[crosvm pid]/stat
+fn get_guest_time(pid: u32) -> Result<i64> {
+    let file = read_to_string(format!("/proc/{}/stat", pid))?;
+    let data_list: Vec<_> = file.split_whitespace().collect();
+
+    // Information about guest_time is at 43th place of the file split with the whitespace.
+    // Example of /proc/[pid]/stat :
+    // 6603 (kworker/104:1H-kblockd) I 2 0 0 0 -1 69238880 0 0 0 0 0 88 0 0 0 -20 1 0 1845 0 0
+    // 18446744073709551615 0 0 0 0 0 0 0 2147483647 0 0 0 0 17 104 0 0 0 0 0 0 0 0 0 0 0 0 0
+    if data_list.len() < 43 {
+        bail!("Failed to parse command result for getting guest time : {}", file);
+    }
+
+    let guest_time_ticks = data_list[42].parse::<i64>()?;
+    // SAFETY : It just returns an integer about CPU tick information.
+    let ticks_per_sec = unsafe { sysconf(_SC_CLK_TCK) } as i64;
+    Ok(guest_time_ticks * MILLIS_PER_SEC / ticks_per_sec)
+}
+
+// Get rss from /proc/[crosvm pid]/smaps
+fn get_rss(pid: u32) -> Result<Rss> {
+    let file = read_to_string(format!("/proc/{}/smaps", pid))?;
+    let lines: Vec<_> = file.split('\n').collect();
+
+    let mut rss_vm_total = 0i64;
+    let mut rss_crosvm_total = 0i64;
+    let mut is_vm = false;
+    for line in lines {
+        if line.contains("crosvm_guest") {
+            is_vm = true;
+        } else if line.contains("Rss:") {
+            let data_list: Vec<_> = line.split_whitespace().collect();
+            if data_list.len() < 2 {
+                bail!("Failed to parse command result for getting rss :\n{}", line);
+            }
+            let rss = data_list[1].parse::<i64>()?;
+
+            if is_vm {
+                rss_vm_total += rss;
+                is_vm = false;
+            }
+            rss_crosvm_total += rss;
+        }
+    }
+
+    Ok(Rss { vm: rss_vm_total, crosvm: rss_crosvm_total })
+}
+
 fn death_reason(result: &Result<ExitStatus, io::Error>, mut failure_reason: &str) -> DeathReason {
     if let Some(position) = failure_reason.find('|') {
         // Separator indicates extra context information is present after the failure name.
@@ -457,12 +593,22 @@
         .arg("info,disk=off")
         .arg("run")
         .arg("--disable-sandbox")
-        .arg("--no-balloon")
         .arg("--cid")
         .arg(config.cid.to_string());
 
+    if system_properties::read_bool("hypervisor.memory_reclaim.supported", false)? {
+        command.arg("--balloon-page-reporting");
+    } else {
+        command.arg("--no-balloon");
+    }
+
     if config.protected {
-        command.arg("--protected-vm");
+        match system_properties::read(SYSPROP_CUSTOM_PVMFW_PATH)? {
+            Some(pvmfw_path) if !pvmfw_path.is_empty() => {
+                command.arg("--protected-vm-with-firmware").arg(pvmfw_path)
+            }
+            _ => command.arg("--protected-vm"),
+        };
 
         // 3 virtio-console devices + vsock = 4.
         let virtio_pci_device_count = 4 + config.disks.len();
@@ -546,6 +692,7 @@
     debug!("Preserving FDs {:?}", preserved_fds);
     command.preserved_fds(preserved_fds);
 
+    command.arg("--params").arg("crashkernel=17M");
     print_crosvm_args(&command);
 
     let result = SharedChild::spawn(&mut command)?;
diff --git a/virtualizationservice/src/main.rs b/virtualizationservice/src/main.rs
index cea2747..714bcfd 100644
--- a/virtualizationservice/src/main.rs
+++ b/virtualizationservice/src/main.rs
@@ -24,8 +24,8 @@
 use crate::aidl::{VirtualizationService, BINDER_SERVICE_IDENTIFIER, TEMPORARY_DIRECTORY};
 use android_logger::{Config, FilterBuilder};
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::BnVirtualizationService;
+use anyhow::{bail, Context, Error};
 use binder::{register_lazy_service, BinderFeatures, ProcessState};
-use anyhow::Error;
 use log::{info, Level};
 use std::fs::{remove_dir_all, remove_file, read_dir};
 
@@ -44,6 +44,7 @@
             ),
     );
 
+    remove_memlock_rlimit().expect("Failed to remove memlock rlimit");
     clear_temporary_files().expect("Failed to delete old temporary files");
 
     let service = VirtualizationService::init();
@@ -53,6 +54,18 @@
     ProcessState::join_thread_pool();
 }
 
+/// Set this PID's RLIMIT_MEMLOCK to RLIM_INFINITY to allow crosvm (a child process) to mlock()
+/// arbitrary amounts of memory. This is necessary for spawning protected VMs.
+fn remove_memlock_rlimit() -> Result<(), Error> {
+    let lim = libc::rlimit { rlim_cur: libc::RLIM_INFINITY, rlim_max: libc::RLIM_INFINITY };
+    // SAFETY - borrowing the new limit struct only
+    match unsafe { libc::setrlimit(libc::RLIMIT_MEMLOCK, &lim) } {
+        0 => Ok(()),
+        -1 => Err(std::io::Error::last_os_error()).context("setrlimit failed"),
+        n => bail!("Unexpected return value from setrlimit(): {}", n),
+    }
+}
+
 /// Remove any files under `TEMPORARY_DIRECTORY`.
 fn clear_temporary_files() -> Result<(), Error> {
     for dir_entry in read_dir(TEMPORARY_DIRECTORY)? {
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index 4190fbb..f6e8a7b 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -27,7 +27,9 @@
 use microdroid_metadata::{ApexPayload, ApkPayload, Metadata, PayloadConfig, PayloadMetadata};
 use microdroid_payload_config::{ApexConfig, VmPayloadConfig};
 use once_cell::sync::OnceCell;
-use packagemanager_aidl::aidl::android::content::pm::IPackageManagerNative::IPackageManagerNative;
+use packagemanager_aidl::aidl::android::content::pm::{
+    IPackageManagerNative::IPackageManagerNative, StagedApexInfo::StagedApexInfo,
+};
 use regex::Regex;
 use serde::Deserialize;
 use serde_xml_rs::from_reader;
@@ -48,7 +50,7 @@
 const PACKAGE_MANAGER_NATIVE_SERVICE: &str = "package_native";
 
 /// Represents the list of APEXes
-#[derive(Clone, Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
 struct ApexInfoList {
     #[serde(rename = "apex-info")]
     list: Vec<ApexInfo>,
@@ -58,6 +60,8 @@
 struct ApexInfo {
     #[serde(rename = "moduleName")]
     name: String,
+    #[serde(rename = "versionCode")]
+    version: u64,
     #[serde(rename = "modulePath")]
     path: PathBuf,
 
@@ -101,6 +105,40 @@
             Ok(apex_info_list)
         })
     }
+
+    // Override apex info with the staged one
+    fn override_staged_apex(&mut self, staged_apex_info: &StagedApexInfo) -> Result<()> {
+        let mut need_to_add: Option<ApexInfo> = None;
+        for apex_info in self.list.iter_mut() {
+            if staged_apex_info.moduleName == apex_info.name {
+                if apex_info.is_active && apex_info.is_factory {
+                    // Copy the entry to the end as factory/non-active after the loop
+                    // to keep the factory version. Typically this step is unncessary,
+                    // but some apexes (like sharedlibs) need to be kept even if it's inactive.
+                    need_to_add.replace(ApexInfo { is_active: false, ..apex_info.clone() });
+                    // And make this one as non-factory. Note that this one is still active
+                    // and overridden right below.
+                    apex_info.is_factory = false;
+                }
+                // Active one is overridden with the staged one.
+                if apex_info.is_active {
+                    apex_info.version = staged_apex_info.versionCode as u64;
+                    apex_info.path = PathBuf::from(&staged_apex_info.diskImagePath);
+                    apex_info.has_classpath_jar = staged_apex_info.hasClassPathJars;
+                    apex_info.last_update_seconds = last_updated(&apex_info.path)?;
+                }
+            }
+        }
+        if let Some(info) = need_to_add {
+            self.list.push(info);
+        }
+        Ok(())
+    }
+}
+
+fn last_updated<P: AsRef<Path>>(path: P) -> Result<u64> {
+    let metadata = metadata(path)?;
+    Ok(metadata.modified()?.duration_since(SystemTime::UNIX_EPOCH)?.as_secs())
 }
 
 impl ApexInfo {
@@ -137,19 +175,11 @@
                     .context("Failed to get service when prefer_staged is set.")?;
             let staged =
                 pm.getStagedApexModuleNames().context("getStagedApexModuleNames failed")?;
-            for apex_info in list.list.iter_mut() {
-                if staged.contains(&apex_info.name) {
-                    if let Some(staged_apex_info) =
-                        pm.getStagedApexInfo(&apex_info.name).context("getStagedApexInfo failed")?
-                    {
-                        apex_info.path = PathBuf::from(staged_apex_info.diskImagePath);
-                        apex_info.has_classpath_jar = staged_apex_info.hasClassPathJars;
-                        let metadata = metadata(&apex_info.path)?;
-                        apex_info.last_update_seconds =
-                            metadata.modified()?.duration_since(SystemTime::UNIX_EPOCH)?.as_secs();
-                        // by definition, staged apex can't be a factory apex.
-                        apex_info.is_factory = false;
-                    }
+            for name in staged {
+                if let Some(staged_apex_info) =
+                    pm.getStagedApexInfo(&name).context("getStagedApexInfo failed")?
+                {
+                    list.override_staged_apex(&staged_apex_info)?;
                 }
             }
         }
@@ -243,8 +273,13 @@
     let apex_list = pm.get_apex_list(vm_payload_config.prefer_staged)?;
 
     // collect APEXes from config
-    let apex_infos =
+    let mut apex_infos =
         collect_apex_infos(&apex_list, &vm_payload_config.apexes, app_config.debugLevel);
+
+    // Pass sorted list of apexes. Sorting key shouldn't use `path` because it will change after
+    // reboot with prefer_staged. `last_update_seconds` is added to distinguish "samegrade"
+    // update.
+    apex_infos.sort_by_key(|info| (&info.name, &info.version, &info.last_update_seconds));
     info!("Microdroid payload APEXes: {:?}", apex_infos.iter().map(|ai| &ai.name));
 
     let metadata_file = make_metadata_file(app_config, &apex_infos, temporary_directory)?;
@@ -385,7 +420,7 @@
 
     if let Some(file) = storage_image {
         writable_partitions.push(Partition {
-            label: "encrypted-storage".to_owned(),
+            label: "encryptedstore".to_owned(),
             image: Some(ParcelFileDescriptor::new(file)),
             writable: true,
         });
@@ -422,6 +457,7 @@
 #[cfg(test)]
 mod tests {
     use super::*;
+    use tempfile::NamedTempFile;
 
     #[test]
     fn test_find_apex_names_in_classpath() {
@@ -560,4 +596,90 @@
             ]
         );
     }
+
+    #[test]
+    fn test_prefer_staged_apex_with_factory_active_apex() {
+        let single_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 1,
+            path: PathBuf::from("foo.apex"),
+            is_factory: true,
+            is_active: true,
+            ..Default::default()
+        };
+        let mut apex_info_list = ApexInfoList { list: vec![single_apex.clone()] };
+
+        let staged = NamedTempFile::new().unwrap();
+        apex_info_list
+            .override_staged_apex(&StagedApexInfo {
+                moduleName: "foo".to_string(),
+                versionCode: 2,
+                diskImagePath: staged.path().to_string_lossy().to_string(),
+                ..Default::default()
+            })
+            .expect("should be ok");
+
+        assert_eq!(
+            apex_info_list,
+            ApexInfoList {
+                list: vec![
+                    ApexInfo {
+                        version: 2,
+                        is_factory: false,
+                        path: staged.path().to_owned(),
+                        last_update_seconds: last_updated(staged.path()).unwrap(),
+                        ..single_apex.clone()
+                    },
+                    ApexInfo { is_active: false, ..single_apex },
+                ],
+            }
+        );
+    }
+
+    #[test]
+    fn test_prefer_staged_apex_with_factory_and_inactive_apex() {
+        let factory_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 1,
+            path: PathBuf::from("foo.apex"),
+            is_factory: true,
+            ..Default::default()
+        };
+        let active_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 2,
+            path: PathBuf::from("foo.downloaded.apex"),
+            is_active: true,
+            ..Default::default()
+        };
+        let mut apex_info_list =
+            ApexInfoList { list: vec![factory_apex.clone(), active_apex.clone()] };
+
+        let staged = NamedTempFile::new().unwrap();
+        apex_info_list
+            .override_staged_apex(&StagedApexInfo {
+                moduleName: "foo".to_string(),
+                versionCode: 3,
+                diskImagePath: staged.path().to_string_lossy().to_string(),
+                ..Default::default()
+            })
+            .expect("should be ok");
+
+        assert_eq!(
+            apex_info_list,
+            ApexInfoList {
+                list: vec![
+                    // factory apex isn't touched
+                    factory_apex,
+                    // update active one
+                    ApexInfo {
+                        version: 3,
+                        path: staged.path().to_owned(),
+                        last_update_seconds: last_updated(staged.path()).unwrap(),
+                        ..active_apex
+                    },
+                ],
+            }
+        );
+    }
 }
diff --git a/vm/Android.bp b/vm/Android.bp
index 7b016d4..95ef082 100644
--- a/vm/Android.bp
+++ b/vm/Android.bp
@@ -17,6 +17,7 @@
         "liblibc",
         "liblog_rust",
         "libmicrodroid_payload_config",
+        "librand",
         "librustutils",
         "libserde_json",
         "libserde",
diff --git a/vm/src/main.rs b/vm/src/main.rs
index 3b887d3..b5046fb 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -27,7 +27,7 @@
 use clap::Parser;
 use create_idsig::command_create_idsig;
 use create_partition::command_create_partition;
-use run::{command_run, command_run_app};
+use run::{command_run, command_run_app, command_run_microdroid};
 use rustutils::system_properties;
 use std::path::{Path, PathBuf};
 
@@ -48,8 +48,13 @@
         instance: PathBuf,
 
         /// Path to VM config JSON within APK (e.g. assets/vm_config.json)
+        #[clap(long)]
         config_path: Option<String>,
 
+        /// Path to VM payload binary within APK (e.g. MicrodroidTestNativeLib.so)
+        #[clap(long)]
+        payload_path: Option<String>,
+
         /// Name of VM
         #[clap(long)]
         name: Option<String>,
@@ -105,6 +110,65 @@
         #[clap(long = "extra-idsig")]
         extra_idsigs: Vec<PathBuf>,
     },
+    /// Run a virtual machine with Microdroid inside
+    RunMicrodroid {
+        /// Path to the directory where VM-related files (e.g. instance.img, apk.idsig, etc.) will
+        /// be stored. If not specified a random directory under /data/local/tmp/microdroid will be
+        /// created and used.
+        #[clap(long)]
+        work_dir: Option<PathBuf>,
+
+        /// Name of VM
+        #[clap(long)]
+        name: Option<String>,
+
+        /// Detach VM from the terminal and run in the background
+        #[clap(short, long)]
+        daemonize: bool,
+
+        /// Path to the file backing the storage.
+        /// Created if the option is used but the path does not exist in the device.
+        #[clap(long)]
+        storage: Option<PathBuf>,
+
+        /// Size of the storage. Used only if --storage is supplied but path does not exist
+        /// Default size is 10*1024*1024
+        #[clap(long)]
+        storage_size: Option<u64>,
+
+        /// Path to file for VM console output.
+        #[clap(long)]
+        console: Option<PathBuf>,
+
+        /// Path to file for VM log output.
+        #[clap(long)]
+        log: Option<PathBuf>,
+
+        /// Path to file where ramdump is recorded on kernel panic
+        #[clap(long)]
+        ramdump: Option<PathBuf>,
+
+        /// Debug level of the VM. Supported values: "none" (default), "app_only", and "full".
+        #[clap(long, default_value = "none", value_parser = parse_debug_level)]
+        debug: DebugLevel,
+
+        /// Run VM in protected mode.
+        #[clap(short, long)]
+        protected: bool,
+
+        /// Memory size (in MiB) of the VM. If unspecified, defaults to the value of `memory_mib`
+        /// in the VM config file.
+        #[clap(short, long)]
+        mem: Option<u32>,
+
+        /// Number of vCPUs in the VM. If unspecified, defaults to 1.
+        #[clap(long)]
+        cpus: Option<u32>,
+
+        /// Comma separated list of task profile names to apply to the VM
+        #[clap(long)]
+        task_profiles: Vec<String>,
+    },
     /// Run a virtual machine
     Run {
         /// Path to VM config JSON
@@ -201,6 +265,7 @@
             storage,
             storage_size,
             config_path,
+            payload_path,
             daemonize,
             console,
             log,
@@ -219,7 +284,8 @@
             &instance,
             storage.as_deref(),
             storage_size,
-            config_path.as_deref().unwrap_or(""),
+            config_path,
+            payload_path,
             daemonize,
             console.as_deref(),
             log.as_deref(),
@@ -231,6 +297,36 @@
             task_profiles,
             &extra_idsigs,
         ),
+        Opt::RunMicrodroid {
+            name,
+            work_dir,
+            storage,
+            storage_size,
+            daemonize,
+            console,
+            log,
+            ramdump,
+            debug,
+            protected,
+            mem,
+            cpus,
+            task_profiles,
+        } => command_run_microdroid(
+            name,
+            service.as_ref(),
+            work_dir,
+            storage.as_deref(),
+            storage_size,
+            daemonize,
+            console.as_deref(),
+            log.as_deref(),
+            ramdump.as_deref(),
+            debug,
+            protected,
+            mem,
+            cpus,
+            task_profiles,
+        ),
         Opt::Run { name, config, daemonize, cpus, task_profiles, console, log } => {
             command_run(
                 name,
diff --git a/vm/src/run.rs b/vm/src/run.rs
index de8f1c0..3f25bba 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -20,15 +20,19 @@
     PartitionType::PartitionType,
     VirtualMachineAppConfig::{DebugLevel::DebugLevel, Payload::Payload, VirtualMachineAppConfig},
     VirtualMachineConfig::VirtualMachineConfig,
+    VirtualMachinePayloadConfig::VirtualMachinePayloadConfig,
     VirtualMachineState::VirtualMachineState,
 };
-use anyhow::{bail, Context, Error};
+use anyhow::{anyhow, bail, Context, Error};
 use binder::ParcelFileDescriptor;
 use microdroid_payload_config::VmPayloadConfig;
+use rand::{distributions::Alphanumeric, Rng};
+use std::fs;
 use std::fs::File;
-use std::io::{self, BufRead, BufReader};
+use std::io;
 use std::os::unix::io::{AsRawFd, FromRawFd};
 use std::path::{Path, PathBuf};
+use std::process::Command;
 use vmclient::{ErrorCode, VmInstance};
 use vmconfig::{open_parcel_file, VmConfig};
 use zip::ZipArchive;
@@ -43,7 +47,8 @@
     instance: &Path,
     storage: Option<&Path>,
     storage_size: Option<u64>,
-    config_path: &str,
+    config_path: Option<String>,
+    payload_path: Option<String>,
     daemonize: bool,
     console_path: Option<&Path>,
     log_path: Option<&Path>,
@@ -57,7 +62,11 @@
 ) -> Result<(), Error> {
     let apk_file = File::open(apk).context("Failed to open APK file")?;
 
-    let extra_apks = parse_extra_apk_list(apk, config_path)?;
+    let extra_apks = match config_path.as_deref() {
+        Some(path) => parse_extra_apk_list(apk, path)?,
+        None => vec![],
+    };
+
     if extra_apks.len() != extra_idsigs.len() {
         bail!(
             "Found {} extra apks, but there are {} extra idsigs",
@@ -97,7 +106,7 @@
                 service,
                 path,
                 storage_size.unwrap_or(10 * 1024 * 1024),
-                PartitionType::RAW,
+                PartitionType::ENCRYPTEDSTORE,
             )?;
         }
         Some(open_parcel_file(path, true)?)
@@ -108,6 +117,19 @@
     let extra_idsig_files: Result<Vec<File>, _> = extra_idsigs.iter().map(File::open).collect();
     let extra_idsig_fds = extra_idsig_files?.into_iter().map(ParcelFileDescriptor::new).collect();
 
+    let payload = if let Some(config_path) = config_path {
+        if payload_path.is_some() {
+            bail!("Only one of --config-path or --payload-path can be defined")
+        }
+        Payload::ConfigPath(config_path)
+    } else if let Some(payload_path) = payload_path {
+        Payload::PayloadConfig(VirtualMachinePayloadConfig { payloadPath: payload_path })
+    } else {
+        bail!("Either --config-path or --payload-path must be defined")
+    };
+
+    let payload_config_str = format!("{:?}!{:?}", apk, payload);
+
     let config = VirtualMachineConfig::AppConfig(VirtualMachineAppConfig {
         name: name.unwrap_or_else(|| String::from("VmRunApp")),
         apk: apk_fd.into(),
@@ -115,21 +137,90 @@
         extraIdsigs: extra_idsig_fds,
         instanceImage: open_parcel_file(instance, true /* writable */)?.into(),
         encryptedStorageImage: storage,
-        payload: Payload::ConfigPath(config_path.to_owned()),
+        payload,
         debugLevel: debug_level,
         protectedVm: protected,
         memoryMib: mem.unwrap_or(0) as i32, // 0 means use the VM default
         numCpus: cpus.unwrap_or(1) as i32,
         taskProfiles: task_profiles,
     });
-    run(
+    run(service, &config, &payload_config_str, daemonize, console_path, log_path, ramdump_path)
+}
+
+const EMPTY_PAYLOAD_APK: &str = "com.android.microdroid.empty_payload";
+
+fn find_empty_payload_apk_path() -> Result<PathBuf, Error> {
+    let output = Command::new("/system/bin/pm")
+        .arg("path")
+        .arg(EMPTY_PAYLOAD_APK)
+        .output()
+        .context("failed to execute pm path")?;
+    let output_str = String::from_utf8(output.stdout).context("failed to parse output")?;
+    match output_str.strip_prefix("package:") {
+        None => Err(anyhow!("Unexpected output {}", output_str)),
+        Some(apk_path) => Ok(PathBuf::from(apk_path.trim())),
+    }
+}
+
+fn create_work_dir() -> Result<PathBuf, Error> {
+    let s: String =
+        rand::thread_rng().sample_iter(&Alphanumeric).take(17).map(char::from).collect();
+    let work_dir = PathBuf::from("/data/local/tmp/microdroid").join(s);
+    println!("creating work dir {}", work_dir.display());
+    fs::create_dir_all(&work_dir).context("failed to mkdir")?;
+    Ok(work_dir)
+}
+
+/// Run a VM with Microdroid
+#[allow(clippy::too_many_arguments)]
+pub fn command_run_microdroid(
+    name: Option<String>,
+    service: &dyn IVirtualizationService,
+    work_dir: Option<PathBuf>,
+    storage: Option<&Path>,
+    storage_size: Option<u64>,
+    daemonize: bool,
+    console_path: Option<&Path>,
+    log_path: Option<&Path>,
+    ramdump_path: Option<&Path>,
+    debug_level: DebugLevel,
+    protected: bool,
+    mem: Option<u32>,
+    cpus: Option<u32>,
+    task_profiles: Vec<String>,
+) -> Result<(), Error> {
+    let apk = find_empty_payload_apk_path()
+        .context(anyhow!("failed to find path for {} apk", EMPTY_PAYLOAD_APK))?;
+    println!("found path for {} apk: {}", EMPTY_PAYLOAD_APK, apk.display());
+
+    let work_dir = work_dir.unwrap_or(create_work_dir()?);
+    let idsig = work_dir.join("apk.idsig");
+    println!("apk.idsig path: {}", idsig.display());
+    let instance_img = work_dir.join("instance.img");
+    println!("instance.img path: {}", instance_img.display());
+
+    let payload_path = "MicrodroidEmptyPayloadJniLib.so";
+    let extra_sig = [];
+    command_run_app(
+        name,
         service,
-        &config,
-        &format!("{:?}!{:?}", apk, config_path),
+        &apk,
+        &idsig,
+        &instance_img,
+        storage,
+        storage_size,
+        /* config_path= */ None,
+        Some(payload_path.to_owned()),
         daemonize,
         console_path,
         log_path,
         ramdump_path,
+        debug_level,
+        protected,
+        mem,
+        cpus,
+        task_profiles,
+        &extra_sig,
     )
 }
 
@@ -187,7 +278,7 @@
 fn run(
     service: &dyn IVirtualizationService,
     config: &VirtualMachineConfig,
-    config_path: &str,
+    payload_config: &str,
     daemonize: bool,
     console_path: Option<&Path>,
     log_path: Option<&Path>,
@@ -221,7 +312,7 @@
 
     println!(
         "Created VM from {} with CID {}, state is {}.",
-        config_path,
+        payload_config,
         vm.cid(),
         state_to_str(vm.state()?)
     );
@@ -265,19 +356,8 @@
 struct Callback {}
 
 impl vmclient::VmCallback for Callback {
-    fn on_payload_started(&self, _cid: i32, stream: Option<&File>) {
-        // Show the output of the payload
-        if let Some(stream) = stream {
-            let mut reader = BufReader::new(stream.try_clone().unwrap());
-            std::thread::spawn(move || loop {
-                let mut s = String::new();
-                match reader.read_line(&mut s) {
-                    Ok(0) => break,
-                    Ok(_) => print!("{}", s),
-                    Err(e) => eprintln!("error reading from virtual machine: {}", e),
-                };
-            });
-        }
+    fn on_payload_started(&self, _cid: i32) {
+        eprintln!("payload started");
     }
 
     fn on_payload_ready(&self, _cid: i32) {
diff --git a/vm_payload/Android.bp b/vm_payload/Android.bp
new file mode 100644
index 0000000..967d1cf
--- /dev/null
+++ b/vm_payload/Android.bp
@@ -0,0 +1,80 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+// The Rust implementation of the C API.
+rust_ffi_static {
+    name: "libvm_payload_impl",
+    crate_name: "vm_payload",
+    visibility: ["//visibility:private"],
+    srcs: ["src/*.rs"],
+    include_dirs: ["include"],
+    prefer_rlib: true,
+    rustlibs: [
+        "android.system.virtualization.payload-rust",
+        "libandroid_logger",
+        "libanyhow",
+        "libbinder_rs",
+        "liblazy_static",
+        "liblibc",
+        "liblog_rust",
+        "librpcbinder_rs",
+        "libvsock",
+    ],
+    // The sanitize section below fixes the fuzzer build in b/256166339.
+    // TODO(b/250854486): Remove the sanitize section once the bug is fixed.
+    sanitize: {
+        address: false,
+    },
+}
+
+// Rust wrappers round the C API for Rust clients.
+// (Yes, this involves going Rust -> C -> Rust.)
+rust_bindgen {
+    name: "libvm_payload_bindgen",
+    wrapper_src: "include-restricted/vm_payload_restricted.h",
+    crate_name: "vm_payload_bindgen",
+    source_stem: "bindings",
+    apex_available: ["com.android.compos"],
+    visibility: ["//packages/modules/Virtualization/compos"],
+    shared_libs: [
+        "libvm_payload#current",
+    ],
+}
+
+// Shared library for clients to link against.
+cc_library_shared {
+    name: "libvm_payload",
+    shared_libs: [
+        "libbinder_ndk",
+        "libbinder_rpc_unstable",
+        "liblog",
+    ],
+    whole_static_libs: ["libvm_payload_impl"],
+    export_static_lib_headers: ["libvm_payload_impl"],
+    installable: false,
+    version_script: "libvm_payload.map.txt",
+    stubs: {
+        symbol_file: "libvm_payload.map.txt",
+        // Implementation is available inside a Microdroid VM.
+        implementation_installable: false,
+    },
+}
+
+// Just the headers. Mostly useful for clients that only want the
+// declaration of AVmPayload_main().
+cc_library_headers {
+    name: "vm_payload_headers",
+    apex_available: ["com.android.compos"],
+    export_include_dirs: ["include"],
+}
+
+// Restricted headers for use by internal clients & associated tests.
+cc_library_headers {
+    name: "vm_payload_restricted_headers",
+    header_libs: ["vm_payload_headers"],
+    export_header_lib_headers: ["vm_payload_headers"],
+    export_include_dirs: ["include-restricted"],
+    apex_available: ["com.android.compos"],
+    visibility: ["//packages/modules/Virtualization:__subpackages__"],
+}
diff --git a/vm_payload/include-restricted/vm_payload_restricted.h b/vm_payload/include-restricted/vm_payload_restricted.h
new file mode 100644
index 0000000..0b78541
--- /dev/null
+++ b/vm_payload/include-restricted/vm_payload_restricted.h
@@ -0,0 +1,54 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <sys/cdefs.h>
+
+#include "vm_payload.h"
+
+// The functions declared here are restricted to VMs created with a config file;
+// they will fail if called in other VMs. The ability to create such VMs
+// requires the android.permission.USE_CUSTOM_VIRTUAL_MACHINE permission, and is
+// therefore not available to privileged or third party apps.
+
+// These functions can be used by tests, if the permission is granted via shell.
+
+__BEGIN_DECLS
+
+/**
+ * Get the VM's DICE attestation chain.
+ *
+ * \param data pointer to size bytes where the chain is written.
+ * \param size number of bytes that can be written to data.
+ *
+ * \return the total size of the chain
+ */
+size_t AVmPayload_getDiceAttestationChain(void *data, size_t size);
+
+/**
+ * Get the VM's DICE attestation CDI.
+ *
+ * \param data pointer to size bytes where the CDI is written.
+ * \param size number of bytes that can be written to data.
+ *
+ * \return the total size of the CDI
+ */
+size_t AVmPayload_getDiceAttestationCdi(void *data, size_t size);
+
+__END_DECLS
diff --git a/microdroid/vm_payload/include/vm_main.h b/vm_payload/include/vm_main.h
similarity index 100%
rename from microdroid/vm_payload/include/vm_main.h
rename to vm_payload/include/vm_main.h
diff --git a/vm_payload/include/vm_payload.h b/vm_payload/include/vm_payload.h
new file mode 100644
index 0000000..7c224f6
--- /dev/null
+++ b/vm_payload/include/vm_payload.h
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdnoreturn.h>
+#include <sys/cdefs.h>
+
+#include "vm_main.h"
+
+__BEGIN_DECLS
+
+struct AIBinder;
+typedef struct AIBinder AIBinder;
+
+/**
+ * Notifies the host that the payload is ready.
+ *
+ * If the host app has set a `VirtualMachineCallback` for the VM, its
+ * `onPayloadReady` method will be called.
+ *
+ * Note that subsequent calls to this function after the first have no effect;
+ * `onPayloadReady` is never called more than once.
+ */
+void AVmPayload_notifyPayloadReady(void);
+
+/**
+ * Runs a binder RPC server, serving the supplied binder service implementation on the given vsock
+ * port.
+ *
+ * If and when the server is ready for connections (it is listening on the port), `on_ready` is
+ * called to allow appropriate action to be taken - e.g. to notify clients that they may now
+ * attempt to connect with `AVmPayload_notifyPayloadReady`.
+ *
+ * Note that this function does not return. The calling thread joins the binder
+ * thread pool to handle incoming messages.
+ *
+ * \param service the service to bind to the given port.
+ * \param port vsock port.
+ * \param on_ready the callback to execute once the server is ready for connections. The callback
+ *                 will be called at most once.
+ * \param param param for the `on_ready` callback.
+ */
+noreturn void AVmPayload_runVsockRpcServer(AIBinder *service, unsigned int port,
+                                           void (*on_ready)(void *param), void *param);
+
+/**
+ * Get a secret that is uniquely bound to this VM instance. The secrets are
+ * 32-byte values and the value associated with an identifier will not change
+ * over the lifetime of the VM instance.
+ *
+ * \param identifier identifier of the secret to return.
+ * \param identifier_size size of the secret identifier.
+ * \param secret pointer to size bytes where the secret is written.
+ * \param size number of bytes of the secret to get, <= 32.
+ */
+void AVmPayload_getVmInstanceSecret(const void *identifier, size_t identifier_size, void *secret,
+                                    size_t size);
+
+/**
+ * Gets the path to the APK contents. It is a directory, under which are
+ * the unzipped contents of the APK containing the payload, all read-only
+ * but accessible to the payload.
+ *
+ * \return the path to the APK contents. The returned string should not be
+ * deleted or freed by the application. The string remains valid for the
+ * lifetime of the VM.
+ */
+const char *AVmPayload_getApkContentsPath(void);
+
+/**
+ * Gets the path to the encrypted persistent storage for the VM, if any. This is
+ * a directory under which any files or directories created will be stored on
+ * behalf of the VM by the host app. All data is encrypted using a key known
+ * only to the VM, so the host cannot decrypt it, but may delete it.
+ *
+ * \return the path to the APK contents, or NULL if no encrypted storage was
+ * requested in the VM configuration. If non-null the returned string should not
+ * be deleted or freed by the application and remains valid for the lifetime of
+ * the VM.
+ */
+const char *AVmPayload_getEncryptedStoragePath(void);
+
+__END_DECLS
diff --git a/vm_payload/libvm_payload.map.txt b/vm_payload/libvm_payload.map.txt
new file mode 100644
index 0000000..a2402d1
--- /dev/null
+++ b/vm_payload/libvm_payload.map.txt
@@ -0,0 +1,12 @@
+LIBVM_PAYLOAD {
+  global:
+    AVmPayload_notifyPayloadReady;       # systemapi
+    AVmPayload_runVsockRpcServer;        # systemapi
+    AVmPayload_getVmInstanceSecret;      # systemapi
+    AVmPayload_getDiceAttestationChain;  # systemapi
+    AVmPayload_getDiceAttestationCdi;    # systemapi
+    AVmPayload_getApkContentsPath;       # systemapi
+    AVmPayload_getEncryptedStoragePath;  # systemapi
+  local:
+    *;
+};
diff --git a/vm_payload/src/api.rs b/vm_payload/src/api.rs
new file mode 100644
index 0000000..a79c0bb
--- /dev/null
+++ b/vm_payload/src/api.rs
@@ -0,0 +1,260 @@
+// Copyright 2022, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! This module handles the interaction with virtual machine payload service.
+
+// We're implementing unsafe functions, but we still want warnings on unsafe usage within them.
+#![warn(unsafe_op_in_unsafe_fn)]
+
+use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
+    IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
+use anyhow::{ensure, bail, Context, Result};
+use binder::{Strong, unstable_api::{AIBinder, new_spibinder}};
+use lazy_static::lazy_static;
+use log::{error, info, Level};
+use rpcbinder::{get_unix_domain_rpc_interface, RpcServer};
+use std::convert::Infallible;
+use std::ffi::CString;
+use std::fmt::Debug;
+use std::os::raw::{c_char, c_void};
+use std::ptr;
+use std::sync::{Mutex, atomic::{AtomicBool, Ordering}};
+
+lazy_static! {
+    static ref VM_APK_CONTENTS_PATH_C: CString =
+        CString::new(VM_APK_CONTENTS_PATH).expect("CString::new failed");
+    static ref PAYLOAD_CONNECTION: Mutex<Option<Strong<dyn IVmPayloadService>>> = Mutex::default();
+}
+
+static ALREADY_NOTIFIED: AtomicBool = AtomicBool::new(false);
+
+/// Return a connection to the payload service in Microdroid Manager. Uses the existing connection
+/// if there is one, otherwise attempts to create a new one.
+fn get_vm_payload_service() -> Result<Strong<dyn IVmPayloadService>> {
+    let mut connection = PAYLOAD_CONNECTION.lock().unwrap();
+    if let Some(strong) = &*connection {
+        Ok(strong.clone())
+    } else {
+        let new_connection: Strong<dyn IVmPayloadService> = get_unix_domain_rpc_interface(
+            VM_PAYLOAD_SERVICE_SOCKET_NAME,
+        )
+        .context(format!("Failed to connect to service: {}", VM_PAYLOAD_SERVICE_SOCKET_NAME))?;
+        *connection = Some(new_connection.clone());
+        Ok(new_connection)
+    }
+}
+
+/// Make sure our logging goes to logcat. It is harmless to call this more than once.
+fn initialize_logging() {
+    android_logger::init_once(
+        android_logger::Config::default().with_tag("vm_payload").with_min_level(Level::Info),
+    );
+}
+
+/// In many cases clients can't do anything useful if API calls fail, and the failure
+/// generally indicates that the VM is exiting or otherwise doomed. So rather than
+/// returning a non-actionable error indication we just log the problem and abort
+/// the process.
+fn unwrap_or_abort<T, E: Debug>(result: Result<T, E>) -> T {
+    result.unwrap_or_else(|e| {
+        let msg = format!("{:?}", e);
+        error!("{msg}");
+        panic!("{msg}")
+    })
+}
+
+/// Notifies the host that the payload is ready.
+/// Panics on failure.
+#[no_mangle]
+pub extern "C" fn AVmPayload_notifyPayloadReady() {
+    initialize_logging();
+
+    if !ALREADY_NOTIFIED.swap(true, Ordering::Relaxed) {
+        unwrap_or_abort(try_notify_payload_ready());
+
+        info!("Notified host payload ready successfully");
+    }
+}
+
+/// Notifies the host that the payload is ready.
+/// Returns a `Result` containing error information if failed.
+fn try_notify_payload_ready() -> Result<()> {
+    get_vm_payload_service()?.notifyPayloadReady().context("Cannot notify payload ready")
+}
+
+/// Runs a binder RPC server, serving the supplied binder service implementation on the given vsock
+/// port.
+///
+/// If and when the server is ready for connections (it is listening on the port), `on_ready` is
+/// called to allow appropriate action to be taken - e.g. to notify clients that they may now
+/// attempt to connect.
+///
+/// The current thread joins the binder thread pool to handle incoming messages.
+/// This function never returns.
+///
+/// Panics on error (including unexpected server exit).
+///
+/// # Safety
+///
+/// If present, the `on_ready` callback must be a valid function pointer, which will be called at
+/// most once, while this function is executing, with the `param` parameter.
+#[no_mangle]
+pub unsafe extern "C" fn AVmPayload_runVsockRpcServer(
+    service: *mut AIBinder,
+    port: u32,
+    on_ready: Option<unsafe extern "C" fn(param: *mut c_void)>,
+    param: *mut c_void,
+) -> Infallible {
+    initialize_logging();
+
+    // SAFETY: try_run_vsock_server has the same requirements as this function
+    unwrap_or_abort(unsafe { try_run_vsock_server(service, port, on_ready, param) })
+}
+
+/// # Safety: Same as `AVmPayload_runVsockRpcServer`.
+unsafe fn try_run_vsock_server(
+    service: *mut AIBinder,
+    port: u32,
+    on_ready: Option<unsafe extern "C" fn(param: *mut c_void)>,
+    param: *mut c_void,
+) -> Result<Infallible> {
+    // SAFETY: AIBinder returned has correct reference count, and the ownership can
+    // safely be taken by new_spibinder.
+    let service = unsafe { new_spibinder(service) };
+    if let Some(service) = service {
+        match RpcServer::new_vsock(service, port) {
+            Ok(server) => {
+                if let Some(on_ready) = on_ready {
+                    // SAFETY: We're calling the callback with the parameter specified within the
+                    // allowed lifetime.
+                    unsafe { on_ready(param) };
+                }
+                server.join();
+                bail!("RpcServer unexpectedly terminated");
+            }
+            Err(err) => {
+                bail!("Failed to start RpcServer: {:?}", err);
+            }
+        }
+    } else {
+        bail!("Failed to convert the given service from AIBinder to SpIBinder.");
+    }
+}
+
+/// Get a secret that is uniquely bound to this VM instance.
+/// Panics on failure.
+///
+/// # Safety
+///
+/// Behavior is undefined if any of the following conditions are violated:
+///
+/// * `identifier` must be [valid] for reads of `identifier_size` bytes.
+/// * `secret` must be [valid] for writes of `size` bytes.
+///
+/// [valid]: ptr#safety
+#[no_mangle]
+pub unsafe extern "C" fn AVmPayload_getVmInstanceSecret(
+    identifier: *const u8,
+    identifier_size: usize,
+    secret: *mut u8,
+    size: usize,
+) {
+    initialize_logging();
+
+    // SAFETY: See the requirements on `identifier` above.
+    let identifier = unsafe { std::slice::from_raw_parts(identifier, identifier_size) };
+    let vm_secret = unwrap_or_abort(try_get_vm_instance_secret(identifier, size));
+
+    // SAFETY: See the requirements on `secret` above; `vm_secret` is known to have length `size`,
+    // and cannot overlap `secret` because we just allocated it.
+    unsafe {
+        ptr::copy_nonoverlapping(vm_secret.as_ptr(), secret, size);
+    }
+}
+
+fn try_get_vm_instance_secret(identifier: &[u8], size: usize) -> Result<Vec<u8>> {
+    let vm_secret = get_vm_payload_service()?
+        .getVmInstanceSecret(identifier, i32::try_from(size)?)
+        .context("Cannot get VM instance secret")?;
+    ensure!(
+        vm_secret.len() == size,
+        "Returned secret has {} bytes, expected {}",
+        vm_secret.len(),
+        size
+    );
+    Ok(vm_secret)
+}
+
+/// Get the VM's attestation chain.
+/// Panics on failure.
+///
+/// # Safety
+///
+/// Behavior is undefined if any of the following conditions are violated:
+///
+/// * `data` must be [valid] for writes of `size` bytes.
+///
+/// [valid]: ptr#safety
+#[no_mangle]
+pub unsafe extern "C" fn AVmPayload_getDiceAttestationChain(data: *mut u8, size: usize) -> usize {
+    initialize_logging();
+
+    let chain = unwrap_or_abort(try_get_dice_attestation_chain());
+    // SAFETY: See the requirements on `data` above. The number of bytes copied doesn't exceed
+    // the length of either buffer, and `chain` cannot overlap `data` because we just allocated it.
+    unsafe { ptr::copy_nonoverlapping(chain.as_ptr(), data, std::cmp::min(chain.len(), size)) };
+    chain.len()
+}
+
+fn try_get_dice_attestation_chain() -> Result<Vec<u8>> {
+    get_vm_payload_service()?.getDiceAttestationChain().context("Cannot get attestation chain")
+}
+
+/// Get the VM's attestation CDI.
+/// Panics on failure.
+///
+/// # Safety
+///
+/// Behavior is undefined if any of the following conditions are violated:
+///
+/// * `data` must be [valid] for writes of `size` bytes.
+///
+/// [valid]: ptr#safety
+#[no_mangle]
+pub unsafe extern "C" fn AVmPayload_getDiceAttestationCdi(data: *mut u8, size: usize) -> usize {
+    initialize_logging();
+
+    let cdi = unwrap_or_abort(try_get_dice_attestation_cdi());
+    // SAFETY: See the requirements on `data` above. The number of bytes copied doesn't exceed
+    // the length of either buffer, and `cdi` cannot overlap `data` because we just allocated it.
+    unsafe { ptr::copy_nonoverlapping(cdi.as_ptr(), data, std::cmp::min(cdi.len(), size)) };
+    cdi.len()
+}
+
+fn try_get_dice_attestation_cdi() -> Result<Vec<u8>> {
+    get_vm_payload_service()?.getDiceAttestationCdi().context("Cannot get attestation CDI")
+}
+
+/// Gets the path to the APK contents.
+#[no_mangle]
+pub extern "C" fn AVmPayload_getApkContentsPath() -> *const c_char {
+    (*VM_APK_CONTENTS_PATH_C).as_ptr()
+}
+
+/// Gets the path to the VM's encrypted storage.
+#[no_mangle]
+pub extern "C" fn AVmPayload_getEncryptedStoragePath() -> *const c_char {
+    // TODO(b/254454578): Return a real path if storage is present
+    ptr::null()
+}
diff --git a/microdroid/vm_payload/src/lib.rs b/vm_payload/src/lib.rs
similarity index 93%
rename from microdroid/vm_payload/src/lib.rs
rename to vm_payload/src/lib.rs
index be6cf93..5c3ee31 100644
--- a/microdroid/vm_payload/src/lib.rs
+++ b/vm_payload/src/lib.rs
@@ -14,9 +14,9 @@
 
 //! Library for payload to communicate with the Microdroid Manager.
 
-mod vm_payload_service;
+mod api;
 
-pub use vm_payload_service::{
+pub use api::{
     AVmPayload_getDiceAttestationCdi, AVmPayload_getDiceAttestationChain,
     AVmPayload_getVmInstanceSecret, AVmPayload_notifyPayloadReady,
 };
diff --git a/vmbase/example/Android.bp b/vmbase/example/Android.bp
index e9a3f98..505de6b 100644
--- a/vmbase/example/Android.bp
+++ b/vmbase/example/Android.bp
@@ -12,9 +12,13 @@
         "libaarch64_paging",
         "libbuddy_system_allocator",
         "libdice_nostd",
+        "liblibfdt",
         "liblog_rust_nostd",
         "libvmbase",
     ],
+    static_libs: [
+        "libarm-optimized-routines-mem",
+    ],
     apex_available: ["com.android.virt"],
 }
 
@@ -54,10 +58,11 @@
     edition: "2021",
     rustlibs: [
         "android.system.virtualizationservice-rust",
+        "libandroid_logger",
         "libanyhow",
-        "libenv_logger",
         "liblibc",
         "liblog_rust",
+        "libnix",
         "libvmclient",
     ],
     data: [
diff --git a/vmbase/example/idmap.S b/vmbase/example/idmap.S
index 7fc5d5e..71a6ade 100644
--- a/vmbase/example/idmap.S
+++ b/vmbase/example/idmap.S
@@ -44,7 +44,7 @@
 	.fill		509, 8, 0x0			// 509 GiB of remaining VA space
 
 	/* level 2 */
-0:	.quad		.L_BLOCK_RO  | 0x80000000	// DT provided by VMM
+0:	.quad		.L_BLOCK_MEM | 0x80000000	// DT provided by VMM
 	.quad		.L_BLOCK_MEM_XIP | 0x80200000	// 2 MiB of DRAM containing image
 	.quad		.L_BLOCK_MEM | 0x80400000	// 2 MiB of writable DRAM
 	.fill		509, 8, 0x0
diff --git a/vmbase/example/src/main.rs b/vmbase/example/src/main.rs
index 03f0603..bb64651 100644
--- a/vmbase/example/src/main.rs
+++ b/vmbase/example/src/main.rs
@@ -27,9 +27,14 @@
     bionic_tls, dtb_range, print_addresses, rodata_range, stack_chk_guard, text_range,
     writable_region, DEVICE_REGION,
 };
-use aarch64_paging::{idmap::IdMap, paging::Attributes};
+use aarch64_paging::{
+    idmap::IdMap,
+    paging::{Attributes, MemoryRegion},
+};
 use alloc::{vec, vec::Vec};
 use buddy_system_allocator::LockedHeap;
+use core::ffi::CStr;
+use libfdt::Fdt;
 use log::{info, LevelFilter};
 use vmbase::{logger, main, println};
 
@@ -57,6 +62,7 @@
     assert_eq!(arg0, dtb_range().start.0 as u64);
     check_data();
     check_stack_guard();
+    check_fdt();
 
     unsafe {
         HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
@@ -138,6 +144,51 @@
     info!("Data looks good");
 }
 
+fn check_fdt() {
+    info!("Checking FDT...");
+    let fdt = MemoryRegion::from(layout::dtb_range());
+    let fdt = unsafe { core::slice::from_raw_parts_mut(fdt.start().0 as *mut u8, fdt.len()) };
+
+    let reader = Fdt::from_slice(fdt).unwrap();
+    info!("FDT passed verification.");
+    for reg in reader.memory().unwrap().unwrap() {
+        info!("memory @ {reg:#x?}");
+    }
+
+    let compatible = CStr::from_bytes_with_nul(b"ns16550a\0").unwrap();
+
+    for c in reader.compatible_nodes(compatible).unwrap() {
+        let reg = c.reg().unwrap().unwrap().next().unwrap();
+        info!("node compatible with '{}' at {reg:?}", compatible.to_str().unwrap());
+    }
+
+    let writer = Fdt::from_mut_slice(fdt).unwrap();
+    writer.unpack().unwrap();
+    info!("FDT successfully unpacked.");
+
+    let path = CStr::from_bytes_with_nul(b"/memory\0").unwrap();
+    let mut node = writer.node_mut(path).unwrap().unwrap();
+    let name = CStr::from_bytes_with_nul(b"child\0").unwrap();
+    let mut child = node.add_subnode(name).unwrap();
+    info!("Created subnode '{}/{}'.", path.to_str().unwrap(), name.to_str().unwrap());
+
+    let name = CStr::from_bytes_with_nul(b"str-property\0").unwrap();
+    child.appendprop(name, b"property-value\0").unwrap();
+    info!("Appended property '{}'.", name.to_str().unwrap());
+
+    let name = CStr::from_bytes_with_nul(b"pair-property\0").unwrap();
+    let addr = 0x0123_4567u64;
+    let size = 0x89ab_cdefu64;
+    child.appendprop_addrrange(name, addr, size).unwrap();
+    info!("Appended property '{}'.", name.to_str().unwrap());
+
+    let writer = child.fdt();
+    writer.pack().unwrap();
+    info!("FDT successfully packed.");
+
+    info!("FDT checks done.");
+}
+
 fn check_alloc() {
     info!("Allocating a Vec...");
     let mut vector: Vec<u32> = vec![1, 2, 3, 4];
diff --git a/vmbase/example/tests/test.rs b/vmbase/example/tests/test.rs
index 85e0213..57b68ed 100644
--- a/vmbase/example/tests/test.rs
+++ b/vmbase/example/tests/test.rs
@@ -25,8 +25,9 @@
 use log::info;
 use std::{
     fs::File,
-    io,
-    os::unix::io::{AsRawFd, FromRawFd},
+    io::{self, BufRead, BufReader},
+    os::unix::io::FromRawFd,
+    panic, thread,
 };
 use vmclient::{DeathReason, VmInstance};
 
@@ -36,7 +37,14 @@
 /// Runs the vmbase_example VM as an unprotected VM via VirtualizationService.
 #[test]
 fn test_run_example_vm() -> Result<(), Error> {
-    env_logger::init();
+    android_logger::init_once(
+        android_logger::Config::default().with_tag("vmbase").with_min_level(log::Level::Debug),
+    );
+
+    // Redirect panic messages to logcat.
+    panic::set_hook(Box::new(|panic_info| {
+        log::error!("{}", panic_info);
+    }));
 
     // We need to start the thread pool for Binder to work properly, especially link_to_death.
     ProcessState::start_thread_pool();
@@ -62,8 +70,8 @@
         platformVersion: "~1.0".to_string(),
         taskProfiles: vec![],
     });
-    let console = duplicate_stdout()?;
-    let log = duplicate_stdout()?;
+    let console = android_log_fd()?;
+    let log = android_log_fd()?;
     let vm = VmInstance::create(service.as_ref(), &config, Some(console), Some(log), None)
         .context("Failed to create VM")?;
     vm.start().context("Failed to start VM")?;
@@ -76,17 +84,17 @@
     Ok(())
 }
 
-/// Safely duplicate the standard output file descriptor.
-fn duplicate_stdout() -> io::Result<File> {
-    let stdout_fd = io::stdout().as_raw_fd();
-    // Safe because this just duplicates a file descriptor which we know to be valid, and we check
-    // for an error.
-    let dup_fd = unsafe { libc::dup(stdout_fd) };
-    if dup_fd < 0 {
-        Err(io::Error::last_os_error())
-    } else {
-        // Safe because we have just duplicated the file descriptor so we own it, and `from_raw_fd`
-        // takes ownership of it.
-        Ok(unsafe { File::from_raw_fd(dup_fd) })
-    }
+fn android_log_fd() -> io::Result<File> {
+    let (reader_fd, writer_fd) = nix::unistd::pipe()?;
+
+    // SAFETY: These are new FDs with no previous owner.
+    let reader = unsafe { File::from_raw_fd(reader_fd) };
+    let writer = unsafe { File::from_raw_fd(writer_fd) };
+
+    thread::spawn(|| {
+        for line in BufReader::new(reader).lines() {
+            info!("{}", line.unwrap());
+        }
+    });
+    Ok(writer)
 }
diff --git a/vmclient/src/lib.rs b/vmclient/src/lib.rs
index e6f32b4..20b7f02 100644
--- a/vmclient/src/lib.rs
+++ b/vmclient/src/lib.rs
@@ -74,7 +74,7 @@
 pub trait VmCallback {
     /// Called when the payload has been started within the VM. If present, `stream` is connected
     /// to the stdin/stdout of the payload.
-    fn on_payload_started(&self, cid: i32, stream: Option<&File>) {}
+    fn on_payload_started(&self, cid: i32) {}
 
     /// Callend when the payload has notified Virtualization Service that it is ready to serve
     /// clients.
@@ -269,14 +269,10 @@
 impl Interface for VirtualMachineCallback {}
 
 impl IVirtualMachineCallback for VirtualMachineCallback {
-    fn onPayloadStarted(
-        &self,
-        cid: i32,
-        stream: Option<&ParcelFileDescriptor>,
-    ) -> BinderResult<()> {
+    fn onPayloadStarted(&self, cid: i32) -> BinderResult<()> {
         self.state.notify_state(VirtualMachineState::STARTED);
         if let Some(ref callback) = self.client_callback {
-            callback.on_payload_started(cid, stream.map(ParcelFileDescriptor::as_ref));
+            callback.on_payload_started(cid);
         }
         Ok(())
     }