Merge "Migrate pkvm_perf_test.py dex2oat64 and crosvm memory footprint collection during compilation to AOSP."
diff --git a/apex/Android.bp b/apex/Android.bp
index 83985cc..a9fad55 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -63,6 +63,9 @@
     ],
     prebuilts: [
         "com.android.virt.init.rc",
+        "microdroid_initrd_app_debuggable",
+        "microdroid_initrd_full_debuggable",
+        "microdroid_initrd_normal",
         "microdroid.json",
         "microdroid_bootloader",
         "microdroid_bootloader.avbpubkey",
diff --git a/apex/product_packages.mk b/apex/product_packages.mk
index ec295f5..5bc0044 100644
--- a/apex/product_packages.mk
+++ b/apex/product_packages.mk
@@ -32,4 +32,4 @@
 
 PRODUCT_SYSTEM_EXT_PROPERTIES := ro.config.isolated_compilation_enabled=true
 
-PRODUCT_SYSTEM_FSVERITY_GENERATE_METADATA := true
+PRODUCT_FSVERITY_GENERATE_METADATA := true
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index df95323..6863b0c 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -317,22 +317,6 @@
         RunCommand(args, cmd)
 
 
-def SignSuperImg(args, key, super_img, work_dir):
-    # unpack super.img
-    UnpackSuperImg(args, super_img, work_dir)
-
-    system_a_img = os.path.join(work_dir, 'system_a.img')
-    vendor_a_img = os.path.join(work_dir, 'vendor_a.img')
-
-    # re-sign each partition
-    system_a_f = Async(AddHashTreeFooter, args, key, system_a_img)
-    vendor_a_f = Async(AddHashTreeFooter, args, key, vendor_a_img)
-
-    # 3. re-pack super.img
-    partitions = {"system_a": system_a_img, "vendor_a": vendor_a_img}
-    Async(MakeSuperImage, args, partitions, super_img, wait=[system_a_f, vendor_a_f])
-
-
 def ReplaceBootloaderPubkey(args, key, bootloader, bootloader_pubkey):
     if os.path.basename(bootloader) in args.key_overrides:
         key = args.key_overrides[os.path.basename(bootloader)]
@@ -399,13 +383,20 @@
     init_boot_img_f = Async(AddHashFooter, args, key, files['init_boot.img'])
 
     # re-sign super.img
-    super_img_f = Async(SignSuperImg, args, key, files['super.img'], unpack_dir.name)
+    # 1. unpack super.img
+    # 2. resign system and vendor
+    # 3. repack super.img out of resigned system and vendor
+    UnpackSuperImg(args, files['super.img'], unpack_dir.name)
+    system_a_f = Async(AddHashTreeFooter, args, key, system_a_img)
+    vendor_a_f = Async(AddHashTreeFooter, args, key, vendor_a_img)
+    partitions = {"system_a": system_a_img, "vendor_a": vendor_a_img}
+    Async(MakeSuperImage, args, partitions, files['super.img'], wait=[system_a_f, vendor_a_f])
 
     # re-generate vbmeta from re-signed {boot, vendor_boot, init_boot, system_a, vendor_a}.img
     Async(MakeVbmetaImage, args, key, files['vbmeta.img'],
           images=[files['boot.img'], files['vendor_boot.img'],
                   files['init_boot.img'], system_a_img, vendor_a_img],
-          wait=[boot_img_f, vendor_boot_img_f, init_boot_img_f, super_img_f])
+          wait=[boot_img_f, vendor_boot_img_f, init_boot_img_f, system_a_f, vendor_a_f])
 
     # Re-sign bootconfigs and the uboot_env with the same key
     bootconfig_sign_key = key
diff --git a/authfs/Android.bp b/authfs/Android.bp
index 7788702..523da35 100644
--- a/authfs/Android.bp
+++ b/authfs/Android.bp
@@ -16,6 +16,7 @@
         "libauthfs_fsverity_metadata",
         "libbinder_rs",
         "libcfg_if",
+        "libclap",
         "libfsverity_digests_proto_rust",
         "libfuse_rust",
         "liblibc",
@@ -24,7 +25,6 @@
         "libopenssl",
         "libprotobuf",
         "librpcbinder_rs",
-        "libstructopt",
         "libthiserror",
     ],
     prefer_rlib: true,
diff --git a/authfs/src/main.rs b/authfs/src/main.rs
index c09ed71..9ff0ae3 100644
--- a/authfs/src/main.rs
+++ b/authfs/src/main.rs
@@ -33,13 +33,13 @@
 //! the state is not persistent, thus only new file/directory are supported.
 
 use anyhow::{anyhow, bail, Result};
+use clap::Parser;
 use log::error;
 use protobuf::Message;
 use std::convert::TryInto;
 use std::fs::File;
 use std::num::NonZeroU8;
 use std::path::{Path, PathBuf};
-use structopt::StructOpt;
 
 mod common;
 mod file;
@@ -53,22 +53,21 @@
 use fsverity_digests_proto::fsverity_digests::FSVerityDigests;
 use fusefs::{AuthFs, AuthFsEntry, LazyVerifiedReadonlyFile};
 
-#[derive(StructOpt)]
+#[derive(Parser)]
 struct Args {
     /// Mount point of AuthFS.
-    #[structopt(parse(from_os_str))]
     mount_point: PathBuf,
 
     /// CID of the VM where the service runs.
-    #[structopt(long)]
+    #[clap(long)]
     cid: u32,
 
     /// Extra options to FUSE
-    #[structopt(short = "o")]
+    #[clap(short = 'o')]
     extra_options: Option<String>,
 
     /// Number of threads to serve FUSE requests.
-    #[structopt(short = "j")]
+    #[clap(short = 'j')]
     thread_number: Option<NonZeroU8>,
 
     /// A read-only remote file with integrity check. Can be multiple.
@@ -76,21 +75,21 @@
     /// For example, `--remote-ro-file 5:sha256-1234abcd` tells the filesystem to associate the
     /// file $MOUNTPOINT/5 with a remote FD 5, and has a fs-verity digest with sha256 of the hex
     /// value 1234abcd.
-    #[structopt(long, parse(try_from_str = parse_remote_ro_file_option))]
+    #[clap(long, value_parser = parse_remote_ro_file_option)]
     remote_ro_file: Vec<OptionRemoteRoFile>,
 
     /// A read-only remote file without integrity check. Can be multiple.
     ///
     /// For example, `--remote-ro-file-unverified 5` tells the filesystem to associate the file
     /// $MOUNTPOINT/5 with a remote FD 5.
-    #[structopt(long)]
+    #[clap(long)]
     remote_ro_file_unverified: Vec<i32>,
 
     /// A new read-writable remote file with integrity check. Can be multiple.
     ///
     /// For example, `--remote-new-rw-file 5` tells the filesystem to associate the file
     /// $MOUNTPOINT/5 with a remote FD 5.
-    #[structopt(long)]
+    #[clap(long)]
     remote_new_rw_file: Vec<i32>,
 
     /// A read-only directory that represents a remote directory. The directory view is constructed
@@ -107,7 +106,7 @@
     /// include a file like /5/system/framework/framework.jar. "prefix/" tells the filesystem to
     /// strip the path (e.g. "system/") from the mount point to match the expected location of the
     /// remote FD (e.g. a directory FD of "/system" in the remote).
-    #[structopt(long, parse(try_from_str = parse_remote_new_ro_dir_option))]
+    #[clap(long, value_parser = parse_remote_new_ro_dir_option)]
     remote_ro_dir: Vec<OptionRemoteRoDir>,
 
     /// A new directory that is assumed empty in the backing filesystem. New files created in this
@@ -116,14 +115,15 @@
     ///
     /// For example, `--remote-new-rw-dir 5` tells the filesystem to associate $MOUNTPOINT/5
     /// with a remote dir FD 5.
-    #[structopt(long)]
+    #[clap(long)]
     remote_new_rw_dir: Vec<i32>,
 
     /// Enable debugging features.
-    #[structopt(long)]
+    #[clap(long)]
     debug: bool,
 }
 
+#[derive(Clone)]
 struct OptionRemoteRoFile {
     /// ID to refer to the remote file.
     remote_fd: i32,
@@ -132,6 +132,7 @@
     digest: String,
 }
 
+#[derive(Clone)]
 struct OptionRemoteRoDir {
     /// ID to refer to the remote dir.
     remote_dir_fd: i32,
@@ -305,7 +306,7 @@
 }
 
 fn try_main() -> Result<()> {
-    let args = Args::from_args_safe()?;
+    let args = Args::parse();
 
     let log_level = if args.debug { log::Level::Debug } else { log::Level::Info };
     android_logger::init_once(
diff --git a/avmd/Android.bp b/avmd/Android.bp
index dc6a896..7237f5f 100644
--- a/avmd/Android.bp
+++ b/avmd/Android.bp
@@ -40,6 +40,9 @@
     name: "avmdtool_tests",
     srcs: ["tests/*_test.rs"],
     test_suites: ["general-tests"],
+    rustlibs: [
+        "libtempfile",
+    ],
     compile_multilib: "first",
     data_bins: ["avmdtool"],
     data: ["tests/data/*"],
diff --git a/avmd/tests/avmdtool_test.rs b/avmd/tests/avmdtool_test.rs
index d93cb6f..4647f06 100644
--- a/avmd/tests/avmdtool_test.rs
+++ b/avmd/tests/avmdtool_test.rs
@@ -16,19 +16,50 @@
 
 use std::fs;
 use std::process::Command;
+use tempfile::TempDir;
 
 #[test]
 fn test_dump() {
-    // test.avmd is generated with
-    // ```
-    // avmdtool create /tmp/test.amvd \
-    // --apex-payload microdroid vbmeta ./libs/apexutil/tests/data/test.apex \
-    // --apk microdroid_manager apk \
-    // ./libs/apkverify/tests/data/v3-only-with-rsa-pkcs1-sha256-4096.apk \
-    // --apk microdroid_manager extra-apk ./libs/apkverify/tests/data/v3-only-with-stamp.apk
-    //```
-    let output =
-        Command::new("./avmdtool").args(["dump", "tests/data/test.avmd"]).output().unwrap();
+    let filename = "tests/data/test.avmd";
+    assert!(
+        fs::metadata(filename).is_ok(),
+        "File '{}' does not exist. You can re-create it with:
+    avmdtool create {} \\
+    --apex-payload microdroid vbmeta tests/data/test.apex \\
+    --apk microdroid_manager apk \\
+    tests/data/v3-only-with-rsa-pkcs1-sha256-4096.apk \\
+    --apk microdroid_manager extra-apk tests/data/v3-only-with-stamp.apk",
+        filename,
+        filename
+    );
+    let output = Command::new("./avmdtool").args(["dump", filename]).output().unwrap();
     assert!(output.status.success());
     assert_eq!(output.stdout, fs::read("tests/data/test.avmd.dump").unwrap());
 }
+
+#[test]
+fn test_create() {
+    let test_dir = TempDir::new().unwrap();
+    let test_file_path = test_dir.path().join("tmp_test.amvd");
+    let output = Command::new("./avmdtool")
+        .args([
+            "create",
+            test_file_path.to_str().unwrap(),
+            "--apex-payload",
+            "microdroid",
+            "vbmeta",
+            "tests/data/test.apex",
+            "--apk",
+            "microdroid_manager",
+            "apk",
+            "tests/data/v3-only-with-rsa-pkcs1-sha256-4096.apk",
+            "--apk",
+            "microdroid_manager",
+            "extra-apk",
+            "tests/data/v3-only-with-stamp.apk",
+        ])
+        .output()
+        .unwrap();
+    assert!(output.status.success());
+    assert_eq!(fs::read(test_file_path).unwrap(), fs::read("tests/data/test.avmd").unwrap());
+}
diff --git a/avmd/tests/data/test.apex b/avmd/tests/data/test.apex
new file mode 100644
index 0000000..fd79365
--- /dev/null
+++ b/avmd/tests/data/test.apex
Binary files differ
diff --git a/avmd/tests/data/v3-only-with-rsa-pkcs1-sha256-4096.apk b/avmd/tests/data/v3-only-with-rsa-pkcs1-sha256-4096.apk
new file mode 100644
index 0000000..0c9391c
--- /dev/null
+++ b/avmd/tests/data/v3-only-with-rsa-pkcs1-sha256-4096.apk
Binary files differ
diff --git a/avmd/tests/data/v3-only-with-stamp.apk b/avmd/tests/data/v3-only-with-stamp.apk
new file mode 100644
index 0000000..5f65214
--- /dev/null
+++ b/avmd/tests/data/v3-only-with-stamp.apk
Binary files differ
diff --git a/compos/Android.bp b/compos/Android.bp
index 6aa9d3d..ea7c4d6 100644
--- a/compos/Android.bp
+++ b/compos/Android.bp
@@ -13,7 +13,6 @@
         "libandroid_logger",
         "libanyhow",
         "libbinder_rs",
-        "libclap",
         "libcompos_common",
         "liblibc",
         "liblog_rust",
diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 946bc5b..5ea5c06 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -33,7 +33,7 @@
 use std::num::NonZeroU32;
 use std::path::{Path, PathBuf};
 use std::thread;
-use vmclient::{DeathReason, VmInstance, VmWaitError};
+use vmclient::{DeathReason, ErrorCode, VmInstance, VmWaitError};
 
 /// This owns an instance of the CompOS VM.
 pub struct ComposClient(VmInstance);
@@ -45,9 +45,6 @@
     pub debug_mode: bool,
     /// Number of vCPUs to have in the VM. If None, defaults to 1.
     pub cpus: Option<NonZeroU32>,
-    /// Comma separated list of host CPUs where vCPUs are assigned to. If None, any host CPU can be
-    /// used to run any vCPU.
-    pub cpu_set: Option<String>,
     /// List of task profiles to apply to the VM
     pub task_profiles: Vec<String>,
     /// If present, overrides the path to the VM config JSON file
@@ -112,7 +109,6 @@
             protectedVm: protected_vm,
             memoryMib: parameters.memory_mib.unwrap_or(0), // 0 means use the default
             numCpus: parameters.cpus.map_or(1, NonZeroU32::get) as i32,
-            cpuAffinity: parameters.cpu_set.clone(),
             taskProfiles: parameters.task_profiles.clone(),
         });
 
@@ -241,8 +237,8 @@
         log::warn!("VM payload finished, cid = {}, exit code = {}", cid, exit_code);
     }
 
-    fn on_error(&self, cid: i32, error_code: i32, message: &str) {
-        log::warn!("VM error, cid = {}, error code = {}, message = {}", cid, error_code, message);
+    fn on_error(&self, cid: i32, error_code: ErrorCode, message: &str) {
+        log::warn!("VM error, cid = {}, error code = {:?}, message = {}", cid, error_code, message);
     }
 
     fn on_died(&self, cid: i32, death_reason: DeathReason) {
diff --git a/compos/common/lib.rs b/compos/common/lib.rs
index efbde06..a5b1ea8 100644
--- a/compos/common/lib.rs
+++ b/compos/common/lib.rs
@@ -61,6 +61,3 @@
 
 /// Number of CPUs to run dex2oat (actually the entire compos VM) with
 pub const DEX2OAT_THREADS_PROP_NAME: &str = "dalvik.vm.boot-dex2oat-threads";
-
-/// Set of host-side CPUs to run dex2oat (actually the entire compos VM) on
-pub const DEX2OAT_CPU_SET_PROP_NAME: &str = "dalvik.vm.boot-dex2oat-cpu-set";
diff --git a/compos/composd/src/instance_manager.rs b/compos/composd/src/instance_manager.rs
index 75671d7..451222e 100644
--- a/compos/composd/src/instance_manager.rs
+++ b/compos/composd/src/instance_manager.rs
@@ -23,8 +23,8 @@
 use binder::Strong;
 use compos_common::compos_client::VmParameters;
 use compos_common::{
-    CURRENT_INSTANCE_DIR, DEX2OAT_CPU_SET_PROP_NAME, DEX2OAT_THREADS_PROP_NAME,
-    PREFER_STAGED_VM_CONFIG_PATH, TEST_INSTANCE_DIR,
+    CURRENT_INSTANCE_DIR, DEX2OAT_THREADS_PROP_NAME, PREFER_STAGED_VM_CONFIG_PATH,
+    TEST_INSTANCE_DIR,
 };
 use rustutils::system_properties;
 use std::num::NonZeroU32;
@@ -92,15 +92,8 @@
             NonZeroU32::new(num_cpus::get() as u32)
         }
     };
-    let cpu_set = system_properties::read(DEX2OAT_CPU_SET_PROP_NAME)?;
     let task_profiles = vec!["SCHED_SP_COMPUTE".to_string()];
-    Ok(VmParameters {
-        cpus,
-        cpu_set,
-        task_profiles,
-        memory_mib: Some(VM_MEMORY_MIB),
-        ..Default::default()
-    })
+    Ok(VmParameters { cpus, task_profiles, memory_mib: Some(VM_MEMORY_MIB), ..Default::default() })
 }
 
 // Ensures we only run one instance at a time.
diff --git a/compos/composd_cmd/composd_cmd.rs b/compos/composd_cmd/composd_cmd.rs
index d5feed8..b6d82aa 100644
--- a/compos/composd_cmd/composd_cmd.rs
+++ b/compos/composd_cmd/composd_cmd.rs
@@ -31,30 +31,32 @@
     },
 };
 use anyhow::{bail, Context, Result};
