Merge "Better error messages"
diff --git a/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java b/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java
index c25de71..b884b9e 100644
--- a/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java
+++ b/compos/benchmark/src/java/com/android/compos/benchmark/ComposBenchmark.java
@@ -17,6 +17,7 @@
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.TruthJUnit.assume;
 
 import static org.junit.Assert.assertTrue;
@@ -136,9 +137,7 @@
             Long compileStartTime = System.nanoTime();
             String output = executeCommand(command);
             Long compileEndTime = System.nanoTime();
-            Pattern pattern = Pattern.compile("All Ok");
-            Matcher matcher = pattern.matcher(output);
-            assertTrue(matcher.find());
+            assertThat(output).containsMatch("All Ok");
             double elapsedSec = (compileEndTime - compileStartTime) / NANOS_IN_SEC;
             Log.i(TAG, "Compile time in guest took " + elapsedSec + "s");
             getMetricsRunnable.stop();
diff --git a/javalib/Android.bp b/javalib/Android.bp
index f5f91f2..a124af7 100644
--- a/javalib/Android.bp
+++ b/javalib/Android.bp
@@ -60,3 +60,30 @@
         "//build/orchestrator/apis",
     ],
 }
+
+java_api_contribution {
+    name: "framework-virtualization-system-stubs",
+    api_surface: "system",
+    api_file: "api/system-current.txt",
+    visibility: [
+        "//build/orchestrator/apis",
+    ],
+}
+
+java_api_contribution {
+    name: "framework-virtualization-test-stubs",
+    api_surface: "test",
+    api_file: "api/test-current.txt",
+    visibility: [
+        "//build/orchestrator/apis",
+    ],
+}
+
+java_api_contribution {
+    name: "framework-virtualization-module-lib-stubs",
+    api_surface: "module-lib",
+    api_file: "api/module-lib-current.txt",
+    visibility: [
+        "//build/orchestrator/apis",
+    ],
+}
diff --git a/pvmfw/avb/fuzz/Android.bp b/pvmfw/avb/fuzz/Android.bp
new file mode 100644
index 0000000..451fd8a
--- /dev/null
+++ b/pvmfw/avb/fuzz/Android.bp
@@ -0,0 +1,34 @@
+// 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.
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+rust_fuzz {
+    name: "avb_kernel_without_footer_verify_fuzzer",
+    srcs: ["without_footer_verify_fuzzer.rs"],
+    rustlibs: [
+        "libpvmfw_avb_nostd",
+    ],
+    fuzz_config: {
+        cc: [
+            "android-kvm@google.com",
+        ],
+        fuzz_on_haiku_device: true,
+        fuzz_on_haiku_host: true,
+    },
+}
+
+// TODO(b/260574387): Add avb_kernel_with_footer_verify_fuzzer
diff --git a/pvmfw/avb/fuzz/without_footer_verify_fuzzer.rs b/pvmfw/avb/fuzz/without_footer_verify_fuzzer.rs
new file mode 100644
index 0000000..fc8fa85
--- /dev/null
+++ b/pvmfw/avb/fuzz/without_footer_verify_fuzzer.rs
@@ -0,0 +1,28 @@
+// 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.
+
+#![allow(missing_docs)]
+#![no_main]
+
+use libfuzzer_sys::fuzz_target;
+use pvmfw_avb::verify_payload;
+
+fuzz_target!(|kernel: &[u8]| {
+    // This fuzzer is mostly supposed to catch the memory corruption in
+    // AVB footer parsing. It is unlikely that the randomly generated
+    // kernel can pass the kernel verification, so the value of `initrd`
+    // is not so important as we won't reach initrd verification with
+    // this fuzzer.
+    let _ = verify_payload(kernel, /*initrd=*/ None, &[0u8; 64]);
+});
diff --git a/pvmfw/avb/src/lib.rs b/pvmfw/avb/src/lib.rs
index f0ee4ed..a1e3ee0 100644
--- a/pvmfw/avb/src/lib.rs
+++ b/pvmfw/avb/src/lib.rs
@@ -19,6 +19,7 @@
 #![feature(mixed_integer_ops)]
 
 mod error;
+mod ops;
 mod partition;
 mod utils;
 mod verify;
