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 ffb2e14..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);
@@ -919,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/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/avb/Android.bp b/libs/avb/Android.bp
index 436f672..1d257bc 100644
--- a/libs/avb/Android.bp
+++ b/libs/avb/Android.bp
@@ -14,6 +14,7 @@
         "--size_t-is-usize",
         "--default-enum-style rust",
         "--allowlist-function=.*",
+        "--allowlist-var=AVB.*",
         "--use-core",
         "--raw-line=#![no_std]",
         "--ctypes-prefix=core::ffi",
diff --git a/microdroid/initrd/src/main.rs b/microdroid/initrd/src/main.rs
index 3b0a7d2..c5515af 100644
--- a/microdroid/initrd/src/main.rs
+++ b/microdroid/initrd/src/main.rs
@@ -72,9 +72,9 @@
 
 // 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 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)))?;
diff --git a/pvmfw/avb/Android.bp b/pvmfw/avb/Android.bp
index 837f747..fb950b7 100644
--- a/pvmfw/avb/Android.bp
+++ b/pvmfw/avb/Android.bp
@@ -9,6 +9,7 @@
     prefer_rlib: true,
     rustlibs: [
         "libavb_bindgen",
+        "libtinyvec_nostd",
     ],
     whole_static_libs: [
         "libavb",
@@ -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",
     ],
@@ -78,18 +80,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..b0598de
--- /dev/null
+++ b/pvmfw/avb/src/descriptor.rs
@@ -0,0 +1,204 @@
+// 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;
+
+const DIGEST_SIZE: usize = 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,
+    /// TODO(b/265897559): Pass this digest to DICE.
+    #[allow(dead_code)]
+    pub(crate) digest: [u8; DIGEST_SIZE],
+}
+
+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; DIGEST_SIZE];
+        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..065eca5 100644
--- a/pvmfw/avb/src/lib.rs
+++ b/pvmfw/avb/src/lib.rs
@@ -18,6 +18,7 @@
 // For usize.checked_add_signed(isize), available in Rust 1.66.0
 #![feature(mixed_integer_ops)]
 
+mod descriptor;
 mod error;
 mod ops;
 mod partition;
diff --git a/pvmfw/avb/src/partition.rs b/pvmfw/avb/src/partition.rs
index 82d89f8..636bfd3 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, 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..14b0e7e 100644
--- a/pvmfw/avb/src/verify.rs
+++ b/pvmfw/avb/src/verify.rs
@@ -14,19 +14,12 @@
 
 //! This module handles the pvmfw payload verification.
 
-use crate::error::{AvbIOError, AvbSlotVerifyError};
+use crate::descriptor::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;
 
 /// This enum corresponds to the `DebugLevel` in `VirtualMachineConfig`.
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -37,104 +30,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 +49,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,
@@ -191,13 +96,17 @@
     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)? };
+    // TODO(b/265897559): Pass the digest in kernel descriptor to DICE.
+    let _kernel_descriptor = hash_descriptors.find(PartitionName::Kernel)?;
 
     if initrd.is_none() {
-        verify_vbmeta_has_no_initrd_descriptor(&vbmeta_image)?;
+        verify_vbmeta_has_only_one_hash_descriptor(&hash_descriptors)?;
         return Ok(DebugLevel::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) =
diff --git a/pvmfw/avb/tests/api_test.rs b/pvmfw/avb/tests/api_test.rs
index b5300f9..4f00f1e 100644
--- a/pvmfw/avb/tests/api_test.rs
+++ b/pvmfw/avb/tests/api_test.rs
@@ -25,6 +25,8 @@
 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;
@@ -62,12 +64,22 @@
 }
 
 #[test]
-fn payload_with_non_initrd_descriptor_passes_verification_with_no_initrd() -> Result<()> {
+fn payload_with_non_initrd_descriptor_fails_verification_with_no_initrd() -> Result<()> {
     assert_payload_verification_eq(
         &fs::read(TEST_IMG_WITH_NON_INITRD_HASHDESC_PATH)?,
         /*initrd=*/ None,
         &load_trusted_public_key()?,
-        Ok(DebugLevel::None),
+        Err(AvbSlotVerifyError::InvalidMetadata),
+    )
+}
+
+#[test]
+fn payload_with_non_initrd_descriptor_fails_verification_with_initrd() -> Result<()> {
+    assert_payload_verification_with_initrd_eq(
+        &fs::read(TEST_IMG_WITH_INITRD_AND_NON_INITRD_DESC_PATH)?,
+        &load_latest_initrd_normal()?,
+        &load_trusted_public_key()?,
+        Err(AvbSlotVerifyError::InvalidMetadata),
     )
 }
 
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/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 24c36b3..eabdfe8 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -20,7 +20,6 @@
 
 extern crate alloc;
 
-mod avb;
 mod config;
 mod dice;
 mod entry;
@@ -38,7 +37,6 @@
 use alloc::boxed::Box;
 
 use crate::{
-    avb::PUBLIC_KEY,
     dice::derive_next_bcc,
     entry::RebootReason,
     fdt::add_dice_node,
@@ -52,6 +50,7 @@
 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;
 
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/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 5cd0cb1..c6915e5 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);
@@ -674,7 +676,7 @@
         VirtualMachineConfig config =
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidExitNativeLib.so")
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .build();
 
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_delete", config);
@@ -708,7 +710,7 @@
                 newVmConfigBuilder()
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
                         .setApkPath(getContext().getPackageCodePath())
-                        .setMemoryMib(minMemoryRequired())
+                        .setMemoryBytes(minMemoryRequired())
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
 
@@ -744,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);
@@ -765,7 +767,7 @@
             VirtualMachineConfig lowMemConfig =
                     newVmConfigBuilder()
                             .setPayloadBinaryName("MicrodroidTestNativeLib.so")
-                            .setMemoryMib(memMib)
+                            .setMemoryBytes(memMib)
                             .setDebugLevel(DEBUG_LEVEL_NONE)
                             .setVmOutputCaptured(false)
                             .build();
@@ -1230,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";
@@ -1284,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);
@@ -1307,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);
@@ -1331,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);
@@ -1366,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);
@@ -1482,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) {
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/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index ca42999..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)),
@@ -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/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