+use clap::Parser;
 use compos_common::timeouts::TIMEOUTS;
 use std::sync::{Arc, Condvar, Mutex};
 use std::time::Duration;
 
+#[derive(Parser)]
+enum Actions {
+    /// Compile classpath for real. Output can be used after a reboot.
+    StagedApexCompile {},
+
+    /// Compile classpath in a debugging VM. Output is ignored.
+    TestCompile {
+        /// If any APEX is staged, prefer the staged version.
+        #[clap(long)]
+        prefer_staged: bool,
+    },
+}
+
 fn main() -> Result<()> {
-    #[rustfmt::skip]
-    let app = clap::App::new("composd_cmd")
-        .subcommand(
-            clap::SubCommand::with_name("staged-apex-compile"))
-        .subcommand(
-            clap::SubCommand::with_name("test-compile")
-                .arg(clap::Arg::with_name("prefer-staged").long("prefer-staged")),
-        );
-    let args = app.get_matches();
+    let action = Actions::parse();
 
     ProcessState::start_thread_pool();
 
-    match args.subcommand() {
-        Some(("staged-apex-compile", _)) => run_staged_apex_compile()?,
-        Some(("test-compile", sub_matches)) => {
-            let prefer_staged = sub_matches.is_present("prefer-staged");
-            run_test_compile(prefer_staged)?;
-        }
-        _ => panic!("Unrecognized subcommand"),
+    match action {
+        Actions::StagedApexCompile {} => run_staged_apex_compile()?,
+        Actions::TestCompile { prefer_staged } => run_test_compile(prefer_staged)?,
     }
 
     println!("All Ok!");
diff --git a/compos/verify/verify.rs b/compos/verify/verify.rs
index 2ece8f5..3abdc74 100644
--- a/compos/verify/verify.rs
+++ b/compos/verify/verify.rs
@@ -20,6 +20,7 @@
 use android_logger::LogId;
 use anyhow::{bail, Context, Result};
 use binder::ProcessState;
+use clap::{Parser, ValueEnum};
 use compos_common::compos_client::{ComposClient, VmParameters};
 use compos_common::odrefresh::{
     CURRENT_ARTIFACTS_SUBDIR, ODREFRESH_OUTPUT_ROOT_DIR, PENDING_ARTIFACTS_SUBDIR,
@@ -37,6 +38,24 @@
 
 const MAX_FILE_SIZE_BYTES: u64 = 100 * 1024;
 
+#[derive(Parser)]
+struct Args {
+    /// Type of the VM instance
+    #[clap(long, value_enum)]
+    instance: Instance,
+
+    /// Starts the VM in debug mode
+    #[clap(long, action)]
+    debug: bool,
+}
+
+#[derive(ValueEnum, Clone)]
+enum Instance {
+    Current,
+    Pending,
+    Test,
+}
+
 fn main() {
     android_logger::init_once(
         android_logger::Config::default()
@@ -57,23 +76,11 @@
 }
 
 fn try_main() -> Result<()> {
-    let matches = clap::App::new("compos_verify")
-        .arg(
-            clap::Arg::with_name("instance")
-                .long("instance")
-                .takes_value(true)
-                .required(true)
-                .possible_values(&["current", "pending", "test"]),
-        )
-        .arg(clap::Arg::with_name("debug").long("debug"))
-        .get_matches();
-
-    let debug_mode = matches.is_present("debug");
-    let (instance_dir, artifacts_dir) = match matches.value_of("instance").unwrap() {
-        "current" => (CURRENT_INSTANCE_DIR, CURRENT_ARTIFACTS_SUBDIR),
-        "pending" => (CURRENT_INSTANCE_DIR, PENDING_ARTIFACTS_SUBDIR),
-        "test" => (TEST_INSTANCE_DIR, TEST_ARTIFACTS_SUBDIR),
-        _ => unreachable!("Unexpected instance name"),
+    let args = Args::parse();
+    let (instance_dir, artifacts_dir) = match args.instance {
+        Instance::Current => (CURRENT_INSTANCE_DIR, CURRENT_ARTIFACTS_SUBDIR),
+        Instance::Pending => (CURRENT_INSTANCE_DIR, PENDING_ARTIFACTS_SUBDIR),
+        Instance::Test => (TEST_INSTANCE_DIR, TEST_ARTIFACTS_SUBDIR),
     };
 
     let instance_dir = Path::new(COMPOS_DATA_ROOT).join(instance_dir);
@@ -104,7 +111,7 @@
         instance_image,
         &idsig,
         &idsig_manifest_apk,
-        &VmParameters { debug_mode, ..Default::default() },
+        &VmParameters { debug_mode: args.debug, ..Default::default() },
     )?;
 
     let service = vm_instance.connect_service()?;
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 828ac9f..ae84c08 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -45,11 +45,15 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.lang.ref.WeakReference;
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.Files;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.WeakHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -65,6 +69,11 @@
  * @hide
  */
 public class VirtualMachine {
+    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";
 
@@ -159,6 +168,8 @@
 
     private final ExecutorService mExecutorService = Executors.newCachedThreadPool();
 
+    @NonNull private final Context mContext;
+
     static {
         System.loadLibrary("virtualmachine_jni");
     }
@@ -166,6 +177,7 @@
     private VirtualMachine(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
+        mContext = context;
         mPackageName = context.getPackageName();
         mName = name;
         mConfig = config;
@@ -231,6 +243,18 @@
             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);
+            }
+
+            instancesMap.put(name, new WeakReference<>(vm));
+        }
+
         return vm;
     }
 
@@ -249,7 +273,23 @@
             throw new VirtualMachineException(e);
         }
 
-        VirtualMachine vm = new VirtualMachine(context, name, config);
+        VirtualMachine vm;
+        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();
+            } else {
+                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
@@ -544,6 +584,11 @@
         mInstanceFilePath.delete();
         mIdsigFilePath.delete();
         vmRootDir.delete();
+
+        synchronized (sInstancesLock) {
+            Map<String, WeakReference<VirtualMachine>> instancesMap = sInstances.get(mContext);
+            if (instancesMap != null) instancesMap.remove(mName);
+        }
     }
 
     /**
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index b7c7a88..4ecd942 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -37,7 +37,6 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
-import java.util.regex.Pattern;
 
 /**
  * Represents a configuration of a virtual machine. A configuration consists of hardware
@@ -57,7 +56,6 @@
     private static final String KEY_PROTECTED_VM = "protectedVm";
     private static final String KEY_MEMORY_MIB = "memoryMib";
     private static final String KEY_NUM_CPUS = "numCpus";
-    private static final String KEY_CPU_AFFINITY = "cpuAffinity";
 
     // Paths to the APK file of this application.
     @NonNull private final String mApkPath;
@@ -106,13 +104,6 @@
     private final int mNumCpus;
 
     /**
-     * Comma-separated list of CPUs or CPU ranges to run vCPUs on (e.g. 0,1-3,5), or
-     * colon-separated list of assignments of vCPU to host CPU assignments (e.g. 0=0:1=1:2=2).
-     * Default is no mask which means a vCPU can run on any host CPU.
-     */
-    private final String mCpuAffinity;
-
-    /**
      * Path within the APK to the payload config file that defines software aspects of this config.
      */
     @NonNull private final String mPayloadConfigPath;
@@ -124,8 +115,7 @@
             DebugLevel debugLevel,
             boolean protectedVm,
             int memoryMib,
-            int numCpus,
-            String cpuAffinity) {
+            int numCpus) {
         mApkPath = apkPath;
         mCerts = certs;
         mPayloadConfigPath = payloadConfigPath;
@@ -133,7 +123,6 @@
         mProtectedVm = protectedVm;
         mMemoryMib = memoryMib;
         mNumCpus = numCpus;
-        mCpuAffinity = cpuAffinity;
     }
 
     /** Loads a config from a stream, for example a file. */
@@ -166,9 +155,8 @@
         final boolean protectedVm = b.getBoolean(KEY_PROTECTED_VM);
         final int memoryMib = b.getInt(KEY_MEMORY_MIB);
         final int numCpus = b.getInt(KEY_NUM_CPUS);
-        final String cpuAffinity = b.getString(KEY_CPU_AFFINITY);
         return new VirtualMachineConfig(apkPath, certs, payloadConfigPath, debugLevel, protectedVm,
-                memoryMib, numCpus, cpuAffinity);
+                memoryMib, numCpus);
     }
 
     /** Persists this config to a stream, for example a file. */
@@ -249,7 +237,6 @@
         parcel.protectedVm = mProtectedVm;
         parcel.memoryMib = mMemoryMib;
         parcel.numCpus = mNumCpus;
-        parcel.cpuAffinity = mCpuAffinity;
         // 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];
@@ -268,7 +255,6 @@
         private boolean mProtectedVm;
         private int mMemoryMib;
         private int mNumCpus;
-        private String mCpuAffinity;
 
         /**
          * Creates a builder for the given context (APK), and the payload config file in APK.
@@ -281,7 +267,6 @@
             mDebugLevel = DebugLevel.NONE;
             mProtectedVm = false;
             mNumCpus = 1;
-            mCpuAffinity = null;
         }
 
         /**
@@ -326,19 +311,6 @@
         }
 
         /**
-         * Sets on which host CPUs the vCPUs can run. The format is a comma-separated list of CPUs
-         * or CPU ranges to run vCPUs on. e.g. "0,1-3,5" to choose host CPUs 0, 1, 2, 3, and 5.
-         * Or this can be a colon-separated list of assignments of vCPU to host CPU assignments.
-         * e.g. "0=0:1=1:2=2" to map vCPU 0 to host CPU 0, and so on.
-         *
-         * @hide
-         */
-        public Builder cpuAffinity(String affinity) {
-            mCpuAffinity = affinity;
-            return this;
-        }
-
-        /**
          * Builds an immutable {@link VirtualMachineConfig}
          *
          * @hide
@@ -365,13 +337,6 @@
                         + "range [1, " + availableCpus + "]");
             }
 
-            if (mCpuAffinity != null
-                    && !Pattern.matches("[\\d]+(-[\\d]+)?(,[\\d]+(-[\\d]+)?)*", mCpuAffinity)
-                    && !Pattern.matches("[\\d]+=[\\d]+(:[\\d]+=[\\d]+)*", mCpuAffinity)) {
-                throw new IllegalArgumentException("CPU affinity [" + mCpuAffinity + "]"
-                        + " is invalid");
-            }
-
             if (mProtectedVm
                     && !HypervisorProperties.hypervisor_protected_vm_supported().orElse(false)) {
                 throw new UnsupportedOperationException(
@@ -384,7 +349,7 @@
 
             return new VirtualMachineConfig(
                     apkPath, certs, mPayloadConfigPath, mDebugLevel, mProtectedVm, mMemoryMib,
-                    mNumCpus, mCpuAffinity);
+                    mNumCpus);
         }
     }
 }
diff --git a/libs/apkverify/Android.bp b/libs/apkverify/Android.bp
index ab9265d..9bb8f8e 100644
--- a/libs/apkverify/Android.bp
+++ b/libs/apkverify/Android.bp
@@ -13,9 +13,11 @@
         "libbyteorder",
         "libbytes",
         "liblog_rust",
+        "libnum_traits",
         "libopenssl",
         "libzip",
     ],
+    proc_macros: ["libnum_derive"],
 }
 
 rust_library {
@@ -30,6 +32,7 @@
     name: "libapkverify.test",
     defaults: ["libapkverify.defaults"],
     test_suites: ["general-tests"],
+    data: ["tests/data/*"],
 }
 
 rust_test {
diff --git a/libs/apkverify/src/algorithms.rs b/libs/apkverify/src/algorithms.rs
new file mode 100644
index 0000000..edfa946
--- /dev/null
+++ b/libs/apkverify/src/algorithms.rs
@@ -0,0 +1,167 @@
+/*
+ * 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.
+ */
+
+//! Algorithms used for APK Signature Scheme.
+
+use anyhow::{ensure, Result};
+use num_derive::FromPrimitive;
+use openssl::hash::MessageDigest;
+use openssl::pkey::{self, PKey};
+use openssl::rsa::Padding;
+use openssl::sign::Verifier;
+use std::cmp::Ordering;
+
+/// [Signature Algorithm IDs]: https://source.android.com/docs/security/apksigning/v2#signature-algorithm-ids
+///
+/// Some of the algorithms are not implemented. See b/197052981.
+#[derive(Clone, Debug, Eq, FromPrimitive)]
+#[repr(u32)]
+pub enum SignatureAlgorithmID {
+    RsaPssWithSha256 = 0x0101,
+    RsaPssWithSha512 = 0x0102,
+    RsaPkcs1V15WithSha256 = 0x0103,
+    RsaPkcs1V15WithSha512 = 0x0104,
+    EcdsaWithSha256 = 0x0201,
+    EcdsaWithSha512 = 0x0202,
+    DsaWithSha256 = 0x0301,
+    VerityRsaPkcs1V15WithSha256 = 0x0421,
+    VerityEcdsaWithSha256 = 0x0423,
+    VerityDsaWithSha256 = 0x0425,
+}
+
+impl Ord for SignatureAlgorithmID {
+    /// Ranks the signature algorithm according to the corresponding content
+    /// digest algorithm's rank.
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.to_content_digest_algorithm().cmp(&other.to_content_digest_algorithm())
+    }
+}
+
+impl PartialOrd for SignatureAlgorithmID {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl PartialEq for SignatureAlgorithmID {
+    fn eq(&self, other: &Self) -> bool {
+        self.cmp(other) == Ordering::Equal
+    }
+}
+
+impl SignatureAlgorithmID {
+    pub(crate) fn new_verifier<'a>(
+        &self,
+        public_key: &'a PKey<pkey::Public>,
+    ) -> Result<Verifier<'a>> {
+        ensure!(
+            !matches!(
+                self,
+                SignatureAlgorithmID::DsaWithSha256 | SignatureAlgorithmID::VerityDsaWithSha256
+            ),
+            "TODO(b/197052981): Algorithm '{:#?}' is not implemented.",
+            self
+        );
+        ensure!(public_key.id() == self.pkey_id(), "Public key has the wrong ID");
+        let mut verifier = Verifier::new(self.new_message_digest(), public_key)?;
+        if public_key.id() == pkey::Id::RSA {
+            verifier.set_rsa_padding(self.rsa_padding())?;
+        }
+        Ok(verifier)
+    }
+
+    /// Returns the message digest corresponding to the signature algorithm
+    /// according to the spec [Signature Algorithm IDs].
+    pub(crate) fn new_message_digest(&self) -> MessageDigest {
+        match self {
+            SignatureAlgorithmID::RsaPssWithSha256
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::EcdsaWithSha256
+            | SignatureAlgorithmID::DsaWithSha256
+            | SignatureAlgorithmID::VerityRsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::VerityEcdsaWithSha256
+            | SignatureAlgorithmID::VerityDsaWithSha256 => MessageDigest::sha256(),
+            SignatureAlgorithmID::RsaPssWithSha512
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha512
+            | SignatureAlgorithmID::EcdsaWithSha512 => MessageDigest::sha512(),
+        }
+    }
+
+    fn pkey_id(&self) -> pkey::Id {
+        match self {
+            SignatureAlgorithmID::RsaPssWithSha256
+            | SignatureAlgorithmID::RsaPssWithSha512
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha512
+            | SignatureAlgorithmID::VerityRsaPkcs1V15WithSha256 => pkey::Id::RSA,
+            SignatureAlgorithmID::EcdsaWithSha256
+            | SignatureAlgorithmID::EcdsaWithSha512
+            | SignatureAlgorithmID::VerityEcdsaWithSha256 => pkey::Id::EC,
+            SignatureAlgorithmID::DsaWithSha256 | SignatureAlgorithmID::VerityDsaWithSha256 => {
+                pkey::Id::DSA
+            }
+        }
+    }
+
+    fn rsa_padding(&self) -> Padding {
+        match self {
+            SignatureAlgorithmID::RsaPssWithSha256 | SignatureAlgorithmID::RsaPssWithSha512 => {
+                Padding::PKCS1_PSS
+            }
+            SignatureAlgorithmID::RsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::VerityRsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha512 => Padding::PKCS1,
+            SignatureAlgorithmID::EcdsaWithSha256
+            | SignatureAlgorithmID::EcdsaWithSha512
+            | SignatureAlgorithmID::VerityEcdsaWithSha256
+            | SignatureAlgorithmID::DsaWithSha256
+            | SignatureAlgorithmID::VerityDsaWithSha256 => Padding::NONE,
+        }
+    }
+
+    fn to_content_digest_algorithm(&self) -> ContentDigestAlgorithm {
+        match self {
+            SignatureAlgorithmID::RsaPssWithSha256
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::EcdsaWithSha256
+            | SignatureAlgorithmID::DsaWithSha256 => ContentDigestAlgorithm::ChunkedSha256,
+            SignatureAlgorithmID::RsaPssWithSha512
+            | SignatureAlgorithmID::RsaPkcs1V15WithSha512
+            | SignatureAlgorithmID::EcdsaWithSha512 => ContentDigestAlgorithm::ChunkedSha512,
+            SignatureAlgorithmID::VerityRsaPkcs1V15WithSha256
+            | SignatureAlgorithmID::VerityEcdsaWithSha256
+            | SignatureAlgorithmID::VerityDsaWithSha256 => {
+                ContentDigestAlgorithm::VerityChunkedSha256
+            }
+        }
+    }
+}
+
+/// The rank of the content digest algorithm in this enum is used to help pick
+/// v4 apk digest.
+/// According to APK Signature Scheme v4, [apk digest] is the first available
+/// content digest of the highest rank (rank N).
+///
+/// This rank was also used for step 3a of the v3 signature verification.
+///
+/// [apk digest]: https://source.android.com/docs/security/features/apksigning/v4#apk-digest
+/// [v3 verification]: https://source.android.com/docs/security/apksigning/v3#v3-verification
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+enum ContentDigestAlgorithm {
+    ChunkedSha256 = 1,
+    VerityChunkedSha256,
+    ChunkedSha512,
+}
diff --git a/libs/apkverify/src/lib.rs b/libs/apkverify/src/lib.rs
index c5aa9e5..040c304 100644
--- a/libs/apkverify/src/lib.rs
+++ b/libs/apkverify/src/lib.rs
@@ -16,6 +16,7 @@
 
 //! Verifies APK/APEX signing with v2/v3 scheme
 