diff --git a/pvmfw/avb/src/ops.rs b/pvmfw/avb/src/ops.rs
new file mode 100644
index 0000000..903fecb
--- /dev/null
+++ b/pvmfw/avb/src/ops.rs
@@ -0,0 +1,330 @@
+// 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 `AvbOps`.
+
+use crate::error::{
+    slot_verify_result_to_verify_payload_result, to_avb_io_result, AvbIOError, AvbSlotVerifyError,
+};
+use crate::partition::PartitionName;
+use crate::utils::{self, as_ref, is_not_null, to_nonnull, write};
+use avb_bindgen::{
+    avb_slot_verify, avb_slot_verify_data_free, AvbHashtreeErrorMode, AvbIOResult, AvbOps,
+    AvbSlotVerifyData, AvbSlotVerifyFlags, AvbVBMetaData,
+};
+use core::{
+    ffi::{c_char, c_void, CStr},
+    mem::MaybeUninit,
+    ptr, slice,
+};
+
+const NULL_BYTE: &[u8] = b"\0";
+
+pub(crate) struct Payload<'a> {
+    kernel: &'a [u8],
+    initrd: Option<&'a [u8]>,
+    trusted_public_key: &'a [u8],
+}
+
+impl<'a> AsRef<Payload<'a>> for AvbOps {
+    fn as_ref(&self) -> &Payload<'a> {
+        let payload = self.user_data as *const Payload;
+        // SAFETY: It is safe to cast the `AvbOps.user_data` to Payload as we have saved a
+        // pointer to a valid value of Payload in user_data when creating AvbOps.
+        unsafe { &*payload }
+    }
+}
+
+impl<'a> Payload<'a> {
+    pub(crate) fn new(
+        kernel: &'a [u8],
+        initrd: Option<&'a [u8]>,
+        trusted_public_key: &'a [u8],
+    ) -> Self {
+        Self { kernel, initrd, trusted_public_key }
+    }
+
+    fn get_partition(&self, partition_name: *const c_char) -> Result<&[u8], AvbIOError> {
+        match partition_name.try_into()? {
+            PartitionName::Kernel => Ok(self.kernel),
+            PartitionName::InitrdNormal | PartitionName::InitrdDebug => {
+                self.initrd.ok_or(AvbIOError::NoSuchPartition)
+            }
+        }
+    }
+}
+
+/// `Ops` wraps the class `AvbOps` in libavb. It provides pvmfw customized
+/// operations used in the verification.
+pub(crate) struct Ops(AvbOps);
+
+impl<'a> From<&mut Payload<'a>> for Ops {
+    fn from(payload: &mut Payload<'a>) -> Self {
+        let avb_ops = AvbOps {
+            user_data: payload as *mut _ as *mut c_void,
+            ab_ops: ptr::null_mut(),
+            atx_ops: ptr::null_mut(),
+            read_from_partition: Some(read_from_partition),
+            get_preloaded_partition: Some(get_preloaded_partition),
+            write_to_partition: None,
+            validate_vbmeta_public_key: Some(validate_vbmeta_public_key),
+            read_rollback_index: Some(read_rollback_index),
+            write_rollback_index: None,
+            read_is_device_unlocked: Some(read_is_device_unlocked),
+            get_unique_guid_for_partition: Some(get_unique_guid_for_partition),
+            get_size_of_partition: Some(get_size_of_partition),
+            read_persistent_value: None,
+            write_persistent_value: None,
+            validate_public_key_for_partition: None,
+        };
+        Self(avb_ops)
+    }
+}
+
+impl Ops {
+    pub(crate) fn verify_partition(
+        &mut self,
+        partition_name: &CStr,
+    ) -> Result<AvbSlotVerifyDataWrap, AvbSlotVerifyError> {
+        let requested_partitions = [partition_name.as_ptr(), ptr::null()];
+        let ab_suffix = CStr::from_bytes_with_nul(NULL_BYTE).unwrap();
+        let mut out_data = MaybeUninit::uninit();
+        // SAFETY: It is safe to call `avb_slot_verify()` as the pointer arguments (`ops`,
+        // `requested_partitions` and `ab_suffix`) passed to the method are all valid and
+        // initialized.
+        let result = unsafe {
+            avb_slot_verify(
+                &mut self.0,
+                requested_partitions.as_ptr(),
+                ab_suffix.as_ptr(),
+                AvbSlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NONE,
+                AvbHashtreeErrorMode::AVB_HASHTREE_ERROR_MODE_RESTART_AND_INVALIDATE,
+                out_data.as_mut_ptr(),
+            )
+        };
+        slot_verify_result_to_verify_payload_result(result)?;
+        // SAFETY: This is safe because `out_data` has been properly initialized after
+        // calling `avb_slot_verify` and it returns OK.
+        let out_data = unsafe { out_data.assume_init() };
+        out_data.try_into()
+    }
+}
+
+extern "C" fn read_is_device_unlocked(
+    _ops: *mut AvbOps,
+    out_is_unlocked: *mut bool,
+) -> AvbIOResult {
+    to_avb_io_result(write(out_is_unlocked, false))
+}
+
+extern "C" fn get_preloaded_partition(
+    ops: *mut AvbOps,
+    partition: *const c_char,
+    num_bytes: usize,
+    out_pointer: *mut *mut u8,
+    out_num_bytes_preloaded: *mut usize,
+) -> AvbIOResult {
+    to_avb_io_result(try_get_preloaded_partition(
+        ops,
+        partition,
+        num_bytes,
+        out_pointer,
+        out_num_bytes_preloaded,
+    ))
+}
+
+fn try_get_preloaded_partition(
+    ops: *mut AvbOps,
+    partition: *const c_char,
+    num_bytes: usize,
+    out_pointer: *mut *mut u8,
+    out_num_bytes_preloaded: *mut usize,
+) -> utils::Result<()> {
+    let ops = as_ref(ops)?;
+    let partition = ops.as_ref().get_partition(partition)?;
+    write(out_pointer, partition.as_ptr() as *mut u8)?;
+    write(out_num_bytes_preloaded, partition.len().min(num_bytes))
+}
+
+extern "C" fn read_from_partition(
+    ops: *mut AvbOps,
+    partition: *const c_char,
+    offset: i64,
+    num_bytes: usize,
+    buffer: *mut c_void,
+    out_num_read: *mut usize,
+) -> AvbIOResult {
+    to_avb_io_result(try_read_from_partition(
+        ops,
+        partition,
+        offset,
+        num_bytes,
+        buffer,
+        out_num_read,
+    ))
+}
+
+fn try_read_from_partition(
+    ops: *mut AvbOps,
+    partition: *const c_char,
+    offset: i64,
+    num_bytes: usize,
+    buffer: *mut c_void,
+    out_num_read: *mut usize,
+) -> utils::Result<()> {
+    let ops = as_ref(ops)?;
+    let partition = ops.as_ref().get_partition(partition)?;
+    let buffer = to_nonnull(buffer)?;
+    // SAFETY: It is safe to copy the requested number of bytes to `buffer` as `buffer`
+    // is created to point to the `num_bytes` of bytes in memory.
+    let buffer_slice = unsafe { slice::from_raw_parts_mut(buffer.as_ptr() as *mut u8, num_bytes) };
+    copy_data_to_dst(partition, offset, buffer_slice)?;
+    write(out_num_read, buffer_slice.len())
+}
+
+fn copy_data_to_dst(src: &[u8], offset: i64, dst: &mut [u8]) -> utils::Result<()> {
+    let start = to_copy_start(offset, src.len()).ok_or(AvbIOError::InvalidValueSize)?;
+    let end = start.checked_add(dst.len()).ok_or(AvbIOError::InvalidValueSize)?;
+    dst.copy_from_slice(src.get(start..end).ok_or(AvbIOError::RangeOutsidePartition)?);
+    Ok(())
+}
+
+fn to_copy_start(offset: i64, len: usize) -> Option<usize> {
+    usize::try_from(offset)
+        .ok()
+        .or_else(|| isize::try_from(offset).ok().and_then(|v| len.checked_add_signed(v)))
+}
+
+extern "C" fn get_size_of_partition(
+    ops: *mut AvbOps,
+    partition: *const c_char,
+    out_size_num_bytes: *mut u64,
+) -> AvbIOResult {
+    to_avb_io_result(try_get_size_of_partition(ops, partition, out_size_num_bytes))
+}
+
+fn try_get_size_of_partition(
+    ops: *mut AvbOps,
+    partition: *const c_char,
+    out_size_num_bytes: *mut u64,
+) -> utils::Result<()> {
+    let ops = as_ref(ops)?;
+    let partition = ops.as_ref().get_partition(partition)?;
+    let partition_size =
+        u64::try_from(partition.len()).map_err(|_| AvbIOError::InvalidValueSize)?;
+    write(out_size_num_bytes, partition_size)
+}
+
+extern "C" fn read_rollback_index(
+    _ops: *mut AvbOps,
+    _rollback_index_location: usize,
+    _out_rollback_index: *mut u64,
+) -> AvbIOResult {
+    // TODO(b/256148034): Write -1 to out_rollback_index
+    // so that we won't compare the current rollback index with uninitialized number
+    // in avb_slot_verify.
+
+    // Rollback protection is not yet implemented, but
+    // this method is required by `avb_slot_verify()`.
+    AvbIOResult::AVB_IO_RESULT_OK
+}
+
+extern "C" fn get_unique_guid_for_partition(
+    _ops: *mut AvbOps,
+    _partition: *const c_char,
+    _guid_buf: *mut c_char,
+    _guid_buf_size: usize,
+) -> AvbIOResult {
+    // TODO(b/256148034): Check if it's possible to throw an error here instead of having
+    // an empty method.
+    // This method is required by `avb_slot_verify()`.
+    AvbIOResult::AVB_IO_RESULT_OK
+}
+
+extern "C" fn validate_vbmeta_public_key(
+    ops: *mut AvbOps,
+    public_key_data: *const u8,
+    public_key_length: usize,
+    public_key_metadata: *const u8,
+    public_key_metadata_length: usize,
+    out_is_trusted: *mut bool,
+) -> AvbIOResult {
+    to_avb_io_result(try_validate_vbmeta_public_key(
+        ops,
+        public_key_data,
+        public_key_length,
+        public_key_metadata,
+        public_key_metadata_length,
+        out_is_trusted,
+    ))
+}
+
+fn try_validate_vbmeta_public_key(
+    ops: *mut AvbOps,
+    public_key_data: *const u8,
+    public_key_length: usize,
+    _public_key_metadata: *const u8,
+    _public_key_metadata_length: usize,
+    out_is_trusted: *mut bool,
+) -> utils::Result<()> {
+    // The public key metadata is not used when we build the VBMeta.
+    is_not_null(public_key_data)?;
+    // SAFETY: It is safe to create a slice with the given pointer and length as
+    // `public_key_data` is a valid pointer and it points to an array of length
+    // `public_key_length`.
+    let public_key = unsafe { slice::from_raw_parts(public_key_data, public_key_length) };
+    let ops = as_ref(ops)?;
+    let trusted_public_key = ops.as_ref().trusted_public_key;
+    write(out_is_trusted, public_key == trusted_public_key)
+}
+
+pub(crate) struct AvbSlotVerifyDataWrap(*mut AvbSlotVerifyData);
+
+impl TryFrom<*mut AvbSlotVerifyData> for AvbSlotVerifyDataWrap {
+    type Error = AvbSlotVerifyError;
+
+    fn try_from(data: *mut AvbSlotVerifyData) -> Result<Self, Self::Error> {
+        is_not_null(data).map_err(|_| AvbSlotVerifyError::Io)?;
+        Ok(Self(data))
+    }
+}
+
+impl Drop for AvbSlotVerifyDataWrap {
+    fn drop(&mut self) {
+        // SAFETY: This is safe because `self.0` is checked nonnull when the
+        // instance is created. We can free this pointer when the instance is
+        // no longer needed.
+        unsafe {
+            avb_slot_verify_data_free(self.0);
+        }
+    }
+}
+
+impl AsRef<AvbSlotVerifyData> for AvbSlotVerifyDataWrap {
+    fn as_ref(&self) -> &AvbSlotVerifyData {
+        // This is safe because `self.0` is checked nonnull when the instance is created.
+        as_ref(self.0).unwrap()
+    }
+}
+
+impl AvbSlotVerifyDataWrap {
+    pub(crate) fn vbmeta_images(&self) -> Result<&[AvbVBMetaData], AvbSlotVerifyError> {
+        let data = self.as_ref();
+        is_not_null(data.vbmeta_images).map_err(|_| AvbSlotVerifyError::Io)?;
+        // SAFETY: It is safe as the raw pointer `data.vbmeta_images` is a nonnull pointer.
+        let vbmeta_images =
+            unsafe { slice::from_raw_parts(data.vbmeta_images, data.num_vbmeta_images) };
+        Ok(vbmeta_images)
+    }
+}
diff --git a/pvmfw/avb/src/verify.rs b/pvmfw/avb/src/verify.rs
index b39bb5a..d6a0cb2 100644
--- a/pvmfw/avb/src/verify.rs
+++ b/pvmfw/avb/src/verify.rs
@@ -14,24 +14,20 @@
 
 //! This module handles the pvmfw payload verification.
 
