Merge "Set future_updatable: false for virt and compos APEXes"
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 72926ff..14452a3 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -19,7 +19,7 @@
       "name": "compos_key_tests"
     },
     {
-      "name": "composd_verify.test"
+      "name": "compos_verify.test"
     },
     {
       "name": "initrd_bootconfig.test"
diff --git a/apex/Android.bp b/apex/Android.bp
index 7d14da2..e39b459 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -3,15 +3,8 @@
 }
 
 microdroid_filesystem_images = [
-    "microdroid_boot",
-    "microdroid_bootconfig_debuggable",
-    "microdroid_bootconfig_normal",
-    "microdroid_init_boot",
     "microdroid_super",
-    "microdroid_uboot_env",
     "microdroid_vbmeta",
-    "microdroid_vbmeta_bootconfig",
-    "microdroid_vendor_boot",
 ]
 
 soong_config_module_type {
@@ -98,8 +91,6 @@
         "microdroid_initrd_debuggable",
         "microdroid_initrd_normal",
         "microdroid.json",
-        "microdroid_bootloader",
-        "microdroid_bootloader.avbpubkey",
         "microdroid_kernel",
     ],
     host_required: [
@@ -146,7 +137,10 @@
         },
     },
     required: [
+        // sign_virt_apex should be runnable from outside the source tree,
+        // therefore, any required tool should be listed in build/make/core/Makefile as well.
         "img2simg",
+        "initrd_bootconfig",
         "lpmake",
         "lpunpack",
         "simg2img",
@@ -167,6 +161,7 @@
         // sign_virt_apex
         "avbtool",
         "img2simg",
+        "initrd_bootconfig",
         "lpmake",
         "lpunpack",
         "sign_virt_apex",
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index 557c8aa..3f3600d 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -24,7 +24,7 @@
 
 sign_virt_apex uses external tools which are assumed to be available via PATH.
 - avbtool (--avbtool can override the tool)
-- lpmake, lpunpack, simg2img, img2simg
+- lpmake, lpunpack, simg2img, img2simg, initrd_bootconfig
 """
 import argparse
 import hashlib
@@ -102,6 +102,11 @@
     parser.add_argument(
         'input_dir',
         help='the directory having files to be packaged')
+    parser.add_argument(
+        '--do_not_update_bootconfigs',
+        action='store_true',
+        help='This will NOT update the vbmeta related bootconfigs while signing the apex.\
+            Used for testing only!!')
     args = parser.parse_args(argv)
     # preprocess --key_override into a map
     args.key_overrides = {}
@@ -200,23 +205,19 @@
     return info, descriptors
 
 
-# Look up a list of (key, value) with a key. Returns the value of the first matching pair.
+# Look up a list of (key, value) with a key. Returns the list of value(s) with the matching key.
+# The order of those values is maintained.
 def LookUp(pairs, key):
-    for k, v in pairs:
-        if key == k:
-            return v
-    return None
+    return [v for (k, v) in pairs if k == key]
 
 
-def AddHashFooter(args, key, image_path):
+def AddHashFooter(args, key, image_path, partition_name, additional_descriptors=None):
     if os.path.basename(image_path) in args.key_overrides:
         key = args.key_overrides[os.path.basename(image_path)]
-    info, descriptors = AvbInfo(args, image_path)
+    info, _ = AvbInfo(args, image_path)
     if info:
-        descriptor = LookUp(descriptors, 'Hash descriptor')
         image_size = ReadBytesSize(info['Image size'])
         algorithm = info['Algorithm']
-        partition_name = descriptor['Partition Name']
         partition_size = str(image_size)
 
         cmd = ['avbtool', 'add_hash_footer',
@@ -227,6 +228,9 @@
                '--image', image_path]
         if args.signing_args:
             cmd.extend(shlex.split(args.signing_args))
+        if additional_descriptors:
+            for image in additional_descriptors:
+                cmd.extend(['--include_descriptors_from_image', image])
         RunCommand(args, cmd)
 
 
@@ -235,7 +239,7 @@
         key = args.key_overrides[os.path.basename(image_path)]
     info, descriptors = AvbInfo(args, image_path)
     if info:
-        descriptor = LookUp(descriptors, 'Hashtree descriptor')
+        descriptor = LookUp(descriptors, 'Hashtree descriptor')[0]
         image_size = ReadBytesSize(info['Image size'])
         algorithm = info['Algorithm']
         partition_name = descriptor['Partition Name']
@@ -254,6 +258,82 @@
         RunCommand(args, cmd)
 
 
+def UpdateVbmetaBootconfig(args, initrds, vbmeta_img):
+    # Update the bootconfigs in ramdisk
+    def detach_bootconfigs(initrd_bc, initrd, bc):
+        cmd = ['initrd_bootconfig', 'detach', initrd_bc, initrd, bc]
+        RunCommand(args, cmd)
+
+    def attach_bootconfigs(initrd_bc, initrd, bc):
+        cmd = ['initrd_bootconfig', 'attach',
+               initrd, bc, '--output', initrd_bc]
+        RunCommand(args, cmd)
+
+    # Validate that avb version used while signing the apex is the same as used by build server
+    def validate_avb_version(bootconfigs):
+        cmd = ['avbtool', 'version']
+        stdout, _ = RunCommand(args, cmd)
+        avb_version_curr = stdout.split(" ")[1].strip()
+        avb_version_curr = avb_version_curr[0:avb_version_curr.rfind('.')]
+
+        avb_version_bc = re.search(
+            r"androidboot.vbmeta.avb_version = \"([^\"]*)\"", bootconfigs).group(1)
+        if avb_version_curr != avb_version_bc:
+            raise Exception(f'AVB version mismatch between current & one & \
+                used to build bootconfigs:{avb_version_curr}&{avb_version_bc}')
+
+    def calc_vbmeta_digest():
+        cmd = ['avbtool', 'calculate_vbmeta_digest', '--image',
+               vbmeta_img, '--hash_algorithm', 'sha256']
+        stdout, _ = RunCommand(args, cmd)
+        return stdout.strip()
+
+    def calc_vbmeta_size():
+        cmd = ['avbtool', 'info_image', '--image', vbmeta_img]
+        stdout, _ = RunCommand(args, cmd)
+        size = 0
+        for line in stdout.split("\n"):
+            line = line.split(":")
+            if line[0] in ['Header Block', 'Authentication Block', 'Auxiliary Block']:
+                size += int(line[1].strip()[0:-6])
+        return size
+
+    def update_vbmeta_digest(bootconfigs):
+        # Update androidboot.vbmeta.digest in bootconfigs
+        result = re.search(
+            r"androidboot.vbmeta.digest = \"[^\"]*\"", bootconfigs)
+        if not result:
+            raise ValueError("Failed to find androidboot.vbmeta.digest")
+
+        return bootconfigs.replace(result.group(),
+                                   f'androidboot.vbmeta.digest = "{calc_vbmeta_digest()}"')
+
+    def update_vbmeta_size(bootconfigs):
+        # Update androidboot.vbmeta.size in bootconfigs
+        result = re.search(r"androidboot.vbmeta.size = [0-9]+", bootconfigs)
+        if not result:
+            raise ValueError("Failed to find androidboot.vbmeta.size")
+        return bootconfigs.replace(result.group(),
+                                   f'androidboot.vbmeta.size = {calc_vbmeta_size()}')
+
+    with tempfile.TemporaryDirectory() as work_dir:
+        tmp_initrd = os.path.join(work_dir, 'initrd')
+        tmp_bc = os.path.join(work_dir, 'bc')
+
+        for initrd in initrds:
+            detach_bootconfigs(initrd, tmp_initrd, tmp_bc)
+            bc_file = open(tmp_bc, "rt", encoding="utf-8")
+            bc_data = bc_file.read()
+            validate_avb_version(bc_data)
+            bc_data = update_vbmeta_digest(bc_data)
+            bc_data = update_vbmeta_size(bc_data)
+            bc_file.close()
+            bc_file = open(tmp_bc, "wt", encoding="utf-8")
+            bc_file.write(bc_data)
+            bc_file.flush()
+            attach_bootconfigs(initrd, tmp_initrd, tmp_bc)
+
+
 def MakeVbmetaImage(args, key, vbmeta_img, images=None, chained_partitions=None):
     if os.path.basename(vbmeta_img) in args.key_overrides:
         key = args.key_overrides[os.path.basename(vbmeta_img)]
@@ -318,43 +398,21 @@
         RunCommand(args, cmd)
 
 
-def ReplaceBootloaderPubkey(args, key, bootloader, bootloader_pubkey):
-    if os.path.basename(bootloader) in args.key_overrides:
-        key = args.key_overrides[os.path.basename(bootloader)]
-    # read old pubkey before replacement
-    with open(bootloader_pubkey, 'rb') as f:
-        old_pubkey = f.read()
-
-    # replace bootloader pubkey (overwrite the old one with the new one)
-    ExtractAvbPubkey(args, key, bootloader_pubkey)
-
-    # read new pubkey
-    with open(bootloader_pubkey, 'rb') as f:
-        new_pubkey = f.read()
-
-    assert len(old_pubkey) == len(new_pubkey)
-
-    # replace pubkey embedded in bootloader
-    with open(bootloader, 'r+b') as bl_f:
-        pos = bl_f.read().find(old_pubkey)
-        assert pos != -1
-        bl_f.seek(pos)
-        bl_f.write(new_pubkey)
-
+def GenVbmetaImage(args, image, output, partition_name):
+    cmd = ['avbtool', 'add_hash_footer', '--dynamic_partition_size',
+           '--do_not_append_vbmeta_image',
+           '--partition_name', partition_name,
+           '--image', image,
+           '--output_vbmeta_image', output]
+    RunCommand(args, cmd)
 
 # dict of (key, file) for re-sign/verification. keys are un-versioned for readability.
 virt_apex_files = {
-    'bootloader.pubkey': 'etc/microdroid_bootloader.avbpubkey',
-    'bootloader': 'etc/fs/microdroid_bootloader',
-    'boot.img': 'etc/fs/microdroid_boot.img',
-    'vendor_boot.img': 'etc/fs/microdroid_vendor_boot.img',
-    'init_boot.img': 'etc/fs/microdroid_init_boot.img',
-    'super.img': 'etc/fs/microdroid_super.img',
+    'kernel': 'etc/fs/microdroid_kernel',
     'vbmeta.img': 'etc/fs/microdroid_vbmeta.img',
-    'vbmeta_bootconfig.img': 'etc/fs/microdroid_vbmeta_bootconfig.img',
-    'bootconfig.normal': 'etc/fs/microdroid_bootconfig.normal',
-    'bootconfig.debuggable': 'etc/fs/microdroid_bootconfig.debuggable',
-    'uboot_env.img': 'etc/fs/uboot_env.img'
+    'super.img': 'etc/fs/microdroid_super.img',
+    'initrd_normal.img': 'etc/microdroid_initrd_normal.img',
+    'initrd_debuggable.img': 'etc/microdroid_initrd_debuggable.img',
 }
 
 
@@ -371,17 +429,6 @@
     system_a_img = os.path.join(unpack_dir.name, 'system_a.img')
     vendor_a_img = os.path.join(unpack_dir.name, 'vendor_a.img')
 
-    # Key(pubkey) embedded in bootloader should match with the one used to make VBmeta below
-    # while it's okay to use different keys for other image files.
-    replace_f = Async(ReplaceBootloaderPubkey, args,
-                      key, files['bootloader'], files['bootloader.pubkey'])
-
-    # re-sign bootloader, boot.img, vendor_boot.img, and init_boot.img
-    Async(AddHashFooter, args, key, files['bootloader'], wait=[replace_f])
-    Async(AddHashFooter, args, key, files['boot.img'])
-    Async(AddHashFooter, args, key, files['vendor_boot.img'])
-    Async(AddHashFooter, args, key, files['init_boot.img'])
-
     # re-sign super.img
     # 1. unpack super.img
     # 2. resign system and vendor
@@ -390,27 +437,34 @@
     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])
+    Async(MakeSuperImage, args, partitions,
+          files['super.img'], wait=[system_a_f, vendor_a_f])
 
     # re-generate vbmeta from re-signed {system_a, vendor_a}.img
-    Async(MakeVbmetaImage, args, key, files['vbmeta.img'],
-          images=[system_a_img, vendor_a_img],
-          wait=[system_a_f, vendor_a_f])
+    vbmeta_f = Async(MakeVbmetaImage, args, key, files['vbmeta.img'],
+                     images=[system_a_img, vendor_a_img],
+                     wait=[system_a_f, vendor_a_f])
 
-    # Re-sign bootconfigs and the uboot_env with the same key
-    bootconfig_sign_key = key
-    Async(AddHashFooter, args, bootconfig_sign_key, files['bootconfig.normal'])
-    Async(AddHashFooter, args, bootconfig_sign_key, files['bootconfig.debuggable'])
-    Async(AddHashFooter, args, bootconfig_sign_key, files['uboot_env.img'])
+    vbmeta_bc_f = None
+    if not args.do_not_update_bootconfigs:
+        vbmeta_bc_f = Async(UpdateVbmetaBootconfig, args,
+                            [files['initrd_normal.img'],
+                                files['initrd_debuggable.img']], files['vbmeta.img'],
+                            wait=[vbmeta_f])
 
-    # Re-sign vbmeta_bootconfig with chained_partitions to "bootconfig" and
-    # "uboot_env". Note that, for now, `key` and `bootconfig_sign_key` are the
-    # same, but technically they can be different. Vbmeta records pubkeys which
-    # signed chained partitions.
-    Async(MakeVbmetaImage, args, key, files['vbmeta_bootconfig.img'], chained_partitions={
-        'bootconfig': bootconfig_sign_key,
-        'uboot_env': bootconfig_sign_key,
-    })
+    # Re-sign kernel. Note kernel's vbmeta contain addition descriptor from ramdisk(s)
+    initrd_normal_hashdesc = tempfile.NamedTemporaryFile(delete=False).name
+    initrd_debug_hashdesc = tempfile.NamedTemporaryFile(delete=False).name
+    initrd_n_f = Async(GenVbmetaImage, args, files['initrd_normal.img'],
+                       initrd_normal_hashdesc, "initrd_normal",
+                       wait=[vbmeta_bc_f] if vbmeta_bc_f is not None else [])
+    initrd_d_f = Async(GenVbmetaImage, args, files['initrd_debuggable.img'],
+                       initrd_debug_hashdesc, "initrd_debug",
+                       wait=[vbmeta_bc_f] if vbmeta_bc_f is not None else [])
+    Async(AddHashFooter, args, key, files['kernel'], partition_name="boot",
+          additional_descriptors=[
+              initrd_normal_hashdesc, initrd_debug_hashdesc],
+          wait=[initrd_n_f, initrd_d_f])
 
 
 def VerifyVirtApex(args):
@@ -430,27 +484,16 @@
             pubkey = f.read()
             pubkey_digest = hashlib.sha1(pubkey).hexdigest()
 
-    def contents(file):
-        with open(file, 'rb') as f:
-            return f.read()
-
-    def check_equals_pubkey(file):
-        assert contents(file) == pubkey, f'pubkey mismatch: {file}'
-
-    def check_contains_pubkey(file):
-        assert contents(file).find(pubkey) != -1, f'pubkey missing: {file}'
-
     def check_avb_pubkey(file):
         info, _ = AvbInfo(args, file)
         assert info is not None, f'no avbinfo: {file}'
         assert info['Public key (sha1)'] == pubkey_digest, f'pubkey mismatch: {file}'
 
     for f in files.values():
-        if f == files['bootloader.pubkey']:
-            Async(check_equals_pubkey, f)
-        elif f == files['bootloader']:
-            Async(check_contains_pubkey, f)
-        elif f == files['super.img']:
+        if f in (files['initrd_normal.img'], files['initrd_debuggable.img']):
+            # TODO(b/245277660): Verify that ramdisks contain the correct vbmeta digest
+            continue
+        if f == files['super.img']:
             Async(check_avb_pubkey, system_a_img)
             Async(check_avb_pubkey, vendor_a_img)
         else:
@@ -467,7 +510,7 @@
             SignVirtApex(args)
         # ensure all tasks are completed without exceptions
         AwaitAll(tasks)
-    except: # pylint: disable=bare-except
+    except:  # pylint: disable=bare-except
         traceback.print_exc()
         sys.exit(1)
 
diff --git a/authfs/src/fsverity/editor.rs b/authfs/src/fsverity/editor.rs
index 1e298be..4af6e80 100644
--- a/authfs/src/fsverity/editor.rs
+++ b/authfs/src/fsverity/editor.rs
@@ -204,7 +204,7 @@
             let mut merkle_tree = self.merkle_tree.write().unwrap();
 
             let offset_in_buf = (output_offset - offset) as usize;
-            let source = &buf[offset_in_buf as usize..offset_in_buf as usize + current_size];
+            let source = &buf[offset_in_buf..offset_in_buf + current_size];
             let output_chunk_index = (output_offset / CHUNK_SIZE) as usize;
             let offset_from_alignment = (output_offset % CHUNK_SIZE) as usize;
 
diff --git a/authfs/src/fusefs.rs b/authfs/src/fusefs.rs
index cfb53ca..544a94e 100644
--- a/authfs/src/fusefs.rs
+++ b/authfs/src/fusefs.rs
@@ -23,7 +23,7 @@
     SetattrValid, ZeroCopyReader, ZeroCopyWriter,
 };
 use fuse::sys::OpenOptions as FuseOpenOptions;
-use log::{debug, error, warn};
+use log::{error, trace, warn};
 use std::collections::{btree_map, BTreeMap};
 use std::convert::{TryFrom, TryInto};
 use std::ffi::{CStr, CString, OsStr};
@@ -1062,7 +1062,7 @@
             | SetattrValid::MTIME
             | SetattrValid::MTIME_NOW,
     ) {
-        debug!("Ignoring ctime/atime/mtime change as authfs does not maintain timestamp currently");
+        trace!("Ignoring ctime/atime/mtime change as authfs does not maintain timestamp currently");
     }
     Ok(())
 }
diff --git a/authfs/testdata/README.md b/authfs/testdata/README.md
index cf641a9..2df6753 100644
--- a/authfs/testdata/README.md
+++ b/authfs/testdata/README.md
@@ -2,7 +2,7 @@
 =================
 With a key pair, fs-verity signature can be generated by simply running
 `fsverity_metadata_generator` command line tool, which uses
-[fsverity-util](https://git.kernel.org/pub/scm/linux/kernel/git/ebiggers/fsverity-utils.git).
+[fsverity-utils](https://git.kernel.org/pub/scm/fs/fsverity/fsverity-utils.git)
 
 ```
 fsverity_metadata_generator --fsverity-path {fsverity_path} --key key.pem --key-format pem \
diff --git a/compos/tests/java/android/compos/test/ComposTestCase.java b/compos/tests/java/android/compos/test/ComposTestCase.java
index fe1c4f0..4e3d0a8 100644
--- a/compos/tests/java/android/compos/test/ComposTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposTestCase.java
@@ -27,6 +27,7 @@
 
 import com.android.microdroid.test.host.CommandRunner;
 import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
+import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.FileInputStreamSource;
 import com.android.tradefed.result.LogDataType;
@@ -94,10 +95,7 @@
     public void tearDown() throws Exception {
         killVmAndReconnectAdb();
 
-        archiveLogThenDelete(mTestLogs, getDevice(), COMPOS_APEXDATA_DIR + "/vm_console.log",
-                "vm_console.log-" + mTestName.getMethodName());
-        archiveLogThenDelete(mTestLogs, getDevice(), COMPOS_APEXDATA_DIR + "/vm.log",
-                "vm.log-" + mTestName.getMethodName());
+        archiveVmLogsThenDelete("teardown");
 
         CommandRunner android = new CommandRunner(getDevice());
 
@@ -115,6 +113,19 @@
         }
     }
 
+    private void archiveVmLogsThenDelete(String suffix) throws DeviceNotAvailableException {
+        archiveLogThenDelete(
+                mTestLogs,
+                getDevice(),
+                COMPOS_APEXDATA_DIR + "/vm_console.log",
+                "vm_console.log-" + suffix + "-" + mTestName.getMethodName());
+        archiveLogThenDelete(
+                mTestLogs,
+                getDevice(),
+                COMPOS_APEXDATA_DIR + "/vm.log",
+                "vm.log-" + suffix + "-" + mTestName.getMethodName());
+    }
+
     @Test
     public void testOdrefreshSpeed() throws Exception {
         getDevice().setProperty(SYSTEM_SERVER_COMPILER_FILTER_PROP_NAME, "speed");
@@ -159,6 +170,10 @@
         }
         killVmAndReconnectAdb();
 
+        // These logs are potentially useful, capture them before they are overwritten by
+        // compos_verify.
+        archiveVmLogsThenDelete("compile");
+
         // Expect the BCC extracted from the BCC to be well-formed.
         assertVmBccIsValid();
 
diff --git a/javalib/api/system-current.txt b/javalib/api/system-current.txt
index fe9943d..b455c85 100644
--- a/javalib/api/system-current.txt
+++ b/javalib/api/system-current.txt
@@ -58,8 +58,8 @@
   public final class VirtualMachineConfig {
     method @Nullable public String getApkPath();
     method @NonNull public int getDebugLevel();
-    method @IntRange(from=0) public long getEncryptedStorageKib();
-    method @IntRange(from=0) public int getMemoryMib();
+    method @IntRange(from=0) public long getEncryptedStorageBytes();
+    method @IntRange(from=0) public long getMemoryBytes();
     method @IntRange(from=1) public int getNumCpus();
     method @Nullable public String getPayloadBinaryName();
     method public boolean isCompatibleWith(@NonNull android.system.virtualmachine.VirtualMachineConfig);
@@ -75,8 +75,8 @@
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig build();
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setApkPath(@NonNull String);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setDebugLevel(int);
-    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setEncryptedStorageKib(@IntRange(from=1) long);
-    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setMemoryMib(@IntRange(from=1) int);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setEncryptedStorageBytes(@IntRange(from=1) long);
+    method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setMemoryBytes(@IntRange(from=1) long);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setNumCpus(@IntRange(from=1) int);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setPayloadBinaryName(@NonNull String);
     method @NonNull public android.system.virtualmachine.VirtualMachineConfig.Builder setProtectedVm(boolean);
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 7c7f4b5..ba7174e 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -476,7 +476,7 @@
                 try {
                     service.initializeWritablePartition(
                             ParcelFileDescriptor.open(vm.mEncryptedStoreFilePath, MODE_READ_WRITE),
-                            config.getEncryptedStorageKib() * 1024L,
+                            config.getEncryptedStorageBytes(),
                             PartitionType.ENCRYPTEDSTORE);
                 } catch (FileNotFoundException e) {
                     throw new VirtualMachineException("encrypted storage image missing", e);
@@ -621,15 +621,26 @@
             }
             virtualMachine = mVirtualMachine;
         }
+
+        int status;
         if (virtualMachine == null) {
-            return mVmRootPath.exists() ? STATUS_STOPPED : STATUS_DELETED;
+            status = STATUS_STOPPED;
         } else {
             try {
-                return stateToStatus(virtualMachine.getState());
+                status = stateToStatus(virtualMachine.getState());
             } catch (RemoteException e) {
                 throw e.rethrowAsRuntimeException();
             }
         }
+        if (status == STATUS_STOPPED && !mVmRootPath.exists()) {
+            // A VM can quite happily keep running if its backing files have been deleted.
+            // But once it stops, it's gone forever.
+            synchronized (mLock) {
+                dropVm();
+            }
+            return STATUS_DELETED;
+        }
+        return status;
     }
 
     private int stateToStatus(@VirtualMachineState int state) {
@@ -908,7 +919,7 @@
      * Stops this virtual machine. Stopping a virtual machine is like pulling the plug on a real
      * computer; the machine halts immediately. Software running on the virtual machine is not
      * notified of the event. Writes to {@linkplain
-     * VirtualMachineConfig.Builder#setEncryptedStorageKib encrypted storage} might not be
+     * VirtualMachineConfig.Builder#setEncryptedStorageBytes encrypted storage} might not be
      * persisted, and the instance might be left in an inconsistent state.
      *
      * <p>For a graceful shutdown, you could request the payload to call {@code exit()}, e.g. via a
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index b358f9e..cb9bad0 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -60,7 +60,7 @@
     private static final String[] EMPTY_STRING_ARRAY = {};
 
     // These define the schema of the config file persisted on disk.
-    private static final int VERSION = 4;
+    private static final int VERSION = 5;
     private static final String KEY_VERSION = "version";
     private static final String KEY_PACKAGENAME = "packageName";
     private static final String KEY_APKPATH = "apkPath";
@@ -68,9 +68,9 @@
     private static final String KEY_PAYLOADBINARYNAME = "payloadBinaryPath";
     private static final String KEY_DEBUGLEVEL = "debugLevel";
     private static final String KEY_PROTECTED_VM = "protectedVm";
-    private static final String KEY_MEMORY_MIB = "memoryMib";
+    private static final String KEY_MEMORY_BYTES = "memoryBytes";
     private static final String KEY_NUM_CPUS = "numCpus";
-    private static final String KEY_ENCRYPTED_STORAGE_KIB = "encryptedStorageKib";
+    private static final String KEY_ENCRYPTED_STORAGE_BYTES = "encryptedStorageBytes";
     private static final String KEY_VM_OUTPUT_CAPTURED = "vmOutputCaptured";
 
     /** @hide */
@@ -111,9 +111,10 @@
     private final boolean mProtectedVm;
 
     /**
-     * The amount of RAM to give the VM, in MiB. If this is 0 or negative the default will be used.
+     * The amount of RAM to give the VM, in bytes. If this is 0 or negative the default will be
+     * used.
      */
-    private final int mMemoryMib;
+    private final long mMemoryBytes;
 
     /**
      * Number of vCPUs in the VM. Defaults to 1 when not specified.
@@ -128,8 +129,8 @@
     /** Name of the payload binary file within the APK that will be executed within the VM. */
     @Nullable private final String mPayloadBinaryName;
 
-    /** The size of storage in KiB. 0 indicates that encryptedStorage is not required */
-    private final long mEncryptedStorageKib;
+    /** The size of storage in bytes. 0 indicates that encryptedStorage is not required */
+    private final long mEncryptedStorageBytes;
 
     /** Whether the app can read console and log output. */
     private final boolean mVmOutputCaptured;
@@ -141,9 +142,9 @@
             @Nullable String payloadBinaryName,
             @DebugLevel int debugLevel,
             boolean protectedVm,
-            int memoryMib,
+            long memoryBytes,
             int numCpus,
-            long encryptedStorageKib,
+            long encryptedStorageBytes,
             boolean vmOutputCaptured) {
         // This is only called from Builder.build(); the builder handles parameter validation.
         mPackageName = packageName;
@@ -152,9 +153,9 @@
         mPayloadBinaryName = payloadBinaryName;
         mDebugLevel = debugLevel;
         mProtectedVm = protectedVm;
-        mMemoryMib = memoryMib;
+        mMemoryBytes = memoryBytes;
         mNumCpus = numCpus;
-        mEncryptedStorageKib = encryptedStorageKib;
+        mEncryptedStorageBytes = encryptedStorageBytes;
         mVmOutputCaptured = vmOutputCaptured;
     }
 
@@ -220,14 +221,14 @@
         }
         builder.setDebugLevel(debugLevel);
         builder.setProtectedVm(b.getBoolean(KEY_PROTECTED_VM));
-        int memoryMib = b.getInt(KEY_MEMORY_MIB);
-        if (memoryMib != 0) {
-            builder.setMemoryMib(memoryMib);
+        long memoryBytes = b.getLong(KEY_MEMORY_BYTES);
+        if (memoryBytes != 0) {
+            builder.setMemoryBytes(memoryBytes);
         }
         builder.setNumCpus(b.getInt(KEY_NUM_CPUS));
-        long encryptedStorageKib = b.getLong(KEY_ENCRYPTED_STORAGE_KIB);
-        if (encryptedStorageKib != 0) {
-            builder.setEncryptedStorageKib(encryptedStorageKib);
+        long encryptedStorageBytes = b.getLong(KEY_ENCRYPTED_STORAGE_BYTES);
+        if (encryptedStorageBytes != 0) {
+            builder.setEncryptedStorageBytes(encryptedStorageBytes);
         }
         builder.setVmOutputCaptured(b.getBoolean(KEY_VM_OUTPUT_CAPTURED));
 
@@ -258,11 +259,11 @@
         b.putInt(KEY_DEBUGLEVEL, mDebugLevel);
         b.putBoolean(KEY_PROTECTED_VM, mProtectedVm);
         b.putInt(KEY_NUM_CPUS, mNumCpus);
-        if (mMemoryMib > 0) {
-            b.putInt(KEY_MEMORY_MIB, mMemoryMib);
+        if (mMemoryBytes > 0) {
+            b.putLong(KEY_MEMORY_BYTES, mMemoryBytes);
         }
-        if (mEncryptedStorageKib > 0) {
-            b.putLong(KEY_ENCRYPTED_STORAGE_KIB, mEncryptedStorageKib);
+        if (mEncryptedStorageBytes > 0) {
+            b.putLong(KEY_ENCRYPTED_STORAGE_BYTES, mEncryptedStorageBytes);
         }
         b.putBoolean(KEY_VM_OUTPUT_CAPTURED, mVmOutputCaptured);
         b.writeToStream(output);
@@ -335,8 +336,8 @@
      */
     @SystemApi
     @IntRange(from = 0)
-    public int getMemoryMib() {
-        return mMemoryMib;
+    public long getMemoryBytes() {
+        return mMemoryBytes;
     }
 
     /**
@@ -357,26 +358,26 @@
      */
     @SystemApi
     public boolean isEncryptedStorageEnabled() {
-        return mEncryptedStorageKib > 0;
+        return mEncryptedStorageBytes > 0;
     }
 
     /**
-     * Returns the size of encrypted storage (in KiB) available in the VM, or 0 if encrypted storage
-     * is not enabled
+     * Returns the size of encrypted storage (in bytes) available in the VM, or 0 if encrypted
+     * storage is not enabled
      *
      * @hide
      */
     @SystemApi
     @IntRange(from = 0)
-    public long getEncryptedStorageKib() {
-        return mEncryptedStorageKib;
+    public long getEncryptedStorageBytes() {
+        return mEncryptedStorageBytes;
     }
 
     /**
      * Returns whether the app can read the VM console or log output. If not, the VM output is
      * automatically forwarded to the host logcat.
      *
-     * @see #setVmOutputCaptured
+     * @see Builder#setVmOutputCaptured
      * @hide
      */
     @SystemApi
@@ -398,7 +399,7 @@
     public boolean isCompatibleWith(@NonNull VirtualMachineConfig other) {
         return this.mDebugLevel == other.mDebugLevel
                 && this.mProtectedVm == other.mProtectedVm
-                && this.mEncryptedStorageKib == other.mEncryptedStorageKib
+                && this.mEncryptedStorageBytes == other.mEncryptedStorageBytes
                 && this.mVmOutputCaptured == other.mVmOutputCaptured
                 && Objects.equals(this.mPayloadConfigPath, other.mPayloadConfigPath)
                 && Objects.equals(this.mPayloadBinaryName, other.mPayloadBinaryName)
@@ -453,13 +454,23 @@
                 break;
         }
         vsConfig.protectedVm = mProtectedVm;
-        vsConfig.memoryMib = mMemoryMib;
+        vsConfig.memoryMib = bytesToMebiBytes(mMemoryBytes);
         vsConfig.numCpus = mNumCpus;
         // Don't allow apps to set task profiles ... at least for now.
         vsConfig.taskProfiles = EMPTY_STRING_ARRAY;
         return vsConfig;
     }
 
+    private int bytesToMebiBytes(long mMemoryBytes) {
+        long oneMebi = 1024 * 1024;
+        // We can't express requests for more than 2 exabytes, but then they're not going to succeed
+        // anyway.
+        if (mMemoryBytes > (Integer.MAX_VALUE - 1) * oneMebi) {
+            return Integer.MAX_VALUE;
+        }
+        return (int) ((mMemoryBytes + oneMebi - 1) / oneMebi);
+    }
+
     /**
      * A builder used to create a {@link VirtualMachineConfig}.
      *
@@ -474,9 +485,9 @@
         @DebugLevel private int mDebugLevel = DEBUG_LEVEL_NONE;
         private boolean mProtectedVm;
         private boolean mProtectedVmSet;
-        private int mMemoryMib;
+        private long mMemoryBytes;
         private int mNumCpus = 1;
-        private long mEncryptedStorageKib;
+        private long mEncryptedStorageBytes;
         private boolean mVmOutputCaptured = false;
 
         /**
@@ -543,9 +554,9 @@
                     mPayloadBinaryName,
                     mDebugLevel,
                     mProtectedVm,
-                    mMemoryMib,
+                    mMemoryBytes,
                     mNumCpus,
-                    mEncryptedStorageKib,
+                    mEncryptedStorageBytes,
                     mVmOutputCaptured);
         }
 
@@ -660,18 +671,18 @@
         }
 
         /**
-         * Sets the amount of RAM to give the VM, in mebibytes. If not explicitly set then a default
+         * Sets the amount of RAM to give the VM, in bytes. If not explicitly set then a default
          * size will be used.
          *
          * @hide
          */
         @SystemApi
         @NonNull
-        public Builder setMemoryMib(@IntRange(from = 1) int memoryMib) {
-            if (memoryMib <= 0) {
+        public Builder setMemoryBytes(@IntRange(from = 1) long memoryBytes) {
+            if (memoryBytes <= 0) {
                 throw new IllegalArgumentException("Memory size must be positive");
             }
-            mMemoryMib = memoryMib;
+            mMemoryBytes = memoryBytes;
             return this;
         }
 
@@ -699,8 +710,8 @@
         }
 
         /**
-         * Sets the size (in KiB) of encrypted storage available to the VM. If not set, no encrypted
-         * storage is provided.
+         * Sets the size (in bytes) of encrypted storage available to the VM. If not set, no
+         * encrypted storage is provided.
          *
          * <p>The storage is encrypted with a key deterministically derived from the VM identity
          *
@@ -716,11 +727,11 @@
          */
         @SystemApi
         @NonNull
-        public Builder setEncryptedStorageKib(@IntRange(from = 1) long encryptedStorageKib) {
-            if (encryptedStorageKib <= 0) {
+        public Builder setEncryptedStorageBytes(@IntRange(from = 1) long encryptedStorageBytes) {
+            if (encryptedStorageBytes <= 0) {
                 throw new IllegalArgumentException("Encrypted Storage size must be positive");
             }
-            mEncryptedStorageKib = encryptedStorageKib;
+            mEncryptedStorageBytes = encryptedStorageBytes;
             return this;
         }
 
diff --git a/libs/apexutil/src/lib.rs b/libs/apexutil/src/lib.rs
index 698d93a..8a934e2 100644
--- a/libs/apexutil/src/lib.rs
+++ b/libs/apexutil/src/lib.rs
@@ -144,7 +144,7 @@
         let res = verify("tests/data/test.apex").unwrap();
         // The expected hex is generated when we ran the method the first time.
         assert_eq!(
-            hex::encode(&res.root_digest),
+            hex::encode(res.root_digest),
             "fe11ab17da0a3a738b54bdc3a13f6139cbdf91ec32f001f8d4bbbf8938e04e39"
         );
     }
@@ -153,7 +153,7 @@
     fn payload_vbmeta_has_valid_image_hash() {
         let result = get_payload_vbmeta_image_hash("tests/data/test.apex").unwrap();
         assert_eq!(
-            hex::encode(&result),
+            hex::encode(result),
             "296e32a76544de9da01713e471403ab4667705ad527bb4f1fac0cf61e7ce122d"
         );
     }
diff --git a/libs/apkverify/src/algorithms.rs b/libs/apkverify/src/algorithms.rs
index 6315606..442b47c 100644
--- a/libs/apkverify/src/algorithms.rs
+++ b/libs/apkverify/src/algorithms.rs
@@ -34,11 +34,14 @@
 /// [SignatureAlgorithm.java]: (tools/apksig/src/main/java/com/android/apksig/internal/apk/SignatureAlgorithm.java)
 ///
 /// Some of the algorithms are not implemented. See b/197052981.
-#[derive(Serialize, Deserialize, Clone, Copy, Debug, Eq, PartialEq, FromPrimitive, ToPrimitive)]
+#[derive(
+    Serialize, Deserialize, Clone, Copy, Debug, Default, Eq, PartialEq, FromPrimitive, ToPrimitive,
+)]
 #[repr(u32)]
 pub enum SignatureAlgorithmID {
     /// RSASSA-PSS with SHA2-256 digest, SHA2-256 MGF1, 32 bytes of salt, trailer: 0xbc, content
     /// digested using SHA2-256 in 1 MB chunks.
+    #[default]
     RsaPssWithSha256 = 0x0101,
 
     /// RSASSA-PSS with SHA2-512 digest, SHA2-512 MGF1, 64 bytes of salt, trailer: 0xbc, content
@@ -77,12 +80,6 @@
     VerityDsaWithSha256 = 0x0425,
 }
 
-impl Default for SignatureAlgorithmID {
-    fn default() -> Self {
-        SignatureAlgorithmID::RsaPssWithSha256
-    }
-}
-
 impl ReadFromBytes for Option<SignatureAlgorithmID> {
     fn read_from_bytes(buf: &mut Bytes) -> Result<Self> {
         Ok(SignatureAlgorithmID::from_u32(buf.get_u32_le()))
diff --git a/libs/apkverify/src/ziputil.rs b/libs/apkverify/src/ziputil.rs
index eb2826a..cc8bc58 100644
--- a/libs/apkverify/src/ziputil.rs
+++ b/libs/apkverify/src/ziputil.rs
@@ -46,7 +46,7 @@
     reader = archive.into_inner();
     // the current position should point EOCD offset
     let eocd_offset = reader.seek(SeekFrom::Current(0))? as u32;
-    let mut eocd = vec![0u8; eocd_size as usize];
+    let mut eocd = vec![0u8; eocd_size];
     reader.read_exact(&mut eocd)?;
     ensure!(
         (&eocd[0..]).get_u32_le() == EOCD_SIGNATURE,
diff --git a/libs/avb/Android.bp b/libs/avb/Android.bp
index 436f672..c173f1c 100644
--- a/libs/avb/Android.bp
+++ b/libs/avb/Android.bp
@@ -14,12 +14,13 @@
         "--size_t-is-usize",
         "--default-enum-style rust",
         "--allowlist-function=.*",
+        "--allowlist-var=AVB.*",
         "--use-core",
         "--raw-line=#![no_std]",
         "--ctypes-prefix=core::ffi",
     ],
     static_libs: [
-        "libavb",
+        "libavb_non_debug",
     ],
     shared_libs: [
         "libcrypto",
diff --git a/libs/devicemapper/src/lib.rs b/libs/devicemapper/src/lib.rs
index 9069eb2..4cf4e99 100644
--- a/libs/devicemapper/src/lib.rs
+++ b/libs/devicemapper/src/lib.rs
@@ -226,7 +226,7 @@
 
     let context = Context::new(0);
     let now = SystemTime::now().duration_since(UNIX_EPOCH)?;
-    let ts = Timestamp::from_unix(&context, now.as_secs(), now.subsec_nanos());
+    let ts = Timestamp::from_unix(context, now.as_secs(), now.subsec_nanos());
     let uuid = Uuid::new_v1(ts, node_id)?;
     Ok(String::from(uuid.to_hyphenated().encode_lower(&mut Uuid::encode_buffer())))
 }
diff --git a/libs/devicemapper/src/loopdevice.rs b/libs/devicemapper/src/loopdevice.rs
index 5533e17..16b5b65 100644
--- a/libs/devicemapper/src/loopdevice.rs
+++ b/libs/devicemapper/src/loopdevice.rs
@@ -172,13 +172,13 @@
 
     fn is_direct_io(dev: &Path) -> bool {
         let dio = Path::new("/sys/block").join(dev.file_name().unwrap()).join("loop/dio");
-        "1" == fs::read_to_string(&dio).unwrap().trim()
+        "1" == fs::read_to_string(dio).unwrap().trim()
     }
 
     // kernel exposes /sys/block/loop*/ro which gives the read-only value
     fn is_direct_io_writable(dev: &Path) -> bool {
         let ro = Path::new("/sys/block").join(dev.file_name().unwrap()).join("ro");
-        "0" == fs::read_to_string(&ro).unwrap().trim()
+        "0" == fs::read_to_string(ro).unwrap().trim()
     }
 
     #[test]
diff --git a/libs/dice/src/lib.rs b/libs/dice/src/lib.rs
index 9bbacc6..5332092 100644
--- a/libs/dice/src/lib.rs
+++ b/libs/dice/src/lib.rs
@@ -23,15 +23,12 @@
 use core::ptr;
 use core::result;
 
+pub use open_dice_cbor_bindgen::DiceMode;
+
 use open_dice_cbor_bindgen::DiceConfigType_kDiceConfigTypeDescriptor as DICE_CONFIG_TYPE_DESCRIPTOR;
 use open_dice_cbor_bindgen::DiceConfigType_kDiceConfigTypeInline as DICE_CONFIG_TYPE_INLINE;
 use open_dice_cbor_bindgen::DiceHash;
 use open_dice_cbor_bindgen::DiceInputValues;
-use open_dice_cbor_bindgen::DiceMode;
-use open_dice_cbor_bindgen::DiceMode_kDiceModeDebug as DICE_MODE_DEBUG;
-use open_dice_cbor_bindgen::DiceMode_kDiceModeMaintenance as DICE_MODE_MAINTENANCE;
-use open_dice_cbor_bindgen::DiceMode_kDiceModeNormal as DICE_MODE_NORMAL;
-use open_dice_cbor_bindgen::DiceMode_kDiceModeNotInitialized as DICE_MODE_NOT_INITIALIZED;
 use open_dice_cbor_bindgen::DiceResult;
 use open_dice_cbor_bindgen::DiceResult_kDiceResultBufferTooSmall as DICE_RESULT_BUFFER_TOO_SMALL;
 use open_dice_cbor_bindgen::DiceResult_kDiceResultInvalidInput as DICE_RESULT_INVALID_INPUT;
@@ -90,26 +87,6 @@
     }
 }
 
-/// DICE mode values.
-#[derive(Clone, Copy, Debug)]
-pub enum Mode {
-    /// At least one security mechanism has not been configured. Also acts as a catch-all.
-    /// Invalid mode values should be treated like this mode.
-    NotInitialized = DICE_MODE_NOT_INITIALIZED as _,
-    /// Indicates the device is operating normally under secure configuration.
-    Normal = DICE_MODE_NORMAL as _,
-    /// Indicates at least one criteria for Normal mode is not met.
-    Debug = DICE_MODE_DEBUG as _,
-    /// Indicates a recovery or maintenance mode of some kind.
-    Maintenance = DICE_MODE_MAINTENANCE as _,
-}
-
-impl From<Mode> for DiceMode {
-    fn from(mode: Mode) -> Self {
-        mode as Self
-    }
-}
-
 /// DICE configuration input type.
 #[derive(Debug)]
 pub enum ConfigType<'a> {
@@ -132,7 +109,7 @@
         config: &ConfigType,
         auth_hash: Option<&Hash>,
         auth_descriptor: Option<&[u8]>,
-        mode: Mode,
+        mode: DiceMode,
         hidden: Option<&Hidden>,
     ) -> Self {
         const ZEROED_INLINE_CONFIG: InlineConfig = [0; INLINE_CONFIG_SIZE];
@@ -157,7 +134,7 @@
             authority_hash: auth_hash.map_or([0; mem::size_of::<Hash>()], |h| *h),
             authority_descriptor,
             authority_descriptor_size,
-            mode: mode.into(),
+            mode,
             hidden: hidden.map_or([0; mem::size_of::<Hidden>()], |h| *h),
         })
     }
diff --git a/libs/fdtpci/src/lib.rs b/libs/fdtpci/src/lib.rs
index 1ddda9f..e32e16d 100644
--- a/libs/fdtpci/src/lib.rs
+++ b/libs/fdtpci/src/lib.rs
@@ -91,7 +91,7 @@
 }
 
 /// Information about the PCI bus parsed from the device tree.
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct PciInfo {
     /// The MMIO range used by the memory-mapped PCI CAM.
     pub cam_range: Range<usize>,
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
index 8fd1879..29d7abe 100644
--- a/libs/libfdt/src/lib.rs
+++ b/libs/libfdt/src/lib.rs
@@ -388,6 +388,23 @@
         fdt_err_expect_zero(ret)
     }
 
+    /// Create or change a property name-value pair to the given node.
+    pub fn setprop(&mut self, name: &CStr, value: &[u8]) -> Result<()> {
+        // SAFETY - New value size is constrained to the DT totalsize
+        //          (validated by underlying libfdt).
+        let ret = unsafe {
+            libfdt_bindgen::fdt_setprop(
+                self.fdt.as_mut_ptr(),
+                self.offset,
+                name.as_ptr(),
+                value.as_ptr().cast::<c_void>(),
+                value.len().try_into().map_err(|_| FdtError::BadValue)?,
+            )
+        };
+
+        fdt_err_expect_zero(ret)
+    }
+
     /// Get reference to the containing device tree.
     pub fn fdt(&mut self) -> &mut Fdt {
         self.fdt
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 529686d..9264692 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -240,37 +240,6 @@
     ],
 }
 
-bootimg {
-    name: "microdroid_boot",
-    // We don't have kernel for arm and x86. But Soong demands one when it builds for
-    // arm or x86 target. Satisfy that by providing an empty file as the kernel.
-    kernel_prebuilt: "empty_kernel",
-    arch: {
-        arm64: {
-            kernel_prebuilt: ":microdroid_kernel_prebuilts-6.1-arm64",
-        },
-        x86_64: {
-            kernel_prebuilt: ":microdroid_kernel_prebuilts-6.1-x86_64",
-        },
-    },
-
-    dtb_prebuilt: "dummy_dtb.img",
-    header_version: "4",
-    partition_name: "boot",
-    use_avb: true,
-    avb_private_key: ":microdroid_sign_key",
-}
-
-bootimg {
-    name: "microdroid_init_boot",
-    ramdisk_module: "microdroid_ramdisk",
-    kernel_prebuilt: "empty_kernel",
-    header_version: "4",
-    partition_name: "init_boot",
-    use_avb: true,
-    avb_private_key: ":microdroid_sign_key",
-}
-
 android_filesystem {
     name: "microdroid_ramdisk",
     deps: [
@@ -289,25 +258,6 @@
     type: "compressed_cpio",
 }
 
-bootimg {
-    name: "microdroid_vendor_boot",
-    ramdisk_module: "microdroid_vendor_ramdisk",
-    dtb_prebuilt: "dummy_dtb.img",
-    header_version: "4",
-    vendor_boot: true,
-    arch: {
-        arm64: {
-            bootconfig: ":microdroid_bootconfig_arm64_gen",
-        },
-        x86_64: {
-            bootconfig: ":microdroid_bootconfig_x86_64_gen",
-        },
-    },
-    partition_name: "vendor_boot",
-    use_avb: true,
-    avb_private_key: ":microdroid_sign_key",
-}
-
 android_filesystem {
     name: "microdroid_vendor_ramdisk",
     deps: [
@@ -347,43 +297,6 @@
     cmd: "cat $(in) > $(out)",
 }
 
-vbmeta {
-    name: "microdroid_vbmeta_bootconfig",
-    partition_name: "vbmeta",
-    private_key: ":microdroid_sign_key",
-    chained_partitions: [
-        {
-            name: "bootconfig",
-            private_key: ":microdroid_sign_key",
-        },
-        {
-            name: "uboot_env",
-            private_key: ":microdroid_sign_key",
-        },
-    ],
-}
-
-// python -c "import hashlib; print(hashlib.sha256(b'bootconfig').hexdigest())"
-bootconfig_salt = "e158851fbebb402e1f18ea9372ea2f76b4dea23eceb5c4b92e5b27ade8537f5b"
-
-avb_add_hash_footer {
-    name: "microdroid_bootconfig_normal",
-    src: "bootconfig.normal",
-    filename: "microdroid_bootconfig.normal",
-    partition_name: "bootconfig",
-    private_key: ":microdroid_sign_key",
-    salt: bootconfig_salt,
-}
-
-avb_add_hash_footer {
-    name: "microdroid_bootconfig_debuggable",
-    src: "bootconfig.debuggable",
-    filename: "microdroid_bootconfig.debuggable",
-    partition_name: "bootconfig",
-    private_key: ":microdroid_sign_key",
-    salt: bootconfig_salt,
-}
-
 prebuilt_etc {
     name: "microdroid_fstab",
     src: "fstab.microdroid",
@@ -391,89 +304,9 @@
     installable: false,
 }
 
-prebuilt_etc {
-    name: "microdroid_bootloader",
-    src: ":microdroid_bootloader_signed",
-    arch: {
-        x86_64: {
-            // For unknown reason, the signed bootloader doesn't work on x86_64. Until the problem
-            // is fixed, let's use the unsigned bootloader for the architecture.
-            // TODO(b/185115783): remove this
-            src: ":microdroid_bootloader_pubkey_replaced",
-        },
-    },
-    relative_install_path: "fs",
-    filename: "microdroid_bootloader",
-}
-
 // python -c "import hashlib; print(hashlib.sha256(b'bootloader').hexdigest())"
 bootloader_salt = "3b4a12881d11f33cff968a24d7c53723a8232cde9a8d91e29fdbd6a95ae6adf0"
 
-avb_add_hash_footer {
-    name: "microdroid_bootloader_signed",
-    src: ":microdroid_bootloader_pubkey_replaced",
-    filename: "microdroid_bootloader",
-    partition_name: "bootloader",
-    private_key: ":microdroid_sign_key",
-    salt: bootloader_salt,
-}
-
-// Replace avbpubkey of prebuilt bootloader with the avbpubkey of the signing key
-genrule {
-    name: "microdroid_bootloader_pubkey_replaced",
-    tools: ["replace_bytes"],
-    srcs: [
-        ":microdroid_crosvm_bootloader", // input (bootloader)
-        ":microdroid_crosvm_bootloader.avbpubkey", // old bytes (old pubkey)
-        ":microdroid_bootloader_avbpubkey_gen", // new bytes (new pubkey)
-    ],
-    out: ["bootloader-pubkey-replaced"],
-    // 1. Copy the input to the output (replace_bytes modifies the file in-place)
-    // 2. Replace embedded pubkey with new one.
-    cmd: "cp $(location :microdroid_crosvm_bootloader) $(out) && " +
-        "$(location replace_bytes) $(out) " +
-        "$(location :microdroid_crosvm_bootloader.avbpubkey) " +
-        "$(location :microdroid_bootloader_avbpubkey_gen)",
-}
-
-// Apex keeps a copy of avbpubkey embedded in bootloader so that embedded avbpubkey can be replaced
-// while re-signing bootloader.
-prebuilt_etc {
-    name: "microdroid_bootloader.avbpubkey",
-    src: ":microdroid_bootloader_avbpubkey_gen",
-}
-
-// Generate avbpukey from the signing key
-genrule {
-    name: "microdroid_bootloader_avbpubkey_gen",
-    tools: ["avbtool"],
-    srcs: [":microdroid_sign_key"],
-    out: ["bootloader.pubkey"],
-    cmd: "$(location avbtool) extract_public_key " +
-        "--key $(location :microdroid_sign_key) " +
-        "--output $(out)",
-}
-
-// python -c "import hashlib; print(hashlib.sha256(b'uboot_env').hexdigest())"
-uboot_env_salt = "cbf2d76827ece5ca8d176a40c94ac6355edcf6511b4b887364a8c0e05850df10"
-
-avb_add_hash_footer {
-    name: "microdroid_uboot_env",
-    src: ":microdroid_uboot_env_gen",
-    filename: "uboot_env.img",
-    partition_name: "uboot_env",
-    private_key: ":microdroid_sign_key",
-    salt: uboot_env_salt,
-}
-
-genrule {
-    name: "microdroid_uboot_env_gen",
-    tools: ["mkenvimage_slim"],
-    srcs: ["uboot-env.txt"],
-    out: ["output.img"],
-    cmd: "$(location mkenvimage_slim) -output_path $(out) -input_path $(location uboot-env.txt)",
-}
-
 // Note that keys can be different for filesystem images even though we're using the same key
 // for microdroid. However, the key signing VBmeta should match with the pubkey embedded in
 // bootloader.
diff --git a/microdroid/init.rc b/microdroid/init.rc
index bc42791..ce0cab4 100644
--- a/microdroid/init.rc
+++ b/microdroid/init.rc
@@ -149,6 +149,10 @@
     # Mark boot completed. This will notify microdroid_manager to run payload.
     setprop dev.bootcomplete 1
 
+on property:tombstone_transmit.start=1
+    mkdir /data/tombstones 0771 system system
+    start tombstone_transmit
+
 service tombstone_transmit /system/bin/tombstone_transmit.microdroid -cid 2 -port 2000 -remove_tombstones_after_transmitting
     user system
     group system
@@ -175,4 +179,3 @@
     group shell log readproc
     seclabel u:r:shell:s0
     setenv HOSTNAME console
-
diff --git a/microdroid/initrd/Android.bp b/microdroid/initrd/Android.bp
index 7a95ce6..ff6314b 100644
--- a/microdroid/initrd/Android.bp
+++ b/microdroid/initrd/Android.bp
@@ -68,7 +68,7 @@
         ":microdroid_bootconfig_debuggable_src",
     ] + bootconfigs_arm64,
     out: ["microdroid_initrd_debuggable_arm64"],
-    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
 genrule {
@@ -79,7 +79,7 @@
         ":microdroid_bootconfig_debuggable_src",
     ] + bootconfigs_x86_64,
     out: ["microdroid_initrd_debuggable_x86_64"],
-    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
 genrule {
@@ -90,7 +90,7 @@
         ":microdroid_bootconfig_normal_src",
     ] + bootconfigs_arm64,
     out: ["microdroid_initrd_normal_arm64"],
-    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
 genrule {
@@ -101,7 +101,7 @@
         ":microdroid_bootconfig_normal_src",
     ] + bootconfigs_x86_64,
     out: ["microdroid_initrd_normal_x86_64"],
-    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
 prebuilt_etc {
diff --git a/microdroid/initrd/src/main.rs b/microdroid/initrd/src/main.rs
index 74e4ba6..c5515af 100644
--- a/microdroid/initrd/src/main.rs
+++ b/microdroid/initrd/src/main.rs
@@ -12,36 +12,95 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! Append bootconfig to initrd image
-use anyhow::Result;
+//! Attach/Detach bootconfigs to initrd image
+use anyhow::{bail, Result};
 use clap::Parser;
+use std::cmp::min;
 use std::fs::File;
-use std::io::{Read, Write};
+use std::io::{Read, Seek, SeekFrom, Write};
+use std::mem::size_of;
 use std::path::PathBuf;
 
 const FOOTER_ALIGNMENT: usize = 4;
 const ZEROS: [u8; 4] = [0u8; 4_usize];
+const BOOTCONFIG_MAGIC: &str = "#BOOTCONFIG\n";
+// Footer includes [size(le32)][checksum(le32)][#BOOTCONFIG\n] at the end of bootconfigs.
+const INITRD_FOOTER_LEN: usize = 2 * std::mem::size_of::<u32>() + BOOTCONFIG_MAGIC.len();
 
 #[derive(Parser, Debug)]
-struct Args {
-    /// Initrd (without bootconfig)
-    initrd: PathBuf,
-    /// Bootconfig
-    bootconfigs: Vec<PathBuf>,
-    /// Output
-    #[clap(long = "output")]
-    output: PathBuf,
+enum Opt {
+    /// Append bootconfig(s) to initrd image
+    Attach {
+        /// Initrd (without bootconfigs) <- Input
+        initrd: PathBuf,
+        /// Bootconfigs <- Input
+        bootconfigs: Vec<PathBuf>,
+        /// Initrd (with bootconfigs) <- Output
+        #[clap(long = "output")]
+        output: PathBuf,
+    },
+
+    /// Detach the initrd & bootconfigs - this is required for cases when we update
+    /// bootconfigs in sign_virt_apex
+    Detach {
+        /// Initrd (with bootconfigs) <- Input
+        initrd_with_bootconfigs: PathBuf,
+        /// Initrd (without bootconfigs) <- Output
+        initrd: PathBuf,
+        /// Bootconfigs <- Output
+        bootconfigs: PathBuf,
+    },
 }
 
 fn get_checksum(file_path: &PathBuf) -> Result<u32> {
     File::open(file_path)?.bytes().map(|x| Ok(x? as u32)).sum()
 }
 
+// Copy n bytes of file_in to file_out. Note: copying starts from the current cursors of files.
+// On successful return, the files' cursors would have moved forward by k bytes.
+fn copyfile2file(file_in: &mut File, file_out: &mut File, n: usize) -> Result<()> {
+    let mut buf = vec![0; 1024];
+    let mut copied: usize = 0;
+    while copied < n {
+        let k = min(n - copied, buf.len());
+        file_in.read_exact(&mut buf[..k])?;
+        file_out.write_all(&buf[..k])?;
+        copied += k;
+    }
+    Ok(())
+}
+
+// Note: attaching & then detaching bootconfigs can lead to extra padding in bootconfigs
+fn detach_bootconfig(initrd_bc: PathBuf, initrd: PathBuf, bootconfig: PathBuf) -> Result<()> {
+    let mut initrd_bc = File::open(initrd_bc)?;
+    let mut bootconfig = File::create(bootconfig)?;
+    let mut initrd = File::create(initrd)?;
+    let initrd_bc_size: usize = initrd_bc.metadata()?.len().try_into()?;
+
+    initrd_bc.seek(SeekFrom::End(-(BOOTCONFIG_MAGIC.len() as i64)))?;
+    let mut magic_buf = [0; BOOTCONFIG_MAGIC.len()];
+    initrd_bc.read_exact(&mut magic_buf)?;
+    if magic_buf != BOOTCONFIG_MAGIC.as_bytes() {
+        bail!("BOOTCONFIG_MAGIC not found in initrd. Bootconfigs might not be attached correctly");
+    }
+    let mut size_buf = [0; size_of::<u32>()];
+    initrd_bc.seek(SeekFrom::End(-(INITRD_FOOTER_LEN as i64)))?;
+    initrd_bc.read_exact(&mut size_buf)?;
+    let bc_size: usize = u32::from_le_bytes(size_buf) as usize;
+
+    let initrd_size: usize = initrd_bc_size - bc_size - INITRD_FOOTER_LEN;
+
+    initrd_bc.seek(SeekFrom::Start(0))?;
+    copyfile2file(&mut initrd_bc, &mut initrd, initrd_size)?;
+    copyfile2file(&mut initrd_bc, &mut bootconfig, bc_size)?;
+    Ok(())
+}
+
 // 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 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;
@@ -59,14 +118,21 @@
     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.write_all(BOOTCONFIG_MAGIC.as_bytes())?;
     output_file.flush()?;
     Ok(())
 }
 
 fn try_main() -> Result<()> {
-    let args = Args::parse();
-    attach_bootconfig(args.initrd, args.bootconfigs, args.output)?;
+    let args = Opt::parse();
+    match args {
+        Opt::Attach { initrd, bootconfigs, output } => {
+            attach_bootconfig(initrd, bootconfigs, output)?
+        }
+        Opt::Detach { initrd_with_bootconfigs, initrd, bootconfigs } => {
+            detach_bootconfig(initrd_with_bootconfigs, initrd, bootconfigs)?
+        }
+    };
     Ok(())
 }
 
@@ -82,6 +148,6 @@
     #[test]
     fn verify_args() {
         // Check that the command parsing has been configured in a valid way.
-        Args::command().debug_assert();
+        Opt::command().debug_assert();
     }
 }
diff --git a/microdroid_manager/src/dice.rs b/microdroid_manager/src/dice.rs
index 499835f..e740ed4 100644
--- a/microdroid_manager/src/dice.rs
+++ b/microdroid_manager/src/dice.rs
@@ -17,7 +17,7 @@
 use anyhow::{bail, Context, Error, Result};
 use byteorder::{NativeEndian, ReadBytesExt};
 use diced_open_dice_cbor::{
-    Config, ContextImpl, InputValuesOwned, Mode, OpenDiceCborContext, CDI_SIZE, HASH_SIZE,
+    Config, ContextImpl, DiceMode, InputValuesOwned, OpenDiceCborContext, CDI_SIZE, HASH_SIZE,
     HIDDEN_SIZE,
 };
 use keystore2_crypto::ZVec;
@@ -148,7 +148,7 @@
             Config::Descriptor(config_desc),
             authority_hash,
             None,
-            if debug { Mode::Debug } else { Mode::Normal },
+            if debug { DiceMode::kDiceModeDebug } else { DiceMode::kDiceModeNormal },
             hidden,
         );
         let (cdi_attest, cdi_seal, bcc) = match &self {
diff --git a/microdroid_manager/src/ioutil.rs b/microdroid_manager/src/ioutil.rs
index 8ac3712..d36e349 100644
--- a/microdroid_manager/src/ioutil.rs
+++ b/microdroid_manager/src/ioutil.rs
@@ -76,7 +76,7 @@
         });
 
         let test_file = test_dir.path().join("test.txt");
-        let mut file = wait_for_file(&test_file, Duration::from_secs(5))?;
+        let mut file = wait_for_file(test_file, Duration::from_secs(5))?;
         let mut buffer = String::new();
         file.read_to_string(&mut buffer)?;
         assert_eq!("test", buffer);
@@ -87,7 +87,7 @@
     fn test_wait_for_file_fails() {
         let test_dir = tempfile::TempDir::new().unwrap();
         let test_file = test_dir.path().join("test.txt");
-        let file = wait_for_file(&test_file, Duration::from_secs(1));
+        let file = wait_for_file(test_file, Duration::from_secs(1));
         assert!(file.is_err());
         assert_eq!(
             io::ErrorKind::NotFound,
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 24a12f7..13bc9e3 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -53,6 +53,7 @@
 use std::borrow::Cow::{Borrowed, Owned};
 use std::convert::TryInto;
 use std::env;
+use std::ffi::CString;
 use std::fs::{self, create_dir, OpenOptions};
 use std::io::Write;
 use std::os::unix::process::CommandExt;
@@ -216,13 +217,21 @@
 
     match try_run_payload(&service) {
         Ok(code) => {
-            info!("notifying payload finished");
-            service.notifyPayloadFinished(code)?;
             if code == 0 {
                 info!("task successfully finished");
             } else {
                 error!("task exited with exit code: {}", code);
             }
+            if let Err(e) = post_payload_work() {
+                error!(
+                    "Failed to run post payload work. It is possible that certain tasks
+                    like syncing encrypted store might be incomplete. Error: {:?}",
+                    e
+                );
+            };
+
+            info!("notifying payload finished");
+            service.notifyPayloadFinished(code)?;
             Ok(())
         }
         Err(err) => {
@@ -233,6 +242,28 @@
     }
 }
 
+fn post_payload_work() -> Result<()> {
+    // Sync the encrypted storage filesystem (flushes the filesystem caches).
+    if Path::new(ENCRYPTEDSTORE_BACKING_DEVICE).exists() {
+        let mountpoint = CString::new(ENCRYPTEDSTORE_MOUNTPOINT).unwrap();
+
+        let ret = unsafe {
+            let dirfd = libc::open(
+                mountpoint.as_ptr(),
+                libc::O_DIRECTORY | libc::O_RDONLY | libc::O_CLOEXEC,
+            );
+            ensure!(dirfd >= 0, "Unable to open {:?}", mountpoint);
+            let ret = libc::syncfs(dirfd);
+            libc::close(dirfd);
+            ret
+        };
+        if ret != 0 {
+            error!("failed to sync encrypted storage.");
+            return Err(anyhow!(std::io::Error::last_os_error()));
+        }
+    }
+    Ok(())
+}
 fn dice_derivation(
     dice: DiceDriver,
     verified_data: &MicrodroidData,
@@ -426,7 +457,8 @@
 
     // Start tombstone_transmit if enabled
     if config.export_tombstones {
-        control_service("start", "tombstone_transmit")?;
+        system_properties::write("tombstone_transmit.start", "1")
+            .context("set tombstone_transmit.start")?;
     } else {
         control_service("stop", "tombstoned")?;
     }
@@ -436,10 +468,6 @@
 
     register_vm_payload_service(allow_restricted_apis, service.clone(), dice_context)?;
 
-    if config.export_tombstones {
-        wait_for_tombstone_transmit_done()?;
-    }
-
     // Wait for encryptedstore to finish mounting the storage (if enabled) before setting
     // microdroid_manager.init_done. Reason is init stops uneventd after that.
     // Encryptedstore, however requires ueventd
@@ -451,6 +479,12 @@
     wait_for_property_true("dev.bootcomplete").context("failed waiting for dev.bootcomplete")?;
     system_properties::write("microdroid_manager.init_done", "1")
         .context("set microdroid_manager.init_done")?;
+
+    // Wait for tombstone_transmit to init
+    if config.export_tombstones {
+        wait_for_tombstone_transmit_done()?;
+    }
+
     info!("boot completed, time to run payload");
     exec_task(task, service).context("Failed to run payload")
 }
diff --git a/microdroid_manager/src/swap.rs b/microdroid_manager/src/swap.rs
index d7916db..2f4d176 100644
--- a/microdroid_manager/src/swap.rs
+++ b/microdroid_manager/src/swap.rs
@@ -66,7 +66,7 @@
     f.write_all(Uuid::new_v4().as_bytes())?;
 
     // Write the magic signature string.
-    f.seek(SeekFrom::Start((pagesize - 10) as u64))?;
+    f.seek(SeekFrom::Start(pagesize - 10))?;
     f.write_all("SWAPSPACE2".as_bytes())?;
 
     Ok(())
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index f5e214e..21f84a5 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -18,6 +18,7 @@
         "libfdtpci",
         "liblibfdt",
         "liblog_rust_nostd",
+        "libonce_cell_nostd",
         "libpvmfw_avb_nostd",
         "libpvmfw_embedded_key",
         "libtinyvec_nostd",
diff --git a/pvmfw/README.md b/pvmfw/README.md
index f46c718..1e4b605 100644
--- a/pvmfw/README.md
+++ b/pvmfw/README.md
@@ -214,7 +214,9 @@
 
 [AVB]: https://source.android.com/docs/security/features/verifiedboot/boot-flow
 [BccHandover]: https://pigweed.googlesource.com/open-dice/+/825e3beb6c/src/android/bcc.c#260
+[BccHandoverMainFlow]: https://pigweed.googlesource.com/open-dice/+/825e3beb6c/src/android/bcc.c#199
 [CDDL]: https://datatracker.ietf.org/doc/rfc8610
+[dice-mode]: https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md#Mode-Value-Details
 [dice-dt]: https://www.kernel.org/doc/Documentation/devicetree/bindings/reserved-memory/google%2Copen-dice.yaml
 [Layering]: https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md#layering-details
 [Trusty-BCC]: https://android.googlesource.com/trusty/lib/+/1696be0a8f3a7103/lib/hwbcc/common/swbcc.c#554
diff --git a/pvmfw/avb/Android.bp b/pvmfw/avb/Android.bp
index 837f747..0c1e392 100644
--- a/pvmfw/avb/Android.bp
+++ b/pvmfw/avb/Android.bp
@@ -9,9 +9,10 @@
     prefer_rlib: true,
     rustlibs: [
         "libavb_bindgen",
+        "libtinyvec_nostd",
     ],
     whole_static_libs: [
-        "libavb",
+        "libavb_non_debug",
     ],
 }
 
@@ -37,6 +38,7 @@
         ":microdroid_initrd_debuggable",
         ":test_image_with_one_hashdesc",
         ":test_image_with_non_initrd_hashdesc",
+        ":test_image_with_initrd_and_non_initrd_desc",
         ":test_image_with_prop_desc",
         ":unsigned_test_image",
     ],
@@ -44,7 +46,9 @@
     rustlibs: [
         "libanyhow",
         "libavb_bindgen",
+        "libhex",
         "libpvmfw_avb_nostd",
+        "libopenssl",
     ],
     enabled: false,
     arch: {
@@ -78,18 +82,40 @@
     src: ":unsigned_test_image",
     partition_name: "boot",
     private_key: ":pvmfw_sign_key",
-    salt: "1111",
+    salt: "3322",
     include_descriptors_from_images: [
         ":test_non_initrd_hashdesc",
     ],
 }
 
 avb_add_hash_footer {
+    name: "test_image_with_initrd_and_non_initrd_desc",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "3241",
+    include_descriptors_from_images: [
+        ":microdroid_initrd_normal_hashdesc",
+        ":test_non_initrd_hashdesc",
+    ],
+    enabled: false,
+    arch: {
+        // microdroid_initrd_normal_hashdesc is only available in these architectures.
+        arm64: {
+            enabled: true,
+        },
+        x86_64: {
+            enabled: true,
+        },
+    },
+}
+
+avb_add_hash_footer {
     name: "test_image_with_prop_desc",
     src: ":unsigned_test_image",
     partition_name: "boot",
     private_key: ":pvmfw_sign_key",
-    salt: "1111",
+    salt: "2134",
     props: [
         {
             name: "mock_prop",
diff --git a/pvmfw/avb/src/descriptor.rs b/pvmfw/avb/src/descriptor.rs
new file mode 100644
index 0000000..c54d416
--- /dev/null
+++ b/pvmfw/avb/src/descriptor.rs
@@ -0,0 +1,203 @@
+// Copyright 2023, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Structs and functions relating to the descriptors.
+
+#![warn(unsafe_op_in_unsafe_fn)]
+
+use crate::error::{AvbIOError, AvbSlotVerifyError};
+use crate::partition::PartitionName;
+use crate::utils::{self, is_not_null, to_nonnull, to_usize, usize_checked_add};
+use avb_bindgen::{
+    avb_descriptor_foreach, avb_hash_descriptor_validate_and_byteswap, AvbDescriptor,
+    AvbHashDescriptor, AvbVBMetaData, AVB_SHA256_DIGEST_SIZE,
+};
+use core::{
+    ffi::c_void,
+    mem::{size_of, MaybeUninit},
+    ops::Range,
+    slice,
+};
+use tinyvec::ArrayVec;
+
+/// Digest type for kernel and initrd.
+pub type Digest = [u8; AVB_SHA256_DIGEST_SIZE as usize];
+
+/// `HashDescriptors` can have maximum one `HashDescriptor` per known partition.
+#[derive(Default)]
+pub(crate) struct HashDescriptors(
+    ArrayVec<[HashDescriptor; PartitionName::NUM_OF_KNOWN_PARTITIONS]>,
+);
+
+impl HashDescriptors {
+    /// Builds `HashDescriptors` from `AvbVBMetaData`.
+    /// Returns an error if the given `AvbVBMetaData` contains non-hash descriptor, hash
+    /// descriptor of unknown `PartitionName` or duplicated hash descriptors.
+    ///
+    /// # Safety
+    ///
+    /// Behavior is undefined if any of the following conditions are violated:
+    /// * `vbmeta.vbmeta_data` must be non-null and points to a valid VBMeta.
+    /// * `vbmeta.vbmeta_data` must be valid for reading `vbmeta.vbmeta_size` bytes.
+    pub(crate) unsafe fn from_vbmeta(vbmeta: AvbVBMetaData) -> Result<Self, AvbSlotVerifyError> {
+        is_not_null(vbmeta.vbmeta_data).map_err(|_| AvbSlotVerifyError::Io)?;
+        let mut descriptors = Self::default();
+        // SAFETY: It is safe as the raw pointer `vbmeta.vbmeta_data` is a non-null pointer and
+        // points to a valid VBMeta structure.
+        if !unsafe {
+            avb_descriptor_foreach(
+                vbmeta.vbmeta_data,
+                vbmeta.vbmeta_size,
+                Some(check_and_save_descriptor),
+                &mut descriptors as *mut _ as *mut c_void,
+            )
+        } {
+            return Err(AvbSlotVerifyError::InvalidMetadata);
+        }
+        Ok(descriptors)
+    }
+
+    fn push(&mut self, descriptor: HashDescriptor) -> utils::Result<()> {
+        if self.0.iter().any(|d| d.partition_name_eq(&descriptor)) {
+            return Err(AvbIOError::Io);
+        }
+        self.0.push(descriptor);
+        Ok(())
+    }
+
+    pub(crate) fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Finds the `HashDescriptor` for the given `PartitionName`.
+    /// Throws an error if no corresponding descriptor found.
+    pub(crate) fn find(
+        &self,
+        partition_name: PartitionName,
+    ) -> Result<&HashDescriptor, AvbSlotVerifyError> {
+        self.0
+            .iter()
+            .find(|d| d.partition_name == partition_name)
+            .ok_or(AvbSlotVerifyError::InvalidMetadata)
+    }
+}
+
+/// # Safety
+///
+/// Behavior is undefined if any of the following conditions are violated:
+/// * The `descriptor` pointer must be non-null and points to a valid `AvbDescriptor` struct.
+/// * The `user_data` pointer must be non-null and points to a valid `HashDescriptors` struct.
+unsafe extern "C" fn check_and_save_descriptor(
+    descriptor: *const AvbDescriptor,
+    user_data: *mut c_void,
+) -> bool {
+    // SAFETY: It is safe because the caller must ensure that the `descriptor` pointer and
+    // the `user_data` are non-null and valid.
+    unsafe { try_check_and_save_descriptor(descriptor, user_data).is_ok() }
+}
+
+/// # Safety
+///
+/// Behavior is undefined if any of the following conditions are violated:
+/// * The `descriptor` pointer must be non-null and points to a valid `AvbDescriptor` struct.
+/// * The `user_data` pointer must be non-null and points to a valid `HashDescriptors` struct.
+unsafe fn try_check_and_save_descriptor(
+    descriptor: *const AvbDescriptor,
+    user_data: *mut c_void,
+) -> utils::Result<()> {
+    is_not_null(descriptor)?;
+    // SAFETY: It is safe because the caller ensures that `descriptor` is a non-null pointer
+    // pointing to a valid struct.
+    let desc = unsafe { AvbHashDescriptorWrap::from_descriptor_ptr(descriptor)? };
+    // SAFETY: It is safe because the caller ensures that `descriptor` is a non-null pointer
+    // pointing to a valid struct.
+    let data = unsafe { slice::from_raw_parts(descriptor as *const u8, desc.len()?) };
+    let mut descriptors = to_nonnull(user_data as *mut HashDescriptors)?;
+    // SAFETY: It is safe because the caller ensures that `user_data` is a non-null pointer
+    // pointing to a valid struct.
+    let descriptors = unsafe { descriptors.as_mut() };
+    descriptors.push(HashDescriptor::new(&desc, data)?)
+}
+
+#[derive(Default)]
+pub(crate) struct HashDescriptor {
+    partition_name: PartitionName,
+    pub(crate) digest: Digest,
+}
+
+impl HashDescriptor {
+    fn new(desc: &AvbHashDescriptorWrap, data: &[u8]) -> utils::Result<Self> {
+        let partition_name = data
+            .get(desc.partition_name_range()?)
+            .ok_or(AvbIOError::RangeOutsidePartition)?
+            .try_into()?;
+        let partition_digest =
+            data.get(desc.digest_range()?).ok_or(AvbIOError::RangeOutsidePartition)?;
+        let mut digest = [0u8; size_of::<Digest>()];
+        digest.copy_from_slice(partition_digest);
+        Ok(Self { partition_name, digest })
+    }
+
+    fn partition_name_eq(&self, other: &HashDescriptor) -> bool {
+        self.partition_name == other.partition_name
+    }
+}
+
+/// `AvbHashDescriptor` contains the metadata for the given descriptor.
+struct AvbHashDescriptorWrap(AvbHashDescriptor);
+
+impl AvbHashDescriptorWrap {
+    /// # Safety
+    ///
+    /// Behavior is undefined if any of the following conditions are violated:
+    /// * The `descriptor` pointer must be non-null and point to a valid `AvbDescriptor`.
+    unsafe fn from_descriptor_ptr(descriptor: *const AvbDescriptor) -> utils::Result<Self> {
+        is_not_null(descriptor)?;
+        // SAFETY: It is safe as the raw pointer `descriptor` is non-null and points to
+        // a valid `AvbDescriptor`.
+        let desc = unsafe {
+            let mut desc = MaybeUninit::uninit();
+            if !avb_hash_descriptor_validate_and_byteswap(
+                descriptor as *const AvbHashDescriptor,
+                desc.as_mut_ptr(),
+            ) {
+                return Err(AvbIOError::Io);
+            }
+            desc.assume_init()
+        };
+        Ok(Self(desc))
+    }
+
+    fn len(&self) -> utils::Result<usize> {
+        usize_checked_add(
+            size_of::<AvbDescriptor>(),
+            to_usize(self.0.parent_descriptor.num_bytes_following)?,
+        )
+    }
+
+    fn partition_name_end(&self) -> utils::Result<usize> {
+        usize_checked_add(size_of::<AvbHashDescriptor>(), to_usize(self.0.partition_name_len)?)
+    }
+
+    fn partition_name_range(&self) -> utils::Result<Range<usize>> {
+        let start = size_of::<AvbHashDescriptor>();
+        Ok(start..(self.partition_name_end()?))
+    }
+
+    fn digest_range(&self) -> utils::Result<Range<usize>> {
+        let start = usize_checked_add(self.partition_name_end()?, to_usize(self.0.salt_len)?)?;
+        let end = usize_checked_add(start, to_usize(self.0.digest_len)?)?;
+        Ok(start..end)
+    }
+}
diff --git a/pvmfw/avb/src/lib.rs b/pvmfw/avb/src/lib.rs
index a1e3ee0..8fea162 100644
--- a/pvmfw/avb/src/lib.rs
+++ b/pvmfw/avb/src/lib.rs
@@ -18,11 +18,13 @@
 // For usize.checked_add_signed(isize), available in Rust 1.66.0
 #![feature(mixed_integer_ops)]
 
+mod descriptor;
 mod error;
 mod ops;
 mod partition;
 mod utils;
 mod verify;
 
+pub use descriptor::Digest;
 pub use error::AvbSlotVerifyError;
-pub use verify::{verify_payload, DebugLevel};
+pub use verify::{verify_payload, DebugLevel, VerifiedBootData};
diff --git a/pvmfw/avb/src/partition.rs b/pvmfw/avb/src/partition.rs
index 82d89f8..bc63003 100644
--- a/pvmfw/avb/src/partition.rs
+++ b/pvmfw/avb/src/partition.rs
@@ -18,14 +18,18 @@
 use crate::utils::is_not_null;
 use core::ffi::{c_char, CStr};
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 pub(crate) enum PartitionName {
+    /// The default `PartitionName` is needed to build the default `HashDescriptor`.
+    #[default]
     Kernel,
     InitrdNormal,
     InitrdDebug,
 }
 
 impl PartitionName {
+    pub(crate) const NUM_OF_KNOWN_PARTITIONS: usize = 3;
+
     const KERNEL_PARTITION_NAME: &[u8] = b"boot\0";
     const INITRD_NORMAL_PARTITION_NAME: &[u8] = b"initrd_normal\0";
     const INITRD_DEBUG_PARTITION_NAME: &[u8] = b"initrd_debug\0";
diff --git a/pvmfw/avb/src/utils.rs b/pvmfw/avb/src/utils.rs
index 1b35c22..a24d61f 100644
--- a/pvmfw/avb/src/utils.rs
+++ b/pvmfw/avb/src/utils.rs
@@ -22,7 +22,7 @@
 
 pub(crate) fn write<T>(ptr: *mut T, value: T) -> Result<()> {
     let ptr = to_nonnull(ptr)?;
-    // SAFETY: It is safe as the raw pointer `ptr` is a nonnull pointer.
+    // SAFETY: It is safe as the raw pointer `ptr` is a non-null pointer.
     unsafe {
         *ptr.as_ptr() = value;
     }
@@ -31,7 +31,7 @@
 
 pub(crate) fn as_ref<'a, T>(ptr: *mut T) -> Result<&'a T> {
     let ptr = to_nonnull(ptr)?;
-    // SAFETY: It is safe as the raw pointer `ptr` is a nonnull pointer.
+    // SAFETY: It is safe as the raw pointer `ptr` is a non-null pointer.
     unsafe { Ok(ptr.as_ref()) }
 }
 
diff --git a/pvmfw/avb/src/verify.rs b/pvmfw/avb/src/verify.rs
index a062061..1a79c83 100644
--- a/pvmfw/avb/src/verify.rs
+++ b/pvmfw/avb/src/verify.rs
@@ -14,22 +14,26 @@
 
 //! This module handles the pvmfw payload verification.
 
-use crate::error::{AvbIOError, AvbSlotVerifyError};
+use crate::descriptor::{Digest, HashDescriptors};
+use crate::error::AvbSlotVerifyError;
 use crate::ops::{Ops, Payload};
 use crate::partition::PartitionName;
-use crate::utils::{is_not_null, to_usize, usize_checked_add, write};
-use avb_bindgen::{
-    avb_descriptor_foreach, avb_hash_descriptor_validate_and_byteswap, AvbDescriptor,
-    AvbHashDescriptor, AvbPartitionData, AvbVBMetaData,
-};
-use core::{
-    ffi::{c_char, c_void},
-    mem::{size_of, MaybeUninit},
-    slice,
-};
+use avb_bindgen::{AvbPartitionData, AvbVBMetaData};
+use core::ffi::c_char;
+
+/// Verified data returned when the payload verification succeeds.
+#[derive(Debug)]
+pub struct VerifiedBootData {
+    /// DebugLevel of the VM.
+    pub debug_level: DebugLevel,
+    /// Kernel digest.
+    pub kernel_digest: Digest,
+    /// Initrd digest if initrd exists.
+    pub initrd_digest: Option<Digest>,
+}
 
 /// This enum corresponds to the `DebugLevel` in `VirtualMachineConfig`.
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum DebugLevel {
     /// Not debuggable at all.
     None,
@@ -37,104 +41,6 @@
     Full,
 }
 
-extern "C" fn search_initrd_hash_descriptor(
-    descriptor: *const AvbDescriptor,
-    user_data: *mut c_void,
-) -> bool {
-    try_search_initrd_hash_descriptor(descriptor, user_data).is_ok()
-}
-
-fn try_search_initrd_hash_descriptor(
-    descriptor: *const AvbDescriptor,
-    user_data: *mut c_void,
-) -> Result<(), AvbIOError> {
-    let hash_desc = AvbHashDescriptorRef::try_from(descriptor)?;
-    if matches!(
-        hash_desc.partition_name()?.try_into(),
-        Ok(PartitionName::InitrdDebug) | Ok(PartitionName::InitrdNormal),
-    ) {
-        write(user_data as *mut bool, true)?;
-    }
-    Ok(())
-}
-
-/// `hash_desc` only contains the metadata like fields length and flags of the descriptor.
-/// The data itself is contained in `ptr`.
-struct AvbHashDescriptorRef {
-    hash_desc: AvbHashDescriptor,
-    ptr: *const AvbDescriptor,
-}
-
-impl TryFrom<*const AvbDescriptor> for AvbHashDescriptorRef {
-    type Error = AvbIOError;
-
-    fn try_from(descriptor: *const AvbDescriptor) -> Result<Self, Self::Error> {
-        is_not_null(descriptor)?;
-        // SAFETY: It is safe as the raw pointer `descriptor` is a nonnull pointer and
-        // we have validated that it is of hash descriptor type.
-        let hash_desc = unsafe {
-            let mut desc = MaybeUninit::uninit();
-            if !avb_hash_descriptor_validate_and_byteswap(
-                descriptor as *const AvbHashDescriptor,
-                desc.as_mut_ptr(),
-            ) {
-                return Err(AvbIOError::Io);
-            }
-            desc.assume_init()
-        };
-        Ok(Self { hash_desc, ptr: descriptor })
-    }
-}
-
-impl AvbHashDescriptorRef {
-    fn check_is_in_range(&self, index: usize) -> Result<(), AvbIOError> {
-        let parent_desc = self.hash_desc.parent_descriptor;
-        let total_len = usize_checked_add(
-            size_of::<AvbDescriptor>(),
-            to_usize(parent_desc.num_bytes_following)?,
-        )?;
-        if index <= total_len {
-            Ok(())
-        } else {
-            Err(AvbIOError::Io)
-        }
-    }
-
-    /// Returns the non null-terminated partition name.
-    fn partition_name(&self) -> Result<&[u8], AvbIOError> {
-        let partition_name_offset = size_of::<AvbHashDescriptor>();
-        let partition_name_len = to_usize(self.hash_desc.partition_name_len)?;
-        self.check_is_in_range(usize_checked_add(partition_name_offset, partition_name_len)?)?;
-        let desc = self.ptr as *const u8;
-        // SAFETY: The descriptor has been validated as nonnull and the partition name is
-        // contained within the image.
-        unsafe { Ok(slice::from_raw_parts(desc.add(partition_name_offset), partition_name_len)) }
-    }
-}
-
-fn verify_vbmeta_has_no_initrd_descriptor(
-    vbmeta_image: &AvbVBMetaData,
-) -> Result<(), AvbSlotVerifyError> {
-    is_not_null(vbmeta_image.vbmeta_data).map_err(|_| AvbSlotVerifyError::Io)?;
-    let mut has_unexpected_descriptor = false;
-    // SAFETY: It is safe as the raw pointer `vbmeta_image.vbmeta_data` is a nonnull pointer.
-    if !unsafe {
-        avb_descriptor_foreach(
-            vbmeta_image.vbmeta_data,
-            vbmeta_image.vbmeta_size,
-            Some(search_initrd_hash_descriptor),
-            &mut has_unexpected_descriptor as *mut _ as *mut c_void,
-        )
-    } {
-        return Err(AvbSlotVerifyError::InvalidMetadata);
-    }
-    if has_unexpected_descriptor {
-        Err(AvbSlotVerifyError::InvalidMetadata)
-    } else {
-        Ok(())
-    }
-}
-
 fn verify_only_one_vbmeta_exists(
     vbmeta_images: &[AvbVBMetaData],
 ) -> Result<(), AvbSlotVerifyError> {
@@ -154,6 +60,16 @@
     }
 }
 
+fn verify_vbmeta_has_only_one_hash_descriptor(
+    hash_descriptors: &HashDescriptors,
+) -> Result<(), AvbSlotVerifyError> {
+    if hash_descriptors.len() == 1 {
+        Ok(())
+    } else {
+        Err(AvbSlotVerifyError::InvalidMetadata)
+    }
+}
+
 fn verify_loaded_partition_has_expected_length(
     loaded_partitions: &[AvbPartitionData],
     partition_name: PartitionName,
@@ -182,7 +98,7 @@
     kernel: &[u8],
     initrd: Option<&[u8]>,
     trusted_public_key: &[u8],
-) -> Result<DebugLevel, AvbSlotVerifyError> {
+) -> Result<VerifiedBootData, AvbSlotVerifyError> {
     let mut payload = Payload::new(kernel, initrd, trusted_public_key);
     let mut ops = Ops::from(&mut payload);
     let kernel_verify_result = ops.verify_partition(PartitionName::Kernel.as_cstr())?;
@@ -191,13 +107,20 @@
     verify_only_one_vbmeta_exists(vbmeta_images)?;
     let vbmeta_image = vbmeta_images[0];
     verify_vbmeta_is_from_kernel_partition(&vbmeta_image)?;
+    // SAFETY: It is safe because the `vbmeta_image` is collected from `AvbSlotVerifyData`,
+    // which is returned by `avb_slot_verify()` when the verification succeeds. It is
+    // guaranteed by libavb to be non-null and to point to a valid VBMeta structure.
+    let hash_descriptors = unsafe { HashDescriptors::from_vbmeta(vbmeta_image)? };
+    let kernel_descriptor = hash_descriptors.find(PartitionName::Kernel)?;
 
     if initrd.is_none() {
-        verify_vbmeta_has_no_initrd_descriptor(&vbmeta_image)?;
-        return Ok(DebugLevel::None);
+        verify_vbmeta_has_only_one_hash_descriptor(&hash_descriptors)?;
+        return Ok(VerifiedBootData {
+            debug_level: DebugLevel::None,
+            kernel_digest: kernel_descriptor.digest,
+            initrd_digest: None,
+        });
     }
-    // TODO(b/256148034): Check the vbmeta doesn't have hash descriptors other than
-    // boot, initrd_normal, initrd_debug.
 
     let initrd = initrd.unwrap();
     let (debug_level, initrd_verify_result, initrd_partition_name) =
@@ -214,5 +137,10 @@
         initrd_partition_name,
         initrd.len(),
     )?;
-    Ok(debug_level)
+    let initrd_descriptor = hash_descriptors.find(initrd_partition_name)?;
+    Ok(VerifiedBootData {
+        debug_level,
+        kernel_digest: kernel_descriptor.digest,
+        initrd_digest: Some(initrd_descriptor.digest),
+    })
 }
diff --git a/pvmfw/avb/tests/api_test.rs b/pvmfw/avb/tests/api_test.rs
index 0572789..261d8a8 100644
--- a/pvmfw/avb/tests/api_test.rs
+++ b/pvmfw/avb/tests/api_test.rs
@@ -16,15 +16,17 @@
 
 mod utils;
 
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use avb_bindgen::{AvbFooter, AvbVBMetaImageHeader};
-use pvmfw_avb::{AvbSlotVerifyError, DebugLevel};
+use pvmfw_avb::{verify_payload, AvbSlotVerifyError, DebugLevel};
 use std::{fs, mem::size_of, ptr};
 use utils::*;
 
 const TEST_IMG_WITH_ONE_HASHDESC_PATH: &str = "test_image_with_one_hashdesc.img";
 const TEST_IMG_WITH_PROP_DESC_PATH: &str = "test_image_with_prop_desc.img";
 const TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH: &str = "test_image_with_non_initrd_hashdesc.img";
+const TEST_IMG_WITH_INITRD_AND_NON_INITRD_DESC_PATH: &str =
+    "test_image_with_initrd_and_non_initrd_desc.img";
 const UNSIGNED_TEST_IMG_PATH: &str = "unsigned_test.img";
 
 const RANDOM_FOOTER_POS: usize = 30;
@@ -33,111 +35,125 @@
 /// the latest payload can be verified successfully.
 #[test]
 fn latest_normal_payload_passes_verification() -> Result<()> {
-    assert_payload_verification_with_initrd_eq(
-        &load_latest_signed_kernel()?,
+    assert_latest_payload_verification_passes(
         &load_latest_initrd_normal()?,
-        &load_trusted_public_key()?,
-        Ok(DebugLevel::None),
+        b"initrd_normal",
+        DebugLevel::None,
     )
 }
 
 #[test]
 fn latest_debug_payload_passes_verification() -> Result<()> {
-    assert_payload_verification_with_initrd_eq(
-        &load_latest_signed_kernel()?,
+    assert_latest_payload_verification_passes(
         &load_latest_initrd_debug()?,
-        &load_trusted_public_key()?,
-        Ok(DebugLevel::Full),
+        b"initrd_debug",
+        DebugLevel::Full,
     )
 }
 
 #[test]
 fn payload_expecting_no_initrd_passes_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_eq(
+    let verified_boot_data = verify_payload(
         &fs::read(TEST_IMG_WITH_ONE_HASHDESC_PATH)?,
         /*initrd=*/ None,
         &load_trusted_public_key()?,
-        Ok(DebugLevel::None),
+    )
+    .map_err(|e| anyhow!("Verification failed. Error: {}", e))?;
+
+    assert_eq!(DebugLevel::None, verified_boot_data.debug_level);
+    let digest = hash(&[&hex::decode("1111")?, &fs::read(UNSIGNED_TEST_IMG_PATH)?]);
+    assert_eq!(digest, verified_boot_data.kernel_digest);
+    assert!(verified_boot_data.initrd_digest.is_none());
+    Ok(())
+}
+
+#[test]
+fn payload_with_non_initrd_descriptor_fails_verification_with_no_initrd() -> Result<()> {
+    assert_payload_verification_fails(
+        &fs::read(TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH)?,
+        /*initrd=*/ None,
+        &load_trusted_public_key()?,
+        AvbSlotVerifyError::InvalidMetadata,
     )
 }
 
 #[test]
-fn payload_with_non_initrd_descriptor_passes_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_eq(
-        &fs::read(TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH)?,
-        /*initrd=*/ None,
+fn payload_with_non_initrd_descriptor_fails_verification_with_initrd() -> Result<()> {
+    assert_payload_verification_with_initrd_fails(
+        &fs::read(TEST_IMG_WITH_INITRD_AND_NON_INITRD_DESC_PATH)?,
+        &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        Ok(DebugLevel::None),
+        AvbSlotVerifyError::InvalidMetadata,
     )
 }
 
 #[test]
 fn payload_with_prop_descriptor_fails_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_eq(
+    assert_payload_verification_fails(
         &fs::read(TEST_IMG_WITH_PROP_DESC_PATH)?,
         /*initrd=*/ None,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::InvalidMetadata),
+        AvbSlotVerifyError::InvalidMetadata,
     )
 }
 
 #[test]
 fn payload_expecting_initrd_fails_verification_with_no_initrd() -> Result<()> {
-    assert_payload_verification_eq(
+    assert_payload_verification_fails(
         &load_latest_signed_kernel()?,
         /*initrd=*/ None,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::InvalidMetadata),
+        AvbSlotVerifyError::InvalidMetadata,
     )
 }
 
 #[test]
 fn payload_with_empty_public_key_fails_verification() -> Result<()> {
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &load_latest_signed_kernel()?,
         &load_latest_initrd_normal()?,
         /*trusted_public_key=*/ &[0u8; 0],
-        Err(AvbSlotVerifyError::PublicKeyRejected),
+        AvbSlotVerifyError::PublicKeyRejected,
     )
 }
 
 #[test]
 fn payload_with_an_invalid_public_key_fails_verification() -> Result<()> {
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &load_latest_signed_kernel()?,
         &load_latest_initrd_normal()?,
         /*trusted_public_key=*/ &[0u8; 512],
-        Err(AvbSlotVerifyError::PublicKeyRejected),
+        AvbSlotVerifyError::PublicKeyRejected,
     )
 }
 
 #[test]
 fn payload_with_a_different_valid_public_key_fails_verification() -> Result<()> {
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &load_latest_signed_kernel()?,
         &load_latest_initrd_normal()?,
         &fs::read(PUBLIC_KEY_RSA2048_PATH)?,
-        Err(AvbSlotVerifyError::PublicKeyRejected),
+        AvbSlotVerifyError::PublicKeyRejected,
     )
 }
 
 #[test]
 fn payload_with_an_invalid_initrd_fails_verification() -> Result<()> {
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &load_latest_signed_kernel()?,
         /*initrd=*/ &fs::read(UNSIGNED_TEST_IMG_PATH)?,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::Verification),
+        AvbSlotVerifyError::Verification,
     )
 }
 
 #[test]
 fn unsigned_kernel_fails_verification() -> Result<()> {
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &fs::read(UNSIGNED_TEST_IMG_PATH)?,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::Io),
+        AvbSlotVerifyError::Io,
     )
 }
 
@@ -146,11 +162,11 @@
     let mut kernel = load_latest_signed_kernel()?;
     kernel[1] = !kernel[1]; // Flip the bits
 
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::Verification),
+        AvbSlotVerifyError::Verification,
     )
 }
 
@@ -176,12 +192,14 @@
 
         // Assert.
         let footer = extract_avb_footer(&kernel)?;
-        assert_eq!(wrong_offset, footer.vbmeta_offset as u64);
-        assert_payload_verification_with_initrd_eq(
+        // footer is unaligned; copy vbmeta_offset to local variable
+        let vbmeta_offset = footer.vbmeta_offset;
+        assert_eq!(wrong_offset, vbmeta_offset);
+        assert_payload_verification_with_initrd_fails(
             &kernel,
             &load_latest_initrd_normal()?,
             &load_trusted_public_key()?,
-            Err(AvbSlotVerifyError::Io),
+            AvbSlotVerifyError::Io,
         )?;
     }
     Ok(())
@@ -193,11 +211,11 @@
     let avb_footer_index = kernel.len() - size_of::<AvbFooter>() + RANDOM_FOOTER_POS;
     kernel[avb_footer_index] = !kernel[avb_footer_index];
 
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::InvalidMetadata),
+        AvbSlotVerifyError::InvalidMetadata,
     )
 }
 
@@ -206,11 +224,11 @@
     let mut initrd = load_latest_initrd_normal()?;
     initrd.extend(b"androidboot.vbmeta.digest=1111");
 
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &load_latest_signed_kernel()?,
         &initrd,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::Verification),
+        AvbSlotVerifyError::Verification,
     )
 }
 
@@ -222,11 +240,11 @@
 
     kernel[vbmeta_index] = !kernel[vbmeta_index]; // Flip the bits
 
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::InvalidMetadata),
+        AvbSlotVerifyError::InvalidMetadata,
     )
 }
 
@@ -245,17 +263,17 @@
     kernel[public_key_offset..(public_key_offset + public_key_size)]
         .copy_from_slice(&empty_public_key);
 
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &kernel,
         &load_latest_initrd_normal()?,
         &empty_public_key,
-        Err(AvbSlotVerifyError::Verification),
+        AvbSlotVerifyError::Verification,
     )?;
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::Verification),
+        AvbSlotVerifyError::Verification,
     )
 }
 
@@ -268,10 +286,10 @@
     let mut kernel = load_latest_signed_kernel()?;
     let footer = extract_avb_footer(&kernel)?;
     let vbmeta_header = extract_vbmeta_header(&kernel, &footer)?;
-    assert_eq!(
-        0, vbmeta_header.flags as u32,
-        "The disable flag should not be set in the latest kernel."
-    );
+
+    // vbmeta_header is unaligned; copy flags to local variable
+    let vbmeta_header_flags = vbmeta_header.flags;
+    assert_eq!(0, vbmeta_header_flags, "The disable flag should not be set in the latest kernel.");
     let flags_addr = ptr::addr_of!(vbmeta_header.flags) as *const u8;
     // SAFETY: It is safe as both raw pointers `flags_addr` and `vbmeta_header` are not null.
     let flags_offset = unsafe { flags_addr.offset_from(ptr::addr_of!(vbmeta_header) as *const u8) };
@@ -283,14 +301,16 @@
 
     // Assert.
     let vbmeta_header = extract_vbmeta_header(&kernel, &footer)?;
+    // vbmeta_header is unaligned; copy flags to local variable
+    let vbmeta_header_flags = vbmeta_header.flags;
     assert_eq!(
-        AVB_VBMETA_IMAGE_FLAGS_VERIFICATION_DISABLED, vbmeta_header.flags as u32,
+        AVB_VBMETA_IMAGE_FLAGS_VERIFICATION_DISABLED, vbmeta_header_flags,
         "VBMeta verification flag should be disabled now."
     );
-    assert_payload_verification_with_initrd_eq(
+    assert_payload_verification_with_initrd_fails(
         &kernel,
         &load_latest_initrd_normal()?,
         &load_trusted_public_key()?,
-        Err(AvbSlotVerifyError::Verification),
+        AvbSlotVerifyError::Verification,
     )
 }
diff --git a/pvmfw/avb/tests/utils.rs b/pvmfw/avb/tests/utils.rs
index 0a2eac6..8756d06 100644
--- a/pvmfw/avb/tests/utils.rs
+++ b/pvmfw/avb/tests/utils.rs
@@ -16,12 +16,13 @@
 
 //! Utility functions used by API tests.
 
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use avb_bindgen::{
     avb_footer_validate_and_byteswap, avb_vbmeta_image_header_to_host_byte_order, AvbFooter,
     AvbVBMetaImageHeader,
 };
-use pvmfw_avb::{verify_payload, AvbSlotVerifyError, DebugLevel};
+use openssl::sha;
+use pvmfw_avb::{verify_payload, AvbSlotVerifyError, DebugLevel, Digest};
 use std::{
     fs,
     mem::{size_of, transmute, MaybeUninit},
@@ -34,22 +35,22 @@
 
 pub const PUBLIC_KEY_RSA2048_PATH: &str = "data/testkey_rsa2048_pub.bin";
 
-pub fn assert_payload_verification_with_initrd_eq(
+pub fn assert_payload_verification_with_initrd_fails(
     kernel: &[u8],
     initrd: &[u8],
     trusted_public_key: &[u8],
-    expected_result: Result<DebugLevel, AvbSlotVerifyError>,
+    expected_error: AvbSlotVerifyError,
 ) -> Result<()> {
-    assert_payload_verification_eq(kernel, Some(initrd), trusted_public_key, expected_result)
+    assert_payload_verification_fails(kernel, Some(initrd), trusted_public_key, expected_error)
 }
 
-pub fn assert_payload_verification_eq(
+pub fn assert_payload_verification_fails(
     kernel: &[u8],
     initrd: Option<&[u8]>,
     trusted_public_key: &[u8],
-    expected_result: Result<DebugLevel, AvbSlotVerifyError>,
+    expected_error: AvbSlotVerifyError,
 ) -> Result<()> {
-    assert_eq!(expected_result, verify_payload(kernel, initrd, trusted_public_key));
+    assert_eq!(expected_error, verify_payload(kernel, initrd, trusted_public_key).unwrap_err());
     Ok(())
 }
 
@@ -95,3 +96,34 @@
     };
     Ok(vbmeta_header)
 }
+
+pub fn assert_latest_payload_verification_passes(
+    initrd: &[u8],
+    initrd_salt: &[u8],
+    expected_debug_level: DebugLevel,
+) -> Result<()> {
+    let kernel = load_latest_signed_kernel()?;
+    let verified_boot_data = verify_payload(&kernel, Some(initrd), &load_trusted_public_key()?)
+        .map_err(|e| anyhow!("Verification failed. Error: {}", e))?;
+
+    assert_eq!(expected_debug_level, verified_boot_data.debug_level);
+
+    let footer = extract_avb_footer(&kernel)?;
+    assert_eq!(
+        hash(&[&hash(&[b"bootloader"]), &kernel[..usize::try_from(footer.original_image_size)?]]),
+        verified_boot_data.kernel_digest,
+        "Kernel digest is not equal to the expected."
+    );
+    assert_eq!(
+        hash(&[&hash(&[initrd_salt]), initrd,]),
+        verified_boot_data.initrd_digest.unwrap(),
+        "initrd digest is not equal to the expected."
+    );
+    Ok(())
+}
+
+pub fn hash(inputs: &[&[u8]]) -> Digest {
+    let mut digester = sha::Sha256::new();
+    inputs.iter().for_each(|input| digester.update(input));
+    digester.finish()
+}
diff --git a/pvmfw/src/avb.rs b/pvmfw/src/avb.rs
deleted file mode 100644
index 1abe73f..0000000
--- a/pvmfw/src/avb.rs
+++ /dev/null
@@ -1,17 +0,0 @@
-// Copyright 2022, The Android Open Source Project
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//! Image verification.
-
-pub use pvmfw_embedded_key::PUBLIC_KEY;
diff --git a/pvmfw/src/debug_policy.rs b/pvmfw/src/debug_policy.rs
new file mode 100644
index 0000000..37e2af8
--- /dev/null
+++ b/pvmfw/src/debug_policy.rs
@@ -0,0 +1,152 @@
+// Copyright 2023, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Support for the debug policy overlay in pvmfw
+
+use alloc::vec;
+use core::ffi::CStr;
+use core::fmt;
+use libfdt::FdtError;
+use log::info;
+
+#[derive(Debug, Clone)]
+pub enum DebugPolicyError {
+    /// The provided baseline FDT was invalid or malformed, so cannot access certain node/prop
+    Fdt(&'static str, FdtError),
+    /// The provided debug policy FDT was invalid or malformed.
+    DebugPolicyFdt(&'static str, FdtError),
+    /// The overlaid result FDT is invalid or malformed, and may be corrupted.
+    OverlaidFdt(&'static str, FdtError),
+}
+
+impl fmt::Display for DebugPolicyError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::Fdt(s, e) => write!(f, "Invalid baseline FDT. {s}: {e}"),
+            Self::DebugPolicyFdt(s, e) => write!(f, "Invalid overlay FDT. {s}: {e}"),
+            Self::OverlaidFdt(s, e) => write!(f, "Invalid overlaid FDT. {s}: {e}"),
+        }
+    }
+}
+
+/// Applies the debug policy device tree overlay to the pVM DT.
+///
+/// # Safety
+///
+/// When an error is returned by this function, the input `Fdt` should be
+/// discarded as it may have have been partially corrupted during the overlay
+/// application process.
+unsafe fn apply_debug_policy(
+    fdt: &mut libfdt::Fdt,
+    debug_policy: &mut [u8],
+) -> Result<(), DebugPolicyError> {
+    let overlay = libfdt::Fdt::from_mut_slice(debug_policy)
+        .map_err(|e| DebugPolicyError::DebugPolicyFdt("Failed to load debug policy overlay", e))?;
+
+    fdt.unpack().map_err(|e| DebugPolicyError::Fdt("Failed to unpack", e))?;
+
+    let fdt = fdt
+        .apply_overlay(overlay)
+        .map_err(|e| DebugPolicyError::DebugPolicyFdt("Failed to apply overlay", e))?;
+
+    fdt.pack().map_err(|e| DebugPolicyError::OverlaidFdt("Failed to re-pack", e))
+}
+
+/// Dsiables ramdump by removing crashkernel from bootargs in /chosen.
+///
+/// # Safety
+///
+/// This may corrupt the input `Fdt` when error happens while editing prop value.
+unsafe fn disable_ramdump(fdt: &mut libfdt::Fdt) -> Result<(), DebugPolicyError> {
+    let chosen_path = CStr::from_bytes_with_nul(b"/chosen\0").unwrap();
+    let bootargs_name = CStr::from_bytes_with_nul(b"bootargs\0").unwrap();
+
+    let chosen = match fdt
+        .node(chosen_path)
+        .map_err(|e| DebugPolicyError::Fdt("Failed to find /chosen", e))?
+    {
+        Some(node) => node,
+        None => return Ok(()),
+    };
+
+    let bootargs = match chosen
+        .getprop_str(bootargs_name)
+        .map_err(|e| DebugPolicyError::Fdt("Failed to find bootargs prop", e))?
+    {
+        Some(value) if !value.to_bytes().is_empty() => value,
+        _ => return Ok(()),
+    };
+
+    // TODO: Improve add 'crashkernel=17MB' only when it's unnecessary.
+    //       Currently 'crashkernel=17MB' in virtualizationservice and passed by
+    //       chosen node, because it's not exactly a debug policy but a
+    //       configuration. However, it's actually microdroid specific
+    //       so we need a way to generalize it.
+    let mut args = vec![];
+    for arg in bootargs.to_bytes().split(|byte| byte.is_ascii_whitespace()) {
+        if arg.is_empty() || arg.starts_with(b"crashkernel=") {
+            continue;
+        }
+        args.push(arg);
+    }
+    let mut new_bootargs = args.as_slice().join(&b" "[..]);
+    new_bootargs.push(b'\0');
+
+    // We've checked existence of /chosen node at the beginning.
+    let mut chosen_mut = fdt.node_mut(chosen_path).unwrap().unwrap();
+    chosen_mut.setprop(bootargs_name, new_bootargs.as_slice()).map_err(|e| {
+        DebugPolicyError::OverlaidFdt("Failed to remove crashkernel. FDT might be corrupted", e)
+    })
+}
+
+/// Returns true only if fdt has ramdump prop in the /avf/guest/common node with value <1>
+fn is_ramdump_enabled(fdt: &libfdt::Fdt) -> Result<bool, DebugPolicyError> {
+    let common = match fdt
+        .node(CStr::from_bytes_with_nul(b"/avf/guest/common\0").unwrap())
+        .map_err(|e| DebugPolicyError::DebugPolicyFdt("Failed to find /avf/guest/common node", e))?
+    {
+        Some(node) => node,
+        None => return Ok(false),
+    };
+
+    match common
+        .getprop_u32(CStr::from_bytes_with_nul(b"ramdump\0").unwrap())
+        .map_err(|e| DebugPolicyError::DebugPolicyFdt("Failed to find ramdump prop", e))?
+    {
+        Some(1) => Ok(true),
+        _ => Ok(false),
+    }
+}
+
+/// Handles debug policies.
+///
+/// # Safety
+///
+/// This may corrupt the input `Fdt` when overlaying debug policy or applying
+/// ramdump configuration.
+pub unsafe fn handle_debug_policy(
+    fdt: &mut libfdt::Fdt,
+    debug_policy: Option<&mut [u8]>,
+) -> Result<(), DebugPolicyError> {
+    if let Some(dp) = debug_policy {
+        apply_debug_policy(fdt, dp)?;
+    }
+
+    // Handles ramdump in the debug policy
+    if is_ramdump_enabled(fdt)? {
+        info!("ramdump is enabled by debug policy");
+        return Ok(());
+    }
+    disable_ramdump(fdt)
+}
diff --git a/pvmfw/src/dice.rs b/pvmfw/src/dice.rs
index b322850..49218b0 100644
--- a/pvmfw/src/dice.rs
+++ b/pvmfw/src/dice.rs
@@ -15,24 +15,41 @@
 //! Support for DICE derivation and BCC generation.
 
 use core::ffi::CStr;
-
+use core::mem::size_of;
 use dice::bcc::format_config_descriptor;
 use dice::bcc::Handover;
 use dice::hash;
 use dice::ConfigType;
+use dice::DiceMode;
 use dice::InputValues;
+use pvmfw_avb::{DebugLevel, Digest, VerifiedBootData};
+
+fn to_dice_mode(debug_level: DebugLevel) -> DiceMode {
+    match debug_level {
+        DebugLevel::None => DiceMode::kDiceModeNormal,
+        DebugLevel::Full => DiceMode::kDiceModeDebug,
+    }
+}
+
+fn to_dice_hash(verified_boot_data: &VerifiedBootData) -> dice::Result<dice::Hash> {
+    let mut digests = [0u8; size_of::<Digest>() * 2];
+    digests[..size_of::<Digest>()].copy_from_slice(&verified_boot_data.kernel_digest);
+    if let Some(initrd_digest) = verified_boot_data.initrd_digest {
+        digests[size_of::<Digest>()..].copy_from_slice(&initrd_digest);
+    }
+    hash(&digests)
+}
 
 /// Derive the VM-specific secrets and certificate through DICE.
 pub fn derive_next_bcc(
     bcc: &Handover,
     next_bcc: &mut [u8],
-    code: &[u8],
-    debug_mode: bool,
+    verified_boot_data: &VerifiedBootData,
     authority: &[u8],
 ) -> dice::Result<usize> {
-    let code_hash = hash(code)?;
+    let code_hash = to_dice_hash(verified_boot_data)?;
     let auth_hash = hash(authority)?;
-    let mode = if debug_mode { dice::Mode::Debug } else { dice::Mode::Normal };
+    let mode = to_dice_mode(verified_boot_data.debug_level);
     let component_name = CStr::from_bytes_with_nul(b"vm_entry\0").unwrap();
     let mut config_descriptor_buffer = [0; 128];
     let config_descriptor_size = format_config_descriptor(
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 4f30902..530449c 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -15,6 +15,7 @@
 //! Low-level entry and exit points of pvmfw.
 
 use crate::config;
+use crate::debug_policy::{handle_debug_policy, DebugPolicyError};
 use crate::fdt;
 use crate::heap;
 use crate::helpers;
@@ -52,6 +53,16 @@
     SecretDerivationError,
 }
 
+impl From<DebugPolicyError> for RebootReason {
+    fn from(error: DebugPolicyError) -> Self {
+        match error {
+            DebugPolicyError::Fdt(_, _) => RebootReason::InvalidFdt,
+            DebugPolicyError::DebugPolicyFdt(_, _) => RebootReason::InvalidConfig,
+            DebugPolicyError::OverlaidFdt(_, _) => RebootReason::InternalError,
+        }
+    }
+}
+
 main!(start);
 
 /// Entry point for pVM firmware.
@@ -61,7 +72,7 @@
     // - can't access MMIO (therefore, no logging)
 
     match main_wrapper(fdt_address as usize, payload_start as usize, payload_size as usize) {
-        Ok(_) => jump_to_payload(fdt_address, payload_start),
+        Ok(entry) => jump_to_payload(fdt_address, entry.try_into().unwrap()),
         Err(_) => reboot(), // TODO(b/220071963) propagate the reason back to the host.
     }
 
@@ -178,42 +189,11 @@
     }
 }
 
-/// Applies the debug policy device tree overlay to the pVM DT.
-///
-/// # Safety
-///
-/// When an error is returned by this function, the input `Fdt` should be discarded as it may have
-/// have been partially corrupted during the overlay application process.
-unsafe fn apply_debug_policy(
-    fdt: &mut libfdt::Fdt,
-    debug_policy: &mut [u8],
-) -> Result<(), RebootReason> {
-    let overlay = libfdt::Fdt::from_mut_slice(debug_policy).map_err(|e| {
-        error!("Failed to load the debug policy overlay: {e}");
-        RebootReason::InvalidConfig
-    })?;
-
-    fdt.unpack().map_err(|e| {
-        error!("Failed to unpack DT for debug policy: {e}");
-        RebootReason::InternalError
-    })?;
-
-    let fdt = fdt.apply_overlay(overlay).map_err(|e| {
-        error!("Failed to apply the debug policy overlay: {e}");
-        RebootReason::InvalidConfig
-    })?;
-
-    fdt.pack().map_err(|e| {
-        error!("Failed to re-pack DT after debug policy: {e}");
-        RebootReason::InternalError
-    })
-}
-
 /// Sets up the environment for main() and wraps its result for start().
 ///
 /// Provide the abstractions necessary for start() to abort the pVM boot and for main() to run with
 /// the assumption that its environment has been properly configured.
-fn main_wrapper(fdt: usize, payload: usize, payload_size: usize) -> Result<(), RebootReason> {
+fn main_wrapper(fdt: usize, payload: usize, payload_size: usize) -> Result<usize, RebootReason> {
     // Limitations in this function:
     // - only access MMIO once (and while) it has been mapped and configured
     // - only perform logging once the logger has been initialized
@@ -283,9 +263,12 @@
     helpers::flushed_zeroize(bcc_slice);
     helpers::flush(slices.fdt.as_slice());
 
-    if let Some(debug_policy) = appended.get_debug_policy() {
-        // SAFETY - As we `?` the result, there is no risk of re-using a bad `slices.fdt`.
-        unsafe { apply_debug_policy(slices.fdt, debug_policy) }?;
+    // SAFETY - As we `?` the result, there is no risk of using a bad `slices.fdt`.
+    unsafe {
+        handle_debug_policy(slices.fdt, appended.get_debug_policy()).map_err(|e| {
+            error!("Unexpected error when handling debug policy: {e:?}");
+            RebootReason::from(e)
+        })?;
     }
 
     info!("Expecting a bug making MMIO_GUARD_UNMAP return NOT_SUPPORTED on success");
@@ -298,7 +281,7 @@
         RebootReason::InternalError
     })?;
 
-    Ok(())
+    Ok(slices.kernel.as_ptr() as usize)
 }
 
 fn jump_to_payload(fdt_address: u64, payload_start: u64) -> ! {
diff --git a/pvmfw/src/heap.rs b/pvmfw/src/heap.rs
index e04451f..435a6ff 100644
--- a/pvmfw/src/heap.rs
+++ b/pvmfw/src/heap.rs
@@ -31,7 +31,7 @@
 static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
 
 /// 128 KiB
-const HEAP_SIZE: usize =  0x20000;
+const HEAP_SIZE: usize = 0x20000;
 static mut HEAP: [u8; HEAP_SIZE] = [0; HEAP_SIZE];
 
 pub unsafe fn init() {
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index a249e8d..be5a16a 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -20,8 +20,8 @@
 
 extern crate alloc;
 
-mod avb;
 mod config;
+mod debug_policy;
 mod dice;
 mod entry;
 mod exceptions;
@@ -38,20 +38,20 @@
 use alloc::boxed::Box;
 
 use crate::{
-    avb::PUBLIC_KEY,
     dice::derive_next_bcc,
     entry::RebootReason,
     fdt::add_dice_node,
     helpers::flush,
     helpers::GUEST_PAGE_SIZE,
     memory::MemoryTracker,
-    virtio::pci::{find_virtio_devices, map_mmio},
+    virtio::pci::{self, find_virtio_devices},
 };
 use ::dice::bcc;
 use fdtpci::{PciError, PciInfo};
 use libfdt::Fdt;
 use log::{debug, error, info, trace};
 use pvmfw_avb::verify_payload;
+use pvmfw_embedded_key::PUBLIC_KEY;
 
 const NEXT_BCC_SIZE: usize = GUEST_PAGE_SIZE;
 
@@ -76,34 +76,14 @@
     // Set up PCI bus for VirtIO devices.
     let pci_info = PciInfo::from_fdt(fdt).map_err(handle_pci_error)?;
     debug!("PCI: {:#x?}", pci_info);
-    map_mmio(&pci_info, memory)?;
-    // Safety: This is the only place where we call make_pci_root, and this main function is only
-    // called once.
-    let mut pci_root = unsafe { pci_info.make_pci_root() };
+    let mut pci_root = pci::initialise(pci_info, memory)?;
     find_virtio_devices(&mut pci_root).map_err(handle_pci_error)?;
 
-    verify_payload(signed_kernel, ramdisk, PUBLIC_KEY).map_err(|e| {
+    let verified_boot_data = verify_payload(signed_kernel, ramdisk, PUBLIC_KEY).map_err(|e| {
         error!("Failed to verify the payload: {e}");
         RebootReason::PayloadVerificationError
     })?;
 
-    let debug_mode = false; // TODO(b/256148034): Derive the DICE mode from the received initrd.
-    const HASH_SIZE: usize = 64;
-    let mut hashes = [0; HASH_SIZE * 2]; // TODO(b/256148034): Extract AvbHashDescriptor digests.
-    hashes[..HASH_SIZE].copy_from_slice(&::dice::hash(signed_kernel).map_err(|_| {
-        error!("Failed to hash the kernel");
-        RebootReason::InternalError
-    })?);
-    // Note: Using signed_kernel currently makes the DICE code input depend on its VBMeta fields.
-    let code_hash = if let Some(rd) = ramdisk {
-        hashes[HASH_SIZE..].copy_from_slice(&::dice::hash(rd).map_err(|_| {
-            error!("Failed to hash the ramdisk");
-            RebootReason::InternalError
-        })?);
-        &hashes[..]
-    } else {
-        &hashes[..HASH_SIZE]
-    };
     let next_bcc = heap::aligned_boxed_slice(NEXT_BCC_SIZE, GUEST_PAGE_SIZE).ok_or_else(|| {
         error!("Failed to allocate the next-stage BCC");
         RebootReason::InternalError
@@ -111,7 +91,7 @@
     // By leaking the slice, its content will be left behind for the next stage.
     let next_bcc = Box::leak(next_bcc);
     let next_bcc_size =
-        derive_next_bcc(bcc, next_bcc, code_hash, debug_mode, PUBLIC_KEY).map_err(|e| {
+        derive_next_bcc(bcc, next_bcc, &verified_boot_data, PUBLIC_KEY).map_err(|e| {
             error!("Failed to derive next-stage DICE secrets: {e:?}");
             RebootReason::SecretDerivationError
         })?;
diff --git a/pvmfw/src/virtio/hal.rs b/pvmfw/src/virtio/hal.rs
index c6c7a99..5f70b33 100644
--- a/pvmfw/src/virtio/hal.rs
+++ b/pvmfw/src/virtio/hal.rs
@@ -1,5 +1,9 @@
+use super::pci::PCI_INFO;
 use crate::memory::{alloc_shared, dealloc_shared, phys_to_virt, virt_to_phys};
-use core::ptr::{copy_nonoverlapping, NonNull};
+use core::{
+    ops::Range,
+    ptr::{copy_nonoverlapping, NonNull},
+};
 use log::debug;
 use virtio_drivers::{BufferDirection, Hal, PhysAddr, PAGE_SIZE};
 
@@ -26,7 +30,21 @@
         0
     }
 
-    fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
+    fn mmio_phys_to_virt(paddr: PhysAddr, size: usize) -> NonNull<u8> {
+        let pci_info = PCI_INFO.get().expect("VirtIO HAL used before PCI_INFO was initialised");
+        // Check that the region is within the PCI MMIO range that we read from the device tree. If
+        // not, the host is probably trying to do something malicious.
+        if !contains_range(
+            &pci_info.bar_range,
+            &(paddr.try_into().expect("PCI MMIO region start was outside of 32-bit address space")
+                ..paddr
+                    .checked_add(size)
+                    .expect("PCI MMIO region end overflowed")
+                    .try_into()
+                    .expect("PCI MMIO region end was outside of 32-bit address space")),
+        ) {
+            panic!("PCI MMIO region was outside of expected BAR range.");
+        }
         phys_to_virt(paddr)
     }
 
@@ -68,3 +86,8 @@
         }
     }
 }
+
+/// Returns true if `inner` is entirely contained within `outer`.
+fn contains_range(outer: &Range<u32>, inner: &Range<u32>) -> bool {
+    inner.start >= outer.start && inner.end <= outer.end
+}
diff --git a/pvmfw/src/virtio/pci.rs b/pvmfw/src/virtio/pci.rs
index f9d36c6..d3b3124 100644
--- a/pvmfw/src/virtio/pci.rs
+++ b/pvmfw/src/virtio/pci.rs
@@ -16,8 +16,10 @@
 
 use super::hal::HalImpl;
 use crate::{entry::RebootReason, memory::MemoryTracker};
+use alloc::boxed::Box;
 use fdtpci::{PciError, PciInfo};
 use log::{debug, error, info};
+use once_cell::race::OnceBox;
 use virtio_drivers::{
     device::blk::VirtIOBlk,
     transport::{
@@ -26,8 +28,29 @@
     },
 };
 
+pub(super) static PCI_INFO: OnceBox<PciInfo> = OnceBox::new();
+
+/// Prepares to use VirtIO PCI devices.
+///
+/// In particular:
+///
+/// 1. Maps the PCI CAM and BAR range in the page table and MMIO guard.
+/// 2. Stores the `PciInfo` for the VirtIO HAL to use later.
+/// 3. Creates and returns a `PciRoot`.
+///
+/// This must only be called once; it will panic if it is called a second time.
+pub fn initialise(pci_info: PciInfo, memory: &mut MemoryTracker) -> Result<PciRoot, RebootReason> {
+    map_mmio(&pci_info, memory)?;
+
+    PCI_INFO.set(Box::new(pci_info.clone())).expect("Tried to set PCI_INFO a second time");
+
+    // Safety: This is the only place where we call make_pci_root, and `PCI_INFO.set` above will
+    // panic if it is called a second time.
+    Ok(unsafe { pci_info.make_pci_root() })
+}
+
 /// Maps the CAM and BAR range in the page table and MMIO guard.
-pub fn map_mmio(pci_info: &PciInfo, memory: &mut MemoryTracker) -> Result<(), RebootReason> {
+fn map_mmio(pci_info: &PciInfo, memory: &mut MemoryTracker) -> Result<(), RebootReason> {
     memory.map_mmio_range(pci_info.cam_range.clone()).map_err(|e| {
         error!("Failed to map PCI CAM: {}", e);
         RebootReason::InternalError
diff --git a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
index 7ee1f01..4611134 100644
--- a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
@@ -55,4 +55,10 @@
 
     /* get the content of the specified file. */
     String readFromFile(String path);
+
+    /**
+     * Request the service to exit, triggering the termination of the VM. This may cause any
+     * requests in flight to fail.
+     */
+    oneway void quit();
 }
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 40114fd..f17000a 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -60,6 +60,7 @@
     private static final String TAG = "MicrodroidBenchmarks";
     private static final String METRIC_NAME_PREFIX = getMetricPrefix() + "microdroid/";
     private static final int IO_TEST_TRIAL_COUNT = 5;
+    private static final long ONE_MEBI = 1024 * 1024;
 
     @Rule public Timeout globalTimeout = Timeout.seconds(300);
 
@@ -95,7 +96,7 @@
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidIdleNativeLib.so")
                         .setDebugLevel(DEBUG_LEVEL_NONE)
-                        .setMemoryMib(mem)
+                        .setMemoryBytes(mem * ONE_MEBI)
                         .build();
 
         // returns true if succeeded at least once.
@@ -148,7 +149,7 @@
                     newVmConfigBuilder()
                             .setPayloadBinaryName("MicrodroidIdleNativeLib.so")
                             .setDebugLevel(DEBUG_LEVEL_NONE)
-                            .setMemoryMib(256)
+                            .setMemoryBytes(256 * ONE_MEBI)
                             .build();
             forceCreateNewVirtualMachine("test_vm_boot_time", normalConfig);
 
@@ -176,7 +177,7 @@
                         newVmConfigBuilder()
                                 .setPayloadBinaryName("MicrodroidIdleNativeLib.so")
                                 .setDebugLevel(DEBUG_LEVEL_NONE)
-                                .setMemoryMib(256)
+                                .setMemoryBytes(256 * ONE_MEBI)
                                 .setNumCpus(numCpus)
                                 .build();
                 forceCreateNewVirtualMachine("test_vm_boot_time_multicore", normalConfig);
@@ -212,7 +213,7 @@
                             .setPayloadBinaryName("MicrodroidIdleNativeLib.so")
                             .setDebugLevel(DEBUG_LEVEL_FULL)
                             .setVmOutputCaptured(true)
-                            .setMemoryMib(256)
+                            .setMemoryBytes(256 * ONE_MEBI)
                             .build();
             forceCreateNewVirtualMachine("test_vm_boot_time_debug", normalConfig);
 
@@ -397,7 +398,7 @@
                 newVmConfigBuilder()
                         .setPayloadConfigPath("assets/vm_config_io.json")
                         .setDebugLevel(DEBUG_LEVEL_NONE)
-                        .setMemoryMib(256)
+                        .setMemoryBytes(256 * ONE_MEBI)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine(vmName, config);
         MemoryUsageListener listener = new MemoryUsageListener(this::executeCommand);
@@ -469,7 +470,7 @@
                 newVmConfigBuilder()
                         .setPayloadConfigPath("assets/vm_config_io.json")
                         .setDebugLevel(DEBUG_LEVEL_NONE)
-                        .setMemoryMib(256)
+                        .setMemoryBytes(256 * ONE_MEBI)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine(vmName, config);
         MemoryReclaimListener listener = new MemoryReclaimListener(this::executeCommand);
diff --git a/tests/hostside/Android.bp b/tests/hostside/Android.bp
index 7679c57..6e0cf5a 100644
--- a/tests/hostside/Android.bp
+++ b/tests/hostside/Android.bp
@@ -29,6 +29,7 @@
         // For re-sign test
         "avbtool",
         "img2simg",
+        "initrd_bootconfig",
         "lpmake",
         "lpunpack",
         "mk_payload",
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 87129bb..0623ff2 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -55,7 +55,6 @@
 import org.json.JSONObject;
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestName;
@@ -173,7 +172,11 @@
                 .isSuccess();
     }
 
-    private void resignVirtApex(File virtApexDir, File signingKey, Map<String, File> keyOverrides) {
+    private void resignVirtApex(
+            File virtApexDir,
+            File signingKey,
+            Map<String, File> keyOverrides,
+            boolean updateBootconfigs) {
         File signVirtApex = findTestFile("sign_virt_apex");
 
         RunUtil runUtil = new RunUtil();
@@ -184,6 +187,9 @@
 
         List<String> command = new ArrayList<>();
         command.add(signVirtApex.getAbsolutePath());
+        if (!updateBootconfigs) {
+            command.add("--do_not_update_bootconfigs");
+        }
         keyOverrides.forEach(
                 (filename, keyFile) ->
                         command.add("--key_override " + filename + "=" + keyFile.getPath()));
@@ -271,7 +277,11 @@
     }
 
     private VmInfo runMicrodroidWithResignedImages(
-            File key, Map<String, File> keyOverrides, boolean isProtected) throws Exception {
+            File key,
+            Map<String, File> keyOverrides,
+            boolean isProtected,
+            boolean updateBootconfigs)
+            throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
 
         File virtApexDir = FileUtil.createTempDir("virt_apex");
@@ -284,7 +294,7 @@
         assertWithMessage("Failed to pull " + VIRT_APEX + "etc")
                 .that(getDevice().pullDir(VIRT_APEX + "etc", virtApexEtcDir)).isTrue();
 
-        resignVirtApex(virtApexDir, key, keyOverrides);
+        resignVirtApex(virtApexDir, key, keyOverrides, updateBootconfigs);
 
         // Push back re-signed virt APEX contents and updated microdroid.json
         getDevice().pushDir(virtApexDir, TEST_ROOT);
@@ -460,7 +470,8 @@
 
         // Act
         VmInfo vmInfo =
-                runMicrodroidWithResignedImages(key, /*keyOverrides=*/ Map.of(), protectedVm);
+                runMicrodroidWithResignedImages(
+                        key, /*keyOverrides=*/ Map.of(), protectedVm, /*updateBootconfigs=*/ true);
 
         // Assert
         vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
@@ -471,16 +482,15 @@
         vmInfo.mProcess.destroy();
     }
 
-    // TODO(b/245277660): Resigning the system/vendor image changes the vbmeta hash.
-    // So, unless vbmeta related bootconfigs are updated the following test will fail
     @Test
-    @Ignore("b/245277660")
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
     public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
-        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, /*isProtected=*/ false);
+        VmInfo vmInfo =
+                runMicrodroidWithResignedImages(
+                        key, keyOverrides, /*isProtected=*/ false, /*updateBootconfigs=*/ true);
         // Device online means that boot must have succeeded.
         adbConnectToMicrodroid(getDevice(), vmInfo.mCid);
         vmInfo.mProcess.destroy();
@@ -491,15 +501,15 @@
     public void testBootFailsWhenVbMetaDigestDoesNotMatchBootconfig() throws Exception {
         // Sign everything with key1 except vbmeta
         File key = findTestFile("test.com.android.virt.pem");
-        File key2 = findTestFile("test2.com.android.virt.pem");
-        Map<String, File> keyOverrides = Map.of("microdroid_vbmeta.img", key2);
         // To be able to stop it, it should be a daemon.
-        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, /*isProtected=*/ false);
+        VmInfo vmInfo =
+                runMicrodroidWithResignedImages(
+                        key, Map.of(), /*isProtected=*/ false, /*updateBootconfigs=*/ false);
         // Wait so that init can print errors to console (time in cuttlefish >> in real device)
         assertThatEventually(
                 100000,
                 () -> getDevice().pullFileContents(CONSOLE_PATH),
-                containsString("init: [libfs_avb]Failed to verify vbmeta digest"));
+                containsString("init: [libfs_avb] Failed to verify vbmeta digest"));
         vmInfo.mProcess.destroy();
     }
 
@@ -839,6 +849,64 @@
         assertThat(ret).contains("Payload binary name must not specify a path");
     }
 
+    @Test
+    @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
+    public void testAllVbmetaUseSHA256() throws Exception {
+        File virtApexDir = FileUtil.createTempDir("virt_apex");
+        // Pull the virt apex's etc/ directory (which contains images)
+        File virtApexEtcDir = new File(virtApexDir, "etc");
+        // We need only etc/ directory for images
+        assertWithMessage("Failed to mkdir " + virtApexEtcDir)
+                .that(virtApexEtcDir.mkdirs())
+                .isTrue();
+        assertWithMessage("Failed to pull " + VIRT_APEX + "etc")
+                .that(getDevice().pullDir(VIRT_APEX + "etc", virtApexEtcDir))
+                .isTrue();
+
+        checkHashAlgorithm(virtApexEtcDir);
+    }
+
+    private String avbInfo(String image_path) throws Exception {
+        File avbtool = findTestFile("avbtool");
+        List<String> command =
+                Arrays.asList(avbtool.getAbsolutePath(), "info_image", "--image", image_path);
+        CommandResult result =
+                new RunUtil().runTimedCmd(5000, "/bin/bash", "-c", String.join(" ", command));
+        String out = result.getStdout();
+        String err = result.getStderr();
+        assertWithMessage(
+                        "Command "
+                                + command
+                                + " failed."
+                                + ":\n\tout: "
+                                + out
+                                + "\n\terr: "
+                                + err
+                                + "\n")
+                .about(command_results())
+                .that(result)
+                .isSuccess();
+        return out;
+    }
+
+    private void checkHashAlgorithm(File virtApexEtcDir) throws Exception {
+        List<String> images =
+                Arrays.asList(
+                        // kernel image (contains descriptors from initrd(s) as well)
+                        "/fs/microdroid_kernel",
+                        // vbmeta partition (contains descriptors from vendor/system images)
+                        "/fs/microdroid_vbmeta.img");
+
+        for (String path : images) {
+            String info = avbInfo(virtApexEtcDir + path);
+            Pattern pattern = Pattern.compile("Hash Algorithm:[ ]*(sha1|sha256)");
+            Matcher m = pattern.matcher(info);
+            while (m.find()) {
+                assertThat(m.group(1)).isEqualTo("sha256");
+            }
+        }
+    }
+
     @Before
     public void setUp() throws Exception {
         testIfDeviceIsCapable(getDevice());
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index e3c9961..5f9b915 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -24,6 +24,7 @@
         "MicrodroidTestNativeLib",
         "MicrodroidIdleNativeLib",
         "MicrodroidEmptyNativeLib",
+        "MicrodroidExitNativeLib",
         "MicrodroidPrivateLinkingNativeLib",
     ],
     jni_uses_platform_apis: true,
@@ -73,6 +74,14 @@
     stl: "none",
 }
 
+// A payload that exits immediately on start
+cc_library_shared {
+    name: "MicrodroidExitNativeLib",
+    srcs: ["src/native/exitbinary.cpp"],
+    header_libs: ["vm_payload_headers"],
+    stl: "libc++_static",
+}
+
 // A payload which tries to link against libselinux, one of private libraries
 cc_library_shared {
     name: "MicrodroidPrivateLinkingNativeLib",
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 7bd5f08..9cafd68 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -119,8 +119,10 @@
         revokePermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
     }
 
-    private static final int MIN_MEM_ARM64 = 150;
-    private static final int MIN_MEM_X86_64 = 196;
+    private static final long ONE_MEBI = 1024 * 1024;
+
+    private static final long MIN_MEM_ARM64 = 150 * ONE_MEBI;
+    private static final long MIN_MEM_X86_64 = 196 * ONE_MEBI;
     private static final String EXAMPLE_STRING = "Literally any string!! :)";
 
     @Test
@@ -131,7 +133,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
@@ -164,7 +166,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_NONE)
                         .setVmOutputCaptured(false)
                         .build();
@@ -195,7 +197,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .build();
 
         SecurityException e =
@@ -214,7 +216,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
 
@@ -245,7 +247,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
 
@@ -294,7 +296,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_vsock", config);
@@ -345,13 +347,13 @@
 
         assertThat(minimal.getApkPath()).isNull();
         assertThat(minimal.getDebugLevel()).isEqualTo(DEBUG_LEVEL_NONE);
-        assertThat(minimal.getMemoryMib()).isEqualTo(0);
+        assertThat(minimal.getMemoryBytes()).isEqualTo(0);
         assertThat(minimal.getNumCpus()).isEqualTo(1);
         assertThat(minimal.getPayloadBinaryName()).isEqualTo("binary.so");
         assertThat(minimal.getPayloadConfigPath()).isNull();
         assertThat(minimal.isProtectedVm()).isEqualTo(isProtectedVm());
         assertThat(minimal.isEncryptedStorageEnabled()).isFalse();
-        assertThat(minimal.getEncryptedStorageKib()).isEqualTo(0);
+        assertThat(minimal.getEncryptedStorageBytes()).isEqualTo(0);
         assertThat(minimal.isVmOutputCaptured()).isEqualTo(false);
 
         // Maximal has everything that can be set to some non-default value. (And has different
@@ -364,20 +366,20 @@
                         .setApkPath("/apk/path")
                         .setNumCpus(maxCpus)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
-                        .setMemoryMib(42)
-                        .setEncryptedStorageKib(1024)
+                        .setMemoryBytes(42)
+                        .setEncryptedStorageBytes(1_000_000)
                         .setVmOutputCaptured(true);
         VirtualMachineConfig maximal = maximalBuilder.build();
 
         assertThat(maximal.getApkPath()).isEqualTo("/apk/path");
         assertThat(maximal.getDebugLevel()).isEqualTo(DEBUG_LEVEL_FULL);
-        assertThat(maximal.getMemoryMib()).isEqualTo(42);
+        assertThat(maximal.getMemoryBytes()).isEqualTo(42);
         assertThat(maximal.getNumCpus()).isEqualTo(maxCpus);
         assertThat(maximal.getPayloadBinaryName()).isNull();
         assertThat(maximal.getPayloadConfigPath()).isEqualTo("config/path");
         assertThat(maximal.isProtectedVm()).isEqualTo(isProtectedVm());
         assertThat(maximal.isEncryptedStorageEnabled()).isTrue();
-        assertThat(maximal.getEncryptedStorageKib()).isEqualTo(1024);
+        assertThat(maximal.getEncryptedStorageBytes()).isEqualTo(1_000_000);
         assertThat(maximal.isVmOutputCaptured()).isEqualTo(true);
 
         assertThat(minimal.isCompatibleWith(maximal)).isFalse();
@@ -403,9 +405,9 @@
         assertThrows(
                 IllegalArgumentException.class, () -> builder.setPayloadBinaryName("dir/file.so"));
         assertThrows(IllegalArgumentException.class, () -> builder.setDebugLevel(-1));
-        assertThrows(IllegalArgumentException.class, () -> builder.setMemoryMib(0));
+        assertThrows(IllegalArgumentException.class, () -> builder.setMemoryBytes(0));
         assertThrows(IllegalArgumentException.class, () -> builder.setNumCpus(0));
-        assertThrows(IllegalArgumentException.class, () -> builder.setEncryptedStorageKib(0));
+        assertThrows(IllegalArgumentException.class, () -> builder.setEncryptedStorageBytes(0));
 
         // Consistency checks enforced at build time.
         Exception e;
@@ -437,7 +439,7 @@
         assertConfigCompatible(baseline, newBaselineBuilder()).isTrue();
 
         // Changes that must always be compatible
-        assertConfigCompatible(baseline, newBaselineBuilder().setMemoryMib(99)).isTrue();
+        assertConfigCompatible(baseline, newBaselineBuilder().setMemoryBytes(99)).isTrue();
         if (maxCpus > 1) {
             assertConfigCompatible(baseline, newBaselineBuilder().setNumCpus(2)).isTrue();
         }
@@ -457,7 +459,7 @@
         // Changes that are currently incompatible for ease of implementation, but this might change
         // in the future.
         assertConfigCompatible(baseline, newBaselineBuilder().setApkPath("/different")).isFalse();
-        assertConfigCompatible(baseline, newBaselineBuilder().setEncryptedStorageKib(100))
+        assertConfigCompatible(baseline, newBaselineBuilder().setEncryptedStorageBytes(100_000))
                 .isFalse();
 
         VirtualMachineConfig.Builder debuggableBuilder =
@@ -506,14 +508,14 @@
 
         assertThat(vm.getName()).isEqualTo("vm_name");
         assertThat(vm.getConfig().getPayloadBinaryName()).isEqualTo("binary.so");
-        assertThat(vm.getConfig().getMemoryMib()).isEqualTo(0);
+        assertThat(vm.getConfig().getMemoryBytes()).isEqualTo(0);
 
-        VirtualMachineConfig compatibleConfig = builder.setMemoryMib(42).build();
+        VirtualMachineConfig compatibleConfig = builder.setMemoryBytes(42).build();
         vm.setConfig(compatibleConfig);
 
         assertThat(vm.getName()).isEqualTo("vm_name");
         assertThat(vm.getConfig().getPayloadBinaryName()).isEqualTo("binary.so");
-        assertThat(vm.getConfig().getMemoryMib()).isEqualTo(42);
+        assertThat(vm.getConfig().getMemoryBytes()).isEqualTo(42);
 
         assertThat(getVirtualMachineManager().get("vm_name")).isSameInstanceAs(vm);
     }
@@ -526,7 +528,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
 
@@ -624,7 +626,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadConfigPath("assets/vm_config.json")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .build();
 
         VirtualMachine vm =
@@ -646,7 +648,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .build();
 
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_delete", config);
@@ -664,6 +666,40 @@
     }
 
     @Test
+    @CddTest(
+            requirements = {
+                "9.17/C-1-1",
+            })
+    public void deleteVmFiles() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadBinaryName("MicrodroidExitNativeLib.so")
+                        .setMemoryBytes(minMemoryRequired())
+                        .build();
+
+        VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_delete", config);
+        vm.run();
+        // If we explicitly stop a VM, that triggers some tidy up; so for this test we start a VM
+        // that immediately stops itself.
+        while (vm.getStatus() == STATUS_RUNNING) {
+            Thread.sleep(100);
+        }
+
+        // Delete the files without telling VMM. This isn't a good idea, but we can't stop an
+        // app doing it, and we should recover from it.
+        for (File f : vm.getRootDir().listFiles()) {
+            Files.delete(f.toPath());
+        }
+        vm.getRootDir().delete();
+
+        VirtualMachineManager vmm = getVirtualMachineManager();
+        assertThat(vmm.get("test_vm_delete")).isNull();
+        assertThat(vm.getStatus()).isEqualTo(STATUS_DELETED);
+    }
+
+    @Test
     @CddTest(requirements = {
             "9.17/C-1-1",
     })
@@ -674,7 +710,7 @@
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
                         .setApkPath(getContext().getPackageCodePath())
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
 
@@ -710,7 +746,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadConfigPath("assets/vm_config_extra_apk.json")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_extra_apk", config);
@@ -731,7 +767,7 @@
             VirtualMachineConfig lowMemConfig =
                     newVmConfigBuilder()
                             .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                            .setMemoryMib(memMib)
+                            .setMemoryBytes(memMib)
                             .setDebugLevel(DEBUG_LEVEL_NONE)
                             .setVmOutputCaptured(false)
                             .build();
@@ -1196,7 +1232,7 @@
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
                         .setDebugLevel(DEBUG_LEVEL_FULL);
         if (encryptedStoreEnabled) {
-            builder.setEncryptedStorageKib(4096);
+            builder.setEncryptedStorageBytes(4_000_000);
         }
         VirtualMachineConfig config = builder.build();
         String vmNameOrig = "test_vm_orig";
@@ -1250,8 +1286,8 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
-                        .setEncryptedStorageKib(4096)
+                        .setMemoryBytes(minMemoryRequired())
+                        .setEncryptedStorageBytes(4_000_000)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
@@ -1273,7 +1309,7 @@
         final VirtualMachineConfig vmConfig =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         final VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_caps", vmConfig);
@@ -1297,8 +1333,8 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
-                        .setEncryptedStorageKib(4096)
+                        .setMemoryBytes(minMemoryRequired())
+                        .setEncryptedStorageBytes(4_000_000)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_a", config);
@@ -1332,7 +1368,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_read_from_assets", config);
@@ -1448,7 +1484,7 @@
         assertThat(e).hasMessageThat().contains(expectedContents);
     }
 
-    private int minMemoryRequired() {
+    private long minMemoryRequired() {
         if (Build.SUPPORTED_ABIS.length > 0) {
             String primaryAbi = Build.SUPPORTED_ABIS[0];
             switch (primaryAbi) {
@@ -1484,15 +1520,34 @@
             throws Exception {
         CompletableFuture<Boolean> payloadStarted = new CompletableFuture<>();
         CompletableFuture<Boolean> payloadReady = new CompletableFuture<>();
+        CompletableFuture<Boolean> payloadFinished = new CompletableFuture<>();
         TestResults testResults = new TestResults();
         VmEventListener listener =
                 new VmEventListener() {
-                    private void testVMService(VirtualMachine vm) {
+                    ITestService mTestService = null;
+
+                    private void initializeTestService(VirtualMachine vm) {
                         try {
-                            ITestService testService =
+                            mTestService =
                                     ITestService.Stub.asInterface(
                                             vm.connectToVsockServer(ITestService.SERVICE_PORT));
-                            testsToRun.runTests(testService, testResults);
+                        } catch (Exception e) {
+                            testResults.mException = e;
+                        }
+                    }
+
+                    private void testVMService(VirtualMachine vm) {
+                        try {
+                            if (mTestService == null) initializeTestService(vm);
+                            testsToRun.runTests(mTestService, testResults);
+                        } catch (Exception e) {
+                            testResults.mException = e;
+                        }
+                    }
+
+                    private void quitVMService(VirtualMachine vm) {
+                        try {
+                            mTestService.quit();
                         } catch (Exception e) {
                             testResults.mException = e;
                         }
@@ -1503,7 +1558,7 @@
                         Log.i(TAG, "onPayloadReady");
                         payloadReady.complete(true);
                         testVMService(vm);
-                        forceStop(vm);
+                        quitVMService(vm);
                     }
 
                     @Override
@@ -1511,10 +1566,19 @@
                         Log.i(TAG, "onPayloadStarted");
                         payloadStarted.complete(true);
                     }
+
+                    @Override
+                    public void onPayloadFinished(VirtualMachine vm, int exitCode) {
+                        Log.i(TAG, "onPayloadFinished: " + exitCode);
+                        payloadFinished.complete(true);
+                        forceStop(vm);
+                    }
                 };
+
         listener.runToFinish(TAG, vm);
         assertThat(payloadStarted.getNow(false)).isTrue();
         assertThat(payloadReady.getNow(false)).isTrue();
+        assertThat(payloadFinished.getNow(false)).isTrue();
         return testResults;
     }
 
diff --git a/tests/testapk/src/native/exitbinary.cpp b/tests/testapk/src/native/exitbinary.cpp
new file mode 100644
index 0000000..4b70727
--- /dev/null
+++ b/tests/testapk/src/native/exitbinary.cpp
@@ -0,0 +1,21 @@
+/*
+ * Copyright 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <vm_main.h>
+
+extern "C" int AVmPayload_main() {
+    return 0;
+}
diff --git a/tests/testapk/src/native/idlebinary.cpp b/tests/testapk/src/native/idlebinary.cpp
index 9499d94..366120c 100644
--- a/tests/testapk/src/native/idlebinary.cpp
+++ b/tests/testapk/src/native/idlebinary.cpp
@@ -20,6 +20,6 @@
 extern "C" int AVmPayload_main() {
     // do nothing; just leave it alive. good night.
     for (;;) {
-        sleep(1000);
+        pause();
     }
 }
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index 4ba502a..365ea75 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -240,9 +240,6 @@
                 return ScopedAStatus::fromExceptionCodeWithMessage(EX_SERVICE_SPECIFIC,
                                                                    msg.c_str());
             }
-            // TODO(b/264520098): Remove sync() once TestService supports quit() method
-            // and Microdroid manager flushes filesystem caches on shutdown.
-            sync();
             return ScopedAStatus::ok();
         }
 
@@ -255,6 +252,8 @@
             }
             return ScopedAStatus::ok();
         }
+
+        ScopedAStatus quit() override { exit(0); }
     };
     auto testService = ndk::SharedRefBase::make<TestService>();
 
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index c827c2e..cab7c71 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -90,6 +90,9 @@
 
 const UNFORMATTED_STORAGE_MAGIC: &str = "UNFORMATTED-STORAGE";
 
+/// crosvm requires all partitions to be a multiple of 4KiB.
+const PARTITION_GRANULARITY_BYTES: u64 = 4096;
+
 lazy_static! {
     pub static ref GLOBAL_SERVICE: Strong<dyn IVirtualizationServiceInternal> =
         wait_for_interface(BINDER_SERVICE_IDENTIFIER)
@@ -172,16 +175,17 @@
     fn initializeWritablePartition(
         &self,
         image_fd: &ParcelFileDescriptor,
-        size: i64,
+        size_bytes: i64,
         partition_type: PartitionType,
     ) -> binder::Result<()> {
         check_manage_access()?;
-        let size = size.try_into().map_err(|e| {
+        let size_bytes = size_bytes.try_into().map_err(|e| {
             Status::new_exception_str(
                 ExceptionCode::ILLEGAL_ARGUMENT,
-                Some(format!("Invalid size {}: {:?}", size, e)),
+                Some(format!("Invalid size {}: {:?}", size_bytes, e)),
             )
         })?;
+        let size_bytes = round_up(size_bytes, PARTITION_GRANULARITY_BYTES);
         let image = clone_file(image_fd)?;
         // initialize the file. Any data in the file will be erased.
         image.set_len(0).map_err(|e| {
@@ -190,7 +194,7 @@
                 Some(format!("Failed to reset a file: {:?}", e)),
             )
         })?;
-        let mut part = QcowFile::new(image, size).map_err(|e| {
+        let mut part = QcowFile::new(image, size_bytes).map_err(|e| {
             Status::new_service_specific_error_str(
                 -1,
                 Some(format!("Failed to create QCOW2 image: {:?}", e)),
@@ -252,7 +256,7 @@
         const NUM_ATTEMPTS: usize = 5;
 
         for _ in 0..NUM_ATTEMPTS {
-            let vm_context = GLOBAL_SERVICE.allocateGlobalVmContext(requester_debug_pid as i32)?;
+            let vm_context = GLOBAL_SERVICE.allocateGlobalVmContext(requester_debug_pid)?;
             let cid = vm_context.getCid()? as Cid;
             let temp_dir: PathBuf = vm_context.getTemporaryDirectory()?.into();
             let service = VirtualMachineService::new_binder(self.state.clone(), cid).as_binder();
@@ -460,6 +464,15 @@
     File::create(ramdump_path).context(format!("Failed to create ramdump file {:?}", &ramdump_path))
 }
 
+fn round_up(input: u64, granularity: u64) -> u64 {
+    if granularity == 0 {
+        return input;
+    }
+    // If the input is absurdly large we round down instead of up; it's going to fail anyway.
+    let result = input.checked_add(granularity - 1).unwrap_or(input);
+    (result / granularity) * granularity
+}
+
 /// Given the configuration for a disk image, assembles the `DiskFile` to pass to crosvm.
 ///
 /// This may involve assembling a composite disk from a set of partition images.
diff --git a/virtualizationmanager/src/crosvm.rs b/virtualizationmanager/src/crosvm.rs
index c1ac20f..8f88daf 100644
--- a/virtualizationmanager/src/crosvm.rs
+++ b/virtualizationmanager/src/crosvm.rs
@@ -78,9 +78,9 @@
     /// triggered.
     static ref BOOT_HANGUP_TIMEOUT: Duration = if nested_virt::is_nested_virtualization().unwrap() {
         // Nested virtualization is slow, so we need a longer timeout.
-        Duration::from_secs(100)
+        Duration::from_secs(300)
     } else {
-        Duration::from_secs(10)
+        Duration::from_secs(30)
     };
 }
 
@@ -575,7 +575,7 @@
 
     let guest_time_ticks = data_list[42].parse::<i64>()?;
     // SAFETY : It just returns an integer about CPU tick information.
-    let ticks_per_sec = unsafe { sysconf(_SC_CLK_TCK) } as i64;
+    let ticks_per_sec = unsafe { sysconf(_SC_CLK_TCK) };
     Ok(guest_time_ticks * MILLIS_PER_SEC / ticks_per_sec)
 }
 
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
index fc4c9e7..d72d5ac 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
@@ -37,7 +37,7 @@
      * The file must be open with both read and write permissions, and should be a new empty file.
      */
     void initializeWritablePartition(
-            in ParcelFileDescriptor imageFd, long size, PartitionType type);
+            in ParcelFileDescriptor imageFd, long sizeBytes, PartitionType type);
 
     /**
      * Create or update an idsig file that digests the given APK file. The idsig file follows the
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 43b7616..e0b78ba 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -148,7 +148,7 @@
                 cid: vm.cid as i32,
                 temporaryDirectory: vm.get_temp_dir().to_string_lossy().to_string(),
                 requesterUid: vm.requester_uid as i32,
-                requesterPid: vm.requester_debug_pid as i32,
+                requesterPid: vm.requester_debug_pid,
             })
             .collect();
         Ok(cids)
diff --git a/vmbase/common.h b/vmbase/common.h
index 788dcf0..0f73b36 100644
--- a/vmbase/common.h
+++ b/vmbase/common.h
@@ -30,9 +30,13 @@
 	movk \reg, :abs_g0_nc:\imm
 .endm
 
+.macro hvc_call func_id:req
+    mov_i x0, \func_id
+    hvc 0
+.endm
+
 .macro reset_or_hang
-	mov_i x0, PSCI_SYSTEM_RESET
-	hvc 0
+	hvc_call PSCI_SYSTEM_RESET
 999:	wfi
 	b 999b
 .endm
diff --git a/vmbase/entry.S b/vmbase/entry.S
index 5f0a2ce..ab46465 100644
--- a/vmbase/entry.S
+++ b/vmbase/entry.S
@@ -63,13 +63,71 @@
 .set .Lsctlrval, .L_SCTLR_ELx_M | .L_SCTLR_ELx_C | .L_SCTLR_ELx_SA | .L_SCTLR_EL1_ITD | .L_SCTLR_EL1_SED
 .set .Lsctlrval, .Lsctlrval | .L_SCTLR_ELx_I | .L_SCTLR_EL1_SPAN | .L_SCTLR_EL1_RES1 | .L_SCTLR_EL1_WXN
 
+/* SMC function IDs */
+.set .L_SMCCC_VERSION_ID, 0x80000000
+.set .L_SMCCC_TRNG_VERSION_ID, 0x84000050
+.set .L_SMCCC_TRNG_FEATURES_ID, 0x84000051
+.set .L_SMCCC_TRNG_RND64_ID, 0xc4000053
+
+/* SMC function versions */
+.set .L_SMCCC_VERSION_1_1, 0x0101
+.set .L_SMCCC_TRNG_VERSION_1_0, 0x0100
+
 /* Bionic-compatible stack protector */
 .section .data.stack_protector, "aw"
 __bionic_tls:
 	.zero	40
 .global __stack_chk_guard
 __stack_chk_guard:
-	.quad	0x23d6d3f3c3b84098 	/* TODO: randomize */
+	.quad	0
+
+/**
+ * This macro stores a random value into a register.
+ * If a TRNG backed is not present or if an error occurs, the value remains unchanged.
+ */
+.macro rnd_reg reg:req
+	mov x20, x0
+	mov x21, x1
+	mov x22, x2
+	mov x23, x3
+
+	/* Verify SMCCC version >=1.1 */
+	hvc_call .L_SMCCC_VERSION_ID
+	cmp w0, 0
+	b.lt 100f
+	cmp w0, .L_SMCCC_VERSION_1_1
+	b.lt 100f
+
+	/* Verify TRNG ABI version 1.x */
+	hvc_call .L_SMCCC_TRNG_VERSION_ID
+	cmp w0, 0
+	b.lt 100f
+	cmp w0, .L_SMCCC_TRNG_VERSION_1_0
+	b.lt 100f
+
+	/* Call TRNG_FEATURES, ensure TRNG_RND is implemented */
+	mov_i x1, .L_SMCCC_TRNG_RND64_ID
+	hvc_call .L_SMCCC_TRNG_FEATURES_ID
+	cmp w0, 0
+	b.lt 100f
+
+	/* Call TRNG_RND, request 64 bits of entropy */
+	mov x1, #64
+	hvc_call .L_SMCCC_TRNG_RND64_ID
+	cmp x0, 0
+	b.lt 100f
+
+	mov \reg, x3
+	b 101f
+
+100:
+	reset_or_hang
+101:
+	mov x0, x20
+	mov x1, x21
+	mov x2, x22
+	mov x3, x23
+.endm
 
 /**
  * This is a generic entry point for an image. It carries out the operations required to prepare the
@@ -162,9 +220,15 @@
 	adr_l x30, __bionic_tls
 	msr tpidr_el0, x30
 
+	/* Randomize stack protector. */
+	rnd_reg x29
+	adr_l x30, __stack_chk_guard
+	str x29, [x30]
+
 	/* Call into Rust code. */
 	bl rust_entry
 
 	/* Loop forever waiting for interrupts. */
 4:	wfi
 	b 4b
+
diff --git a/zipfuse/src/main.rs b/zipfuse/src/main.rs
index 5e9e160..20d6fd6 100644
--- a/zipfuse/src/main.rs
+++ b/zipfuse/src/main.rs
@@ -673,7 +673,7 @@
             |root| {
                 let path = root.join("foo");
 
-                let metadata = fs::metadata(&path);
+                let metadata = fs::metadata(path);
                 assert!(metadata.is_ok());
                 let metadata = metadata.unwrap();