+mod algorithms;
 mod bytes_ext;
 mod sigutil;
 #[allow(dead_code)]
@@ -23,5 +24,5 @@
 mod v3;
 mod ziputil;
 
-// TODO(jooyung) fallback to v2 when v3 not found
+// TODO(b/197052981) fallback to v2 when v3 not found
 pub use v3::{get_public_key_der, pick_v4_apk_digest, verify};
diff --git a/libs/apkverify/src/sigutil.rs b/libs/apkverify/src/sigutil.rs
index 2b2f9da..3832c09 100644
--- a/libs/apkverify/src/sigutil.rs
+++ b/libs/apkverify/src/sigutil.rs
@@ -16,19 +16,25 @@
 
 //! Utilities for Signature Verification
 
-use anyhow::{anyhow, bail, Result};
+// TODO(b/246254355): Remove this once we migrate all the usages of
+// raw signature algorithm id to the enum.
+#![allow(dead_code)]
+
+use anyhow::{anyhow, ensure, Error, Result};
 use byteorder::{LittleEndian, ReadBytesExt};
 use bytes::{Buf, BufMut, Bytes, BytesMut};
+use num_traits::FromPrimitive;
 use openssl::hash::{DigestBytes, Hasher, MessageDigest};
 use std::cmp::min;
-use std::io::{Cursor, Read, Seek, SeekFrom, Take};
+use std::io::{self, Cursor, ErrorKind, Read, Seek, SeekFrom, Take};
 
+use crate::algorithms::SignatureAlgorithmID;
 use crate::ziputil::{set_central_directory_offset, zip_sections};
 
 const APK_SIG_BLOCK_MIN_SIZE: u32 = 32;
 const APK_SIG_BLOCK_MAGIC: u128 = 0x3234206b636f6c4220676953204b5041;
 
-// TODO(jooyung): introduce type
+// TODO(b/246254355): Migrates usages of raw signature algorithm id to the enum.
 pub const SIGNATURE_RSA_PSS_WITH_SHA256: u32 = 0x0101;
 pub const SIGNATURE_RSA_PSS_WITH_SHA512: u32 = 0x0102;
 pub const SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA256: u32 = 0x0103;
@@ -40,15 +46,15 @@
 pub const SIGNATURE_VERITY_ECDSA_WITH_SHA256: u32 = 0x0423;
 pub const SIGNATURE_VERITY_DSA_WITH_SHA256: u32 = 0x0425;
 
-// TODO(jooyung): introduce type
-const CONTENT_DIGEST_CHUNKED_SHA256: u32 = 1;
-const CONTENT_DIGEST_CHUNKED_SHA512: u32 = 2;
-const CONTENT_DIGEST_VERITY_CHUNKED_SHA256: u32 = 3;
-#[allow(unused)]
-const CONTENT_DIGEST_SHA256: u32 = 4;
-
 const CHUNK_SIZE_BYTES: u64 = 1024 * 1024;
 
+/// The [APK structure] has four major sections:
+///
+/// | Zip contents | APK Signing Block | Central directory | EOCD(End of Central Directory) |
+///
+/// This structure contains the offset/size information of all the sections except the Zip contents.
+///
+/// [APK structure]: https://source.android.com/docs/security/apksigning/v2#apk-signing-block
 pub struct ApkSections<R> {
     inner: R,
     signing_block_offset: u32,
@@ -79,7 +85,6 @@
     /// and the additional information relevant for verifying the block against the file.
     pub fn find_signature(&mut self, block_id: u32) -> Result<Bytes> {
         let signing_block = self.bytes(self.signing_block_offset, self.signing_block_size)?;
-        // TODO(jooyung): propagate NotFound error so that verification can fallback to V2
         find_signature_scheme_block(Bytes::from(signing_block), block_id)
     }
 
@@ -91,7 +96,10 @@
     ///    order the chunks appear in the APK.
     /// (see https://source.android.com/security/apksigning/v2#integrity-protected-contents)
     pub fn compute_digest(&mut self, signature_algorithm_id: u32) -> Result<Vec<u8>> {
-        let digester = Digester::new(signature_algorithm_id)?;
+        // TODO(b/246254355): Passes the enum SignatureAlgorithmID directly to this method.
+        let signature_algorithm_id = SignatureAlgorithmID::from_u32(signature_algorithm_id)
+            .ok_or_else(|| anyhow!("Unsupported algorithm ID: {}", signature_algorithm_id))?;
+        let digester = Digester { message_digest: signature_algorithm_id.new_message_digest() };
 
         let mut digests_of_chunks = BytesMut::new();
         let mut chunk_count = 0u32;
@@ -157,34 +165,20 @@
 }
 
 struct Digester {
-    algorithm: MessageDigest,
+    message_digest: MessageDigest,
 }
 
 const CHUNK_HEADER_TOP: &[u8] = &[0x5a];
 const CHUNK_HEADER_MID: &[u8] = &[0xa5];
 
 impl Digester {
-    fn new(signature_algorithm_id: u32) -> Result<Digester> {
-        let digest_algorithm_id = to_content_digest_algorithm(signature_algorithm_id)?;
-        let algorithm = match digest_algorithm_id {
-            CONTENT_DIGEST_CHUNKED_SHA256 => MessageDigest::sha256(),
-            CONTENT_DIGEST_CHUNKED_SHA512 => MessageDigest::sha512(),
-            // TODO(jooyung): implement
-            CONTENT_DIGEST_VERITY_CHUNKED_SHA256 => {
-                bail!("TODO(b/190343842): CONTENT_DIGEST_VERITY_CHUNKED_SHA256: not implemented")
-            }
-            _ => bail!("Unknown digest algorithm: {}", digest_algorithm_id),
-        };
-        Ok(Digester { algorithm })
-    }
-
     // v2/v3 digests are computed after prepending "header" byte and "size" info.
     fn digest(&self, data: &[u8], header: &[u8], size: u32) -> Result<DigestBytes> {
-        let mut ctx = Hasher::new(self.algorithm)?;
-        ctx.update(header)?;
-        ctx.update(&size.to_le_bytes())?;
-        ctx.update(data)?;
-        Ok(ctx.finish()?)
+        let mut hasher = Hasher::new(self.message_digest)?;
+        hasher.update(header)?;
+        hasher.update(&size.to_le_bytes())?;
+        hasher.update(data)?;
+        Ok(hasher.finish()?)
     }
 }
 
@@ -198,30 +192,30 @@
     // * @+8  bytes payload
     // * @-24 bytes uint64:    size in bytes (same as the one above)
     // * @-16 bytes uint128:   magic
-    if central_directory_offset < APK_SIG_BLOCK_MIN_SIZE {
-        bail!(
-            "APK too small for APK Signing Block. ZIP Central Directory offset: {}",
-            central_directory_offset
-        );
-    }
+    ensure!(
+        central_directory_offset >= APK_SIG_BLOCK_MIN_SIZE,
+        "APK too small for APK Signing Block. ZIP Central Directory offset: {}",
+        central_directory_offset
+    );
     reader.seek(SeekFrom::Start((central_directory_offset - 24) as u64))?;
     let size_in_footer = reader.read_u64::<LittleEndian>()? as u32;
-    if reader.read_u128::<LittleEndian>()? != APK_SIG_BLOCK_MAGIC {
-        bail!("No APK Signing Block before ZIP Central Directory")
-    }
+    ensure!(
+        reader.read_u128::<LittleEndian>()? == APK_SIG_BLOCK_MAGIC,
+        "No APK Signing Block before ZIP Central Directory"
+    );
     let total_size = size_in_footer + 8;
     let signing_block_offset = central_directory_offset
         .checked_sub(total_size)
         .ok_or_else(|| anyhow!("APK Signing Block size out of range: {}", size_in_footer))?;
     reader.seek(SeekFrom::Start(signing_block_offset as u64))?;
     let size_in_header = reader.read_u64::<LittleEndian>()? as u32;
-    if size_in_header != size_in_footer {
-        bail!(
-            "APK Signing Block sizes in header and footer do not match: {} vs {}",
-            size_in_header,
-            size_in_footer
-        );
-    }
+    // This corresponds to APK Signature Scheme v3 verification step 1a.
+    ensure!(
+        size_in_header == size_in_footer,
+        "APK Signing Block sizes in header and footer do not match: {} vs {}",
+        size_in_header,
+        size_in_footer
+    );
     Ok((signing_block_offset, total_size))
 }
 
@@ -236,9 +230,11 @@
     let mut entry_count = 0;
     while pairs.has_remaining() {
         entry_count += 1;
-        if pairs.remaining() < 8 {
-            bail!("Insufficient data to read size of APK Signing Block entry #{}", entry_count);
-        }
+        ensure!(
+            pairs.remaining() >= 8,
+            "Insufficient data to read size of APK Signing Block entry #{}",
+            entry_count
+        );
         let length = pairs.get_u64_le();
         let mut pair = pairs.split_to(length as usize);
         let id = pair.get_u32_le();
@@ -246,52 +242,87 @@
             return Ok(pair);
         }
     }
-    // TODO(jooyung): return NotFound error
-    bail!("No APK Signature Scheme block in APK Signing Block with ID: {}", block_id)
+    let context =
+        format!("No APK Signature Scheme block in APK Signing Block with ID: {}", block_id);
+    Err(Error::new(io::Error::from(ErrorKind::NotFound)).context(context))
 }
 
-pub fn is_supported_signature_algorithm(algorithm_id: u32) -> bool {
-    matches!(
-        algorithm_id,
-        SIGNATURE_RSA_PSS_WITH_SHA256
-            | SIGNATURE_RSA_PSS_WITH_SHA512
-            | SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA256
-            | SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA512
-            | SIGNATURE_ECDSA_WITH_SHA256
-            | SIGNATURE_ECDSA_WITH_SHA512
-            | SIGNATURE_DSA_WITH_SHA256
-            | SIGNATURE_VERITY_RSA_PKCS1_V1_5_WITH_SHA256
-            | SIGNATURE_VERITY_ECDSA_WITH_SHA256
-            | SIGNATURE_VERITY_DSA_WITH_SHA256
-    )
-}
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use byteorder::LittleEndian;
+    use std::fs::File;
+    use std::mem::size_of_val;
 
-fn to_content_digest_algorithm(algorithm_id: u32) -> Result<u32> {
-    match algorithm_id {
-        SIGNATURE_RSA_PSS_WITH_SHA256
-        | SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA256
-        | SIGNATURE_ECDSA_WITH_SHA256
-        | SIGNATURE_DSA_WITH_SHA256 => Ok(CONTENT_DIGEST_CHUNKED_SHA256),
-        SIGNATURE_RSA_PSS_WITH_SHA512
-        | SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA512
-        | SIGNATURE_ECDSA_WITH_SHA512 => Ok(CONTENT_DIGEST_CHUNKED_SHA512),
-        SIGNATURE_VERITY_RSA_PKCS1_V1_5_WITH_SHA256
-        | SIGNATURE_VERITY_ECDSA_WITH_SHA256
-        | SIGNATURE_VERITY_DSA_WITH_SHA256 => Ok(CONTENT_DIGEST_VERITY_CHUNKED_SHA256),
-        _ => bail!("Unknown signature algorithm: {}", algorithm_id),
+    use crate::v3::{to_hex_string, APK_SIGNATURE_SCHEME_V3_BLOCK_ID};
+
+    const CENTRAL_DIRECTORY_HEADER_SIGNATURE: u32 = 0x02014b50;
+
+    #[test]
+    fn test_apk_sections() {
+        let apk_file = File::open("tests/data/v3-only-with-ecdsa-sha512-p521.apk").unwrap();
+        let apk_sections = ApkSections::new(apk_file).unwrap();
+        let mut reader = &apk_sections.inner;
+
+        // Checks APK Signing Block.
+        assert_eq!(
+            apk_sections.signing_block_offset + apk_sections.signing_block_size,
+            apk_sections.central_directory_offset
+        );
+        let apk_signature_offset = SeekFrom::Start(
+            apk_sections.central_directory_offset as u64 - size_of_val(&APK_SIG_BLOCK_MAGIC) as u64,
+        );
+        reader.seek(apk_signature_offset).unwrap();
+        assert_eq!(reader.read_u128::<LittleEndian>().unwrap(), APK_SIG_BLOCK_MAGIC);
+
+        // Checks Central directory.
+        assert_eq!(reader.read_u32::<LittleEndian>().unwrap(), CENTRAL_DIRECTORY_HEADER_SIGNATURE);
+        assert_eq!(
+            apk_sections.central_directory_offset + apk_sections.central_directory_size,
+            apk_sections.eocd_offset
+        );
+
+        // Checks EOCD.
+        assert_eq!(
+            reader.metadata().unwrap().len(),
+            (apk_sections.eocd_offset + apk_sections.eocd_size) as u64
+        );
     }
-}
 
-/// Rank the signature algorithm according to the preferences of the v4 signing scheme.
-pub fn rank_signature_algorithm(algo: u32) -> Result<u32> {
-    rank_content_digest_algorithm(to_content_digest_algorithm(algo)?)
-}
+    #[test]
+    fn test_apk_digest() {
+        let apk_file = File::open("tests/data/v3-only-with-dsa-sha256-1024.apk").unwrap();
+        let mut apk_sections = ApkSections::new(apk_file).unwrap();
+        let digest = apk_sections.compute_digest(SIGNATURE_DSA_WITH_SHA256).unwrap();
+        assert_eq!(
+            "0DF2426EA33AEDAF495D88E5BE0C6A1663FF0A81C5ED12D5B2929AE4B4300F2F",
+            to_hex_string(&digest[..])
+        );
+    }
 
-fn rank_content_digest_algorithm(id: u32) -> Result<u32> {
-    match id {
-        CONTENT_DIGEST_CHUNKED_SHA256 => Ok(0),
-        CONTENT_DIGEST_VERITY_CHUNKED_SHA256 => Ok(1),
-        CONTENT_DIGEST_CHUNKED_SHA512 => Ok(2),
-        _ => bail!("Unknown digest algorithm: {}", id),
+    #[test]
+    fn test_apk_sections_cannot_find_signature() {
+        let apk_file = File::open("tests/data/v2-only-two-signers.apk").unwrap();
+        let mut apk_sections = ApkSections::new(apk_file).unwrap();
+        let result = apk_sections.find_signature(APK_SIGNATURE_SCHEME_V3_BLOCK_ID);
+
+        assert!(result.is_err());
+        let error = result.unwrap_err();
+        assert_eq!(error.downcast_ref::<io::Error>().unwrap().kind(), ErrorKind::NotFound);
+        assert!(
+            error.to_string().contains(&APK_SIGNATURE_SCHEME_V3_BLOCK_ID.to_string()),
+            "Error should contain the block ID: {}",
+            error
+        );
+    }
+
+    #[test]
+    fn test_apk_sections_find_signature() {
+        let apk_file = File::open("tests/data/v3-only-with-dsa-sha256-1024.apk").unwrap();
+        let mut apk_sections = ApkSections::new(apk_file).unwrap();
+        let signature = apk_sections.find_signature(APK_SIGNATURE_SCHEME_V3_BLOCK_ID).unwrap();
+
+        let expected_v3_signature_block_size = 1289; // Only for this specific APK
+        assert_eq!(signature.len(), expected_v3_signature_block_size);
     }
 }
diff --git a/libs/apkverify/src/v3.rs b/libs/apkverify/src/v3.rs
index 0c20a2e..557abcd 100644
--- a/libs/apkverify/src/v3.rs
+++ b/libs/apkverify/src/v3.rs
@@ -15,33 +15,28 @@
  */
 
 //! Verifies APK Signature Scheme V3
-
-// TODO(jooyung) remove this
-#![allow(dead_code)]
+//!
+//! [v3 verification]: https://source.android.com/security/apksigning/v3#verification
 
 use anyhow::{anyhow, bail, ensure, Context, Result};
 use bytes::Bytes;
-use openssl::hash::MessageDigest;
+use num_traits::FromPrimitive;
 use openssl::pkey::{self, PKey};
-use openssl::rsa::Padding;
-use openssl::sign::Verifier;
 use openssl::x509::X509;
 use std::fs::File;
 use std::io::{Read, Seek};
 use std::ops::Range;
 use std::path::Path;
 
+use crate::algorithms::SignatureAlgorithmID;
 use crate::bytes_ext::{BytesExt, LengthPrefixed, ReadFromBytes};
 use crate::sigutil::*;
 
 pub const APK_SIGNATURE_SCHEME_V3_BLOCK_ID: u32 = 0xf05368c0;
 
-// TODO(jooyung): get "ro.build.version.sdk"
+// TODO(b/190343842): get "ro.build.version.sdk"
 const SDK_INT: u32 = 31;
 
-/// Data model for Signature Scheme V3
-/// https://source.android.com/security/apksigning/v3#verification
-
 type Signers = LengthPrefixed<Vec<LengthPrefixed<Signer>>>;
 
 struct Signer {
@@ -63,6 +58,7 @@
     certificates: LengthPrefixed<Vec<LengthPrefixed<X509Certificate>>>,
     min_sdk: u32,
     max_sdk: u32,
+    #[allow(dead_code)]
     additional_attributes: LengthPrefixed<Vec<LengthPrefixed<AdditionalAttributes>>>,
 }
 
@@ -74,6 +70,7 @@
 
 #[derive(Debug)]
 struct Signature {
+    /// TODO(b/246254355): Change the type of signature_algorithm_id to SignatureAlgorithmID
     signature_algorithm_id: u32,
     signature: LengthPrefixed<Bytes>,
 }
@@ -128,14 +125,14 @@
     })
 }
 