-use crate::error::{
-    slot_verify_result_to_verify_payload_result, to_avb_io_result, AvbIOError, AvbSlotVerifyError,
-};
+use crate::error::{AvbIOError, AvbSlotVerifyError};
+use crate::ops::{Ops, Payload};
 use crate::partition::PartitionName;
-use crate::utils::{as_ref, is_not_null, to_nonnull, to_usize, usize_checked_add, write};
+use crate::utils::{is_not_null, to_usize, usize_checked_add, write};
 use avb_bindgen::{
-    avb_descriptor_foreach, avb_hash_descriptor_validate_and_byteswap, avb_slot_verify,
-    avb_slot_verify_data_free, AvbDescriptor, AvbHashDescriptor, AvbHashtreeErrorMode, AvbIOResult,
-    AvbOps, AvbSlotVerifyData, AvbSlotVerifyFlags, AvbVBMetaData,
+    avb_descriptor_foreach, avb_hash_descriptor_validate_and_byteswap, AvbDescriptor,
+    AvbHashDescriptor, AvbVBMetaData,
 };
 use core::{
-    ffi::{c_char, c_void, CStr},
+    ffi::{c_char, c_void},
     mem::{size_of, MaybeUninit},
-    ptr, slice,
+    slice,
 };
 
-const NULL_BYTE: &[u8] = b"\0";
-
 /// This enum corresponds to the `DebugLevel` in `VirtualMachineConfig`.
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum DebugLevel {
@@ -41,167 +37,6 @@
     Full,
 }
 
-extern "C" fn read_is_device_unlocked(
-    _ops: *mut AvbOps,
-    out_is_unlocked: *mut bool,
-) -> AvbIOResult {
-    to_avb_io_result(write(out_is_unlocked, false))
-}
-
-extern "C" fn get_preloaded_partition(
-    ops: *mut AvbOps,
-    partition: *const c_char,
-    num_bytes: usize,
-    out_pointer: *mut *mut u8,
-    out_num_bytes_preloaded: *mut usize,
-) -> AvbIOResult {
-    to_avb_io_result(try_get_preloaded_partition(
-        ops,
-        partition,
-        num_bytes,
-        out_pointer,
-        out_num_bytes_preloaded,
-    ))
-}
-
-fn try_get_preloaded_partition(
-    ops: *mut AvbOps,
-    partition: *const c_char,
-    num_bytes: usize,
-    out_pointer: *mut *mut u8,
-    out_num_bytes_preloaded: *mut usize,
-) -> Result<(), AvbIOError> {
-    let ops = as_ref(ops)?;
-    let partition = ops.as_ref().get_partition(partition)?;
-    write(out_pointer, partition.as_ptr() as *mut u8)?;
-    write(out_num_bytes_preloaded, partition.len().min(num_bytes))
-}
-
-extern "C" fn read_from_partition(
-    ops: *mut AvbOps,
-    partition: *const c_char,
-    offset: i64,
-    num_bytes: usize,
-    buffer: *mut c_void,
-    out_num_read: *mut usize,
-) -> AvbIOResult {
-    to_avb_io_result(try_read_from_partition(
-        ops,
-        partition,
-        offset,
-        num_bytes,
-        buffer,
-        out_num_read,
-    ))
-}
-
-fn try_read_from_partition(
-    ops: *mut AvbOps,
-    partition: *const c_char,
-    offset: i64,
-    num_bytes: usize,
-    buffer: *mut c_void,
-    out_num_read: *mut usize,
-) -> Result<(), AvbIOError> {
-    let ops = as_ref(ops)?;
-    let partition = ops.as_ref().get_partition(partition)?;
-    let buffer = to_nonnull(buffer)?;
-    // SAFETY: It is safe to copy the requested number of bytes to `buffer` as `buffer`
-    // is created to point to the `num_bytes` of bytes in memory.
-    let buffer_slice = unsafe { slice::from_raw_parts_mut(buffer.as_ptr() as *mut u8, num_bytes) };
-    copy_data_to_dst(partition, offset, buffer_slice)?;
-    write(out_num_read, buffer_slice.len())
-}
-
-fn copy_data_to_dst(src: &[u8], offset: i64, dst: &mut [u8]) -> Result<(), AvbIOError> {
-    let start = to_copy_start(offset, src.len()).ok_or(AvbIOError::InvalidValueSize)?;
-    let end = start.checked_add(dst.len()).ok_or(AvbIOError::InvalidValueSize)?;
-    dst.copy_from_slice(src.get(start..end).ok_or(AvbIOError::RangeOutsidePartition)?);
-    Ok(())
-}
-
-fn to_copy_start(offset: i64, len: usize) -> Option<usize> {
-    usize::try_from(offset)
-        .ok()
-        .or_else(|| isize::try_from(offset).ok().and_then(|v| len.checked_add_signed(v)))
-}
-
-extern "C" fn get_size_of_partition(
-    ops: *mut AvbOps,
-    partition: *const c_char,
-    out_size_num_bytes: *mut u64,
-) -> AvbIOResult {
-    to_avb_io_result(try_get_size_of_partition(ops, partition, out_size_num_bytes))
-}
-
-fn try_get_size_of_partition(
-    ops: *mut AvbOps,
-    partition: *const c_char,
-    out_size_num_bytes: *mut u64,
-) -> Result<(), AvbIOError> {
-    let ops = as_ref(ops)?;
-    let partition = ops.as_ref().get_partition(partition)?;
-    let partition_size =
-        u64::try_from(partition.len()).map_err(|_| AvbIOError::InvalidValueSize)?;
-    write(out_size_num_bytes, partition_size)
-}
-
-extern "C" fn read_rollback_index(
-    _ops: *mut AvbOps,
-    _rollback_index_location: usize,
-    _out_rollback_index: *mut u64,
-) -> AvbIOResult {
-    // Rollback protection is not yet implemented, but
-    // this method is required by `avb_slot_verify()`.
-    AvbIOResult::AVB_IO_RESULT_OK
-}
-
-extern "C" fn get_unique_guid_for_partition(
-    _ops: *mut AvbOps,
-    _partition: *const c_char,
-    _guid_buf: *mut c_char,
-    _guid_buf_size: usize,
-) -> AvbIOResult {
-    // This method is required by `avb_slot_verify()`.
-    AvbIOResult::AVB_IO_RESULT_OK
-}
-
-extern "C" fn validate_vbmeta_public_key(
-    ops: *mut AvbOps,
-    public_key_data: *const u8,
-    public_key_length: usize,
-    public_key_metadata: *const u8,
-    public_key_metadata_length: usize,
-    out_is_trusted: *mut bool,
-) -> AvbIOResult {
-    to_avb_io_result(try_validate_vbmeta_public_key(
-        ops,
-        public_key_data,
-        public_key_length,
-        public_key_metadata,
-        public_key_metadata_length,
-        out_is_trusted,
-    ))
-}
-
-fn try_validate_vbmeta_public_key(
-    ops: *mut AvbOps,
-    public_key_data: *const u8,
-    public_key_length: usize,
-    _public_key_metadata: *const u8,
-    _public_key_metadata_length: usize,
-    out_is_trusted: *mut bool,
-) -> Result<(), AvbIOError> {
-    is_not_null(public_key_data)?;
-    // SAFETY: It is safe to create a slice with the given pointer and length as
-    // `public_key_data` is a valid pointer and it points to an array of length
-    // `public_key_length`.
-    let public_key = unsafe { slice::from_raw_parts(public_key_data, public_key_length) };
-    let ops = as_ref(ops)?;
-    let trusted_public_key = ops.as_ref().trusted_public_key;
-    write(out_is_trusted, public_key == trusted_public_key)
-}
-
 extern "C" fn search_initrd_hash_descriptor(
     descriptor: *const AvbDescriptor,
     user_data: *mut c_void,
@@ -277,119 +112,6 @@
     }
 }
 