-/// Gets the APK digest.
+/// Gets the v4 [apk_digest].
+///
+/// [apk_digest]: https://source.android.com/docs/security/apksigning/v4#apk-digest
 pub fn pick_v4_apk_digest<R: Read + Seek>(apk: R) -> Result<(u32, Box<[u8]>)> {
     let mut sections = ApkSections::new(apk)?;
     let mut block = sections.find_signature(APK_SIGNATURE_SCHEME_V3_BLOCK_ID)?;
     let signers = block.read::<Signers>()?;
-    if signers.len() != 1 {
-        bail!("should only have one signer");
-    }
+    ensure!(signers.len() == 1, "should only have one signer");
     signers[0].pick_v4_apk_digest()
 }
 
@@ -146,8 +143,8 @@
         Ok(self
             .signatures
             .iter()
-            .filter(|sig| is_supported_signature_algorithm(sig.signature_algorithm_id))
-            .max_by_key(|sig| rank_signature_algorithm(sig.signature_algorithm_id).unwrap())
+            .filter(|sig| SignatureAlgorithmID::from_u32(sig.signature_algorithm_id).is_some())
+            .max_by_key(|sig| SignatureAlgorithmID::from_u32(sig.signature_algorithm_id).unwrap())
             .ok_or_else(|| anyhow!("No supported signatures found"))?)
     }
 
@@ -162,6 +159,7 @@
         Ok((digest.signature_algorithm_id, digest.digest.as_ref().to_vec().into_boxed_slice()))
     }
 
+    /// The steps in this method implements APK Signature Scheme v3 verification step 3.
     fn verify<R: Read + Seek>(&self, sections: &mut ApkSections<R>) -> Result<Box<[u8]>> {
         // 1. Choose the strongest supported signature algorithm ID from signatures.
         let strongest = self.strongest_signature()?;
@@ -217,44 +215,21 @@
             bail!("Public key mismatch between certificate and signature record");
         }
 
-        // TODO(jooyung) 8. If the proof-of-rotation attribute exists for the signer verify that the struct is valid and this signer is the last certificate in the list.
+        // TODO(b/245914104)
+        // 8. If the proof-of-rotation attribute exists for the signer verify that the
+        // struct is valid and this signer is the last certificate in the list.
         Ok(self.public_key.to_vec().into_boxed_slice())
     }
 }
 
-fn verify_signed_data(data: &Bytes, signature: &Signature, key: &PKey<pkey::Public>) -> Result<()> {
-    let (pkey_id, padding, digest) = match signature.signature_algorithm_id {
-        SIGNATURE_RSA_PSS_WITH_SHA256 => {
-            (pkey::Id::RSA, Padding::PKCS1_PSS, MessageDigest::sha256())
-        }
-        SIGNATURE_RSA_PSS_WITH_SHA512 => {
-            (pkey::Id::RSA, Padding::PKCS1_PSS, MessageDigest::sha512())
-        }
-        SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA256 | SIGNATURE_VERITY_RSA_PKCS1_V1_5_WITH_SHA256 => {
-            (pkey::Id::RSA, Padding::PKCS1, MessageDigest::sha256())
-        }
-        SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA512 => {
-            (pkey::Id::RSA, Padding::PKCS1, MessageDigest::sha512())
-        }
-        SIGNATURE_ECDSA_WITH_SHA256 | SIGNATURE_VERITY_ECDSA_WITH_SHA256 => {
-            (pkey::Id::EC, Padding::NONE, MessageDigest::sha256())
-        }
-        // TODO(b/190343842) not implemented signature algorithm
-        SIGNATURE_ECDSA_WITH_SHA512
-        | SIGNATURE_DSA_WITH_SHA256
-        | SIGNATURE_VERITY_DSA_WITH_SHA256 => {
-            bail!(
-                "TODO(b/190343842) not implemented signature algorithm: {:#x}",
-                signature.signature_algorithm_id
-            );
-        }
-        _ => bail!("Unsupported signature algorithm: {:#x}", signature.signature_algorithm_id),
-    };
-    ensure!(key.id() == pkey_id, "Public key has the wrong ID");
-    let mut verifier = Verifier::new(digest, key)?;
-    if pkey_id == pkey::Id::RSA {
-        verifier.set_rsa_padding(padding)?;
-    }
+fn verify_signed_data(
+    data: &Bytes,
+    signature: &Signature,
+    public_key: &PKey<pkey::Public>,
+) -> Result<()> {
+    let mut verifier = SignatureAlgorithmID::from_u32(signature.signature_algorithm_id)
+        .context("Unsupported algorithm")?
+        .new_verifier(public_key)?;
     verifier.update(data)?;
     let verified = verifier.verify(&signature.signature)?;
     ensure!(verified, "Signature is invalid ");
@@ -262,7 +237,7 @@
 }
 
 // ReadFromBytes implementations
-// TODO(jooyung): add derive macro: #[derive(ReadFromBytes)]
+// TODO(b/190343842): add derive macro: #[derive(ReadFromBytes)]
 
 impl ReadFromBytes for Signer {
     fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
@@ -301,6 +276,39 @@
 }
 
 #[inline]
-fn to_hex_string(buf: &[u8]) -> String {
+pub(crate) fn to_hex_string(buf: &[u8]) -> String {
     buf.iter().map(|b| format!("{:02X}", b)).collect()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs::File;
+
+    #[test]
+    fn test_pick_v4_apk_digest_only_with_v3_dsa_sha256() {
+        check_v4_apk_digest(
+            "tests/data/v3-only-with-dsa-sha256-1024.apk",
+            SIGNATURE_DSA_WITH_SHA256,
+            "0DF2426EA33AEDAF495D88E5BE0C6A1663FF0A81C5ED12D5B2929AE4B4300F2F",
+        );
+    }
+
+    #[test]
+    fn test_pick_v4_apk_digest_only_with_v3_pkcs1_sha512() {
+        check_v4_apk_digest(
+            "tests/data/v3-only-with-rsa-pkcs1-sha512-1024.apk",
+            SIGNATURE_RSA_PKCS1_V1_5_WITH_SHA512,
+            "9B9AE02DA60B18999BF541790F00D380006FDF0655C3C482AA0BB0AF17CF7A42\
+             ECF56B973518546C9080B2FEF83027E895ED2882BFC88EA19790BBAB29AF53B3",
+        );
+    }
+
+    fn check_v4_apk_digest(apk_filename: &str, expected_algorithm: u32, expected_digest: &str) {
+        let apk_file = File::open(apk_filename).unwrap();
+        let (signature_algorithm_id, apk_digest) = pick_v4_apk_digest(apk_file).unwrap();
+
+        assert_eq!(expected_algorithm, signature_algorithm_id);
+        assert_eq!(expected_digest, to_hex_string(apk_digest.as_ref()));
+    }
+}
diff --git a/libs/apkverify/src/ziputil.rs b/libs/apkverify/src/ziputil.rs
index ebb66e0..8badff2 100644
--- a/libs/apkverify/src/ziputil.rs
+++ b/libs/apkverify/src/ziputil.rs
@@ -14,17 +14,18 @@
  * limitations under the License.
  */
 
-//! Utilities for zip handling
+//! Utilities for zip handling of APK files.
 
 use anyhow::{bail, Result};
 use bytes::{Buf, BufMut};
 use std::io::{Read, Seek, SeekFrom};
 use zip::ZipArchive;
 
-const EOCD_MIN_SIZE: usize = 22;
+const EOCD_SIZE_WITHOUT_COMMENT: usize = 22;
 const EOCD_CENTRAL_DIRECTORY_SIZE_FIELD_OFFSET: usize = 12;
 const EOCD_CENTRAL_DIRECTORY_OFFSET_FIELD_OFFSET: usize = 16;
-const EOCD_MAGIC: u32 = 0x06054b50;
+/// End of Central Directory signature
+const EOCD_SIGNATURE: u32 = 0x06054b50;
 const ZIP64_MARK: u32 = 0xffffffff;
 
 #[derive(Debug, PartialEq, Eq)]
@@ -39,7 +40,7 @@
 pub fn zip_sections<R: Read + Seek>(mut reader: R) -> Result<(R, ZipSections)> {
     // open a zip to parse EOCD
     let archive = ZipArchive::new(reader)?;
-    let eocd_size = archive.comment().len() + EOCD_MIN_SIZE;
+    let eocd_size = archive.comment().len() + EOCD_SIZE_WITHOUT_COMMENT;
     if archive.offset() != 0 {
         bail!("Invalid ZIP: offset should be 0, but {}.", archive.offset());
     }
@@ -49,7 +50,7 @@
     let eocd_offset = reader.seek(SeekFrom::Current(0))? as u32;
     let mut eocd = vec![0u8; eocd_size as usize];
     reader.read_exact(&mut eocd)?;
-    if (&eocd[0..]).get_u32_le() != EOCD_MAGIC {
+    if (&eocd[0..]).get_u32_le() != EOCD_SIGNATURE {
         bail!("Invalid ZIP: ZipArchive::new() should point EOCD after reading.");
     }
     let (central_directory_size, central_directory_offset) = get_central_directory(&eocd)?;
@@ -72,7 +73,7 @@
 }
 
 fn get_central_directory(buf: &[u8]) -> Result<(u32, u32)> {
-    if buf.len() < EOCD_MIN_SIZE {
+    if buf.len() < EOCD_SIZE_WITHOUT_COMMENT {
         bail!("Invalid EOCD size: {}", buf.len());
     }
     let mut buf = &buf[EOCD_CENTRAL_DIRECTORY_SIZE_FIELD_OFFSET..];
@@ -83,7 +84,7 @@
 
 /// Update EOCD's central_directory_offset field.
 pub fn set_central_directory_offset(buf: &mut [u8], value: u32) -> Result<()> {
-    if buf.len() < EOCD_MIN_SIZE {
+    if buf.len() < EOCD_SIZE_WITHOUT_COMMENT {
         bail!("Invalid EOCD size: {}", buf.len());
     }
     (&mut buf[EOCD_CENTRAL_DIRECTORY_OFFSET_FIELD_OFFSET..]).put_u32_le(value);
@@ -94,6 +95,8 @@
 mod tests {
     use super::*;
     use crate::testing::assert_contains;
+    use byteorder::{LittleEndian, ReadBytesExt};
+    use std::fs::File;
     use std::io::{Cursor, Write};
     use zip::{write::FileOptions, ZipWriter};
 
@@ -107,7 +110,10 @@
     #[test]
     fn test_zip_sections() {
         let (cursor, sections) = zip_sections(create_test_zip()).unwrap();
-        assert_eq!(sections.eocd_offset, (cursor.get_ref().len() - EOCD_MIN_SIZE) as u32);
+        assert_eq!(
+            sections.eocd_offset,
+            (cursor.get_ref().len() - EOCD_SIZE_WITHOUT_COMMENT) as u32
+        );
     }
 
     #[test]
@@ -118,7 +124,7 @@
         // insert garbage between CD and EOCD.
         // by the way, to mock zip-rs, use CD as garbage. This is implementation detail of zip-rs,
         // which reads CD at (eocd_offset - cd_size) instead of at cd_offset from EOCD.
-        let (pre_eocd, eocd) = buf.split_at(buf.len() - EOCD_MIN_SIZE);
+        let (pre_eocd, eocd) = buf.split_at(buf.len() - EOCD_SIZE_WITHOUT_COMMENT);
         let (_, cd_offset) = get_central_directory(eocd).unwrap();
         let cd = &pre_eocd[cd_offset as usize..];
 
@@ -127,4 +133,24 @@
         assert!(res.is_err());
         assert_contains(&res.err().unwrap().to_string(), "Invalid ZIP: offset should be 0");
     }
+
+    #[test]
+    fn test_zip_sections_with_apk() {
+        let apk = File::open("tests/data/v3-only-with-stamp.apk").unwrap();
+        let (mut reader, sections) = zip_sections(apk).unwrap();
+
+        // Checks Central directory.
+        assert_eq!(
+            sections.central_directory_offset + sections.central_directory_size,
+            sections.eocd_offset
+        );
+
+        // Checks EOCD.
+        reader.seek(SeekFrom::Start(sections.eocd_offset as u64)).unwrap();
+        assert_eq!(reader.read_u32::<LittleEndian>().unwrap(), EOCD_SIGNATURE);
+        assert_eq!(
+            reader.metadata().unwrap().len(),
+            (sections.eocd_offset + sections.eocd_size) as u64
+        );
+    }
 }
diff --git a/libs/apkverify/tests/apkverify_test.rs b/libs/apkverify/tests/apkverify_test.rs
index a674ad7..3818259 100644
--- a/libs/apkverify/tests/apkverify_test.rs
+++ b/libs/apkverify/tests/apkverify_test.rs
@@ -25,7 +25,7 @@
 fn test_verify_truncated_cd() {
     use zip::result::ZipError;
     let res = verify("tests/data/v2-only-truncated-cd.apk");
-    // TODO(jooyung): consider making a helper for err assertion
+    // TODO(b/190343842): consider making a helper for err assertion
     assert!(matches!(
         res.unwrap_err().root_cause().downcast_ref::<ZipError>().unwrap(),
         ZipError::InvalidArchive(_),
@@ -37,16 +37,12 @@
     assert!(verify("tests/data/test.apex").is_ok());
 }
 
-// TODO(b/190343842)
 #[test]
 fn test_verify_v3_dsa_sha256() {
     for key_name in KEY_NAMES_DSA.iter() {
         let res = verify(format!("tests/data/v3-only-with-dsa-sha256-{}.apk", key_name));
         assert!(res.is_err());
-        assert_contains(
-            &res.unwrap_err().to_string(),
-            "TODO(b/190343842) not implemented signature algorithm",
-        );
+        assert_contains(&res.unwrap_err().to_string(), "not implemented");
     }
 }
 
@@ -57,16 +53,10 @@
     }
 }
 
-// TODO(b/190343842)
 #[test]
 fn test_verify_v3_ecdsa_sha512() {
     for key_name in KEY_NAMES_ECDSA.iter() {
-        let res = verify(format!("tests/data/v3-only-with-ecdsa-sha512-{}.apk", key_name));
-        assert!(res.is_err());
-        assert_contains(
-            &res.unwrap_err().to_string(),
-            "TODO(b/190343842) not implemented signature algorithm",
-        );
+        assert!(verify(format!("tests/data/v3-only-with-ecdsa-sha512-{}.apk", key_name)).is_ok());
     }
 }
 
@@ -88,7 +78,6 @@
     }
 }
 
-// TODO(b/190343842)
 #[test]
 fn test_verify_v3_sig_does_not_verify() {
     let path_list = [
@@ -101,13 +90,11 @@
         assert!(res.is_err());
         let error_msg = &res.unwrap_err().to_string();
         assert!(
-            error_msg.contains("Signature is invalid")
-                || error_msg.contains("TODO(b/190343842) not implemented signature algorithm")
+            error_msg.contains("Signature is invalid") || error_msg.contains("not implemented")
         );
     }
 }
 
-// TODO(b/190343842)
 #[test]
 fn test_verify_v3_digest_mismatch() {
     let path_list = [
@@ -118,10 +105,7 @@
         let res = verify(path);
         assert!(res.is_err());
         let error_msg = &res.unwrap_err().to_string();
-        assert!(
-            error_msg.contains("Digest mismatch")
-                || error_msg.contains("TODO(b/190343842) not implemented signature algorithm")
-        );
+        assert!(error_msg.contains("Digest mismatch") || error_msg.contains("not implemented"));
     }
 }
 
diff --git a/libs/apkverify/tests/data/v2-only-two-signers.apk b/libs/apkverify/tests/data/v2-only-two-signers.apk
new file mode 100644
index 0000000..697b046
--- /dev/null
+++ b/libs/apkverify/tests/data/v2-only-two-signers.apk
Binary files differ
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 65aeb07..281416b 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -578,3 +578,18 @@
     filename: "event-log-tags",
     installable: false,
 }
+
+filegroup {
+    name: "microdroid_bootconfig_full_debuggable_src",
+    srcs: ["bootconfig.full_debuggable"],
+}
+
+filegroup {
+    name: "microdroid_bootconfig_app_debuggable_src",
+    srcs: ["bootconfig.app_debuggable"],
+}
+
+filegroup {
+    name: "microdroid_bootconfig_normal_src",
+    srcs: ["bootconfig.normal"],
+}
diff --git a/microdroid/initrd/Android.bp b/microdroid/initrd/Android.bp
new file mode 100644
index 0000000..d8e7069
--- /dev/null
+++ b/microdroid/initrd/Android.bp
@@ -0,0 +1,161 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_binary_host {
+    name: "initrd_bootconfig",
+    srcs: ["src/main.rs"],
+    rustlibs: [
+        "libanyhow",
+        "libclap",
+    ],
+    prefer_rlib: true,
+}
+
+python_binary_host {
+    name: "gen_vbmeta_bootconfig",
+    srcs: ["gen_vbmeta_bootconfig.py"],
+}
+
+genrule {
+    name: "microdroid_initrd_gen",
+    srcs: [
+        ":microdroid_ramdisk",
+        ":microdroid_vendor_ramdisk",
+    ],
+    out: ["microdroid_initrd.img"],
+    cmd: "cat $(in) > $(out)",
+}
+
+// This contains vbmeta hashes & related (boot)configs which are passed to kernel/init
+genrule {
+    name: "microdroid_vbmeta_bootconfig_gen",
+    srcs: [":microdroid_vbmeta"],
+    out: ["bootconfig_microdroid_vbmeta"],
+    tools: [
+        "gen_vbmeta_bootconfig",
+        "avbtool",
+    ],
+    cmd: "$(location gen_vbmeta_bootconfig) $(location avbtool) $(in) > $(out)",
+}
+
+bootconfigs_arm64 = [
+    ":microdroid_bootconfig_arm64_gen",
+    ":microdroid_vbmeta_bootconfig_gen",
+]
+
+bootconfigs_x86_64 = [
+    ":microdroid_bootconfig_x86_64_gen",
+    ":microdroid_vbmeta_bootconfig_gen",
+]
+
+genrule {
+    name: "microdroid_initrd_full_debuggable_arm64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_initrd_gen",
+        ":microdroid_bootconfig_full_debuggable_src",
+    ] + bootconfigs_arm64,
+    out: ["microdroid_initrd_full_debuggable_arm64"],
+    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+}
+
+genrule {
+    name: "microdroid_initrd_full_debuggable_x86_64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_initrd_gen",
+        ":microdroid_bootconfig_full_debuggable_src",
+    ] + bootconfigs_x86_64,
+    out: ["microdroid_initrd_full_debuggable_x86_64"],
+    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+}
+
+genrule {
+    name: "microdroid_initrd_app_debuggable_arm64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_initrd_gen",
+        ":microdroid_bootconfig_app_debuggable_src",
+    ] + bootconfigs_arm64,
+    out: ["microdroid_initrd_app_debuggable_arm64"],
+    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+}
+
+genrule {
+    name: "microdroid_initrd_app_debuggable_x86_64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_initrd_gen",
+        ":microdroid_bootconfig_app_debuggable_src",
+    ] + bootconfigs_x86_64,
+    out: ["microdroid_initrd_app_debuggable_x86_64"],
+    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+}
+
+genrule {
+    name: "microdroid_initrd_normal_arm64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_initrd_gen",
+        ":microdroid_bootconfig_normal_src",
+    ] + bootconfigs_arm64,
+    out: ["microdroid_initrd_normal_arm64"],
+    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+}
+
+genrule {
+    name: "microdroid_initrd_normal_x86_64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_initrd_gen",
+        ":microdroid_bootconfig_normal_src",
+    ] + bootconfigs_x86_64,
+    out: ["microdroid_initrd_normal_x86_64"],
+    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+}
+
+prebuilt_etc {
+    name: "microdroid_initrd_full_debuggable",
+    // We don't have ramdisk for architectures other than x86_64 & arm64
+    src: "empty_file",
+    arch: {
+        x86_64: {
+            src: ":microdroid_initrd_full_debuggable_x86_64",
+        },
+        arm64: {
+            src: ":microdroid_initrd_full_debuggable_arm64",
+        },
+    },
+    filename: "microdroid_initrd_full_debuggable.img",
+}
+
+prebuilt_etc {
+    name: "microdroid_initrd_app_debuggable",
+    // We don't have ramdisk for architectures other than x86_64 & arm64
+    src: "empty_file",
+    arch: {
+        x86_64: {
+            src: ":microdroid_initrd_app_debuggable_x86_64",
+        },
+        arm64: {
+            src: ":microdroid_initrd_app_debuggable_arm64",
+        },
+    },
+    filename: "microdroid_initrd_app_debuggable.img",
+}
+
+prebuilt_etc {
+    name: "microdroid_initrd_normal",
+    // We don't have ramdisk for architectures other than x86_64 & arm64
+    src: "empty_file",
+    arch: {
+        x86_64: {
+            src: ":microdroid_initrd_normal_x86_64",
+        },
+        arm64: {
+            src: ":microdroid_initrd_normal_arm64",
+        },
+    },
+    filename: "microdroid_initrd_normal.img",
+}
diff --git a/microdroid/initrd/empty_file b/microdroid/initrd/empty_file
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/microdroid/initrd/empty_file
diff --git a/microdroid/initrd/gen_vbmeta_bootconfig.py b/microdroid/initrd/gen_vbmeta_bootconfig.py
new file mode 100755
index 0000000..f0fc2e8
--- /dev/null
+++ b/microdroid/initrd/gen_vbmeta_bootconfig.py
@@ -0,0 +1,54 @@
+"""Extract information about vbmeta (digest/size) required in (androidboot.*)
+bootconfig. It uses avbtool to find some values such as vbmeta size, digest"""
+#!/usr/bin/env python3
+
+import sys
+import subprocess
+
+def main(args):
+    """Runs avbtool to generate vbmeta related bootconfigs"""
+    avbtool = args[0]
+    vbmeta_img = args[1]
+    hash_algorithm = 'sha256'
+    size = 0
+
+    with subprocess.Popen([avbtool, 'version'],
+                            stdout=subprocess.PIPE,
+                            stderr=subprocess.STDOUT) as proc:
+        stdout, _ = proc.communicate()
+        avb_version = stdout.decode("utf-8").split(" ")[1].strip()
+        avb_version = avb_version[0:avb_version.rfind('.')]
+
+    with subprocess.Popen([avbtool, 'info_image', '--image', vbmeta_img],
+                           stdout=subprocess.PIPE,
+                           stderr=subprocess.STDOUT) as proc:
+        stdout, _ = proc.communicate()
+        for line in stdout.decode("utf-8").split("\n"):
+            line = line.split(":")
+            if line[0] in \
+                ['Header Block', 'Authentication Block', 'Auxiliary Block']:
+                size += int(line[1].strip()[0:-6])
+
+    with subprocess.Popen([avbtool, 'calculate_vbmeta_digest',
+                            '--image', vbmeta_img,
+                            '--hash_algorithm', hash_algorithm],
+                           stdout=subprocess.PIPE,
+                           stderr=subprocess.STDOUT) as proc:
+        stdout, _ = proc.communicate()
+        vbmeta_hash = stdout.decode("utf-8").strip()
+
+    print(f'androidboot.vbmeta.size = {size}')
+    print(f'androidboot.vbmeta.digest = \"{vbmeta_hash}\"')
+    print(f'androidboot.vbmeta.hash_alg = \"{hash_algorithm}\"')
+    print(f'androidboot.vbmeta.avb_version = \"{avb_version}\"')
+    print('androidboot.veritymode = "enforcing"')
+    print('androidboot.vbmeta.invalidate_on_error = "yes"')
+    print('androidboot.vbmeta.device_state = "locked"')
+    print('androidboot.vbmeta.device = "/dev/block/by-name/vbmeta_a"')
+    print('androidboot.slot_suffix = "_a"')
+    print('androidboot.force_normal_boot = "1"')
+    print('androidboot.verifiedbootstate = "green"')
+
+## Main body
+if __name__ == '__main__':
+    main(sys.argv[1:])
diff --git a/microdroid/initrd/src/main.rs b/microdroid/initrd/src/main.rs
new file mode 100644
index 0000000..69c6ae4
--- /dev/null
+++ b/microdroid/initrd/src/main.rs
@@ -0,0 +1,74 @@
+// 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.
+
+//! Append bootconfig to initrd image
+use anyhow::Result;
+use clap::Parser;
+use std::fs::File;
+use std::io::{Read, Write};
+use std::path::PathBuf;
+
+const FOOTER_ALIGNMENT: usize = 4;
+const ZEROS: [u8; 4] = [0u8; 4_usize];
+
+#[derive(Parser, Debug)]
+struct Args {
+    /// Initrd (without bootconfig)
+    initrd: PathBuf,
+    /// Bootconfig
+    bootconfigs: Vec<PathBuf>,
+    /// Output
+    #[clap(long = "output")]
+    output: PathBuf,
+}
+
+fn get_checksum(file_path: &PathBuf) -> Result<u32> {
+    File::open(file_path)?.bytes().map(|x| Ok(x? as u32)).sum()
+}
+
+// Bootconfig is attached to the initrd in the following way:
+// [initrd][bootconfig][padding][size(le32)][checksum(le32)][#BOOTCONFIG\n]
+fn attach_bootconfig(initrd: PathBuf, bootconfigs: Vec<PathBuf>, output: PathBuf) -> Result<()> {
+    let mut output_file = File::create(&output)?;
+    let mut initrd_file = File::open(&initrd)?;
+    let initrd_size: usize = initrd_file.metadata()?.len().try_into()?;
+    let mut bootconfig_size: usize = 0;
+    let mut checksum: u32 = 0;
+
+    std::io::copy(&mut initrd_file, &mut output_file)?;
+    for bootconfig in bootconfigs {
+        let mut bootconfig_file = File::open(&bootconfig)?;
+        std::io::copy(&mut bootconfig_file, &mut output_file)?;
+        bootconfig_size += bootconfig_file.metadata()?.len() as usize;
+        checksum += get_checksum(&bootconfig)?;
+    }
+
+    let padding_size: usize = FOOTER_ALIGNMENT - (initrd_size + bootconfig_size) % FOOTER_ALIGNMENT;
+    output_file.write_all(&ZEROS[..padding_size])?;
+    output_file.write_all(&((padding_size + bootconfig_size) as u32).to_le_bytes())?;
+    output_file.write_all(&checksum.to_le_bytes())?;
+    output_file.write_all(b"#BOOTCONFIG\n")?;
+    output_file.flush()?;
+    Ok(())
+}
+
+fn try_main() -> Result<()> {
+    let args = Args::parse();
+    attach_bootconfig(args.initrd, args.bootconfigs, args.output)?;
+    Ok(())
+}
+
+fn main() {
+    try_main().unwrap()
+}
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index cfb851b..f3fa2f5 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -11,6 +11,7 @@
     rustlibs: [
         "android.hardware.security.dice-V1-rust",
         "android.security.dice-rust",
+        "android.system.virtualizationcommon-rust",
         "android.system.virtualizationservice-rust",
         "android.system.virtualmachineservice-rust",
         "libanyhow",
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 9e1890f..e3ad495 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -23,6 +23,10 @@
     Config::Config, InputValues::InputValues, Mode::Mode,
 };
 use android_security_dice::aidl::android::security::dice::IDiceMaintenance::IDiceMaintenance;