-struct AvbSlotVerifyDataWrap(*mut AvbSlotVerifyData);
-
-impl TryFrom<*mut AvbSlotVerifyData> for AvbSlotVerifyDataWrap {
-    type Error = AvbSlotVerifyError;
-
-    fn try_from(data: *mut AvbSlotVerifyData) -> Result<Self, Self::Error> {
-        is_not_null(data).map_err(|_| AvbSlotVerifyError::Io)?;
-        Ok(Self(data))
-    }
-}
-
-impl Drop for AvbSlotVerifyDataWrap {
-    fn drop(&mut self) {
-        // SAFETY: This is safe because `self.0` is checked nonnull when the
-        // instance is created. We can free this pointer when the instance is
-        // no longer needed.
-        unsafe {
-            avb_slot_verify_data_free(self.0);
-        }
-    }
-}
-
-impl AsRef<AvbSlotVerifyData> for AvbSlotVerifyDataWrap {
-    fn as_ref(&self) -> &AvbSlotVerifyData {
-        // This is safe because `self.0` is checked nonnull when the instance is created.
-        as_ref(self.0).unwrap()
-    }
-}
-
-impl AvbSlotVerifyDataWrap {
-    fn vbmeta_images(&self) -> Result<&[AvbVBMetaData], AvbSlotVerifyError> {
-        let data = self.as_ref();
-        is_not_null(data.vbmeta_images).map_err(|_| AvbSlotVerifyError::Io)?;
-        // SAFETY: It is safe as the raw pointer `data.vbmeta_images` is a nonnull pointer.
-        let vbmeta_images =
-            unsafe { slice::from_raw_parts(data.vbmeta_images, data.num_vbmeta_images) };
-        Ok(vbmeta_images)
-    }
-}
-
-struct Payload<'a> {
-    kernel: &'a [u8],
-    initrd: Option<&'a [u8]>,
-    trusted_public_key: &'a [u8],
-}
-
-impl<'a> AsRef<Payload<'a>> for AvbOps {
-    fn as_ref(&self) -> &Payload<'a> {
-        let payload = self.user_data as *const Payload;
-        // SAFETY: It is safe to cast the `AvbOps.user_data` to Payload as we have saved a
-        // pointer to a valid value of Payload in user_data when creating AvbOps, and
-        // assume that the Payload isn't used beyond the lifetime of the AvbOps that it
-        // belongs to.
-        unsafe { &*payload }
-    }
-}
-
-impl<'a> Payload<'a> {
-    fn get_partition(&self, partition_name: *const c_char) -> Result<&[u8], AvbIOError> {
-        match partition_name.try_into()? {
-            PartitionName::Kernel => Ok(self.kernel),
-            PartitionName::InitrdNormal | PartitionName::InitrdDebug => {
-                self.initrd.ok_or(AvbIOError::NoSuchPartition)
-            }
-        }
-    }
-
-    fn verify_partition(
-        &mut self,
-        partition_name: &CStr,
-    ) -> Result<AvbSlotVerifyDataWrap, AvbSlotVerifyError> {
-        let requested_partitions = [partition_name.as_ptr(), ptr::null()];
-        let mut avb_ops = AvbOps {
-            user_data: self as *mut _ as *mut c_void,
-            ab_ops: ptr::null_mut(),
-            atx_ops: ptr::null_mut(),
-            read_from_partition: Some(read_from_partition),
-            get_preloaded_partition: Some(get_preloaded_partition),
-            write_to_partition: None,
-            validate_vbmeta_public_key: Some(validate_vbmeta_public_key),
-            read_rollback_index: Some(read_rollback_index),
-            write_rollback_index: None,
-            read_is_device_unlocked: Some(read_is_device_unlocked),
-            get_unique_guid_for_partition: Some(get_unique_guid_for_partition),
-            get_size_of_partition: Some(get_size_of_partition),
-            read_persistent_value: None,
-            write_persistent_value: None,
-            validate_public_key_for_partition: None,
-        };
-        let ab_suffix = CStr::from_bytes_with_nul(NULL_BYTE).unwrap();
-        let mut out_data = MaybeUninit::uninit();
-        // SAFETY: It is safe to call `avb_slot_verify()` as the pointer arguments (`ops`,
-        // `requested_partitions` and `ab_suffix`) passed to the method are all valid and
-        // initialized. The last argument `out_data` is allowed to be null so that nothing
-        // will be written to it.
-        let result = unsafe {
-            avb_slot_verify(
-                &mut avb_ops,
-                requested_partitions.as_ptr(),
-                ab_suffix.as_ptr(),
-                AvbSlotVerifyFlags::AVB_SLOT_VERIFY_FLAGS_NONE,
-                AvbHashtreeErrorMode::AVB_HASHTREE_ERROR_MODE_RESTART_AND_INVALIDATE,
-                out_data.as_mut_ptr(),
-            )
-        };
-        slot_verify_result_to_verify_payload_result(result)?;
-        // SAFETY: This is safe because `out_data` has been properly initialized after
-        // calling `avb_slot_verify` and it returns OK.
-        let out_data = unsafe { out_data.assume_init() };
-        out_data.try_into()
-    }
-}
-
 fn verify_vbmeta_has_no_initrd_descriptor(
     vbmeta_image: &AvbVBMetaData,
 ) -> Result<(), AvbSlotVerifyError> {
@@ -428,8 +150,9 @@
     initrd: Option<&[u8]>,
     trusted_public_key: &[u8],
 ) -> Result<DebugLevel, AvbSlotVerifyError> {
-    let mut payload = Payload { kernel, initrd, trusted_public_key };
-    let kernel_verify_result = payload.verify_partition(PartitionName::Kernel.as_cstr())?;
+    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())?;
     let vbmeta_images = kernel_verify_result.vbmeta_images()?;
     if vbmeta_images.len() != 1 {
         // There can only be one VBMeta.
@@ -437,16 +160,16 @@
     }
     let vbmeta_image = vbmeta_images[0];
     verify_vbmeta_is_from_kernel_partition(&vbmeta_image)?;
-    if payload.initrd.is_none() {
+    if initrd.is_none() {
         verify_vbmeta_has_no_initrd_descriptor(&vbmeta_image)?;
         return Ok(DebugLevel::None);
     }
     // TODO(b/256148034): Check the vbmeta doesn't have hash descriptors other than
     // boot, initrd_normal, initrd_debug.
 
-    let debug_level = if payload.verify_partition(PartitionName::InitrdNormal.as_cstr()).is_ok() {
+    let debug_level = if ops.verify_partition(PartitionName::InitrdNormal.as_cstr()).is_ok() {
         DebugLevel::None
-    } else if payload.verify_partition(PartitionName::InitrdDebug.as_cstr()).is_ok() {
+    } else if ops.verify_partition(PartitionName::InitrdDebug.as_cstr()).is_ok() {
         DebugLevel::Full
     } else {
         return Err(AvbSlotVerifyError::Verification);
diff --git a/pvmfw/src/heap.rs b/pvmfw/src/heap.rs
index acc903a..e04451f 100644
--- a/pvmfw/src/heap.rs
+++ b/pvmfw/src/heap.rs
@@ -30,7 +30,9 @@
 #[global_allocator]
 static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
 
-static mut HEAP: [u8; 131072] = [0; 131072];
+/// 128 KiB
+const HEAP_SIZE: usize =  0x20000;
+static mut HEAP: [u8; HEAP_SIZE] = [0; HEAP_SIZE];
 
 pub unsafe fn init() {
     HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index 604aa80..7eecb97 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -321,8 +321,7 @@
         handle_alloc_error(layout);
     };
 
-    let vaddr = buffer.as_ptr() as usize;
-    let paddr = virt_to_phys(vaddr);
+    let paddr = virt_to_phys(buffer);
     // If share_range fails then we will leak the allocation, but that seems better than having it
     // be reused while maybe still partially shared with the host.
     share_range(&(paddr..paddr + layout.size()), granule)?;
@@ -338,7 +337,7 @@
 ///
 /// The memory must have been allocated by `alloc_shared` with the same size, and not yet
 /// deallocated.
-pub unsafe fn dealloc_shared(vaddr: usize, size: usize) -> smccc::Result<()> {
+pub unsafe fn dealloc_shared(vaddr: NonNull<u8>, size: usize) -> smccc::Result<()> {
     let layout = shared_buffer_layout(size)?;
     let granule = layout.align();
 
@@ -346,7 +345,7 @@
     unshare_range(&(paddr..paddr + layout.size()), granule)?;
     // Safe because the memory was allocated by `alloc_shared` above using the same allocator, and
     // the layout is the same as was used then.
-    unsafe { dealloc(vaddr as *mut u8, layout) };
+    unsafe { dealloc(vaddr.as_ptr(), layout) };
 
     Ok(())
 }
@@ -372,8 +371,16 @@
 
 /// Returns the intermediate physical address corresponding to the given virtual address.
 ///
-/// As we use identity mapping for everything, this is just the identity function, but it's useful
-/// to use it to be explicit about where we are converting from virtual to physical address.
-pub fn virt_to_phys(vaddr: usize) -> usize {
-    vaddr
+/// As we use identity mapping for everything, this is just a cast, but it's useful to use it to be
+/// explicit about where we are converting from virtual to physical address.
+pub fn virt_to_phys(vaddr: NonNull<u8>) -> usize {
+    vaddr.as_ptr() as _
+}
+
+/// Returns a pointer for the virtual address corresponding to the given non-zero intermediate
+/// physical address.
+///
+/// Panics if `paddr` is 0.
+pub fn phys_to_virt(paddr: usize) -> NonNull<u8> {
+    NonNull::new(paddr as _).unwrap()
 }
diff --git a/pvmfw/src/virtio/hal.rs b/pvmfw/src/virtio/hal.rs
index c1c8ae6..c6c7a99 100644
--- a/pvmfw/src/virtio/hal.rs
+++ b/pvmfw/src/virtio/hal.rs
@@ -1,23 +1,22 @@
-use crate::memory::{alloc_shared, dealloc_shared, virt_to_phys};
+use crate::memory::{alloc_shared, dealloc_shared, phys_to_virt, virt_to_phys};
 use core::ptr::{copy_nonoverlapping, NonNull};
 use log::debug;
-use virtio_drivers::{BufferDirection, Hal, PhysAddr, VirtAddr, PAGE_SIZE};
+use virtio_drivers::{BufferDirection, Hal, PhysAddr, PAGE_SIZE};
 
 pub struct HalImpl;
 
 impl Hal for HalImpl {
-    fn dma_alloc(pages: usize) -> PhysAddr {
+    fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
         debug!("dma_alloc: pages={}", pages);
         let size = pages * PAGE_SIZE;
-        let vaddr = alloc_shared(size)
-            .expect("Failed to allocate and share VirtIO DMA range with host")
-            .as_ptr() as VirtAddr;
-        virt_to_phys(vaddr)
+        let vaddr =
+            alloc_shared(size).expect("Failed to allocate and share VirtIO DMA range with host");
+        let paddr = virt_to_phys(vaddr);
+        (paddr, vaddr)
     }
 
-    fn dma_dealloc(paddr: PhysAddr, pages: usize) -> i32 {
+    fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
         debug!("dma_dealloc: paddr={:#x}, pages={}", paddr, pages);
-        let vaddr = Self::phys_to_virt(paddr);
         let size = pages * PAGE_SIZE;
         // Safe because the memory was allocated by `dma_alloc` above using the same allocator, and
         // the layout is the same as was used then.
@@ -27,8 +26,8 @@
         0
     }
 
-    fn phys_to_virt(paddr: PhysAddr) -> VirtAddr {
-        paddr
+    fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
+        phys_to_virt(paddr)
     }
 
     fn share(buffer: NonNull<[u8]>, direction: BufferDirection) -> PhysAddr {
@@ -43,11 +42,11 @@
                 copy_nonoverlapping(buffer.as_ptr() as *mut u8, copy.as_ptr(), size);
             }
         }
-        virt_to_phys(copy.as_ptr() as VirtAddr)
+        virt_to_phys(copy)
     }
 
     fn unshare(paddr: PhysAddr, buffer: NonNull<[u8]>, direction: BufferDirection) {
-        let vaddr = Self::phys_to_virt(paddr);
+        let vaddr = phys_to_virt(paddr);
         let size = buffer.len();
         if direction == BufferDirection::DeviceToDriver {
             debug!(
@@ -56,7 +55,7 @@
                 buffer.as_ptr() as *mut u8 as usize
             );
             unsafe {
-                copy_nonoverlapping(vaddr as *const u8, buffer.as_ptr() as *mut u8, size);
+                copy_nonoverlapping(vaddr.as_ptr(), buffer.as_ptr() as *mut u8, size);
             }
         }
 
diff --git a/virtualizationmanager/src/crosvm.rs b/virtualizationmanager/src/crosvm.rs
index b6109c6..c1ac20f 100644
--- a/virtualizationmanager/src/crosvm.rs
+++ b/virtualizationmanager/src/crosvm.rs
@@ -678,7 +678,7 @@
         // Configure the logger for the crosvm process to silence logs from the disk crate which
         // don't provide much information to us (but do spamming us).
         .arg("--log-level")
-        .arg("info,disk=off")
+        .arg("info,disk=warn")
         .arg("run")
         .arg("--disable-sandbox")
         .arg("--cid")
diff --git a/vmbase/example/src/pci.rs b/vmbase/example/src/pci.rs
index 438ff9e..c0a2d2b 100644
--- a/vmbase/example/src/pci.rs
+++ b/vmbase/example/src/pci.rs
@@ -15,7 +15,7 @@
 //! Functions to scan the PCI bus for VirtIO device.
 
 use aarch64_paging::paging::MemoryRegion;
-use alloc::alloc::{alloc, dealloc, Layout};
+use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
 use core::{mem::size_of, ptr::NonNull};
 use fdtpci::PciInfo;
 use log::{debug, info};
@@ -25,7 +25,7 @@
         pci::{bus::PciRoot, virtio_device_type, PciTransport},
         DeviceType, Transport,
     },
-    BufferDirection, Hal, PhysAddr, VirtAddr, PAGE_SIZE,
+    BufferDirection, Hal, PhysAddr, PAGE_SIZE,
 };
 
 /// The standard sector size of a VirtIO block device, in bytes.
@@ -87,32 +87,34 @@
 struct HalImpl;
 
 impl Hal for HalImpl {
-    fn dma_alloc(pages: usize) -> PhysAddr {
+    fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
         debug!("dma_alloc: pages={}", pages);
         let layout = Layout::from_size_align(pages * PAGE_SIZE, PAGE_SIZE).unwrap();
         // Safe because the layout has a non-zero size.
-        let vaddr = unsafe { alloc(layout) } as VirtAddr;
-        virt_to_phys(vaddr)
+        let vaddr = unsafe { alloc(layout) };
+        let vaddr =
+            if let Some(vaddr) = NonNull::new(vaddr) { vaddr } else { handle_alloc_error(layout) };
+        let paddr = virt_to_phys(vaddr);
+        (paddr, vaddr)
     }
 
-    fn dma_dealloc(paddr: PhysAddr, pages: usize) -> i32 {
+    fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
         debug!("dma_dealloc: paddr={:#x}, pages={}", paddr, pages);
-        let vaddr = Self::phys_to_virt(paddr);
         let layout = Layout::from_size_align(pages * PAGE_SIZE, PAGE_SIZE).unwrap();
         // Safe because the memory was allocated by `dma_alloc` above using the same allocator, and
         // the layout is the same as was used then.
         unsafe {
-            dealloc(vaddr as *mut u8, layout);
+            dealloc(vaddr.as_ptr(), layout);
         }
         0
     }
 
-    fn phys_to_virt(paddr: PhysAddr) -> VirtAddr {
-        paddr
+    fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
+        NonNull::new(paddr as _).unwrap()
     }
 
     fn share(buffer: NonNull<[u8]>, _direction: BufferDirection) -> PhysAddr {
-        let vaddr = buffer.as_ptr() as *mut u8 as usize;
+        let vaddr = buffer.cast();
         // Nothing to do, as the host already has access to all memory.
         virt_to_phys(vaddr)
     }
@@ -123,6 +125,6 @@
     }
 }
 
-fn virt_to_phys(vaddr: VirtAddr) -> PhysAddr {
-    vaddr
+fn virt_to_phys(vaddr: NonNull<u8>) -> PhysAddr {
+    vaddr.as_ptr() as _
 }