+use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::ErrorCode::ErrorCode;
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
+    VM_BINDER_SERVICE_PORT, VM_STREAM_SERVICE_PORT, IVirtualMachineService,
+};
 use anyhow::{anyhow, bail, ensure, Context, Error, Result};
 use apkverify::{get_public_key_der, verify};
 use binder::{wait_for_interface, Strong};
@@ -49,10 +53,6 @@
 use std::time::{Duration, SystemTime};
 use vsock::VsockStream;
 
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
-    ERROR_PAYLOAD_CHANGED, ERROR_PAYLOAD_VERIFICATION_FAILED, ERROR_PAYLOAD_INVALID_CONFIG, ERROR_UNKNOWN, VM_BINDER_SERVICE_PORT, VM_STREAM_SERVICE_PORT, IVirtualMachineService,
-};
-
 const WAIT_TIMEOUT: Duration = Duration::from_secs(10);
 const MAIN_APK_PATH: &str = "/dev/block/by-name/microdroid-apk";
 const MAIN_APK_IDSIG_PATH: &str = "/dev/block/by-name/microdroid-apk-idsig";
@@ -88,22 +88,24 @@
     InvalidConfig(String),
 }
 
-fn translate_error(err: &Error) -> (i32, String) {
+fn translate_error(err: &Error) -> (ErrorCode, String) {
     if let Some(e) = err.downcast_ref::<MicrodroidError>() {
         match e {
-            MicrodroidError::PayloadChanged(msg) => (ERROR_PAYLOAD_CHANGED, msg.to_string()),
+            MicrodroidError::PayloadChanged(msg) => (ErrorCode::PAYLOAD_CHANGED, msg.to_string()),
             MicrodroidError::PayloadVerificationFailed(msg) => {
-                (ERROR_PAYLOAD_VERIFICATION_FAILED, msg.to_string())
+                (ErrorCode::PAYLOAD_VERIFICATION_FAILED, msg.to_string())
             }
-            MicrodroidError::InvalidConfig(msg) => (ERROR_PAYLOAD_INVALID_CONFIG, msg.to_string()),
+            MicrodroidError::InvalidConfig(msg) => {
+                (ErrorCode::PAYLOAD_CONFIG_INVALID, msg.to_string())
+            }
 
             // Connection failure won't be reported to VS; return the default value
             MicrodroidError::FailedToConnectToVirtualizationService(msg) => {
-                (ERROR_UNKNOWN, msg.to_string())
+                (ErrorCode::UNKNOWN, msg.to_string())
             }
         }
     } else {
-        (ERROR_UNKNOWN, err.to_string())
+        (ErrorCode::UNKNOWN, err.to_string())
     }
 }
 
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index fb639d8..5bcafd5 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -15,11 +15,11 @@
 }
 
 cc_binary {
-    name: "pvmfw_elf",
-    stem: "pvmfw",
+    name: "pvmfw",
     defaults: ["vmbase_elf_defaults"],
     srcs: [
         "idmap.S",
+        "payload.S",
     ],
     static_libs: [
         "libpvmfw",
@@ -32,8 +32,9 @@
 }
 
 raw_binary {
-    name: "pvmfw",
-    src: ":pvmfw_elf",
+    name: "pvmfw_bin",
+    stem: "pvmfw.bin",
+    src: ":pvmfw",
     enabled: false,
     target: {
         android_arm64: {
diff --git a/pvmfw/payload.S b/pvmfw/payload.S
new file mode 100644
index 0000000..cbda448
--- /dev/null
+++ b/pvmfw/payload.S
@@ -0,0 +1,85 @@
+/*
+ * 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
+ *
+ *     https://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.
+ */
+
+.macro mov_i, reg:req, imm:req
+	movz \reg, :abs_g3:\imm
+	movk \reg, :abs_g2_nc:\imm
+	movk \reg, :abs_g1_nc:\imm
+	movk \reg, :abs_g0_nc:\imm
+.endm
+
+/* Stage 1 instruction access cacheability is unaffected. */
+.set .L_SCTLR_ELx_I, 0x1 << 12
+/* SETEND instruction disabled at EL0 in aarch32 mode. */
+.set .L_SCTLR_EL1_SED, 0x1 << 8
+/* Various IT instructions are disabled at EL0 in aarch32 mode. */
+.set .L_SCTLR_EL1_ITD, 0x1 << 7
+.set .L_SCTLR_EL1_RES1, (0x1 << 11) | (0x1 << 20) | (0x1 << 22) | (0x1 << 28) | (0x1 << 29)
+.set .Lsctlrval, .L_SCTLR_ELx_I | .L_SCTLR_EL1_SED | .L_SCTLR_EL1_ITD | .L_SCTLR_EL1_RES1
+
+/**
+ * Disable the exception vector, caches and page talbe and then jump to the payload at the given
+ * address, passing it the given FDT pointer.
+ *
+ * x0: FDT address to pass to payload
+ * x1: Payload address
+ */
+.global start_payload
+start_payload:
+        /* Move payload address to a higher register and zero out parameters other than x0. */
+        mov x30, x1
+        mov x1, #0
+        mov x2, #0
+        mov x3, #0
+
+        /* Zero out remaining registers to avoid leaking data. */
+        mov x4, #0
+        mov x5, #0
+        mov x6, #0
+        mov x7, #0
+        mov x8, #0
+        mov x9, #0
+        mov x10, #0
+        mov x11, #0
+        mov x12, #0
+        mov x13, #0
+        mov x14, #0
+        mov x15, #0
+        mov x16, #0
+        mov x17, #0
+        mov x18, #0
+        mov x19, #0
+        mov x20, #0
+        mov x21, #0
+        mov x22, #0
+        mov x23, #0
+        mov x24, #0
+        mov x25, #0
+        mov x26, #0
+        mov x27, #0
+        mov x28, #0
+
+        /* Disable the MMU and cache, and set other settings to valid warm reset values. */
+        mov_i x29, .Lsctlrval
+        msr sctlr_el1, x29
+        isb
+        msr ttbr0_el1, xzr
+
+        isb
+        dsb nsh
+
+        /* Jump into the payload. */
+        br x30
diff --git a/pvmfw/src/exceptions.rs b/pvmfw/src/exceptions.rs
index 61f7846..596ecc7 100644
--- a/pvmfw/src/exceptions.rs
+++ b/pvmfw/src/exceptions.rs
@@ -18,62 +18,71 @@
 use vmbase::{console::emergency_write_str, eprintln, power::reboot};
 
 #[no_mangle]
-extern "C" fn sync_exception_current() {
+extern "C" fn sync_exception_current(_elr: u64, _spsr: u64) {
+    let esr = read_esr();
     emergency_write_str("sync_exception_current\n");
-    print_esr();
+    print_esr(esr);
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn irq_current() {
+extern "C" fn irq_current(_elr: u64, _spsr: u64) {
     emergency_write_str("irq_current\n");
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn fiq_current() {
+extern "C" fn fiq_current(_elr: u64, _spsr: u64) {
     emergency_write_str("fiq_current\n");
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn serr_current() {
+extern "C" fn serr_current(_elr: u64, _spsr: u64) {
+    let esr = read_esr();
     emergency_write_str("serr_current\n");
-    print_esr();
+    print_esr(esr);
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn sync_lower() {
+extern "C" fn sync_lower(_elr: u64, _spsr: u64) {
+    let esr = read_esr();
     emergency_write_str("sync_lower\n");
-    print_esr();
+    print_esr(esr);
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn irq_lower() {
+extern "C" fn irq_lower(_elr: u64, _spsr: u64) {
     emergency_write_str("irq_lower\n");
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn fiq_lower() {
+extern "C" fn fiq_lower(_elr: u64, _spsr: u64) {
     emergency_write_str("fiq_lower\n");
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn serr_lower() {
+extern "C" fn serr_lower(_elr: u64, _spsr: u64) {
+    let esr = read_esr();
     emergency_write_str("serr_lower\n");
-    print_esr();
+    print_esr(esr);
     reboot();
 }
 
 #[inline]
-fn print_esr() {
+fn read_esr() -> u64 {
     let mut esr: u64;
     unsafe {
         asm!("mrs {esr}, esr_el1", esr = out(reg) esr);
     }
+    esr
+}
+
+#[inline]
+fn print_esr(esr: u64) {
     eprintln!("esr={:#08x}", esr);
 }
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index 0d71f6e..8526a92 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -30,4 +30,15 @@
         "fdt_address={:#018x}, payload_start={:#018x}, payload_size={:#018x}, x3={:#018x}",
         fdt_address, payload_start, payload_size, arg3,
     );
+
+    println!("Starting payload...");
+    // Safe because this is a function we have implemented in assembly that matches its signature
+    // here.
+    unsafe {
+        start_payload(fdt_address, payload_start);
+    }
+}
+
+extern "C" {
+    fn start_payload(fdt_address: u64, payload_start: u64) -> !;
 }
diff --git a/rialto/tests/test.rs b/rialto/tests/test.rs
index 8a78861..687ce86 100644
--- a/rialto/tests/test.rs
+++ b/rialto/tests/test.rs
@@ -63,7 +63,6 @@
         protectedVm: false,
         memoryMib: 300,
         numCpus: 1,
-        cpuAffinity: None,
         platformVersion: "~1.0".to_string(),
         taskProfiles: vec![],
     });
diff --git a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
index e34c4b0..260f804 100644
--- a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
@@ -20,8 +20,12 @@
 interface IBenchmarkService {
     const int SERVICE_PORT = 5677;
 
-    /** Reads a file and returns the elapsed seconds for the reading. */
-    double readFile(String filename, long fileSizeBytes, boolean isRand);
+    /**
+     * Measures the read rate for reading the given file.
+     *
+     * @return The read rate in MB/s.
+     */
+    double measureReadRate(String filename, long fileSizeBytes, boolean isRand);
 
     /** Returns an entry from /proc/meminfo. */
     long getMemInfoEntry(String name);
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
index bf78202..f1aebd2 100644
--- a/tests/benchmark/Android.bp
+++ b/tests/benchmark/Android.bp
@@ -29,7 +29,7 @@
 cc_library_shared {
     name: "MicrodroidBenchmarkNativeLib",
     srcs: ["src/native/benchmarkbinary.cpp"],
-    static_libs: ["libiovsock_vm"],
+    static_libs: ["libiobenchmark"],
     shared_libs: [
         "android.system.virtualmachineservice-ndk",
         "com.android.microdroid.testservice-ndk",
@@ -41,7 +41,7 @@
 }
 
 cc_library {
-    name: "libiovsock_vm",
+    name: "libiobenchmark",
     srcs: ["src/native/io_vsock.cpp"],
     export_include_dirs: ["src/native/include"],
     shared_libs: [
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 cdaf70c..7bf3c4e 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -268,9 +268,8 @@
         @Override
         public void onPayloadReady(VirtualMachine vm, IBenchmarkService benchmarkService)
                 throws RemoteException {
-            double elapsedSeconds = benchmarkService.readFile(FILENAME, mFileSizeBytes, mIsRand);
-            double fileSizeMb = mFileSizeBytes / SIZE_MB;
-            mReadRates.add(fileSizeMb / elapsedSeconds);
+            double readRate = benchmarkService.measureReadRate(FILENAME, mFileSizeBytes, mIsRand);
+            mReadRates.add(readRate);
         }
     }
 
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 28799c8..58b4cf0 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -42,6 +42,7 @@
 
 namespace {
 constexpr uint64_t kBlockSizeBytes = 4096;
+constexpr uint64_t kNumBytesPerMB = 1024 * 1024;
 
 template <typename T>
 static ndk::ScopedAStatus resultStatus(const T& result) {
@@ -56,9 +57,9 @@
 
 class IOBenchmarkService : public aidl::com::android::microdroid::testservice::BnBenchmarkService {
 public:
-    ndk::ScopedAStatus readFile(const std::string& filename, int64_t fileSizeBytes, bool isRand,
-                                double* out) override {
-        auto res = read_file(filename, fileSizeBytes, isRand);
+    ndk::ScopedAStatus measureReadRate(const std::string& filename, int64_t fileSizeBytes,
+                                       bool isRand, double* out) override {
+        auto res = measure_read_rate(filename, fileSizeBytes, isRand);
         if (res.ok()) {
             *out = res.value();
         }
@@ -90,8 +91,9 @@
     }
 
 private:
-    /** Returns the elapsed seconds for reading the file. */
-    Result<double> read_file(const std::string& filename, int64_t fileSizeBytes, bool is_rand) {
+    /** 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;
         if (is_rand) {
@@ -120,7 +122,9 @@
                 return ErrnoError() << "failed to read";
             }
         }
-        return {((double)clock() - start) / CLOCKS_PER_SEC};
+        double elapsed_seconds = ((double)clock() - start) / CLOCKS_PER_SEC;
+        double read_rate = (double)fileSizeBytes / kNumBytesPerMB / elapsed_seconds;
+        return {read_rate};
     }
 
     Result<size_t> read_meminfo_entry(const std::string& stat) {
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index ca4d1d0..eeadae1 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -19,8 +19,8 @@
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 
 import static com.google.common.truth.Truth.assertWithMessage;
-import static com.google.common.truth.TruthJUnit.assume;
 
+import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
 
 import android.platform.test.annotations.RootPermissionTest;
@@ -34,6 +34,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -66,6 +67,7 @@
     private static final String METRIC_PREFIX = "avf_perf/hostside/";
 
     private final MetricsProcessor mMetricsProcessor = new MetricsProcessor(METRIC_PREFIX);
+    @Rule public TestMetrics mMetrics = new TestMetrics();
 
     @Before
     public void setUp() throws Exception {
@@ -75,8 +77,10 @@
     @After
     public void tearDown() throws Exception {
         // Set PKVM enable and reboot to prevent previous staged session.
-        setPKVMStatusWithRebootToBootloader(true);
-        rebootFromBootloaderAndWaitBootCompleted();
+        if (!isCuttlefish()) {
+            setPKVMStatusWithRebootToBootloader(true);
+            rebootFromBootloaderAndWaitBootCompleted();
+        }
 
         CommandRunner android = new CommandRunner(getDevice());
 
@@ -86,7 +90,7 @@
 
     @Test
     public void testBootEnableAndDisablePKVM() throws Exception {
-        testPKVMStatusSwitchSupported();
+        skipIfPKVMStatusSwitchNotSupported();
 
         List<Double> bootWithPKVMEnableTime = new ArrayList<>(ROUND_COUNT);
         List<Double> bootWithoutPKVMEnableTime = new ArrayList<>(ROUND_COUNT);
@@ -116,7 +120,7 @@
 
     @Test
     public void testBootWithAndWithoutCompOS() throws Exception {
-        assume().withMessage("Skip on CF; too slow").that(isCuttlefish()).isFalse();
+        assumeFalse(isCuttlefish());
 
         List<Double> bootWithCompOsTime = new ArrayList<>(ROUND_COUNT);
         List<Double> bootWithoutCompOsTime = new ArrayList<>(ROUND_COUNT);
@@ -149,7 +153,9 @@
         reportMetric(bootWithoutCompOsTime, "boot_time_without_compos", "s");
     }
 
-    private void testPKVMStatusSwitchSupported() throws Exception {
+    private void skipIfPKVMStatusSwitchNotSupported() throws Exception {
+        assumeFalse(isCuttlefish());
+
         if (!getDevice().isStateBootloaderOrFastbootd()) {
             getDevice().rebootIntoBootloader();
         }
@@ -158,14 +164,16 @@
         CommandResult result;
         result = getDevice().executeFastbootCommand("oem", "pkvm", "status");
         rebootFromBootloaderAndWaitBootCompleted();
-        assumeTrue(!result.getStderr().contains("Invalid oem command"));
+        assumeFalse(result.getStderr().contains("Invalid oem command"));
+        // Skip the test if running on a build with pkvm_enabler. Disabling pKVM
+        // for such builds results in a bootloop.
+        assumeTrue(result.getStderr().contains("misc=auto"));
     }
 
     private void reportMetric(List<Double> data, String name, String unit) {
         Map<String, Double> stats = mMetricsProcessor.computeStats(data, name, unit);
-        TestMetrics metrics = new TestMetrics();
         for (Map.Entry<String, Double> entry : stats.entrySet()) {
-            metrics.addTestMetric(entry.getKey(), Double.toString(entry.getValue()));
+            mMetrics.addTestMetric(entry.getKey(), entry.getValue().toString());
         }
     }
 
@@ -195,11 +203,6 @@
         }
         assertWithMessage("Failed to set PKVM status. Reason: " + result)
             .that(result.toString()).ignoringCase().contains(expectedOutput);
-
-        // Skip the test if running on a build with pkvm_enabler. Disabling
-        // pKVM for such build results in a bootloop.
-        assertWithMessage("Expected build with PKVM status misc=auto. Reason: " + result)
-            .that(result.toString()).ignoringCase().contains("misc=auto");
     }
 
     private void rebootFromBootloaderAndWaitBootCompleted() throws Exception {
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 e5bc45a..a07731e 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
@@ -210,9 +210,7 @@
 
         protected void forceStop(VirtualMachine vm) {
             try {
-                vm.clearCallback();
                 vm.stop();
-                mExecutorService.shutdown();
             } catch (VirtualMachineException e) {
                 throw new RuntimeException(e);
             }
@@ -233,6 +231,7 @@
         @Override
         @CallSuper
         public void onDied(VirtualMachine vm, int reason) {
+            vm.clearCallback();
             mExecutorService.shutdown();
         }
 
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 cff06d5..7a9e2ea 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
@@ -235,11 +235,10 @@
             String configPath,
             boolean debug,
             int memoryMib,
-            Optional<Integer> numCpus,
-            Optional<String> cpuAffinity)
+            Optional<Integer> numCpus)
             throws DeviceNotAvailableException {
         return startMicrodroid(androidDevice, buildInfo, apkName, packageName, null, configPath,
-                debug, memoryMib, numCpus, cpuAffinity);
+                debug, memoryMib, numCpus);
     }
 
     public static String startMicrodroid(
@@ -251,12 +250,11 @@
             String configPath,
             boolean debug,
             int memoryMib,
-            Optional<Integer> numCpus,
-            Optional<String> cpuAffinity)
+            Optional<Integer> numCpus)
             throws DeviceNotAvailableException {
         return startMicrodroid(androidDevice, buildInfo, apkName, null, packageName,
                 extraIdsigPaths, configPath, debug,
-                memoryMib, numCpus, cpuAffinity);
+                memoryMib, numCpus);
     }
 
     private static void forwardFileToLog(CommandRunner android, String path, String tag)
@@ -280,8 +278,7 @@
             String configPath,
             boolean debug,
             int memoryMib,
-            Optional<Integer> numCpus,
-            Optional<String> cpuAffinity)
+            Optional<Integer> numCpus)
             throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(androidDevice);
 
@@ -314,7 +311,6 @@
                 "--console " + consolePath,
                 "--mem " + memoryMib,
                 numCpus.isPresent() ? "--cpus " + numCpus.get() : "",
-                cpuAffinity.isPresent() ? "--cpu-affinity " + cpuAffinity.get() : "",
                 debugFlag,
                 apkPath,
                 outApkIdsigPath,
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidTestCase.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidTestCase.java
index 5ba145b..69218a8 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidTestCase.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidTestCase.java
@@ -86,9 +86,8 @@
     private static final int MIN_MEM_ARM64 = 145;
     private static final int MIN_MEM_X86_64 = 196;
 
-    // Number of vCPUs and their affinity to host CPUs for testing purpose
+    // Number of vCPUs for testing purpose
     private static final int NUM_VCPUS = 3;
-    private static final String CPU_AFFINITY = "0,1,2";
 
     @Rule public TestLogData mTestLogs = new TestLogData();
     @Rule public TestName mTestName = new TestName();
@@ -510,8 +509,7 @@
                         configPath,
                         /* debug */ true,
                         minMemorySize(),
-                        Optional.of(NUM_VCPUS),
-                        Optional.of(CPU_AFFINITY));
+                        Optional.of(NUM_VCPUS));
         // check until microdroid is shut down
         CommandRunner android = new CommandRunner(getDevice());
         android.runWithTimeout(15000, "logcat", "-m", "1", "-e", "'crosvm has exited normally'");
@@ -563,8 +561,7 @@
                         configPath,
                         /* debug */ true,
                         minMemorySize(),
-                        Optional.of(NUM_VCPUS),
-                        Optional.of(CPU_AFFINITY));
+                        Optional.of(NUM_VCPUS));
 
         // Check VmCreationRequested atom and clear the statsd report
         List<StatsLog.EventMetricData> data;
@@ -586,7 +583,6 @@
                 AtomsProto.VmCreationRequested.ConfigType.VIRTUAL_MACHINE_APP_CONFIG,
                 atomVmCreationRequested.getConfigType());
         assertEquals(NUM_VCPUS, atomVmCreationRequested.getNumCpus());
-        assertEquals(CPU_AFFINITY, atomVmCreationRequested.getCpuAffinity());
         assertEquals(minMemorySize(), atomVmCreationRequested.getMemoryMib());
         assertEquals(
                 "com.android.art:com.android.compos:com.android.sdkext",
@@ -639,8 +635,7 @@
                         configPath,
                         /* debug */ true,
                         minMemorySize(),
-                        Optional.of(NUM_VCPUS),
-                        Optional.of(CPU_AFFINITY));
+                        Optional.of(NUM_VCPUS));
         adbConnectToMicrodroid(getDevice(), cid);
         waitForBootComplete();
         // Test writing to /data partition
@@ -790,8 +785,7 @@
                         configPath,
                         /* debug */ true,
                         minMemorySize(),
-                        Optional.of(NUM_VCPUS),
-                        Optional.of(CPU_AFFINITY));
+                        Optional.of(NUM_VCPUS));
         adbConnectToMicrodroid(getDevice(), cid);
         waitForBootComplete();
         rootMicrodroid();
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 0c048b9..c2060cb 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -170,19 +170,25 @@
                     .debugLevel(DebugLevel.NONE)
                     .build();
             VirtualMachine vm = mInner.forceCreateNewVirtualMachine("low_mem", lowMemConfig);
-            final CompletableFuture<Integer> exception = new CompletableFuture<>();
+            final CompletableFuture<Boolean> onPayloadReadyExecuted = new CompletableFuture<>();
+            final CompletableFuture<Boolean> onDiedExecuted = new CompletableFuture<>();
             VmEventListener listener =
                     new VmEventListener() {
                         @Override
+                        public void onPayloadReady(VirtualMachine vm) {
+                            onPayloadReadyExecuted.complete(true);
+                            super.onPayloadReady(vm);
+                        }
+                        @Override
                         public void onDied(VirtualMachine vm,  int reason) {
-                            exception.complete(reason);
+                            onDiedExecuted.complete(true);
                             super.onDied(vm, reason);
                         }
                     };
             listener.runToFinish(TAG, vm);
-            assertThat(exception.getNow(0)).isAnyOf(VirtualMachineCallback.DEATH_REASON_REBOOT,
-                    VirtualMachineCallback.DEATH_REASON_HANGUP,
-                    VirtualMachineCallback.DEATH_REASON_CRASH);
+            // Assert that onDied() was executed but onPayloadReady() was never run
+            assertThat(onDiedExecuted.getNow(false)).isTrue();
+            assertThat(onPayloadReadyExecuted.getNow(false)).isFalse();
         }
     }
 
@@ -200,39 +206,27 @@
 
         VirtualMachineConfig.Builder builder = mInner.newVmConfigBuilder("assets/vm_config.json");
         VirtualMachineConfig normalConfig = builder.debugLevel(DebugLevel.NONE).build();
-        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
-        VmEventListener listener =
-                new VmEventListener() {
-                    @Override
-                    public void onPayloadReady(VirtualMachine vm) {
-                        forceStop(vm);
-                    }
-                };
-        listener.runToFinish(TAG, vm);
+        mInner.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 vmInstanceBackup = File.createTempFile("instance", ".img");
+        Files.copy(vmInstance.toPath(), vmInstanceBackup.toPath(), REPLACE_EXISTING);
+        mInner.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).
-        // For testing, we do that by creating another VM with debug level, and copy the config file
-        // from the new VM directory to the old VM directory.
+        // 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.debugLevel(DebugLevel.FULL).build();
-        VirtualMachine newVm = mInner.forceCreateNewVirtualMachine("test_debug_vm", debugConfig);
-        File vmRoot = new File(getContext().getFilesDir(), "vm");
-        File newVmConfig = new File(new File(vmRoot, "test_debug_vm"), "config.xml");
-        File oldVmConfig = new File(new File(vmRoot, "test_vm"), "config.xml");
-        Files.copy(newVmConfig.toPath(), oldVmConfig.toPath(), REPLACE_EXISTING);
-        newVm.delete();
-        // re-load with the copied-in config file.
-        vm = mInner.getVirtualMachineManager().get("test_vm");
-        final CompletableFuture<Boolean> payloadStarted = new CompletableFuture<>();
-        listener =
-                new VmEventListener() {
-                    @Override
-                    public void onPayloadStarted(VirtualMachine vm, ParcelFileDescriptor stream) {
-                        payloadStarted.complete(true);
-                        forceStop(vm);
-                    }
-                };
-        listener.runToFinish(TAG, vm);
-        assertThat(payloadStarted.getNow(false)).isFalse();
+        mInner.forceCreateNewVirtualMachine("test_vm", debugConfig);
+        Files.copy(vmInstanceBackup.toPath(), vmInstance.toPath(), REPLACE_EXISTING);
+        assertThat(tryBootVm(TAG, "test_vm").payloadStarted).isFalse();
     }
 
     private class VmCdis {
@@ -510,4 +504,21 @@
         assertThat(bootResult.deathReason).isEqualTo(
                 VirtualMachineCallback.DEATH_REASON_MICRODROID_INVALID_PAYLOAD_CONFIG);
     }
+
+    @Test
+    public void sameInstancesShareTheSameVmObject()
+            throws VirtualMachineException, InterruptedException, IOException {
+        VirtualMachineConfig.Builder builder =
+                mInner.newVmConfigBuilder("assets/vm_config.json");
+        VirtualMachineConfig normalConfig = builder.debugLevel(DebugLevel.NONE).build();
+        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        VirtualMachine vm2 = mInner.getVirtualMachineManager().get("test_vm");
+        assertThat(vm).isEqualTo(vm2);
+
+        VirtualMachine newVm = mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        VirtualMachine newVm2 = mInner.getVirtualMachineManager().get("test_vm");
+        assertThat(newVm).isEqualTo(newVm2);
+
+        assertThat(vm).isNotEqualTo(newVm);
+    }
 }
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index b6b7b09..1493cb8 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -20,6 +20,7 @@
     },
     prefer_rlib: true,
     rustlibs: [
+        "android.system.virtualizationcommon-rust",
         "android.system.virtualizationservice-rust",
         "android.system.virtualmachineservice-rust",
         "android.os.permissions_aidl-rust",
diff --git a/virtualizationservice/aidl/Android.bp b/virtualizationservice/aidl/Android.bp
index 30a4b03..4d5326a 100644
--- a/virtualizationservice/aidl/Android.bp
+++ b/virtualizationservice/aidl/Android.bp
@@ -5,6 +5,7 @@
 aidl_interface {
     name: "android.system.virtualizationservice",
     srcs: ["android/system/virtualizationservice/**/*.aidl"],
+    imports: ["android.system.virtualizationcommon"],
     // This is never accessed directly. Apps are expected to use this indirectly via the Java
     // wrapper android.system.virtualmachine.
     unstable: true,
@@ -35,6 +36,7 @@
 aidl_interface {
     name: "android.system.virtualmachineservice",
     srcs: ["android/system/virtualmachineservice/**/*.aidl"],
+    imports: ["android.system.virtualizationcommon"],
     unstable: true,
     backend: {
         rust: {
@@ -46,3 +48,27 @@
         },
     },
 }
+
+aidl_interface {
+    name: "android.system.virtualizationcommon",
+    srcs: ["android/system/virtualizationcommon/**/*.aidl"],
+    unstable: true,
+    backend: {
+        java: {
+            apex_available: ["com.android.virt"],
+        },
+        ndk: {
+            apex_available: [
+                "com.android.virt",
+                "com.android.compos",
+            ],
+        },
+        rust: {
+            enabled: true,
+            apex_available: [
+                "com.android.virt",
+                "com.android.compos",
+            ],
+        },
+    },
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationcommon/ErrorCode.aidl b/virtualizationservice/aidl/android/system/virtualizationcommon/ErrorCode.aidl
new file mode 100644
index 0000000..04b9749
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationcommon/ErrorCode.aidl
@@ -0,0 +1,43 @@
+/*
+ * 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.virtualizationcommon;
+
+/**
+ * Errors reported from within a VM.
+ */
+@Backing(type="int")
+enum ErrorCode {
+    /**
+     * Error code for all other errors not listed below.
+     */
+    UNKNOWN = 0,
+
+    /**
+     * Error code indicating that the payload can't be verified due to various reasons (e.g invalid
+     * merkle tree, invalid formats, etc).
+     */
+    PAYLOAD_VERIFICATION_FAILED = 1,
+
+    /**
+     * Error code indicating that the payload is verified, but has changed since the last boot.
+     */
+    PAYLOAD_CHANGED = 2,
+
+    /**
+     * Error code indicating that the payload config is invalid.
+     */
+    PAYLOAD_CONFIG_INVALID = 3,
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
index 6c8eb4a..8d6ed08 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
@@ -15,6 +15,7 @@
  */
 package android.system.virtualizationservice;
 
+import android.system.virtualizationcommon.ErrorCode;
 import android.system.virtualizationservice.DeathReason;
 
 /**
@@ -44,7 +45,7 @@
     /**
      * Called when an error occurs in the VM.
      */
-    void onError(int cid, int errorCode, in String message);
+    void onError(int cid, ErrorCode errorCode, in String message);
 
     /**
      * Called when the VM dies.
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
index 8eb5497..cf5398d 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
@@ -48,7 +48,7 @@
     }
 
     /** Debug level of the VM */
-    DebugLevel debugLevel;
+    DebugLevel debugLevel = DebugLevel.NONE;
 
     /** Whether the VM should be a protected VM. */
     boolean protectedVm;
@@ -65,14 +65,9 @@
     int numCpus = 1;
 
     /**
-     * Comma-separated list of CPUs or CPU ranges to run vCPUs on (e.g. 0,1-3,5), or
-     * colon-separated list of assignments of vCPU to host CPU assignments (e.g. 0=0:1=1:2=2).
-     * Default is no mask which means a vCPU can run on any host CPU.
-     */
-    @nullable String cpuAffinity;
-
-    /**
      * List of task profile names to apply for the VM
+     *
+     * Note: Specifying a value here requires android.permission.USE_CUSTOM_VIRTUAL_MACHINE.
      */
     String[] taskProfiles;
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
index 672c41a..bed4097 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineDebugInfo.aidl
@@ -38,5 +38,5 @@
     int requesterPid;
 
     /** The current lifecycle state of the VM. */
-    VirtualMachineState state;
+    VirtualMachineState state = VirtualMachineState.NOT_STARTED;
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
index d11de03..993bbb0 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineRawConfig.aidl
@@ -55,13 +55,6 @@
     int numCpus = 1;
 
     /**
-     * Comma-separated list of CPUs or CPU ranges to run vCPUs on (e.g. 0,1-3,5), or
-     * colon-separated list of assignments of vCPU to host CPU assignments (e.g. 0=0:1=1:2=2).
-     * Default is no mask which means a vCPU can run on any host CPU.
-     */
-    @nullable String cpuAffinity;
-
-    /**
      * A version or range of versions of the virtual platform that this config is compatible with.
      * The format follows SemVer.
      */
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
index dff5d46..e8c1724 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -15,6 +15,8 @@
  */
 package android.system.virtualmachineservice;
 
+import android.system.virtualizationcommon.ErrorCode;
+
 /** {@hide} */
 interface IVirtualMachineService {
     /**
@@ -51,28 +53,7 @@
     void notifyPayloadFinished(int exitCode);
 
     /**
-     * Notifies that an error has occurred. See the ERROR_* constants.
+     * Notifies that an error has occurred inside the VM..
      */
-    void notifyError(int errorCode, in String message);
-
-    /**
-     * Error code for all other errors not listed below.
-     */
-    const int ERROR_UNKNOWN = 0;
-
-    /**
-     * Error code indicating that the payload can't be verified due to various reasons (e.g invalid
-     * merkle tree, invalid formats, etc).
-     */
-    const int ERROR_PAYLOAD_VERIFICATION_FAILED = 1;
-
-    /**
-     * Error code indicating that the payload is verified, but has changed since the last boot.
-     */
-    const int ERROR_PAYLOAD_CHANGED = 2;
-
-    /**
-     * Error code indicating that the payload config is invalid.
-     */
-    const int ERROR_PAYLOAD_INVALID_CONFIG = 3;
+    void notifyError(ErrorCode errorCode, in String message);
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 6a4cc93..1eca9fe 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -39,12 +39,11 @@
     self, BinderFeatures, ExceptionCode, Interface, LazyServiceGuard, ParcelFileDescriptor,
     SpIBinder, Status, StatusCode, Strong, ThreadState,
 };
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
-    IVirtualMachineService::{
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
         BnVirtualMachineService, IVirtualMachineService, VM_BINDER_SERVICE_PORT,
         VM_STREAM_SERVICE_PORT, VM_TOMBSTONES_SERVICE_PORT,
-    },
 };
+use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::ErrorCode::ErrorCode;
 use anyhow::{anyhow, bail, Context, Result};
 use rpcbinder::run_rpc_server_with_factory;
 use disk::QcowFile;
@@ -460,7 +459,6 @@
             protected: *is_protected,
             memory_mib: config.memoryMib.try_into().ok().and_then(NonZeroU32::new),
             cpus: config.numCpus.try_into().ok().and_then(NonZeroU32::new),
-            cpu_affinity: config.cpuAffinity.clone(),
             task_profiles: config.taskProfiles.clone(),
             console_fd,
             log_fd,
@@ -597,6 +595,12 @@
     config: &VirtualMachineAppConfig,
     temporary_directory: &Path,
 ) -> Result<VirtualMachineRawConfig> {
+    // Controlling CPUs is reserved for platform apps only, even when using
+    // VirtualMachineAppConfig.
+    if !config.taskProfiles.is_empty() {
+        check_use_custom_virtual_machine()?
+    }
+
     let apk_file = clone_file(config.apk.as_ref().unwrap())?;
     let idsig_file = clone_file(config.idsig.as_ref().unwrap())?;
     let instance_file = clone_file(config.instanceImage.as_ref().unwrap())?;
@@ -626,7 +630,6 @@
     vm_config.name = config.name.clone();
     vm_config.protectedVm = config.protectedVm;
     vm_config.numCpus = config.numCpus;
-    vm_config.cpuAffinity = config.cpuAffinity.clone();
     vm_config.taskProfiles = config.taskProfiles.clone();
 
     // Microdroid requires an additional payload disk image and the bootconfig partition.
@@ -878,7 +881,7 @@
     }
 
     /// Call all registered callbacks to say that the VM encountered an error.
-    pub fn notify_error(&self, cid: Cid, error_code: i32, message: &str) {
+    pub fn notify_error(&self, cid: Cid, error_code: ErrorCode, message: &str) {
         let callbacks = &*self.0.lock().unwrap();
         for callback in callbacks {
             if let Err(e) = callback.onError(cid as i32, error_code, message) {
@@ -1116,7 +1119,7 @@
         }
     }
 
-    fn notifyError(&self, error_code: i32, message: &str) -> binder::Result<()> {
+    fn notifyError(&self, error_code: ErrorCode, message: &str) -> binder::Result<()> {
         let cid = self.cid;
         if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
             info!("VM having CID {} encountered an error", cid);
diff --git a/virtualizationservice/src/atom.rs b/virtualizationservice/src/atom.rs
index 01f3e27..3b29d19 100644
--- a/virtualizationservice/src/atom.rs
+++ b/virtualizationservice/src/atom.rs
@@ -66,7 +66,6 @@
     let vm_identifier;
     let config_type;
     let num_cpus;
-    let cpu_affinity;
     let memory_mib;
     let apexes;
     match config {
@@ -74,7 +73,6 @@
             vm_identifier = &config.name;
             config_type = vm_creation_requested::ConfigType::VirtualMachineAppConfig;
             num_cpus = config.numCpus;
-            cpu_affinity = config.cpuAffinity.clone().unwrap_or_default();
             memory_mib = config.memoryMib;
 
             let vm_payload_config = get_vm_payload_config(config);
@@ -93,7 +91,6 @@
             vm_identifier = &config.name;
             config_type = vm_creation_requested::ConfigType::VirtualMachineRawConfig;
             num_cpus = config.numCpus;
-            cpu_affinity = config.cpuAffinity.clone().unwrap_or_default();
             memory_mib = config.memoryMib;
             apexes = String::new();
         }
@@ -108,7 +105,7 @@
         binder_exception_code,
         config_type,
         num_cpus,
-        cpu_affinity: &cpu_affinity,
+        cpu_affinity: "", // deprecated
         memory_mib,
         apexes: &apexes,
         // TODO(seungjaeyoo) Fill information about task_profile
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index c579be2..82a9e78 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -80,7 +80,6 @@
     pub protected: bool,
     pub memory_mib: Option<NonZeroU32>,
     pub cpus: Option<NonZeroU32>,
-    pub cpu_affinity: Option<String>,
     pub task_profiles: Vec<String>,
     pub console_fd: Option<File>,
     pub log_fd: Option<File>,
@@ -457,10 +456,6 @@
         command.arg("--cpus").arg(cpus.to_string());
     }
 
-    if let Some(cpu_affinity) = config.cpu_affinity {
-        command.arg("--cpu-affinity").arg(cpu_affinity);
-    }
-
     if !config.task_profiles.is_empty() {
         command.arg("--task-profiles").arg(config.task_profiles.join(","));
     }
diff --git a/vm/Android.bp b/vm/Android.bp
index eac640e..7b016d4 100644
--- a/vm/Android.bp
+++ b/vm/Android.bp
@@ -12,6 +12,7 @@
         "android.system.virtualizationservice-rust",
         "libanyhow",
         "libbinder_rs",
+        "libclap",
         "libenv_logger",
         "liblibc",
         "liblog_rust",
@@ -19,7 +20,6 @@
         "librustutils",
         "libserde_json",
         "libserde",
-        "libstructopt",
         "libvmconfig",
         "libvmclient",
         "libzip",
diff --git a/vm/src/main.rs b/vm/src/main.rs
index ee0e2e6..0845897 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -24,122 +24,104 @@
 };
 use anyhow::{Context, Error};
 use binder::ProcessState;
+use clap::Parser;
 use create_idsig::command_create_idsig;
 use create_partition::command_create_partition;
 use run::{command_run, command_run_app};
 use rustutils::system_properties;
 use std::path::{Path, PathBuf};
-use structopt::clap::AppSettings;
-use structopt::StructOpt;
 
 #[derive(Debug)]
 struct Idsigs(Vec<PathBuf>);
 
-#[derive(StructOpt)]
-#[structopt(no_version, global_settings = &[AppSettings::DisableVersion])]
+#[derive(Parser)]
 enum Opt {
     /// Run a virtual machine with a config in APK
     RunApp {
-        /// Name of VM
-        #[structopt(long)]
-        name: Option<String>,
-
         /// Path to VM Payload APK
-        #[structopt(parse(from_os_str))]
         apk: PathBuf,
 
         /// Path to idsig of the APK
-        #[structopt(parse(from_os_str))]
         idsig: PathBuf,
 
         /// Path to the instance image. Created if not exists.
-        #[structopt(parse(from_os_str))]
         instance: PathBuf,
 
         /// Path to VM config JSON within APK (e.g. assets/vm_config.json)
         config_path: String,
 
+        /// Name of VM
+        #[clap(long)]
+        name: Option<String>,
+
         /// Detach VM from the terminal and run in the background
-        #[structopt(short, long)]
+        #[clap(short, long)]
         daemonize: bool,
 
         /// Path to file for VM console output.
-        #[structopt(long)]
+        #[clap(long)]
         console: Option<PathBuf>,
 
         /// Path to file for VM log output.
-        #[structopt(long)]
+        #[clap(long)]
         log: Option<PathBuf>,
 
         /// Path to file where ramdump is recorded on kernel panic
-        #[structopt(long)]
+        #[clap(long)]
         ramdump: Option<PathBuf>,
 
         /// Debug level of the VM. Supported values: "none" (default), "app_only", and "full".
-        #[structopt(long, default_value = "none", parse(try_from_str=parse_debug_level))]
+        #[clap(long, default_value = "none", value_parser = parse_debug_level)]
         debug: DebugLevel,
 
         /// Run VM in protected mode.
-        #[structopt(short, long)]
+        #[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.
-        #[structopt(short, long)]
+        #[clap(short, long)]
         mem: Option<u32>,
 
         /// Number of vCPUs in the VM. If unspecified, defaults to 1.
-        #[structopt(long)]
+        #[clap(long)]
         cpus: Option<u32>,
 
-        /// Host CPUs where vCPUs are run on. If unspecified, vCPU runs on any host CPU.
-        #[structopt(long)]
-        cpu_affinity: Option<String>,
-
         /// Comma separated list of task profile names to apply to the VM
-        #[structopt(long)]
+        #[clap(long)]
         task_profiles: Vec<String>,
 
         /// Paths to extra idsig files.
-        #[structopt(long = "extra-idsig")]
+        #[clap(long = "extra-idsig")]
         extra_idsigs: Vec<PathBuf>,
     },
     /// Run a virtual machine
     Run {
-        /// Name of VM
-        #[structopt(long)]
-        name: Option<String>,
-
         /// Path to VM config JSON
-        #[structopt(parse(from_os_str))]
         config: PathBuf,
 
+        /// Name of VM
+        #[clap(long)]
+        name: Option<String>,
+
         /// Detach VM from the terminal and run in the background
-        #[structopt(short, long)]
+        #[clap(short, long)]
         daemonize: bool,
 
         /// Number of vCPUs in the VM. If unspecified, defaults to 1.
-        #[structopt(long)]
+        #[clap(long)]
         cpus: Option<u32>,
 
-        /// Host CPUs where vCPUs are run on. If unspecified, vCPU runs on any host CPU. The format
-        /// can be either a comma-separated list of CPUs or CPU ranges to run vCPUs on (e.g.
-        /// "0,1-3,5" to choose host CPUs 0, 1, 2, 3, and 5, or a colon-separated list of
-        /// assignments of vCPU-to-host-CPU assignments e.g. "0=0:1=1:2=2" to map vCPU 0 to host
-        /// CPU 0 and so on.
-        #[structopt(long)]
-        cpu_affinity: Option<String>,
-
         /// Comma separated list of task profile names to apply to the VM
-        #[structopt(long)]
+        #[clap(long)]
         task_profiles: Vec<String>,
 
         /// Path to file for VM console output.
-        #[structopt(long)]
+        #[clap(long)]
         console: Option<PathBuf>,
 
         /// Path to file for VM log output.
-        #[structopt(long)]
+        #[clap(long)]
         log: Option<PathBuf>,
     },
     /// Stop a virtual machine running in the background
@@ -154,23 +136,22 @@
     /// Create a new empty partition to be used as a writable partition for a VM
     CreatePartition {
         /// Path at which to create the image file
-        #[structopt(parse(from_os_str))]
         path: PathBuf,
 
         /// The desired size of the partition, in bytes.
         size: u64,
 
         /// Type of the partition
-        #[structopt(short="t", long="type", default_value="raw", parse(try_from_str=parse_partition_type))]
+        #[clap(short = 't', long = "type", default_value = "raw",
+               value_parser = parse_partition_type)]
         partition_type: PartitionType,
     },
     /// Creates or update the idsig file by digesting the input APK file.
     CreateIdsig {
         /// Path to VM Payload APK
-        #[structopt(parse(from_os_str))]
         apk: PathBuf,
+
         /// Path to idsig of the APK
-        #[structopt(parse(from_os_str))]
         path: PathBuf,
     },
 }
@@ -194,7 +175,7 @@
 
 fn main() -> Result<(), Error> {
     env_logger::init();
-    let opt = Opt::from_args();
+    let opt = Opt::parse();
 
     // We need to start the thread pool for Binder to work properly, especially link_to_death.
     ProcessState::start_thread_pool();
@@ -216,7 +197,6 @@
             protected,
             mem,
             cpus,
-            cpu_affinity,
             task_profiles,
             extra_idsigs,
         } => command_run_app(
@@ -234,11 +214,10 @@
             protected,
             mem,
             cpus,
-            cpu_affinity,
             task_profiles,
             &extra_idsigs,
         ),
-        Opt::Run { name, config, daemonize, cpus, cpu_affinity, task_profiles, console, log } => {
+        Opt::Run { name, config, daemonize, cpus, task_profiles, console, log } => {
             command_run(
                 name,
                 service.as_ref(),
@@ -248,7 +227,6 @@
                 log.as_deref(),
                 /* mem */ None,
                 cpus,
-                cpu_affinity,
                 task_profiles,
             )
         }
diff --git a/vm/src/run.rs b/vm/src/run.rs
index 05a9390..44e15f9 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -28,7 +28,7 @@
 use std::io::{self, BufRead, BufReader};
 use std::os::unix::io::{AsRawFd, FromRawFd};
 use std::path::{Path, PathBuf};
-use vmclient::VmInstance;
+use vmclient::{ErrorCode, VmInstance};
 use vmconfig::{open_parcel_file, VmConfig};
 use zip::ZipArchive;
 
@@ -49,7 +49,6 @@
     protected: bool,
     mem: Option<u32>,
     cpus: Option<u32>,
-    cpu_affinity: Option<String>,
     task_profiles: Vec<String>,
     extra_idsigs: &[PathBuf],
 ) -> Result<(), Error> {
@@ -102,7 +101,6 @@
         protectedVm: protected,
         memoryMib: mem.unwrap_or(0) as i32, // 0 means use the VM default
         numCpus: cpus.unwrap_or(1) as i32,
-        cpuAffinity: cpu_affinity,
         taskProfiles: task_profiles,
     });
     run(
@@ -127,7 +125,6 @@
     log_path: Option<&Path>,
     mem: Option<u32>,
     cpus: Option<u32>,
-    cpu_affinity: Option<String>,
     task_profiles: Vec<String>,
 ) -> Result<(), Error> {
     let config_file = File::open(config_path).context("Failed to open config file")?;
@@ -144,7 +141,6 @@
     } else {
         config.name = String::from("VmRun");
     }
-    config.cpuAffinity = cpu_affinity;
     config.taskProfiles = task_profiles;
     run(
         service,
@@ -223,7 +219,7 @@
         if let Some(path) = ramdump_path {
             save_ramdump_if_available(path, &vm)?;
         }
-        println!("{}", death_reason);
+        println!("VM ended: {:?}", death_reason);
     }
 
     Ok(())
@@ -253,15 +249,15 @@
     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);
-            loop {
+            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),
                 };
-            }
+            });
         }
     }
 
@@ -273,8 +269,8 @@
         eprintln!("payload finished with exit code {}", exit_code);
     }
 
-    fn on_error(&self, _cid: i32, error_code: i32, message: &str) {
-        eprintln!("VM encountered an error: code={}, message={}", error_code, message);
+    fn on_error(&self, _cid: i32, error_code: ErrorCode, message: &str) {
+        eprintln!("VM encountered an error: code={:?}, message={}", error_code, message);
     }
 }
 
diff --git a/vmbase/Android.bp b/vmbase/Android.bp
index ab87053..fb525b5 100644
--- a/vmbase/Android.bp
+++ b/vmbase/Android.bp
@@ -49,6 +49,13 @@
         "libarm-optimized-routines-mem",
         "libvmbase_entry",
     ],
+    installable: false,
+    enabled: false,
+    target: {
+        android_arm64: {
+            enabled: true,
+        },
+    },
 }
 
 rust_library_rlib {
diff --git a/vmbase/example/Android.bp b/vmbase/example/Android.bp
index e11bb2f..4e62090 100644
--- a/vmbase/example/Android.bp
+++ b/vmbase/example/Android.bp
@@ -18,8 +18,7 @@
 }
 
 cc_binary {
-    name: "vmbase_example_elf",
-    stem: "vmbase_example",
+    name: "vmbase_example",
     defaults: ["vmbase_elf_defaults"],
     srcs: [
         "idmap.S",
@@ -35,8 +34,9 @@
 }
 
 raw_binary {
-    name: "vmbase_example",
-    src: ":vmbase_example_elf",
+    name: "vmbase_example_bin",
+    stem: "vmbase_example.bin",
+    src: ":vmbase_example",
     enabled: false,
     target: {
         android_arm64: {
@@ -60,7 +60,7 @@
         "libvmclient",
     ],
     data: [
-        ":vmbase_example",
+        ":vmbase_example_bin",
     ],
     test_suites: ["general-tests"],
     enabled: false,
diff --git a/vmbase/example/src/exceptions.rs b/vmbase/example/src/exceptions.rs
index 61f7846..0e637ac 100644
--- a/vmbase/example/src/exceptions.rs
+++ b/vmbase/example/src/exceptions.rs
@@ -18,52 +18,52 @@
 use vmbase::{console::emergency_write_str, eprintln, power::reboot};
 
 #[no_mangle]
-extern "C" fn sync_exception_current() {
+extern "C" fn sync_exception_current(_elr: u64, _spsr: u64) {
     emergency_write_str("sync_exception_current\n");
     print_esr();
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn irq_current() {
+extern "C" fn irq_current(_elr: u64, _spsr: u64) {
     emergency_write_str("irq_current\n");
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn fiq_current() {
+extern "C" fn fiq_current(_elr: u64, _spsr: u64) {
     emergency_write_str("fiq_current\n");
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn serr_current() {
+extern "C" fn serr_current(_elr: u64, _spsr: u64) {
     emergency_write_str("serr_current\n");
     print_esr();
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn sync_lower() {
+extern "C" fn sync_lower(_elr: u64, _spsr: u64) {
     emergency_write_str("sync_lower\n");
     print_esr();
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn irq_lower() {
+extern "C" fn irq_lower(_elr: u64, _spsr: u64) {
     emergency_write_str("irq_lower\n");
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn fiq_lower() {
+extern "C" fn fiq_lower(_elr: u64, _spsr: u64) {
     emergency_write_str("fiq_lower\n");
     reboot();
 }
 
 #[no_mangle]
-extern "C" fn serr_lower() {
+extern "C" fn serr_lower(_elr: u64, _spsr: u64) {
     emergency_write_str("serr_lower\n");
     print_esr();
     reboot();
diff --git a/vmbase/example/tests/test.rs b/vmbase/example/tests/test.rs
index 58fffff..85e0213 100644
--- a/vmbase/example/tests/test.rs
+++ b/vmbase/example/tests/test.rs
@@ -59,7 +59,6 @@
         protectedVm: false,
         memoryMib: 300,
         numCpus: 1,
-        cpuAffinity: None,
         platformVersion: "~1.0".to_string(),
         taskProfiles: vec![],
     });
diff --git a/vmclient/Android.bp b/vmclient/Android.bp
index 213125e..88b0c9a 100644
--- a/vmclient/Android.bp
+++ b/vmclient/Android.bp
@@ -8,6 +8,7 @@
     srcs: ["src/lib.rs"],
     edition: "2021",
     rustlibs: [
+        "android.system.virtualizationcommon-rust",
         "android.system.virtualizationservice-rust",
         "libbinder_rs",
         "liblog_rust",
diff --git a/vmclient/src/death_reason.rs b/vmclient/src/death_reason.rs
index b976f6f..fbf2523 100644
--- a/vmclient/src/death_reason.rs
+++ b/vmclient/src/death_reason.rs
@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::fmt::{self, Debug, Display, Formatter};
 use android_system_virtualizationservice::{
         aidl::android::system::virtualizationservice::{
             DeathReason::DeathReason as AidlDeathReason}};
@@ -96,48 +95,3 @@
         }
     }
 }
-
-impl Display for DeathReason {
-    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
-        let s = match self {
-            Self::VirtualizationServiceDied => "VirtualizationService died.",
-            Self::InfrastructureError => "Error waiting for VM to finish.",
-            Self::Killed => "VM was killed.",
-            Self::Unknown => "VM died for an unknown reason.",
-            Self::Shutdown => "VM shutdown cleanly.",
-            Self::Error => "Error starting VM.",
-            Self::Reboot => "VM tried to reboot, possibly due to a kernel panic.",
-            Self::Crash => "VM crashed.",
-            Self::PvmFirmwarePublicKeyMismatch => {
-                "pVM firmware failed to verify the VM because the public key doesn't match."
-            }
-            Self::PvmFirmwareInstanceImageChanged => {
-                "pVM firmware failed to verify the VM because the instance image changed."
-            }
-            Self::BootloaderPublicKeyMismatch => {
-                "Bootloader failed to verify the VM because the public key doesn't match."
-            }
-            Self::BootloaderInstanceImageChanged => {
-                "Bootloader failed to verify the VM because the instance image changed."
-            }
-            Self::MicrodroidFailedToConnectToVirtualizationService => {
-                "The microdroid failed to connect to VirtualizationService's RPC server."
-            }
-            Self::MicrodroidPayloadHasChanged => "The payload for microdroid is changed.",
-            Self::MicrodroidPayloadVerificationFailed => {
-                "The microdroid failed to verify given payload APK."
-            }
-            Self::MicrodroidInvalidPayloadConfig => {
-                "The VM config for microdroid is invalid (e.g. missing tasks)."
-            }
-            Self::MicrodroidUnknownRuntimeError => {
-                "There was a runtime error while running microdroid manager."
-            }
-            Self::Hangup => "VM hangup.",
-            Self::Unrecognised(reason) => {
-                return write!(f, "Unrecognised death reason {:?}.", reason);
-            }
-        };
-        f.write_str(s)
-    }
-}
diff --git a/vmclient/src/error_code.rs b/vmclient/src/error_code.rs
new file mode 100644
index 0000000..a7c442f
--- /dev/null
+++ b/vmclient/src/error_code.rs
@@ -0,0 +1,47 @@
+// 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 android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::ErrorCode::ErrorCode as AidlErrorCode;
+
+/// Errors reported from within a VM.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum ErrorCode {
+    /// Error code for all other errors not listed below.
+    Unknown,
+
+    /// Error code indicating that the payload can't be verified due to various reasons (e.g invalid
+    /// merkle tree, invalid formats, etc).
+    PayloadVerificationFailed,
+
+    /// Error code indicating that the payload is verified, but has changed since the last boot.
+    PayloadChanged,
+
+    /// Error code indicating that the payload config is invalid.
+    PayloadConfigInvalid,
+
+    /// Payload sent a death reason which was not recognised by the client library.
+    Unrecognised(AidlErrorCode),
+}
+
+impl From<AidlErrorCode> for ErrorCode {
+    fn from(error_code: AidlErrorCode) -> Self {
+        match error_code {
+            AidlErrorCode::UNKNOWN => Self::Unknown,
+            AidlErrorCode::PAYLOAD_VERIFICATION_FAILED => Self::PayloadVerificationFailed,
+            AidlErrorCode::PAYLOAD_CHANGED => Self::PayloadChanged,
+            AidlErrorCode::PAYLOAD_CONFIG_INVALID => Self::PayloadConfigInvalid,
+            _ => Self::Unrecognised(error_code),
+        }
+    }
+}
diff --git a/vmclient/src/errors.rs b/vmclient/src/errors.rs
index 231f81f..a6dca91 100644
--- a/vmclient/src/errors.rs
+++ b/vmclient/src/errors.rs
@@ -22,7 +22,7 @@
     #[error("Timed out waiting for VM.")]
     TimedOut,
     /// The VM died before it was ready.
-    #[error("VM died. ({reason})")]
+    #[error("VM died. ({reason:?})")]
     Died {
         /// The reason why the VM died.
         reason: DeathReason,
diff --git a/vmclient/src/lib.rs b/vmclient/src/lib.rs
index 16b5d5a..e6f32b4 100644
--- a/vmclient/src/lib.rs
+++ b/vmclient/src/lib.rs
@@ -15,12 +15,15 @@
 //! Client library for VirtualizationService.
 
 mod death_reason;
+mod error_code;
 mod errors;
 mod sync;
 
 pub use crate::death_reason::DeathReason;
+pub use crate::error_code::ErrorCode;
 pub use crate::errors::VmWaitError;
 use crate::sync::Monitor;
+use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::ErrorCode::ErrorCode as AidlErrorCode;
 use android_system_virtualizationservice::{
     aidl::android::system::virtualizationservice::{
         DeathReason::DeathReason as AidlDeathReason,
@@ -83,7 +86,7 @@
 
     /// Called when an error has occurred in the VM. The `error_code` and `message` may give
     /// further details.
-    fn on_error(&self, cid: i32, error_code: i32, message: &str) {}
+    fn on_error(&self, cid: i32, error_code: ErrorCode, message: &str) {}
 
     /// Called when the VM has exited, all resources have been freed, and any logs have been
     /// written. `death_reason` gives an indication why the VM exited.
@@ -294,9 +297,10 @@
         Ok(())
     }
 
-    fn onError(&self, cid: i32, error_code: i32, message: &str) -> BinderResult<()> {
+    fn onError(&self, cid: i32, error_code: AidlErrorCode, message: &str) -> BinderResult<()> {
         self.state.notify_state(VirtualMachineState::FINISHED);
         if let Some(ref callback) = self.client_callback {
+            let error_code = error_code.into();
             callback.on_error(cid, error_code, message);
         }
         Ok(())