Merge "Revert "Disallow UDS-rooted BCC""
diff --git a/apex/Android.bp b/apex/Android.bp
index e39b459..1c4d357 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -92,6 +92,8 @@
         "microdroid_initrd_normal",
         "microdroid.json",
         "microdroid_kernel",
+        // rialto_bin is a prebuilt target wrapping the signed bare-metal service VM.
+        "rialto_bin",
     ],
     host_required: [
         "vm_shell",
diff --git a/compos/tests/java/android/compos/test/ComposTestCase.java b/compos/tests/java/android/compos/test/ComposTestCase.java
index 8c7aeeb..3a09475 100644
--- a/compos/tests/java/android/compos/test/ComposTestCase.java
+++ b/compos/tests/java/android/compos/test/ComposTestCase.java
@@ -221,8 +221,11 @@
         // compos.info.signature since it's only generated by CompOS.
         // TODO(b/211458160): Remove cache-info.xml once we can plumb timestamp and isFactory of
         // APEXes to the VM.
-        return runner.run("cd " + path + "; find -type f -exec sha256sum {} \\;"
-                + "| grep -v cache-info.xml | grep -v compos.info"
-                + "| sort -k2");
+        return runner.run(
+                "cd "
+                        + path
+                        + " && find -type f -exec sha256sum {} \\;"
+                        + "| grep -v cache-info.xml | grep -v compos.info"
+                        + "| sort -k2");
     }
 }
diff --git a/demo/java/com/android/microdroid/demo/MainActivity.java b/demo/java/com/android/microdroid/demo/MainActivity.java
index aab4148..f27b23b 100644
--- a/demo/java/com/android/microdroid/demo/MainActivity.java
+++ b/demo/java/com/android/microdroid/demo/MainActivity.java
@@ -184,7 +184,7 @@
                         private void testVmService(VirtualMachine vm) {
                             IBinder binder;
                             try {
-                                binder = vm.connectToVsockServer(ITestService.SERVICE_PORT);
+                                binder = vm.connectToVsockServer(ITestService.PORT);
                             } catch (Exception e) {
                                 if (!Thread.interrupted()) {
                                     mPayloadOutput.postValue(
diff --git a/docs/getting_started/index.md b/docs/getting_started/index.md
index 25eb909..62529b3 100644
--- a/docs/getting_started/index.md
+++ b/docs/getting_started/index.md
@@ -11,8 +11,8 @@
 We support the following device:
 
 * aosp_cf_x86_64_phone (Cuttlefish a.k.a. Cloud Android)
-* oriole (Pixel 6)
-* raven (Pixel 6 Pro)
+* oriole/raven (Pixel 6, and 6 Pro)
+* panther/cheetah (Pixel 7, and 7 Pro)
 
 ### Cuttlefish
 
@@ -30,12 +30,14 @@
 acloud create --local-instance --local-image
 ```
 
-### Pixel 6 and 6 Pro
+### Google Pixel phones
 
-If the device is running Android 12, join the [Android Beta
-Program](https://www.google.com/android/beta) to upgrade to Android 13 Beta.
+If the device is running Android 13 or earlier, join the [Android Beta
+Program](https://developer.android.com/about/versions/14/get#on_pixel) to upgrade to Android 14
+Beta.
 
-Once upgraded to Android 13, execute the following command to enable pKVM.
+Once upgraded to Android 14, and if you are using Pixel 6 or 6 Pro, execute the following command to
+enable pKVM. You don't need to do this for Pixel 7 and 7 Pro.
 
 ```shell
 adb reboot bootloader
@@ -44,41 +46,6 @@
 fastboot reboot
 ```
 
-Due to a bug in Android 13 for these devices, pKVM may stop working after an
-[OTA update](https://source.android.com/devices/tech/ota). To prevent this, it
-is necessary to manually replicate the `pvmfw` partition to the other slot:
-
-```shell
-git -C <android_root>/packages/modules/Virtualization
-show de6b0b2ecf6225a0a7b43241de27e74fc3e6ceb2:pvmfw/pvmfw.img > /tmp/pvmfw.img
-adb reboot bootloader
-fastboot --slot other flash pvmfw /tmp/pvmfw.img
-fastboot reboot
-```
-
-Otherwise, if an OTA has already made pKVM unusable, the working partition
-should be copied to the "current" slot:
-
-```shell
-adb reboot bootloader
-fastboot flash pvmfw /tmp/pvmfw.img
-fastboot reboot
-```
-
-Finally, if the `pvmfw` partition has been corrupted, both slots may be flashed
-using the [`pvmfw.img` pre-built](https://android.googlesource.com/platform/packages/modules/Virtualization/+/08deac98acefd62e222edfa374d5292458cf97eb%5E/pvmfw/pvmfw.img)
-as long as the bootloader remains unlocked. Otherwise, a fresh install of
-Android 13 followed by the manual steps above for flashing the `other` slot
-should be used as a last resort.
-
-Starting in Android 14, `pvmfw.img` can be built using the Android Build system:
-```
-lunch <target>  # where PRODUCT_BUILD_PVMFW_IMAGE=true
-m pvmfwimage    # partition image under ${ANDROID_PRODUCT_OUT}/pvmfw.img
-```
-Note that the result is not intended to be used in Android 13 as not
-backward-compatibility is guaranteed.
-
 ## Running demo app
 
 The instruction is [here](../../demo/README.md).
@@ -140,23 +107,6 @@
 adb shell /apex/com.android.virt/bin/vm run-microdroid --debug full
 ```
 
-The `instance.img` and `apk.idsig` files will be stored in a subdirectory under
-`/data/local/tmp/microdroid`, that `vm` will create.
-
-Atlernatively, you can manually run the demo app on top of Microdroid as follows:
-
-```shell
-UNBUNDLED_BUILD_SDKS_FROM_SOURCE=true TARGET_BUILD_APPS=MicrodroidDemoApp m apps_only dist
-adb shell mkdir -p /data/local/tmp/virt
-adb push out/dist/MicrodroidDemoApp.apk /data/local/tmp/virt/
-adb shell /apex/com.android.virt/bin/vm run-app \
-  --debug full \
-  /data/local/tmp/virt/MicrodroidDemoApp.apk \
-  /data/local/tmp/virt/MicrodroidDemoApp.apk.idsig \
-  /data/local/tmp/virt/instance.img \
-  --payload-path MicrodroidTestNativeLib.so
-```
-
 ## Building and updating CrosVM and VirtualizationService {#building-and-updating}
 
 You can update CrosVM and the VirtualizationService by updating the `com.android.virt` APEX instead
diff --git a/docs/getting_started/pixel6.md b/docs/getting_started/pixel6.md
deleted file mode 100644
index 82391c4..0000000
--- a/docs/getting_started/pixel6.md
+++ /dev/null
@@ -1,129 +0,0 @@
-# Instructions for building custom AVF on Pixel 6 or 6 Pro
-
-This document provides steps for building AVF from AOSP, and then install it to
-Pixel 6 series to better understand AVF and do some experiments.
-
-**WARNING**: Unless Android 13 is released to AOSP (expected to be at Summer
-2022, exact date TBD) by the time when you read this documentation, or you or
-your company have early access to Android Tiramisu source tree, you **CANNOT**
-follow this instruction. In that case, you can only **USE** the AVF that is
-shipped in the Android 13 Beta Image.
-
-This is because AVF in the beta image is signed by Google and therefore it can't
-be updated to a new AVF built in AOSP which can't be signed by the Google key
-that is not shared with AOSP.
-
-## Upgrade to Android 13 Beta Image
-
-First, upgrade your Pixel 6 or Pixel 6 Pro to the Android 13 Beta Image. This
-can be done in two ways:
-
-* Join [Android Beta Program](https://www.google.com/android/beta) and then OTA
-  to Android 13.
-* Manually flash [Android 13 Beta Images](https://developer.android.com/about/versions/13/download#factory-images).
-
-Then enable ADB debugging in "Settings" -> "System" -> "Developer options".
-Finally, enable PKVM.
-
-```shell
-adb reboot bootloader
-fastboot flashing unlock
-fastboot oem pkvm enable
-fastboot reboot
-```
-
-## Building GSI and flashing it
-
-Prepare your Android 13 (Tiramisu) source tree.
-
-```shell
-mkdir tm
-cd tm
-repo init -u <URL> -m <your_tm_branch>
-repo sync -c --no-tags -j 10
-```
-
-Patch GSI so that it includes AVF. Edit
-`build/make/target/product/gsi_release.mk` and add the following line to the
-end (or anywhere in the file that makes sense):
-
-```
-PRODUCT_PACKAGES += com.android.virt
-```
-
-Build GSI.
-
-```shell
-source build/envsetup.sh
-choosecombo 1 aosp_arm64 userdebug
-m
-```
-
-Flash GSI to the Pixel device.
-
-```shell
-adb reboot bootloader
-fastboot reboot fastboot
-fastboot delete-logical-partition product_a
-fastboot flash system out/target/product/generic_arm64/system.img
-fastboot --disable-verification flash vbmeta out/target/product/generic_arm64/vbmeta.img
-fastboot -w reboot
-```
-
-Deleting the logical partition `product_a` is needed because the GSI image is
-bigger than the logical partition `system_a` of the beta image.
-`--disable-verification` when flashing the `vbmeta` partition is critical. Don't
-miss it.
-
-Lastly, check if you are running GSI.
-
-```shell
-adb shell getprop ro.build.product
-adb shell ls /dev/kvm
-adb shell ls /apex/com.android.virt/bin/vm
-```
-
-The result should be as follows.
-
-```
-generic_arm64
-/dev/kvm
-/apex/com.android.virt/bin/vm
-```
-
-## Building and installing AVF from AOSP
-
-Checkout AOSP master branch.
-
-```shell
-mkdir aosp
-cd aosp
-repo init -u https://android.googlesource.com/platform/manifest -b master
-repo sync -c --no-tags -j 10
-```
-
-Then build the `com.android.virt` APEX.
-
-```shell
-source build/envsetup.sh
-banchan com.android.virt aosp_arm64
-UNBUNDLED_BUILD_SDKS_FROM_SOURCE=true m apps_only dist
-```
-
-Install the newly built AVF to the device
-
-```shell
-adb install out/dist/com.android.virt.apex
-adb reboot
-```
-
-If this doesn't work for some reason, try this:
-
-```
-adb root
-adb shell setenforce 0
-adb push out/dist/com.android.virt.apex /data/local/
-adb shell cmd -w apexservice deactivatePackage /system/system_ext/apex/com.android.virt.apex
-adb shell cmd -w apexservice activatePackage /data/local/com.android.virt.apex
-// Don't adb reboot
-```
diff --git a/libs/hyp/Android.bp b/libs/hyp/Android.bp
index e4353c8..1d572e5 100644
--- a/libs/hyp/Android.bp
+++ b/libs/hyp/Android.bp
@@ -8,7 +8,10 @@
     srcs: ["src/lib.rs"],
     prefer_rlib: true,
     rustlibs: [
+        "libbitflags",
+        "libonce_cell_nostd",
         "libpsci",
+        "libuuid_nostd",
     ],
     no_stdlibs: true,
     stdlibs: [
diff --git a/libs/hyp/src/error.rs b/libs/hyp/src/error.rs
index 4e25e7f..408150e 100644
--- a/libs/hyp/src/error.rs
+++ b/libs/hyp/src/error.rs
@@ -16,6 +16,7 @@
 
 use crate::KvmError;
 use core::{fmt, result};
+use uuid::Uuid;
 
 /// Result type with hypervisor error.
 pub type Result<T> = result::Result<T, Error>;
@@ -27,6 +28,8 @@
     MmioGuardNotsupported,
     /// Failed to invoke a certain KVM HVC function.
     KvmError(KvmError, u32),
+    /// Unsupported Hypervisor.
+    UnsupportedHypervisorUuid(Uuid),
     /// The MMIO_GUARD granule used by the hypervisor is not supported.
     UnsupportedMmioGuardGranule(usize),
 }
@@ -38,6 +41,9 @@
             Self::KvmError(e, function_id) => {
                 write!(f, "Failed to invoke the HVC function with function ID {function_id}: {e}")
             }
+            Self::UnsupportedHypervisorUuid(u) => {
+                write!(f, "Unsupported Hypervisor UUID {u}")
+            }
             Self::UnsupportedMmioGuardGranule(g) => {
                 write!(f, "Unsupported MMIO guard granule: {g}")
             }
diff --git a/libs/hyp/src/hypervisor/common.rs b/libs/hyp/src/hypervisor/common.rs
index 87d35b2..accef72 100644
--- a/libs/hyp/src/hypervisor/common.rs
+++ b/libs/hyp/src/hypervisor/common.rs
@@ -15,6 +15,15 @@
 //! This module regroups some common traits shared by all the hypervisors.
 
 use crate::error::Result;
+use bitflags::bitflags;
+
+bitflags! {
+    /// Capabilities that Hypervisor backends can declare support for.
+    pub struct HypervisorCap: u32 {
+        /// Capability for guest to share its memory with host at runtime.
+        const DYNAMIC_MEM_SHARE = 0b1;
+    }
+}
 
 /// Trait for the hypervisor.
 pub trait Hypervisor {
@@ -43,4 +52,7 @@
 
     /// Returns the memory protection granule size in bytes.
     fn memory_protection_granule(&self) -> Result<usize>;
+
+    /// Check if required capabilities are supported.
+    fn has_cap(&self, cap: HypervisorCap) -> bool;
 }
diff --git a/libs/hyp/src/hypervisor/gunyah.rs b/libs/hyp/src/hypervisor/gunyah.rs
new file mode 100644
index 0000000..b335c87
--- /dev/null
+++ b/libs/hyp/src/hypervisor/gunyah.rs
@@ -0,0 +1,40 @@
+use super::common::{Hypervisor, HypervisorCap};
+use crate::error::Result;
+use crate::util::SIZE_4KB;
+use uuid::{uuid, Uuid};
+
+pub(super) struct GunyahHypervisor;
+
+impl GunyahHypervisor {
+    pub const UUID: Uuid = uuid!("c1d58fcd-a453-5fdb-9265-ce36673d5f14");
+}
+
+impl Hypervisor for GunyahHypervisor {
+    fn mmio_guard_init(&self) -> Result<()> {
+        Ok(())
+    }
+
+    fn mmio_guard_map(&self, _addr: usize) -> Result<()> {
+        Ok(())
+    }
+
+    fn mmio_guard_unmap(&self, _addr: usize) -> Result<()> {
+        Ok(())
+    }
+
+    fn mem_share(&self, _base_ipa: u64) -> Result<()> {
+        unimplemented!();
+    }
+
+    fn mem_unshare(&self, _base_ipa: u64) -> Result<()> {
+        unimplemented!();
+    }
+
+    fn memory_protection_granule(&self) -> Result<usize> {
+        Ok(SIZE_4KB)
+    }
+
+    fn has_cap(&self, _cap: HypervisorCap) -> bool {
+        false
+    }
+}
diff --git a/libs/hyp/src/hypervisor/kvm.rs b/libs/hyp/src/hypervisor/kvm.rs
index c0c1ac9..772160e 100644
--- a/libs/hyp/src/hypervisor/kvm.rs
+++ b/libs/hyp/src/hypervisor/kvm.rs
@@ -14,7 +14,7 @@
 
 //! Wrappers around calls to the KVM hypervisor.
 
-use super::common::Hypervisor;
+use super::common::{Hypervisor, HypervisorCap};
 use crate::error::{Error, Result};
 use crate::util::{page_address, SIZE_4KB};
 use core::fmt::{self, Display, Formatter};
@@ -22,6 +22,7 @@
     error::{positive_or_error_64, success_or_error_32, success_or_error_64},
     hvc64,
 };
+use uuid::{uuid, Uuid};
 
 /// Error from a KVM HVC call.
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
@@ -71,6 +72,13 @@
 
 pub(super) struct KvmHypervisor;
 
+impl KvmHypervisor {
+    // Based on ARM_SMCCC_VENDOR_HYP_UID_KVM_REG values listed in Linux kernel source:
+    // https://github.com/torvalds/linux/blob/master/include/linux/arm-smccc.h
+    pub(super) const UUID: Uuid = uuid!("28b46fb6-2ec5-11e9-a9ca-4b564d003a74");
+    const CAPABILITIES: HypervisorCap = HypervisorCap::DYNAMIC_MEM_SHARE;
+}
+
 impl Hypervisor for KvmHypervisor {
     fn mmio_guard_init(&self) -> Result<()> {
         mmio_guard_enroll()?;
@@ -122,6 +130,10 @@
         let granule = checked_hvc64(ARM_SMCCC_KVM_FUNC_HYP_MEMINFO, args)?;
         Ok(granule.try_into().unwrap())
     }
+
+    fn has_cap(&self, cap: HypervisorCap) -> bool {
+        Self::CAPABILITIES.contains(cap)
+    }
 }
 
 fn mmio_guard_granule() -> Result<usize> {
diff --git a/libs/hyp/src/hypervisor/mod.rs b/libs/hyp/src/hypervisor/mod.rs
index a694029..05c82dc 100644
--- a/libs/hyp/src/hypervisor/mod.rs
+++ b/libs/hyp/src/hypervisor/mod.rs
@@ -14,28 +14,82 @@
 
 //! Wrappers around hypervisor back-ends.
 
+extern crate alloc;
+
 mod common;
+mod gunyah;
 mod kvm;
 
+use crate::error::{Error, Result};
+use alloc::boxed::Box;
 pub use common::Hypervisor;
+pub use common::HypervisorCap;
+use gunyah::GunyahHypervisor;
 pub use kvm::KvmError;
 use kvm::KvmHypervisor;
-
-static HYPERVISOR: HypervisorBackend = HypervisorBackend::Kvm;
+use once_cell::race::OnceBox;
+use psci::smccc::hvc64;
+use uuid::Uuid;
 
 enum HypervisorBackend {
     Kvm,
+    Gunyah,
 }
 
 impl HypervisorBackend {
     fn get_hypervisor(&self) -> &'static dyn Hypervisor {
         match self {
             Self::Kvm => &KvmHypervisor,
+            Self::Gunyah => &GunyahHypervisor,
         }
     }
 }
 
+impl TryFrom<Uuid> for HypervisorBackend {
+    type Error = Error;
+
+    fn try_from(uuid: Uuid) -> Result<HypervisorBackend> {
+        match uuid {
+            GunyahHypervisor::UUID => Ok(HypervisorBackend::Gunyah),
+            KvmHypervisor::UUID => Ok(HypervisorBackend::Kvm),
+            u => Err(Error::UnsupportedHypervisorUuid(u)),
+        }
+    }
+}
+
+const ARM_SMCCC_VENDOR_HYP_CALL_UID_FUNC_ID: u32 = 0x8600ff01;
+
+fn query_vendor_hyp_call_uid() -> Uuid {
+    let args = [0u64; 17];
+    let res = hvc64(ARM_SMCCC_VENDOR_HYP_CALL_UID_FUNC_ID, args);
+
+    // KVM's UUID of "28b46fb6-2ec5-11e9-a9ca-4b564d003a74" is generated by
+    // Uuid::from_u128() from an input value of
+    // 0x28b46fb6_2ec511e9_a9ca4b56_4d003a74. ARM's SMC calling convention
+    // (Document number ARM DEN 0028E) describes the UUID register mapping such
+    // that W0 contains bytes 0..3 of UUID, with byte 0 in lower order bits. In
+    // the KVM example, byte 0 of KVM's UUID (0x28) will be returned in the low
+    // 8-bits of W0, while byte 15 (0x74) will be returned in bits 31-24 of W3.
+    //
+    // `uuid` value derived below thus need to be byte-reversed before
+    // being used in Uuid::from_u128(). Alternately use Uuid::from_u128_le()
+    // to achieve the same.
+
+    let uuid = ((res[3] as u32 as u128) << 96)
+        | ((res[2] as u32 as u128) << 64)
+        | ((res[1] as u32 as u128) << 32)
+        | (res[0] as u32 as u128);
+
+    Uuid::from_u128_le(uuid)
+}
+
+fn detect_hypervisor() -> HypervisorBackend {
+    query_vendor_hyp_call_uid().try_into().expect("Unknown hypervisor")
+}
+
 /// Gets the hypervisor singleton.
 pub fn get_hypervisor() -> &'static dyn Hypervisor {
-    HYPERVISOR.get_hypervisor()
+    static HYPERVISOR: OnceBox<HypervisorBackend> = OnceBox::new();
+
+    HYPERVISOR.get_or_init(|| Box::new(detect_hypervisor())).get_hypervisor()
 }
diff --git a/libs/hyp/src/lib.rs b/libs/hyp/src/lib.rs
index 6db6ba8..694f957 100644
--- a/libs/hyp/src/lib.rs
+++ b/libs/hyp/src/lib.rs
@@ -21,4 +21,4 @@
 mod util;
 
 pub use error::{Error, Result};
-pub use hypervisor::{get_hypervisor, Hypervisor, KvmError};
+pub use hypervisor::{get_hypervisor, Hypervisor, HypervisorCap, KvmError};
diff --git a/libs/libfdt/src/lib.rs b/libs/libfdt/src/lib.rs
index 61b69f5..9785941 100644
--- a/libs/libfdt/src/lib.rs
+++ b/libs/libfdt/src/lib.rs
@@ -780,11 +780,11 @@
 
     /// Return a shared pointer to the device tree.
     pub fn as_ptr(&self) -> *const c_void {
-        self as *const _ as *const c_void
+        self.buffer.as_ptr().cast::<_>()
     }
 
     fn as_mut_ptr(&mut self) -> *mut c_void {
-        self as *mut _ as *mut c_void
+        self.buffer.as_mut_ptr().cast::<_>()
     }
 
     fn capacity(&self) -> usize {
@@ -792,8 +792,9 @@
     }
 
     fn header(&self) -> &libfdt_bindgen::fdt_header {
+        let p = self.as_ptr().cast::<_>();
         // SAFETY - A valid FDT (verified by constructor) must contain a valid fdt_header.
-        unsafe { &*(&self as *const _ as *const libfdt_bindgen::fdt_header) }
+        unsafe { &*p }
     }
 
     fn totalsize(&self) -> usize {
diff --git a/libs/vsutil/Android.bp b/libs/vsutil/Android.bp
deleted file mode 100644
index ccb36a0..0000000
--- a/libs/vsutil/Android.bp
+++ /dev/null
@@ -1,26 +0,0 @@
-package {
-    default_applicable_licenses: ["Android-Apache-2.0"],
-}
-
-rust_library {
-    name: "libvsutil",
-    crate_name: "vsutil",
-    srcs: ["src/lib.rs"],
-    // Only build on targets which crosvm builds on.
-    enabled: false,
-    target: {
-        android64: {
-            compile_multilib: "64",
-            enabled: true,
-        },
-        linux_bionic_arm64: {
-            enabled: true,
-        },
-    },
-    rustlibs: [
-        "android.system.virtualizationservice-rust",
-        "libbinder_rs",
-        "libdisk",
-    ],
-    apex_available: ["com.android.virt"],
-}
diff --git a/libs/vsutil/src/file.rs b/libs/vsutil/src/file.rs
deleted file mode 100644
index 4adaae1..0000000
--- a/libs/vsutil/src/file.rs
+++ /dev/null
@@ -1,28 +0,0 @@
-// 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.
-
-//! Functions to process files.
-
-use binder::{self, ExceptionCode, ParcelFileDescriptor, Status};
-use std::fs::File;
-
-/// Converts a `&ParcelFileDescriptor` to a `File` by cloning the file.
-pub fn clone_file(fd: &ParcelFileDescriptor) -> binder::Result<File> {
-    fd.as_ref().try_clone().map_err(|e| {
-        Status::new_exception_str(
-            ExceptionCode::BAD_PARCELABLE,
-            Some(format!("Failed to clone File from ParcelFileDescriptor: {:?}", e)),
-        )
-    })
-}
diff --git a/libs/vsutil/src/lib.rs b/libs/vsutil/src/lib.rs
deleted file mode 100644
index cac888a..0000000
--- a/libs/vsutil/src/lib.rs
+++ /dev/null
@@ -1,22 +0,0 @@
-// 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.
-
-//! This library provides utility functions used by both various
-//! Virtualization services.
-
-mod file;
-mod partition;
-
-pub use file::clone_file;
-pub use partition::init_writable_partition;
diff --git a/libs/vsutil/src/partition.rs b/libs/vsutil/src/partition.rs
deleted file mode 100644
index a4e72de..0000000
--- a/libs/vsutil/src/partition.rs
+++ /dev/null
@@ -1,93 +0,0 @@
-// 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.
-
-//! Functions to process partitions.
-
-use crate::file::clone_file;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
-    PartitionType::PartitionType,
-};
-use binder::{self, ExceptionCode, ParcelFileDescriptor, Status};
-use disk::QcowFile;
-use std::io::{Error, ErrorKind, Write};
-
-/// crosvm requires all partitions to be a multiple of 4KiB.
-const PARTITION_GRANULE_BYTES: u64 = 4096;
-
-/// Initialize an empty partition image of the given size to be used as a writable partition.
-pub fn init_writable_partition(
-    image_fd: &ParcelFileDescriptor,
-    size_bytes: i64,
-    partition_type: PartitionType,
-) -> binder::Result<()> {
-    let size_bytes = size_bytes.try_into().map_err(|e| {
-        Status::new_exception_str(
-            ExceptionCode::ILLEGAL_ARGUMENT,
-            Some(format!("Invalid size {}: {:?}", size_bytes, e)),
-        )
-    })?;
-    let size_bytes = round_up(size_bytes, PARTITION_GRANULE_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| {
-        Status::new_service_specific_error_str(-1, Some(format!("Failed to reset a file: {:?}", 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)),
-        )
-    })?;
-
-    match partition_type {
-        PartitionType::RAW => Ok(()),
-        PartitionType::ANDROID_VM_INSTANCE => format_as_android_vm_instance(&mut part),
-        PartitionType::ENCRYPTEDSTORE => format_as_encryptedstore(&mut part),
-        _ => Err(Error::new(
-            ErrorKind::Unsupported,
-            format!("Unsupported partition type {:?}", partition_type),
-        )),
-    }
-    .map_err(|e| {
-        Status::new_service_specific_error_str(
-            -1,
-            Some(format!("Failed to initialize partition as {:?}: {:?}", partition_type, e)),
-        )
-    })
-}
-
-fn round_up(input: u64, granule: u64) -> u64 {
-    if granule == 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(granule - 1).unwrap_or(input);
-    (result / granule) * granule
-}
-
-fn format_as_android_vm_instance(part: &mut dyn Write) -> std::io::Result<()> {
-    const ANDROID_VM_INSTANCE_MAGIC: &str = "Android-VM-instance";
-    const ANDROID_VM_INSTANCE_VERSION: u16 = 1;
-
-    part.write_all(ANDROID_VM_INSTANCE_MAGIC.as_bytes())?;
-    part.write_all(&ANDROID_VM_INSTANCE_VERSION.to_le_bytes())?;
-    part.flush()
-}
-
-fn format_as_encryptedstore(part: &mut dyn Write) -> std::io::Result<()> {
-    const UNFORMATTED_STORAGE_MAGIC: &str = "UNFORMATTED-STORAGE";
-
-    part.write_all(UNFORMATTED_STORAGE_MAGIC.as_bytes())?;
-    part.flush()
-}
diff --git a/microdroid/README.md b/microdroid/README.md
index 71be7d0..5e3f586 100644
--- a/microdroid/README.md
+++ b/microdroid/README.md
@@ -8,8 +8,7 @@
 ## Prerequisites
 
 Any 64-bit target (either x86\_64 or arm64) is supported. 32-bit target is not
-supported. Note that we currently don't support user builds; only userdebug
-builds are supported.
+supported.
 
 The only remaining requirement is that `com.android.virt` APEX has to be
 pre-installed. To do this, add the following line in your product makefile.
@@ -18,10 +17,10 @@
 $(call inherit-product, packages/modules/Virtualization/apex/product_packages.mk)
 ```
 
-Build the target after adding the line, and flash it. This step needs to be done
-only once for the target.
+Build the target product after adding the line, and flash it. This step needs
+to be done only once for the target.
 
-If you are using `aosp_oriole` (Pixel 6) or `aosp_cf_x86_64_phone` (Cuttlefish),
+If you are using Pixel 6 and beyond or Cuttlefish (`aosp_cf_x86_64_phone`)
 adding above line is not necessary as it's already done.
 
 ## Building and installing microdroid
@@ -41,8 +40,11 @@
 
 ## Building an app
 
-An app in microdroid is a shared library file embedded in an APK. The shared
-library should have an entry point `AVmPayload_main` as shown below:
+A [vm
+payload](https://android.googlesource.com/platform/packages/modules/Virtualization/+/refs/heads/master/vm_payload/)
+is a shared library file that gets executed in microdroid. It is packaged as
+part of an Android application.  The library should have an entry point
+`AVmPayload_main` as shown below:
 
 ```C++
 extern "C" int AVmPayload_main() {
@@ -54,53 +56,22 @@
 
 ```
 cc_library_shared {
-  name: "MyMicrodroidApp",
+  name: "MyMicrodroidPayload",
   srcs: ["**/*.cpp"],
   sdk_version: "current",
 }
 ```
 
-Then you need a configuration file in JSON format that defines what to load and
-execute in microdroid. The name of the file can be anything and you may have
-multiple configuration files if needed.
-
-```json
-{
-  "os": { "name": "microdroid" },
-  "task": {
-    "type": "microdroid_launcher",
-    "command": "MyMicrodroidApp.so"
-  }
-}
-```
-
-The value of `task.command` should match with the name of the shared library
-defined above. If your app requires APEXes to be imported, you can declare the
-list in `apexes` key like following.
-
-```json
-{
-  "os": ...,
-  "task": ...,
-  "apexes": [
-    {"name": "com.android.awesome_apex"}
-  ]
-}
-```
-
-Embed the shared library and the VM configuration file in an APK:
+Embed the shared library file in an APK:
 
 ```
 android_app {
   name: "MyApp",
-  srcs: ["**/*.java"], // if there is any java code
-  jni_libs: ["MyMicrodroidApp"],
+  srcs: ["**/*.java"],
+  jni_libs: ["MyMicrodroidPayload"],
   use_embedded_native_libs: true,
   sdk_version: "current",
 }
-
-// The VM configuration file can be embedded by simply placing it at `./assets`
-// directory.
 ```
 
 Finally, you build the APK.
@@ -109,7 +80,7 @@
 TARGET_BUILD_APPS=MyApp m apps_only dist
 ```
 
-## Running the app on microdroid
+## Running the VM payload on microdroid
 
 First of all, install the APK to the target device.
 
@@ -117,22 +88,16 @@
 adb install out/dist/MyApp.apk
 ```
 
-`ALL_CAP`s below are placeholders. They need to be replaced with correct
-values:
+There are two ways start a VM and run the payload in it.
 
-* `VM_CONFIG_FILE`: the name of the VM config file that you embedded in the APK.
-  (e.g. `vm_config.json`)
-* `PACKAGE_NAME_OF_YOUR_APP`: package name of your app (e.g. `com.acme.app`).
-* `PATH_TO_YOUR_APP`: path to the installed APK on the device. Can be obtained
-  via the following command.
-  ```sh
-  adb shell pm path PACKAGE_NAME_OF_YOUR_APP
-  ```
-  It shall report a cryptic path similar to `/data/app/~~OgZq==/com.acme.app-HudMahQ==/base.apk`.
+* By manually invoking the `vm` tool via `adb shell`.
+* Calling APIs programmatically in the Java app.
+
+### Using `vm` tool
 
 Execute the following commands to launch a VM. The VM will boot to microdroid
-and then automatically execute your app (the shared library
-`MyMicrodroidApp.so`).
+and then automatically execute your payload (the shared library
+`MyMicrodroidPayload.so`).
 
 ```sh
 TEST_ROOT=/data/local/tmp/virt
@@ -142,23 +107,36 @@
 PATH_TO_YOUR_APP \
 $TEST_ROOT/MyApp.apk.idsig \
 $TEST_ROOT/instance.img \
---config-path assets/VM_CONFIG_FILE
+--payload-binary-name MyMicrodroidPayload.so
 ```
 
-The last command lets you know the CID assigned to the VM. The console output
-from the VM is stored to `$TEST_ROOT/console.txt` and logcat is stored to
-`$TEST_ROOT/log.txt` file for debugging purpose. If you omit `--log` or
-`--console` option, they will be emitted to the current console.
+`ALL_CAP`s below are placeholders. They need to be replaced with correct
+values:
 
-Stopping the VM can be done as follows:
+* `PACKAGE_NAME_OF_YOUR_APP`: package name of your app (e.g. `com.acme.app`).
+* `PATH_TO_YOUR_APP`: path to the installed APK on the device. Can be obtained
+  via the following command.
+  ```sh
+  adb shell pm path PACKAGE_NAME_OF_YOUR_APP
+  ```
+  It shall report a cryptic path similar to `/data/app/~~OgZq==/com.acme.app-HudMahQ==/base.apk`.
 
-```sh
-adb shell /apex/com.android.virt/bin/vm stop $CID
-```
+The console output from the VM is stored to `$TEST_ROOT/console.txt` and logcat
+is stored to `$TEST_ROOT/log.txt` file for debugging purpose. If you omit
+`--log` or `--console` option, the console output will be emitted to the
+current console and the logcat logs are sent to the main logcat in Android.
 
-, where `$CID` is the reported CID value. This works only when the `vm` was
-invoked with the `--daemonize` flag. If the flag was not used, press Ctrl+C on
-the console where the `vm run-app` command was invoked.
+Stopping the VM can be done by pressing `Ctrl+C`.
+
+### Using the APIs
+
+Use the [Android Virtualization Framework Java
+APIs](https://android.googlesource.com/platform/packages/modules/Virtualization/+/refs/heads/master/javalib/api/system-current.txt)
+in your app to create a microdroid VM and run payload in it. The APIs currently
+are @SystemApi, thus available only to privileged apps.
+
+If you are looking for an example usage of the APIs, you may refer to the [demo
+app](https://android.googlesource.com/platform/packages/modules/Virtualization/+/refs/heads/master/demo/).
 
 ## Debuggable microdroid
 
diff --git a/microdroid/init.rc b/microdroid/init.rc
index 29f8970..871db94 100644
--- a/microdroid/init.rc
+++ b/microdroid/init.rc
@@ -146,6 +146,7 @@
     capabilities CHOWN DAC_OVERRIDE DAC_READ_SEARCH FOWNER SYS_ADMIN
 
 service ueventd /system/bin/ueventd
+    user root
     class core
     critical
     seclabel u:r:ueventd:s0
@@ -161,6 +162,7 @@
     setenv HOSTNAME console
 
 service init_debug_policy /system/bin/init_debug_policy
+    user root
     oneshot
     disabled
     stdio_to_kmsg
diff --git a/microdroid/payload/mk_payload.cc b/microdroid/payload/mk_payload.cc
deleted file mode 100644
index d31333f..0000000
--- a/microdroid/payload/mk_payload.cc
+++ /dev/null
@@ -1,305 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <sys/stat.h>
-#include <sys/types.h>
-#include <unistd.h>
-
-#include <fstream>
-#include <iostream>
-#include <optional>
-#include <string>
-#include <vector>
-
-#include <android-base/file.h>
-#include <android-base/result.h>
-#include <image_aggregator.h>
-#include <json/json.h>
-
-#include "microdroid/metadata.h"
-
-using android::base::Dirname;
-using android::base::ErrnoError;
-using android::base::Error;
-using android::base::Result;
-using android::base::unique_fd;
-using android::microdroid::ApexPayload;
-using android::microdroid::ApkPayload;
-using android::microdroid::Metadata;
-using android::microdroid::WriteMetadata;
-
-using cuttlefish::AlignToPartitionSize;
-using cuttlefish::CreateCompositeDisk;
-using cuttlefish::kLinuxFilesystem;
-using cuttlefish::MultipleImagePartition;
-
-Result<uint32_t> GetFileSize(const std::string& path) {
-    struct stat st;
-    if (lstat(path.c_str(), &st) == -1) {
-        return ErrnoError() << "Can't lstat " << path;
-    }
-    return static_cast<uint32_t>(st.st_size);
-}
-
-std::string RelativeTo(const std::string& path, const std::string& dirname) {
-    bool is_absolute = !path.empty() && path[0] == '/';
-    if (is_absolute || dirname == ".") {
-        return path;
-    } else {
-        return dirname + "/" + path;
-    }
-}
-
-// Returns `append` is appended to the end of filename preserving the extension.
-std::string AppendFileName(const std::string& filename, const std::string& append) {
-    size_t pos = filename.find_last_of('.');
-    if (pos == std::string::npos) {
-        return filename + append;
-    } else {
-        return filename.substr(0, pos) + append + filename.substr(pos);
-    }
-}
-
-struct ApexConfig {
-    std::string name; // the apex name
-    std::string path; // the path to the apex file
-                      // absolute or relative to the config file
-};
-
-struct ApkConfig {
-    std::string name;
-    std::string path;
-    std::string idsig_path;
-};
-
-struct Config {
-    std::string dirname; // config file's direname to resolve relative paths in the config
-
-    std::vector<ApexConfig> apexes;
-    std::optional<ApkConfig> apk;
-    // This is a path in the guest side
-    std::optional<std::string> payload_config_path;
-};
-
-#define DO(expr) \
-    if (auto res = (expr); !res.ok()) return res.error()
-
-Result<void> ParseJson(const Json::Value& value, std::string& s) {
-    if (!value.isString()) {
-        return Error() << "should be a string: " << value;
-    }
-    s = value.asString();
-    return {};
-}
-
-template <typename T>
-Result<void> ParseJson(const Json::Value& value, std::optional<T>& s) {
-    if (value.isNull()) {
-        s.reset();
-        return {};
-    }
-    s.emplace();
-    return ParseJson(value, *s);
-}
-
-template <typename T>
-Result<void> ParseJson(const Json::Value& values, std::vector<T>& parsed) {
-    for (const Json::Value& value : values) {
-        T t;
-        DO(ParseJson(value, t));
-        parsed.push_back(std::move(t));
-    }
-    return {};
-}
-
-Result<void> ParseJson(const Json::Value& value, ApexConfig& apex_config) {
-    DO(ParseJson(value["name"], apex_config.name));
-    DO(ParseJson(value["path"], apex_config.path));
-    return {};
-}
-
-Result<void> ParseJson(const Json::Value& value, ApkConfig& apk_config) {
-    DO(ParseJson(value["name"], apk_config.name));
-    DO(ParseJson(value["path"], apk_config.path));
-    DO(ParseJson(value["idsig_path"], apk_config.idsig_path));
-    return {};
-}
-
-Result<void> ParseJson(const Json::Value& value, Config& config) {
-    DO(ParseJson(value["apexes"], config.apexes));
-    DO(ParseJson(value["apk"], config.apk));
-    DO(ParseJson(value["payload_config_path"], config.payload_config_path));
-    return {};
-}
-
-Result<Config> LoadConfig(const std::string& config_file) {
-    std::ifstream in(config_file);
-    Json::CharReaderBuilder builder;
-    Json::Value root;
-    Json::String errs;
-    if (!parseFromStream(builder, in, &root, &errs)) {
-        return Error() << errs;
-    }
-
-    Config config;
-    config.dirname = Dirname(config_file);
-    DO(ParseJson(root, config));
-    return config;
-}
-
-#undef DO
-
-Result<void> MakeMetadata(const Config& config, const std::string& filename) {
-    Metadata metadata;
-    metadata.set_version(1);
-
-    for (const auto& apex_config : config.apexes) {
-        auto* apex = metadata.add_apexes();
-        apex->set_name(apex_config.name);
-        apex->set_partition_name(apex_config.name);
-        apex->set_is_factory(true);
-    }
-
-    if (config.apk.has_value()) {
-        auto* apk = metadata.mutable_apk();
-        apk->set_name(config.apk->name);
-        apk->set_payload_partition_name("microdroid-apk");
-        apk->set_idsig_partition_name("microdroid-apk-idsig");
-    }
-
-    if (config.payload_config_path.has_value()) {
-        *metadata.mutable_config_path() = config.payload_config_path.value();
-    }
-
-    std::ofstream out(filename);
-    return WriteMetadata(metadata, out);
-}
-
-// fill zeros to align |file_path|'s size to BLOCK_SIZE(4096) boundary.
-// return true when the filler is needed.
-Result<bool> ZeroFiller(const std::string& file_path, const std::string& filler_path) {
-    auto file_size = GetFileSize(file_path);
-    if (!file_size.ok()) {
-        return file_size.error();
-    }
-    auto disk_size = AlignToPartitionSize(*file_size);
-    if (disk_size <= *file_size) {
-        return false;
-    }
-    unique_fd fd(TEMP_FAILURE_RETRY(open(filler_path.c_str(), O_CREAT | O_WRONLY | O_TRUNC, 0600)));
-    if (fd.get() == -1) {
-        return ErrnoError() << "open(" << filler_path << ") failed.";
-    }
-    if (ftruncate(fd.get(), disk_size - *file_size) == -1) {
-        return ErrnoError() << "ftruncate(" << filler_path << ") failed.";
-    }
-    return true;
-}
-
-Result<void> MakePayload(const Config& config, const std::string& metadata_file,
-                         const std::string& output_file) {
-    std::vector<MultipleImagePartition> partitions;
-
-    int filler_count = 0;
-    auto add_partition = [&](auto partition_name, auto file_path) -> Result<void> {
-        std::vector<std::string> image_files{file_path};
-
-        std::string filler_path =
-                AppendFileName(output_file, "-filler-" + std::to_string(filler_count++));
-        if (auto ret = ZeroFiller(file_path, filler_path); !ret.ok()) {
-            return ret.error();
-        } else if (*ret) {
-            image_files.push_back(filler_path);
-        }
-        partitions.push_back(MultipleImagePartition{
-                .label = partition_name,
-                .image_file_paths = image_files,
-                .type = kLinuxFilesystem,
-                .read_only = true,
-        });
-        return {};
-    };
-
-    // put metadata at the first partition
-    partitions.push_back(MultipleImagePartition{
-            .label = "payload-metadata",
-            .image_file_paths = {metadata_file},
-            .type = kLinuxFilesystem,
-            .read_only = true,
-    });
-    // put apexes at the subsequent partitions
-    for (size_t i = 0; i < config.apexes.size(); i++) {
-        const auto& apex_config = config.apexes[i];
-        std::string apex_path = RelativeTo(apex_config.path, config.dirname);
-        if (auto ret = add_partition("microdroid-apex-" + std::to_string(i), apex_path);
-            !ret.ok()) {
-            return ret.error();
-        }
-    }
-    // put apk and its idsig
-    if (config.apk.has_value()) {
-        std::string apk_path = RelativeTo(config.apk->path, config.dirname);
-        if (auto ret = add_partition("microdroid-apk", apk_path); !ret.ok()) {
-            return ret.error();
-        }
-        std::string idsig_path = RelativeTo(config.apk->idsig_path, config.dirname);
-        if (auto ret = add_partition("microdroid-apk-idsig", idsig_path); !ret.ok()) {
-            return ret.error();
-        }
-    }
-
-    const std::string gpt_header = AppendFileName(output_file, "-header");
-    const std::string gpt_footer = AppendFileName(output_file, "-footer");
-    CreateCompositeDisk(partitions, gpt_header, gpt_footer, output_file);
-    return {};
-}
-
-int main(int argc, char** argv) {
-    if (argc < 3 || argc > 4) {
-        std::cerr << "Usage: " << argv[0] << " [--metadata-only] <config> <output>\n";
-        return 1;
-    }
-    int arg_index = 1;
-    bool metadata_only = false;
-    if (strcmp(argv[arg_index], "--metadata-only") == 0) {
-        metadata_only = true;
-        arg_index++;
-    }
-
-    auto config = LoadConfig(argv[arg_index++]);
-    if (!config.ok()) {
-        std::cerr << "bad config: " << config.error() << '\n';
-        return 1;
-    }
-
-    const std::string output_file(argv[arg_index++]);
-    const std::string metadata_file =
-            metadata_only ? output_file : AppendFileName(output_file, "-metadata");
-
-    if (const auto res = MakeMetadata(*config, metadata_file); !res.ok()) {
-        std::cerr << res.error() << '\n';
-        return 1;
-    }
-    if (metadata_only) {
-        return 0;
-    }
-    if (const auto res = MakePayload(*config, metadata_file, output_file); !res.ok()) {
-        std::cerr << res.error() << '\n';
-        return 1;
-    }
-
-    return 0;
-}
diff --git a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
index 3859785..50d437f 100644
--- a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
+++ b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
@@ -67,4 +67,13 @@
      * @throws SecurityException if the use of test APIs is not permitted.
      */
     byte[] getDiceAttestationCdi();
+
+    /**
+     * Requests a certificate using the provided certificate signing request (CSR).
+     *
+     * TODO(b/271275206): Define the format of the CSR properly.
+     * @param csr the certificate signing request.
+     * @return the X.509 encoded certificate.
+     */
+    byte[] requestCertificate(in byte[] csr);
 }
diff --git a/microdroid_manager/src/vm_payload_service.rs b/microdroid_manager/src/vm_payload_service.rs
index 96f51f0..11e6967 100644
--- a/microdroid_manager/src/vm_payload_service.rs
+++ b/microdroid_manager/src/vm_payload_service.rs
@@ -67,6 +67,11 @@
         self.check_restricted_apis_allowed()?;
         Ok(self.dice.cdi_attest().to_vec())
     }
+
+    fn requestCertificate(&self, csr: &[u8]) -> binder::Result<Vec<u8>> {
+        self.check_restricted_apis_allowed()?;
+        self.virtual_machine_service.requestCertificate(csr)
+    }
 }
 
 impl Interface for VmPayloadService {}
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index 8be5f7d..0ae2203 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -8,6 +8,8 @@
     defaults: ["vmbase_ffi_defaults"],
     srcs: ["src/main.rs"],
     edition: "2021",
+    // Require unsafe blocks for inside unsafe functions.
+    flags: ["-Dunsafe_op_in_unsafe_fn"],
     features: [
         "legacy",
     ],
@@ -15,6 +17,8 @@
         "libaarch64_paging",
         "libbssl_ffi_nostd",
         "libbuddy_system_allocator",
+        "libciborium_nostd",
+        "libciborium_io_nostd",
         "libdiced_open_dice_nostd",
         "libfdtpci",
         "libhyp",
@@ -30,7 +34,9 @@
         "libuuid_nostd",
         "libvirtio_drivers",
         "libvmbase",
+        "libzerocopy_nostd",
         "libzeroize_nostd",
+        "libspin_nostd",
     ],
 }
 
diff --git a/pvmfw/avb/Android.bp b/pvmfw/avb/Android.bp
index 7ed4895..90f3971 100644
--- a/pvmfw/avb/Android.bp
+++ b/pvmfw/avb/Android.bp
@@ -7,6 +7,8 @@
     crate_name: "pvmfw_avb",
     srcs: ["src/lib.rs"],
     prefer_rlib: true,
+    // Require unsafe blocks for inside unsafe functions.
+    flags: ["-Dunsafe_op_in_unsafe_fn"],
     rustlibs: [
         "libavb_bindgen_nostd",
         "libtinyvec_nostd",
diff --git a/pvmfw/avb/src/descriptor.rs b/pvmfw/avb/src/descriptor.rs
index c54d416..cd623ac 100644
--- a/pvmfw/avb/src/descriptor.rs
+++ b/pvmfw/avb/src/descriptor.rs
@@ -14,8 +14,6 @@
 
 //! 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};
diff --git a/pvmfw/src/config.rs b/pvmfw/src/config.rs
index f62a580..b90b136 100644
--- a/pvmfw/src/config.rs
+++ b/pvmfw/src/config.rs
@@ -19,10 +19,11 @@
 use core::mem;
 use core::ops::Range;
 use core::result;
+use zerocopy::{FromBytes, LayoutVerified};
 
 /// Configuration data header.
 #[repr(C, packed)]
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, FromBytes)]
 struct Header {
     /// Magic number; must be `Header::MAGIC`.
     magic: u32,
@@ -40,6 +41,8 @@
 pub enum Error {
     /// Reserved region can't fit configuration header.
     BufferTooSmall,
+    /// Header has the wrong alignment
+    HeaderMisaligned,
     /// Header doesn't contain the expect magic value.
     InvalidMagic,
     /// Version of the header isn't supported.
@@ -58,6 +61,7 @@
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Self::BufferTooSmall => write!(f, "Reserved region is smaller than config header"),
+            Self::HeaderMisaligned => write!(f, "Reserved region is misaligned"),
             Self::InvalidMagic => write!(f, "Wrong magic number"),
             Self::UnsupportedVersion(x, y) => write!(f, "Version {x}.{y} not supported"),
             Self::InvalidFlags(v) => write!(f, "Flags value {v:#x} is incorrect or reserved"),
@@ -167,7 +171,7 @@
 }
 
 #[repr(packed)]
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, FromBytes)]
 struct HeaderEntry {
     offset: u32,
     size: u32,
@@ -187,7 +191,9 @@
     pub unsafe fn new(data: &'a mut [u8]) -> Result<Self> {
         let header = data.get(..Header::PADDED_SIZE).ok_or(Error::BufferTooSmall)?;
 
-        let header = &*(header.as_ptr() as *const Header);
+        let (header, _) =
+            LayoutVerified::<_, Header>::new_from_prefix(header).ok_or(Error::HeaderMisaligned)?;
+        let header = header.into_ref();
 
         if header.magic != Header::MAGIC {
             return Err(Error::InvalidMagic);
@@ -206,11 +212,13 @@
             header.get_body_range(Entry::Bcc)?.ok_or(Error::MissingEntry(Entry::Bcc))?;
         let dp_range = header.get_body_range(Entry::DebugPolicy)?;
 
+        let body_size = header.body_size();
+        let total_size = header.total_size();
         let body = data
             .get_mut(Header::PADDED_SIZE..)
             .ok_or(Error::BufferTooSmall)?
-            .get_mut(..header.body_size())
-            .ok_or_else(|| Error::InvalidSize(header.total_size()))?;
+            .get_mut(..body_size)
+            .ok_or(Error::InvalidSize(total_size))?;
 
         Ok(Self { body, bcc_range, dp_range })
     }
diff --git a/pvmfw/src/crypto.rs b/pvmfw/src/crypto.rs
index 0785a7a..d607bee 100644
--- a/pvmfw/src/crypto.rs
+++ b/pvmfw/src/crypto.rs
@@ -248,13 +248,14 @@
 ///
 /// # Safety
 ///
-/// The caller needs to ensure that the pointer points to a valid C string and that the C lifetime
-/// of the string is compatible with a static Rust lifetime.
+/// The caller needs to ensure that the pointer is null or points to a valid C string and that the
+/// C lifetime of the string is compatible with a static Rust lifetime.
 unsafe fn as_static_cstr(p: *const c_char) -> Option<&'static CStr> {
     if p.is_null() {
         None
     } else {
-        Some(CStr::from_ptr(p))
+        // Safety: Safe given the requirements of this function.
+        Some(unsafe { CStr::from_ptr(p) })
     }
 }
 
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 1309d73..999baee 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -19,13 +19,16 @@
 use crate::fdt;
 use crate::heap;
 use crate::helpers;
-use crate::memory::MemoryTracker;
+use crate::helpers::RangeExt as _;
+use crate::memory::{MemoryTracker, MEMORY};
 use crate::mmu;
 use crate::rand;
 use core::arch::asm;
+use core::mem::size_of;
 use core::num::NonZeroUsize;
+use core::ops::Range;
 use core::slice;
-use hyp::get_hypervisor;
+use hyp::{get_hypervisor, HypervisorCap};
 use log::debug;
 use log::error;
 use log::info;
@@ -62,7 +65,7 @@
     // - can't access MMIO (therefore, no logging)
 
     match main_wrapper(fdt_address as usize, payload_start as usize, payload_size as usize) {
-        Ok(entry) => jump_to_payload(fdt_address, entry.try_into().unwrap()),
+        Ok((entry, bcc)) => jump_to_payload(fdt_address, entry.try_into().unwrap(), bcc),
         Err(_) => reboot(), // TODO(b/220071963) propagate the reason back to the host.
     }
 
@@ -109,6 +112,18 @@
             RebootReason::InvalidFdt
         })?;
 
+        if !get_hypervisor().has_cap(HypervisorCap::DYNAMIC_MEM_SHARE) {
+            let range = info.swiotlb_info.fixed_range().ok_or_else(|| {
+                error!("Pre-shared pool range not specified in swiotlb node");
+                RebootReason::InvalidFdt
+            })?;
+
+            memory.init_shared_pool(range).map_err(|e| {
+                error!("Failed to initialize pre-shared pool {e}");
+                RebootReason::InvalidFdt
+            })?;
+        }
+
         let kernel_range = if let Some(r) = info.kernel_range {
             memory.alloc_range(&r).map_err(|e| {
                 error!("Failed to obtain the kernel range with DT range: {e}");
@@ -158,7 +173,11 @@
 ///
 /// Provide the abstractions necessary for start() to abort the pVM boot and for main() to run with
 /// the assumption that its environment has been properly configured.
-fn main_wrapper(fdt: usize, payload: usize, payload_size: usize) -> Result<usize, RebootReason> {
+fn main_wrapper(
+    fdt: usize,
+    payload: usize,
+    payload_size: usize,
+) -> Result<(usize, Range<usize>), RebootReason> {
     // Limitations in this function:
     // - only access MMIO once (and while) it has been mapped and configured
     // - only perform logging once the logger has been initialized
@@ -217,8 +236,8 @@
     unsafe { page_table.activate() };
     debug!("... Success!");
 
-    let mut memory = MemoryTracker::new(page_table);
-    let slices = MemorySlices::new(fdt, payload, payload_size, &mut memory)?;
+    MEMORY.lock().replace(MemoryTracker::new(page_table));
+    let slices = MemorySlices::new(fdt, payload, payload_size, MEMORY.lock().as_mut().unwrap())?;
 
     rand::init().map_err(|e| {
         error!("Failed to initialize rand: {e}");
@@ -226,13 +245,19 @@
     })?;
 
     // This wrapper allows main() to be blissfully ignorant of platform details.
-    crate::main(slices.fdt, slices.kernel, slices.ramdisk, bcc_slice, debug_policy, &mut memory)?;
+    let next_bcc = crate::main(
+        slices.fdt,
+        slices.kernel,
+        slices.ramdisk,
+        bcc_slice,
+        debug_policy,
+        MEMORY.lock().as_mut().unwrap(),
+    )?;
 
     helpers::flushed_zeroize(bcc_slice);
-    helpers::flush(slices.fdt.as_slice());
 
     info!("Expecting a bug making MMIO_GUARD_UNMAP return NOT_SUPPORTED on success");
-    memory.mmio_unmap_all().map_err(|e| {
+    MEMORY.lock().as_mut().unwrap().mmio_unmap_all().map_err(|e| {
         error!("Failed to unshare MMIO ranges: {e}");
         RebootReason::InternalError
     })?;
@@ -240,11 +265,13 @@
         error!("Failed to unshare the UART: {e}");
         RebootReason::InternalError
     })?;
+    MEMORY.lock().take().unwrap();
 
-    Ok(slices.kernel.as_ptr() as usize)
+    Ok((slices.kernel.as_ptr() as usize, next_bcc))
 }
 
-fn jump_to_payload(fdt_address: u64, payload_start: u64) -> ! {
+fn jump_to_payload(fdt_address: u64, payload_start: u64, bcc: Range<usize>) -> ! {
+    const ASM_STP_ALIGN: usize = size_of::<u64>() * 2;
     const SCTLR_EL1_RES1: u64 = (0b11 << 28) | (0b101 << 20) | (0b1 << 11);
     // Stage 1 instruction access cacheability is unaffected.
     const SCTLR_EL1_I: u64 = 0b1 << 12;
@@ -255,12 +282,67 @@
 
     const SCTLR_EL1_VAL: u64 = SCTLR_EL1_RES1 | SCTLR_EL1_ITD | SCTLR_EL1_SED | SCTLR_EL1_I;
 
+    let scratch = layout::scratch_range();
+
+    assert_ne!(scratch.len(), 0, "scratch memory is empty.");
+    assert_eq!(scratch.start % ASM_STP_ALIGN, 0, "scratch memory is misaligned.");
+    assert_eq!(scratch.end % ASM_STP_ALIGN, 0, "scratch memory is misaligned.");
+
+    assert!(bcc.is_within(&scratch));
+    assert_eq!(bcc.start % ASM_STP_ALIGN, 0, "Misaligned guest BCC.");
+    assert_eq!(bcc.end % ASM_STP_ALIGN, 0, "Misaligned guest BCC.");
+
+    let stack = mmu::stack_range();
+
+    assert_ne!(stack.len(), 0, "stack region is empty.");
+    assert_eq!(stack.start % ASM_STP_ALIGN, 0, "Misaligned stack region.");
+    assert_eq!(stack.end % ASM_STP_ALIGN, 0, "Misaligned stack region.");
+
+    // Zero all memory that could hold secrets and that can't be safely written to from Rust.
     // Disable the exception vector, caches and page table and then jump to the payload at the
     // given address, passing it the given FDT pointer.
     //
     // SAFETY - We're exiting pvmfw by passing the register values we need to a noreturn asm!().
     unsafe {
         asm!(
+            "cmp {scratch}, {bcc}",
+            "b.hs 1f",
+
+            // Zero .data & .bss until BCC.
+            "0: stp xzr, xzr, [{scratch}], 16",
+            "cmp {scratch}, {bcc}",
+            "b.lo 0b",
+
+            "1:",
+            // Skip BCC.
+            "mov {scratch}, {bcc_end}",
+            "cmp {scratch}, {scratch_end}",
+            "b.hs 1f",
+
+            // Keep zeroing .data & .bss.
+            "0: stp xzr, xzr, [{scratch}], 16",
+            "cmp {scratch}, {scratch_end}",
+            "b.lo 0b",
+
+            "1:",
+            // Flush d-cache over .data & .bss (including BCC).
+            "0: dc cvau, {cache_line}",
+            "add {cache_line}, {cache_line}, {dcache_line_size}",
+            "cmp {cache_line}, {scratch_end}",
+            "b.lo 0b",
+
+            "mov {cache_line}, {stack}",
+            // Zero stack region.
+            "0: stp xzr, xzr, [{stack}], 16",
+            "cmp {stack}, {stack_end}",
+            "b.lo 0b",
+
+            // Flush d-cache over stack region.
+            "0: dc cvau, {cache_line}",
+            "add {cache_line}, {cache_line}, {dcache_line_size}",
+            "cmp {cache_line}, {stack_end}",
+            "b.lo 0b",
+
             "msr sctlr_el1, {sctlr_el1_val}",
             "isb",
             "mov x1, xzr",
@@ -293,13 +375,21 @@
             "mov x28, xzr",
             "mov x29, xzr",
             "msr ttbr0_el1, xzr",
-            "isb",
+            // Ensure that CMOs have completed before entering payload.
             "dsb nsh",
             "br x30",
             sctlr_el1_val = in(reg) SCTLR_EL1_VAL,
+            bcc = in(reg) u64::try_from(bcc.start).unwrap(),
+            bcc_end = in(reg) u64::try_from(bcc.end).unwrap(),
+            cache_line = in(reg) u64::try_from(scratch.start).unwrap(),
+            scratch = in(reg) u64::try_from(scratch.start).unwrap(),
+            scratch_end = in(reg) u64::try_from(scratch.end).unwrap(),
+            stack = in(reg) u64::try_from(stack.start).unwrap(),
+            stack_end = in(reg) u64::try_from(stack.end).unwrap(),
+            dcache_line_size = in(reg) u64::try_from(helpers::min_dcache_line_size()).unwrap(),
             in("x0") fdt_address,
             in("x30") payload_start,
-            options(nomem, noreturn, nostack),
+            options(noreturn),
         );
     };
 }
@@ -309,7 +399,9 @@
     // pvmfw is contained in a 2MiB region so the payload can't be larger than the 2MiB alignment.
     let size = helpers::align_up(base, helpers::SIZE_2MB).unwrap() - base;
 
-    slice::from_raw_parts_mut(base as *mut u8, size)
+    // SAFETY: This region is mapped and the linker script prevents it from overlapping with other
+    // objects.
+    unsafe { slice::from_raw_parts_mut(base as *mut u8, size) }
 }
 
 enum AppendedConfigType {
@@ -328,8 +420,13 @@
 impl<'a> AppendedPayload<'a> {
     /// SAFETY - 'data' should respect the alignment of config::Header.
     unsafe fn new(data: &'a mut [u8]) -> Option<Self> {
-        match Self::guess_config_type(data) {
-            AppendedConfigType::Valid => Some(Self::Config(config::Config::new(data).unwrap())),
+        // Safety: This fn has the same constraint as us.
+        match unsafe { Self::guess_config_type(data) } {
+            AppendedConfigType::Valid => {
+                // Safety: This fn has the same constraint as us.
+                let config = unsafe { config::Config::new(data) };
+                Some(Self::Config(config.unwrap()))
+            }
             AppendedConfigType::NotFound if cfg!(feature = "legacy") => {
                 const BCC_SIZE: usize = helpers::SIZE_4KB;
                 warn!("Assuming the appended data at {:?} to be a raw BCC", data.as_ptr());
@@ -339,11 +436,14 @@
         }
     }
 
+    /// SAFETY - 'data' should respect the alignment of config::Header.
     unsafe fn guess_config_type(data: &mut [u8]) -> AppendedConfigType {
         // This function is necessary to prevent the borrow checker from getting confused
         // about the ownership of data in new(); see https://users.rust-lang.org/t/78467.
         let addr = data.as_ptr();
-        match config::Config::new(data) {
+
+        // Safety: This fn has the same constraint as us.
+        match unsafe { config::Config::new(data) } {
             Err(config::Error::InvalidMagic) => {
                 warn!("No configuration data found at {addr:?}");
                 AppendedConfigType::NotFound
diff --git a/pvmfw/src/exceptions.rs b/pvmfw/src/exceptions.rs
index 42f4c3b..462a9cc 100644
--- a/pvmfw/src/exceptions.rs
+++ b/pvmfw/src/exceptions.rs
@@ -15,40 +15,110 @@
 //! Exception handlers.
 
 use crate::{helpers::page_4kb_of, read_sysreg};
+use core::fmt;
 use vmbase::console;
-use vmbase::{console::emergency_write_str, eprintln, power::reboot};
+use vmbase::logger;
+use vmbase::{eprintln, power::reboot};
 
-const ESR_32BIT_EXT_DABT: usize = 0x96000010;
 const UART_PAGE: usize = page_4kb_of(console::BASE_ADDRESS);
 
+#[derive(Debug)]
+enum HandleExceptionError {
+    UnknownException,
+}
+
+impl fmt::Display for HandleExceptionError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::UnknownException => write!(f, "An unknown exception occurred, not handled."),
+        }
+    }
+}
+
+#[derive(Debug, PartialEq, Copy, Clone)]
+enum Esr {
+    DataAbortTranslationFault,
+    DataAbortPermissionFault,
+    DataAbortSyncExternalAbort,
+    Unknown(usize),
+}
+
+impl Esr {
+    const EXT_DABT_32BIT: usize = 0x96000010;
+    const TRANSL_FAULT_BASE_32BIT: usize = 0x96000004;
+    const TRANSL_FAULT_ISS_MASK_32BIT: usize = !0x143;
+    const PERM_FAULT_BASE_32BIT: usize = 0x9600004C;
+    const PERM_FAULT_ISS_MASK_32BIT: usize = !0x103;
+}
+
+impl From<usize> for Esr {
+    fn from(esr: usize) -> Self {
+        if esr == Self::EXT_DABT_32BIT {
+            Self::DataAbortSyncExternalAbort
+        } else if esr & Self::TRANSL_FAULT_ISS_MASK_32BIT == Self::TRANSL_FAULT_BASE_32BIT {
+            Self::DataAbortTranslationFault
+        } else if esr & Self::PERM_FAULT_ISS_MASK_32BIT == Self::PERM_FAULT_BASE_32BIT {
+            Self::DataAbortPermissionFault
+        } else {
+            Self::Unknown(esr)
+        }
+    }
+}
+
+impl fmt::Display for Esr {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::DataAbortSyncExternalAbort => write!(f, "Synchronous external abort"),
+            Self::DataAbortTranslationFault => write!(f, "Translation fault"),
+            Self::DataAbortPermissionFault => write!(f, "Permission fault"),
+            Self::Unknown(v) => write!(f, "Unknown exception esr={v:#08x}"),
+        }
+    }
+}
+
+fn handle_exception(_esr: Esr, _far: usize) -> Result<(), HandleExceptionError> {
+    Err(HandleExceptionError::UnknownException)
+}
+
+#[inline]
+fn handling_uart_exception(esr: Esr, far: usize) -> bool {
+    esr == Esr::DataAbortSyncExternalAbort && page_4kb_of(far) == UART_PAGE
+}
+
 #[no_mangle]
 extern "C" fn sync_exception_current(_elr: u64, _spsr: u64) {
-    let esr = read_sysreg!("esr_el1");
+    // Disable logging in exception handler to prevent unsafe writes to UART.
+    let _guard = logger::suppress();
+    let esr: Esr = read_sysreg!("esr_el1").into();
     let far = read_sysreg!("far_el1");
-    // Don't print to the UART if we're handling the exception it could raise.
-    if esr != ESR_32BIT_EXT_DABT || page_4kb_of(far) != UART_PAGE {
-        emergency_write_str("sync_exception_current\n");
-        eprintln!("esr={esr:#08x}");
+
+    if let Err(e) = handle_exception(esr, far) {
+        // Don't print to the UART if we are handling an exception it could raise.
+        if !handling_uart_exception(esr, far) {
+            eprintln!("sync_exception_current");
+            eprintln!("{e}");
+            eprintln!("{esr}, far={far:#08x}");
+        }
+        reboot()
     }
-    reboot();
 }
 
 #[no_mangle]
 extern "C" fn irq_current(_elr: u64, _spsr: u64) {
-    emergency_write_str("irq_current\n");
+    eprintln!("irq_current");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn fiq_current(_elr: u64, _spsr: u64) {
-    emergency_write_str("fiq_current\n");
+    eprintln!("fiq_current");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn serr_current(_elr: u64, _spsr: u64) {
     let esr = read_sysreg!("esr_el1");
-    emergency_write_str("serr_current\n");
+    eprintln!("serr_current");
     eprintln!("esr={esr:#08x}");
     reboot();
 }
@@ -56,27 +126,27 @@
 #[no_mangle]
 extern "C" fn sync_lower(_elr: u64, _spsr: u64) {
     let esr = read_sysreg!("esr_el1");
-    emergency_write_str("sync_lower\n");
+    eprintln!("sync_lower");
     eprintln!("esr={esr:#08x}");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn irq_lower(_elr: u64, _spsr: u64) {
-    emergency_write_str("irq_lower\n");
+    eprintln!("irq_lower");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn fiq_lower(_elr: u64, _spsr: u64) {
-    emergency_write_str("fiq_lower\n");
+    eprintln!("fiq_lower");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn serr_lower(_elr: u64, _spsr: u64) {
     let esr = read_sysreg!("esr_el1");
-    emergency_write_str("serr_lower\n");
+    eprintln!("serr_lower");
     eprintln!("esr={esr:#08x}");
     reboot();
 }
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index c68fc6d..5ecb038 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -17,6 +17,7 @@
 use crate::bootargs::BootArgsIterator;
 use crate::cstr;
 use crate::helpers::flatten;
+use crate::helpers::RangeExt as _;
 use crate::helpers::GUEST_PAGE_SIZE;
 use crate::helpers::SIZE_4KB;
 use crate::memory::BASE_ADDR;
@@ -40,6 +41,7 @@
 use log::debug;
 use log::error;
 use log::info;
+use log::warn;
 use tinyvec::ArrayVec;
 
 /// Extract from /config the address range containing the pre-loaded kernel. Absence of /config is
@@ -435,40 +437,79 @@
 }
 
 #[derive(Debug)]
-struct SwiotlbInfo {
-    size: u64,
-    align: u64,
+pub struct SwiotlbInfo {
+    addr: Option<usize>,
+    size: usize,
+    align: usize,
+}
+
+impl SwiotlbInfo {
+    pub fn fixed_range(&self) -> Option<Range<usize>> {
+        self.addr.map(|addr| addr..addr + self.size)
+    }
 }
 
 fn read_swiotlb_info_from(fdt: &Fdt) -> libfdt::Result<SwiotlbInfo> {
     let node =
         fdt.compatible_nodes(cstr!("restricted-dma-pool"))?.next().ok_or(FdtError::NotFound)?;
-    let size = node.getprop_u64(cstr!("size"))?.ok_or(FdtError::NotFound)?;
-    let align = node.getprop_u64(cstr!("alignment"))?.ok_or(FdtError::NotFound)?;
-    Ok(SwiotlbInfo { size, align })
+    let align =
+        node.getprop_u64(cstr!("alignment"))?.ok_or(FdtError::NotFound)?.try_into().unwrap();
+
+    let (addr, size) = if let Some(mut reg) = node.reg()? {
+        let reg = reg.next().ok_or(FdtError::NotFound)?;
+        let size = reg.size.ok_or(FdtError::NotFound)?;
+        reg.addr.checked_add(size).ok_or(FdtError::BadValue)?;
+        (Some(reg.addr.try_into().unwrap()), size.try_into().unwrap())
+    } else {
+        let size = node.getprop_u64(cstr!("size"))?.ok_or(FdtError::NotFound)?.try_into().unwrap();
+        (None, size)
+    };
+
+    Ok(SwiotlbInfo { addr, size, align })
 }
 
-fn validate_swiotlb_info(swiotlb_info: &SwiotlbInfo) -> Result<(), RebootReason> {
+fn validate_swiotlb_info(
+    swiotlb_info: &SwiotlbInfo,
+    memory: &Range<usize>,
+) -> Result<(), RebootReason> {
     let size = swiotlb_info.size;
     let align = swiotlb_info.align;
 
-    if size == 0 || (size % GUEST_PAGE_SIZE as u64) != 0 {
+    if size == 0 || (size % GUEST_PAGE_SIZE) != 0 {
         error!("Invalid swiotlb size {:#x}", size);
         return Err(RebootReason::InvalidFdt);
     }
 
-    if (align % GUEST_PAGE_SIZE as u64) != 0 {
+    if (align % GUEST_PAGE_SIZE) != 0 {
         error!("Invalid swiotlb alignment {:#x}", align);
         return Err(RebootReason::InvalidFdt);
     }
+
+    if let Some(range) = swiotlb_info.fixed_range() {
+        if !range.is_within(memory) {
+            error!("swiotlb range {range:#x?} not part of memory range {memory:#x?}");
+            return Err(RebootReason::InvalidFdt);
+        }
+    }
+
     Ok(())
 }
 
 fn patch_swiotlb_info(fdt: &mut Fdt, swiotlb_info: &SwiotlbInfo) -> libfdt::Result<()> {
     let mut node =
         fdt.root_mut()?.next_compatible(cstr!("restricted-dma-pool"))?.ok_or(FdtError::NotFound)?;
-    node.setprop_inplace(cstr!("size"), &swiotlb_info.size.to_be_bytes())?;
     node.setprop_inplace(cstr!("alignment"), &swiotlb_info.align.to_be_bytes())?;
+
+    if let Some(range) = swiotlb_info.fixed_range() {
+        node.appendprop_addrrange(
+            cstr!("reg"),
+            range.start.try_into().unwrap(),
+            range.len().try_into().unwrap(),
+        )?;
+    } else {
+        node.setprop_inplace(cstr!("size"), &swiotlb_info.size.to_be_bytes())?;
+    }
+
     Ok(())
 }
 
@@ -540,7 +581,7 @@
     num_cpus: usize,
     pci_info: PciInfo,
     serial_info: SerialInfo,
-    swiotlb_info: SwiotlbInfo,
+    pub swiotlb_info: SwiotlbInfo,
 }
 
 impl DeviceTreeInfo {
@@ -603,7 +644,7 @@
         error!("Failed to read swiotlb info from DT: {e}");
         RebootReason::InvalidFdt
     })?;
-    validate_swiotlb_info(&swiotlb_info)?;
+    validate_swiotlb_info(&swiotlb_info, &memory_range)?;
 
     Ok(DeviceTreeInfo {
         kernel_range,
@@ -681,20 +722,26 @@
     debug_policy: Option<&mut [u8]>,
     debuggable: bool,
 ) -> libfdt::Result<()> {
-    fdt.unpack()?;
+    if let Some(debug_policy) = debug_policy {
+        let backup = Vec::from(fdt.as_slice());
+        fdt.unpack()?;
+        let backup_fdt = Fdt::from_slice(backup.as_slice()).unwrap();
+        if apply_debug_policy(fdt, backup_fdt, debug_policy)? {
+            info!("Debug policy applied.");
+        } else {
+            // apply_debug_policy restored fdt to backup_fdt so unpack it again.
+            fdt.unpack()?;
+        }
+    } else {
+        info!("No debug policy found.");
+        fdt.unpack()?;
+    }
 
     patch_dice_node(fdt, bcc.as_ptr() as usize, bcc.len())?;
 
     set_or_clear_chosen_flag(fdt, cstr!("avf,strict-boot"), strict_boot)?;
     set_or_clear_chosen_flag(fdt, cstr!("avf,new-instance"), new_instance)?;
 
-    if let Some(debug_policy) = debug_policy {
-        apply_debug_policy(fdt, debug_policy)?;
-        info!("Debug policy applied.");
-    } else {
-        info!("No debug policy found.");
-    }
-
     if debuggable {
         if let Some(bootargs) = read_bootargs_from(fdt)? {
             filter_out_dangerous_bootargs(fdt, &bootargs)?;
@@ -734,27 +781,33 @@
     Ok(())
 }
 
-fn apply_debug_policy(fdt: &mut Fdt, debug_policy: &mut [u8]) -> libfdt::Result<()> {
-    let backup_fdt = Vec::from(fdt.as_slice());
-
-    let overlay = match Fdt::from_mut_slice(debug_policy) {
+/// Apply the debug policy overlay to the guest DT.
+///
+/// Returns Ok(true) on success, Ok(false) on recovered failure and Err(_) on corruption of the DT.
+fn apply_debug_policy(
+    fdt: &mut Fdt,
+    backup_fdt: &Fdt,
+    debug_policy: &[u8],
+) -> libfdt::Result<bool> {
+    let mut debug_policy = Vec::from(debug_policy);
+    let overlay = match Fdt::from_mut_slice(debug_policy.as_mut_slice()) {
         Ok(overlay) => overlay,
         Err(e) => {
-            info!("Corrupted debug policy found: {e}. Not applying.");
-            return Ok(());
+            warn!("Corrupted debug policy found: {e}. Not applying.");
+            return Ok(false);
         }
     };
-    let backup_overlay = Vec::from(overlay.as_slice());
 
-    // SAFETY - on failure, the corrupted fdts are discarded and are restored using the backups.
+    // SAFETY - on failure, the corrupted DT is restored using the backup.
     if let Err(e) = unsafe { fdt.apply_overlay(overlay) } {
-        error!("Failed to apply debug policy: {e}. Recovering...");
+        warn!("Failed to apply debug policy: {e}. Recovering...");
         fdt.copy_from_slice(backup_fdt.as_slice())?;
-        overlay.copy_from_slice(backup_overlay.as_slice())?;
         // A successful restoration is considered success because an invalid debug policy
         // shouldn't DOS the pvmfw
+        Ok(false)
+    } else {
+        Ok(true)
     }
-    Ok(())
 }
 
 fn read_common_debug_policy(fdt: &Fdt, debug_feature_name: &CStr) -> libfdt::Result<bool> {
diff --git a/pvmfw/src/gpt.rs b/pvmfw/src/gpt.rs
index 6af3047..c3ccb5a 100644
--- a/pvmfw/src/gpt.rs
+++ b/pvmfw/src/gpt.rs
@@ -25,6 +25,7 @@
 use static_assertions::const_assert_eq;
 use uuid::Uuid;
 use virtio_drivers::device::blk::SECTOR_SIZE;
+use zerocopy::FromBytes;
 
 pub enum Error {
     /// VirtIO error during read operation.
@@ -101,8 +102,10 @@
     fn new(mut device: VirtIOBlk) -> Result<Self> {
         let mut blk = [0; Self::LBA_SIZE];
         device.read_block(Header::LBA, &mut blk).map_err(Error::FailedRead)?;
-        let (header_bytes, _) = blk.split_at(size_of::<Header>());
-        let header = Header::from_bytes(header_bytes).ok_or(Error::InvalidHeader)?;
+        let header = Header::read_from_prefix(blk.as_slice()).unwrap();
+        if !header.is_valid() {
+            return Err(Error::InvalidHeader);
+        }
         let entries_count = usize::try_from(header.entries_count()).unwrap();
 
         Ok(Self { device, entries_count })
@@ -151,6 +154,7 @@
 type Lba = u64;
 
 /// Structure as defined in release 2.10 of the UEFI Specification (5.3.2 GPT Header).
+#[derive(FromBytes)]
 #[repr(C, packed)]
 struct Header {
     signature: u64,
@@ -162,7 +166,7 @@
     backup_lba: Lba,
     first_lba: Lba,
     last_lba: Lba,
-    disk_guid: Uuid,
+    disk_guid: u128,
     entries_lba: Lba,
     entries_count: u32,
     entry_size: u32,
@@ -176,20 +180,6 @@
     const LBA: usize = 1;
     const ENTRIES_LBA: usize = 2;
 
-    fn from_bytes(bytes: &[u8]) -> Option<&Self> {
-        let bytes = bytes.get(..size_of::<Self>())?;
-        // SAFETY - We assume that bytes is properly aligned for Header and have verified above
-        // that it holds enough bytes. All potential values of the slice will produce a valid
-        // Header.
-        let header = unsafe { &*bytes.as_ptr().cast::<Self>() };
-
-        if header.is_valid() {
-            Some(header)
-        } else {
-            None
-        }
-    }
-
     fn is_valid(&self) -> bool {
         self.signature() == Self::SIGNATURE
             && self.header_size() == size_of::<Self>().try_into().unwrap()
diff --git a/pvmfw/src/heap.rs b/pvmfw/src/heap.rs
index eea2e98..151049e 100644
--- a/pvmfw/src/heap.rs
+++ b/pvmfw/src/heap.rs
@@ -27,15 +27,22 @@
 
 use buddy_system_allocator::LockedHeap;
 
-#[global_allocator]
-static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
-
 /// 128 KiB
 const HEAP_SIZE: usize = 0x20000;
 static mut HEAP: [u8; HEAP_SIZE] = [0; HEAP_SIZE];
 
+#[global_allocator]
+static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
+
+/// SAFETY: Must be called no more than once.
 pub unsafe fn init() {
-    HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
+    // SAFETY: Nothing else accesses this memory, and we hand it over to the heap to manage and
+    // never touch it again. The heap is locked, so there cannot be any races.
+    let (start, size) = unsafe { (HEAP.as_mut_ptr() as usize, HEAP.len()) };
+
+    let mut heap = HEAP_ALLOCATOR.lock();
+    // SAFETY: We are supplying a valid memory range, and we only do this once.
+    unsafe { heap.init(start, size) };
 }
 
 /// Allocate an aligned but uninitialized slice of heap.
@@ -53,7 +60,7 @@
 
 #[no_mangle]
 unsafe extern "C" fn malloc(size: usize) -> *mut c_void {
-    malloc_(size, false).map_or(ptr::null_mut(), |p| p.cast::<c_void>().as_ptr())
+    allocate(size, false).map_or(ptr::null_mut(), |p| p.cast::<c_void>().as_ptr())
 }
 
 #[no_mangle]
@@ -61,31 +68,69 @@
     let Some(size) = nmemb.checked_mul(size) else {
         return ptr::null_mut()
     };
-    malloc_(size, true).map_or(ptr::null_mut(), |p| p.cast::<c_void>().as_ptr())
+    allocate(size, true).map_or(ptr::null_mut(), |p| p.cast::<c_void>().as_ptr())
 }
 
 #[no_mangle]
+/// SAFETY: ptr must be null or point to a currently-allocated block returned by allocate (either
+/// directly or via malloc or calloc). Note that this function is called directly from C, so we have
+/// to trust that the C code is doing the right thing; there are checks below which will catch some
+/// errors.
 unsafe extern "C" fn free(ptr: *mut c_void) {
-    if let Some(ptr) = NonNull::new(ptr).map(|p| p.cast::<usize>().as_ptr().offset(-1)) {
-        if let Some(size) = NonZeroUsize::new(*ptr) {
-            if let Some(layout) = malloc_layout(size) {
-                HEAP_ALLOCATOR.dealloc(ptr as *mut u8, layout);
-            }
+    let Some(ptr) = NonNull::new(ptr) else { return };
+    // SAFETY: The contents of the HEAP slice may change, but the address range never does.
+    let heap_range = unsafe { HEAP.as_ptr_range() };
+    assert!(
+        heap_range.contains(&(ptr.as_ptr() as *const u8)),
+        "free() called on a pointer that is not part of the HEAP: {ptr:?}"
+    );
+    let (ptr, size) = unsafe {
+        // SAFETY: ptr is non-null and was allocated by allocate, which prepends a correctly aligned
+        // usize.
+        let ptr = ptr.cast::<usize>().as_ptr().offset(-1);
+        (ptr, *ptr)
+    };
+    let size = NonZeroUsize::new(size).unwrap();
+    let layout = malloc_layout(size).unwrap();
+    // SAFETY: If our precondition is satisfied, then this is a valid currently-allocated block.
+    unsafe { HEAP_ALLOCATOR.dealloc(ptr as *mut u8, layout) }
+}
+
+/// Allocate a block of memory suitable to return from `malloc()` etc. Returns a valid pointer
+/// to a suitable aligned region of size bytes, optionally zeroed (and otherwise uninitialized), or
+/// None if size is 0 or allocation fails. The block can be freed by passing the returned pointer to
+/// `free()`.
+fn allocate(size: usize, zeroed: bool) -> Option<NonNull<usize>> {
+    let size = NonZeroUsize::new(size)?.checked_add(mem::size_of::<usize>())?;
+    let layout = malloc_layout(size)?;
+    // SAFETY: layout is known to have non-zero size.
+    let ptr = unsafe {
+        if zeroed {
+            HEAP_ALLOCATOR.alloc_zeroed(layout)
+        } else {
+            HEAP_ALLOCATOR.alloc(layout)
         }
+    };
+    let ptr = NonNull::new(ptr)?.cast::<usize>().as_ptr();
+    // SAFETY: ptr points to a newly allocated block of memory which is properly aligned
+    // for a usize and is big enough to hold a usize as well as the requested number of
+    // bytes.
+    unsafe {
+        *ptr = size.get();
+        NonNull::new(ptr.offset(1))
     }
 }
 
-unsafe fn malloc_(size: usize, zeroed: bool) -> Option<NonNull<usize>> {
-    let size = NonZeroUsize::new(size)?.checked_add(mem::size_of::<usize>())?;
-    let layout = malloc_layout(size)?;
-    let ptr =
-        if zeroed { HEAP_ALLOCATOR.alloc_zeroed(layout) } else { HEAP_ALLOCATOR.alloc(layout) };
-    let ptr = NonNull::new(ptr)?.cast::<usize>().as_ptr();
-    *ptr = size.get();
-    NonNull::new(ptr.offset(1))
+fn malloc_layout(size: NonZeroUsize) -> Option<Layout> {
+    // We want at least 8 byte alignment, and we need to be able to store a usize.
+    const ALIGN: usize = const_max_size(mem::size_of::<usize>(), mem::size_of::<u64>());
+    Layout::from_size_align(size.get(), ALIGN).ok()
 }
 
-fn malloc_layout(size: NonZeroUsize) -> Option<Layout> {
-    const ALIGN: usize = mem::size_of::<u64>();
-    Layout::from_size_align(size.get(), ALIGN).ok()
+const fn const_max_size(a: usize, b: usize) -> usize {
+    if a > b {
+        a
+    } else {
+        b
+    }
 }
diff --git a/pvmfw/src/helpers.rs b/pvmfw/src/helpers.rs
index 9c739d1..a6f0dd5 100644
--- a/pvmfw/src/helpers.rs
+++ b/pvmfw/src/helpers.rs
@@ -15,12 +15,15 @@
 //! Miscellaneous helper functions.
 
 use core::arch::asm;
+use core::ops::Range;
 use zeroize::Zeroize;
 
 pub const SIZE_4KB: usize = 4 << 10;
 pub const SIZE_2MB: usize = 2 << 20;
+pub const SIZE_4MB: usize = 4 << 20;
 
 pub const GUEST_PAGE_SIZE: usize = SIZE_4KB;
+pub const PVMFW_PAGE_SIZE: usize = SIZE_4KB;
 
 /// Read a value from a system register.
 #[macro_export]
@@ -109,7 +112,8 @@
 }
 
 #[inline]
-fn min_dcache_line_size() -> usize {
+/// Read the number of words in the smallest cache line of all the data caches and unified caches.
+pub fn min_dcache_line_size() -> usize {
     const DMINLINE_SHIFT: usize = 16;
     const DMINLINE_MASK: usize = 0xf;
     let ctr_el0 = read_sysreg!("ctr_el0");
@@ -161,6 +165,18 @@
     unsafe { core::slice::from_raw_parts(original.as_ptr().cast(), len) }
 }
 
+/// Trait to check containment of one range within another.
+pub(crate) trait RangeExt {
+    /// Returns true if `self` is contained within the `other` range.
+    fn is_within(&self, other: &Self) -> bool;
+}
+
+impl<T: PartialOrd> RangeExt for Range<T> {
+    fn is_within(&self, other: &Self) -> bool {
+        self.start >= other.start && self.end <= other.end
+    }
+}
+
 /// Create &CStr out of &str literal
 #[macro_export]
 macro_rules! cstr {
diff --git a/pvmfw/src/instance.rs b/pvmfw/src/instance.rs
index a974543..95ddefd 100644
--- a/pvmfw/src/instance.rs
+++ b/pvmfw/src/instance.rs
@@ -26,13 +26,14 @@
 use crate::virtio::pci::VirtIOBlkIterator;
 use core::fmt;
 use core::mem::size_of;
-use core::slice;
 use diced_open_dice::DiceMode;
 use diced_open_dice::Hash;
 use diced_open_dice::Hidden;
 use log::trace;
 use uuid::Uuid;
 use virtio_drivers::transport::pci::bus::PciRoot;
+use zerocopy::AsBytes;
+use zerocopy::FromBytes;
 
 pub enum Error {
     /// Unexpected I/O error while accessing the underlying disk.
@@ -121,7 +122,7 @@
             let aead = AeadCtx::new_aes_256_gcm_randnonce(&key).map_err(Error::FailedOpen)?;
             let decrypted = aead.open(&mut entry, payload).map_err(Error::FailedOpen)?;
 
-            let body: &EntryBody = decrypted.as_ref();
+            let body = EntryBody::read_from(decrypted).unwrap();
             if body.code_hash != dice_inputs.code_hash {
                 Err(Error::RecordedCodeHashMismatch)
             } else if body.auth_hash != dice_inputs.auth_hash {
@@ -134,22 +135,21 @@
         }
         PvmfwEntry::New { header_index } => {
             let salt = rand::random_array().map_err(Error::FailedSaltGeneration)?;
-            let entry_body = EntryBody::new(dice_inputs, &salt);
-            let body = entry_body.as_ref();
+            let body = EntryBody::new(dice_inputs, &salt);
 
             let key = key.map_err(Error::FailedSeal)?;
             let aead = AeadCtx::new_aes_256_gcm_randnonce(&key).map_err(Error::FailedSeal)?;
             // We currently only support single-blk entries.
-            assert!(body.len() + aead.aead().unwrap().max_overhead() < blk.len());
-            let encrypted = aead.seal(&mut blk, body).map_err(Error::FailedSeal)?;
+            let plaintext = body.as_bytes();
+            assert!(plaintext.len() + aead.aead().unwrap().max_overhead() < blk.len());
+            let encrypted = aead.seal(&mut blk, plaintext).map_err(Error::FailedSeal)?;
             let payload_size = encrypted.len();
             let payload_index = header_index + 1;
             instance_img.write_block(payload_index, &blk).map_err(Error::FailedIo)?;
 
             let header = EntryHeader::new(PvmfwEntry::UUID, payload_size);
-            let (blk_header, blk_rest) = blk.split_at_mut(size_of::<EntryHeader>());
-            blk_header.copy_from_slice(header.as_ref());
-            blk_rest.fill(0);
+            header.write_to_prefix(blk.as_mut_slice()).unwrap();
+            blk[header.as_bytes().len()..].fill(0);
             instance_img.write_block(header_index, &blk).map_err(Error::FailedIo)?;
 
             Ok((true, salt))
@@ -157,6 +157,7 @@
     }
 }
 
+#[derive(FromBytes)]
 #[repr(C, packed)]
 struct Header {
     magic: [u8; Header::MAGIC.len()],
@@ -174,23 +175,6 @@
     fn version(&self) -> u16 {
         u16::from_le(self.version)
     }
-
-    fn from_bytes(bytes: &[u8]) -> Option<&Self> {
-        let header: &Self = bytes.as_ref();
-
-        if header.is_valid() {
-            Some(header)
-        } else {
-            None
-        }
-    }
-}
-
-impl AsRef<Header> for [u8] {
-    fn as_ref(&self) -> &Header {
-        // SAFETY - Assume that the alignement and size match Header.
-        unsafe { &*self.as_ptr().cast::<Header>() }
-    }
 }
 
 fn find_instance_img(pci_root: &mut PciRoot) -> Result<Partition> {
@@ -223,12 +207,15 @@
     let header_index = indices.next().ok_or(Error::MissingInstanceImageHeader)?;
     partition.read_block(header_index, &mut blk).map_err(Error::FailedIo)?;
     // The instance.img header is only used for discovery/validation.
-    let _ = Header::from_bytes(&blk).ok_or(Error::InvalidInstanceImageHeader)?;
+    let header = Header::read_from_prefix(blk.as_slice()).unwrap();
+    if !header.is_valid() {
+        return Err(Error::InvalidInstanceImageHeader);
+    }
 
     while let Some(header_index) = indices.next() {
         partition.read_block(header_index, &mut blk).map_err(Error::FailedIo)?;
 
-        let header: &EntryHeader = blk[..size_of::<EntryHeader>()].as_ref();
+        let header = EntryHeader::read_from_prefix(blk.as_slice()).unwrap();
         match (header.uuid(), header.payload_size()) {
             (uuid, _) if uuid.is_nil() => return Ok(PvmfwEntry::New { header_index }),
             (PvmfwEntry::UUID, payload_size) => {
@@ -250,7 +237,8 @@
 /// Marks the start of an instance.img entry.
 ///
 /// Note: Virtualization/microdroid_manager/src/instance.rs uses the name "partition".
-#[repr(C)]
+#[derive(AsBytes, FromBytes)]
+#[repr(C, packed)]
 struct EntryHeader {
     uuid: u128,
     payload_size: u64,
@@ -270,22 +258,7 @@
     }
 }
 
-impl AsRef<EntryHeader> for [u8] {
-    fn as_ref(&self) -> &EntryHeader {
-        assert_eq!(self.len(), size_of::<EntryHeader>());
-        // SAFETY - The size of the slice was checked and any value may be considered valid.
-        unsafe { &*self.as_ptr().cast::<EntryHeader>() }
-    }
-}
-
-impl AsRef<[u8]> for EntryHeader {
-    fn as_ref(&self) -> &[u8] {
-        let s = self as *const Self;
-        // SAFETY - Transmute the (valid) bytes into a slice.
-        unsafe { slice::from_raw_parts(s.cast::<u8>(), size_of::<Self>()) }
-    }
-}
-
+#[derive(AsBytes, FromBytes)]
 #[repr(C)]
 struct EntryBody {
     code_hash: Hash,
@@ -320,19 +293,3 @@
         }
     }
 }
-
-impl AsRef<EntryBody> for [u8] {
-    fn as_ref(&self) -> &EntryBody {
-        assert_eq!(self.len(), size_of::<EntryBody>());
-        // SAFETY - The size of the slice was checked and members are validated by accessors.
-        unsafe { &*self.as_ptr().cast::<EntryBody>() }
-    }
-}
-
-impl AsRef<[u8]> for EntryBody {
-    fn as_ref(&self) -> &[u8] {
-        let s = self as *const Self;
-        // SAFETY - Transmute the (valid) bytes into a slice.
-        unsafe { slice::from_raw_parts(s.cast::<u8>(), size_of::<Self>()) }
-    }
-}
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index a773f1a..96b707e 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -37,6 +37,8 @@
 mod virtio;
 
 use alloc::boxed::Box;
+use alloc::string::ToString;
+use core::ops::Range;
 
 use crate::dice::PartialInputs;
 use crate::entry::RebootReason;
@@ -46,6 +48,7 @@
 use crate::instance::get_or_generate_instance_salt;
 use crate::memory::MemoryTracker;
 use crate::virtio::pci;
+use ciborium::{de::from_reader, value::Value};
 use diced_open_dice::bcc_handover_main_flow;
 use diced_open_dice::bcc_handover_parse;
 use diced_open_dice::DiceArtifacts;
@@ -58,6 +61,19 @@
 
 const NEXT_BCC_SIZE: usize = GUEST_PAGE_SIZE;
 
+type CiboriumError = ciborium::de::Error<ciborium_io::EndOfFile>;
+
+/// Decodes the provided binary CBOR-encoded value and returns a
+/// ciborium::Value struct wrapped in Result.
+fn value_from_bytes(mut bytes: &[u8]) -> Result<Value, CiboriumError> {
+    let value = from_reader(&mut bytes)?;
+    // Ciborium tries to read one Value, but doesn't care if there is trailing data. We do.
+    if !bytes.is_empty() {
+        return Err(CiboriumError::Semantic(Some(0), "unexpected trailing data".to_string()));
+    }
+    Ok(value)
+}
+
 fn main(
     fdt: &mut Fdt,
     signed_kernel: &[u8],
@@ -65,7 +81,7 @@
     current_bcc_handover: &[u8],
     debug_policy: Option<&mut [u8]>,
     memory: &mut MemoryTracker,
-) -> Result<(), RebootReason> {
+) -> Result<Range<usize>, RebootReason> {
     info!("pVM firmware");
     debug!("FDT: {:?}", fdt.as_ptr());
     debug!("Signed kernel: {:?} ({:#x} bytes)", signed_kernel.as_ptr(), signed_kernel.len());
@@ -81,6 +97,18 @@
     })?;
     trace!("BCC: {bcc_handover:x?}");
 
+    // Minimal BCC verification - check the BCC exists & is valid CBOR.
+    // TODO(alanstokes): Do something more useful.
+    if let Some(bytes) = bcc_handover.bcc() {
+        let _ = value_from_bytes(bytes).map_err(|e| {
+            error!("Invalid BCC: {e:?}");
+            RebootReason::InvalidBcc
+        })?;
+    } else {
+        error!("Missing BCC");
+        return Err(RebootReason::InvalidBcc);
+    }
+
     // Set up PCI bus for VirtIO devices.
     let pci_info = PciInfo::from_fdt(fdt).map_err(handle_pci_error)?;
     debug!("PCI: {:#x?}", pci_info);
@@ -129,7 +157,13 @@
         })?;
 
     info!("Starting payload...");
-    Ok(())
+
+    let bcc_range = {
+        let r = next_bcc.as_ptr_range();
+        (r.start as usize)..(r.end as usize)
+    };
+
+    Ok(bcc_range)
 }
 
 /// Logs the given PCI error and returns the appropriate `RebootReason`.
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index 2d5eb5c..7e8423a 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -16,11 +16,14 @@
 
 #![deny(unsafe_op_in_unsafe_fn)]
 
-use crate::helpers::{self, align_down, align_up, page_4kb_of, SIZE_4KB};
+use crate::helpers::{self, align_down, page_4kb_of, RangeExt, SIZE_4KB, SIZE_4MB};
 use crate::mmu;
 use alloc::alloc::alloc_zeroed;
 use alloc::alloc::dealloc;
 use alloc::alloc::handle_alloc_error;
+use alloc::boxed::Box;
+use buddy_system_allocator::LockedHeap;
+use core::alloc::GlobalAlloc as _;
 use core::alloc::Layout;
 use core::cmp::max;
 use core::cmp::min;
@@ -31,6 +34,9 @@
 use core::result;
 use hyp::get_hypervisor;
 use log::error;
+use log::trace;
+use once_cell::race::OnceBox;
+use spin::mutex::SpinMutex;
 use tinyvec::ArrayVec;
 
 /// Base of the system's contiguous "main" memory.
@@ -40,6 +46,9 @@
 
 pub type MemoryRange = Range<usize>;
 
+pub static MEMORY: SpinMutex<Option<MemoryTracker>> = SpinMutex::new(None);
+unsafe impl Send for MemoryTracker {}
+
 #[derive(Clone, Copy, Debug, Default)]
 enum MemoryType {
     #[default]
@@ -61,8 +70,7 @@
 
     /// True if the instance is fully contained within the passed range.
     pub fn is_within(&self, range: &MemoryRange) -> bool {
-        let our: &MemoryRange = self.as_ref();
-        self.as_ref() == &(max(our.start, range.start)..min(our.end, range.end))
+        self.as_ref().is_within(range)
     }
 }
 
@@ -104,6 +112,8 @@
     FailedToMap,
     /// Error from the interaction with the hypervisor.
     Hypervisor(hyp::Error),
+    /// Failure to set `SHARED_POOL`.
+    SharedPoolSetFailure,
 }
 
 impl fmt::Display for MemoryTrackerError {
@@ -117,6 +127,7 @@
             Self::Overlaps => write!(f, "New region overlaps with tracked regions"),
             Self::FailedToMap => write!(f, "Failed to map the new region"),
             Self::Hypervisor(e) => e.fmt(f),
+            Self::SharedPoolSetFailure => write!(f, "Failed to set SHARED_POOL"),
         }
     }
 }
@@ -129,9 +140,12 @@
 
 type Result<T> = result::Result<T, MemoryTrackerError>;
 
+static SHARED_POOL: OnceBox<LockedHeap<32>> = OnceBox::new();
+
 impl MemoryTracker {
     const CAPACITY: usize = 5;
     const MMIO_CAPACITY: usize = 5;
+    const PVMFW_RANGE: MemoryRange = (BASE_ADDR - SIZE_4MB)..BASE_ADDR;
 
     /// Create a new instance from an active page table, covering the maximum RAM size.
     pub fn new(page_table: mmu::PageTable) -> Self {
@@ -197,7 +211,7 @@
     /// appropriately.
     pub fn map_mmio_range(&mut self, range: MemoryRange) -> Result<()> {
         // MMIO space is below the main memory region.
-        if range.end > self.total.start {
+        if range.end > self.total.start || overlaps(&Self::PVMFW_RANGE, &range) {
             return Err(MemoryTrackerError::OutOfRange);
         }
         if self.mmio_regions.iter().any(|r| overlaps(r, &range)) {
@@ -259,6 +273,31 @@
 
         Ok(())
     }
+
+    /// Initialize a separate heap for shared memory allocations.
+    ///
+    /// Some hypervisors such as Gunyah do not support a MemShare API for guest
+    /// to share its memory with host. Instead they allow host to designate part
+    /// of guest memory as "shared" ahead of guest starting its execution. The
+    /// shared memory region is indicated in swiotlb node. On such platforms use
+    /// a separate heap to allocate buffers that can be shared with host.
+    pub fn init_shared_pool(&mut self, range: Range<usize>) -> Result<()> {
+        let size = NonZeroUsize::new(range.len()).unwrap();
+        let range = self.alloc_mut(range.start, size)?;
+        let shared_pool = LockedHeap::<32>::new();
+
+        // SAFETY - `range` should be a valid region of memory as validated by
+        // `validate_swiotlb_info` and not used by any other rust code.
+        unsafe {
+            shared_pool.lock().init(range.start, range.len());
+        }
+
+        SHARED_POOL
+            .set(Box::new(shared_pool))
+            .map_err(|_| MemoryTrackerError::SharedPoolSetFailure)?;
+
+        Ok(())
+    }
 }
 
 impl Drop for MemoryTracker {
@@ -279,6 +318,7 @@
 /// is not aligned with the memory protection granule then it will be extended on either end to
 /// align.
 fn share_range(range: &MemoryRange, granule: usize) -> hyp::Result<()> {
+    trace!("Sharing memory region {range:#x?}");
     for base in (align_down(range.start, granule)
         .expect("Memory protection granule was not a power of two")..range.end)
         .step_by(granule)
@@ -292,6 +332,7 @@
 /// shared. If the range is not aligned with the memory protection granule then it will be extended
 /// on either end to align.
 fn unshare_range(range: &MemoryRange, granule: usize) -> hyp::Result<()> {
+    trace!("Unsharing memory region {range:#x?}");
     for base in (align_down(range.start, granule)
         .expect("Memory protection granule was not a power of two")..range.end)
         .step_by(granule)
@@ -301,16 +342,27 @@
     Ok(())
 }
 
-/// Allocates a memory range of at least the given size from the global allocator, and shares it
-/// with the host. Returns a pointer to the buffer.
+/// Allocates a memory range of at least the given size that is shared with
+/// host. Returns a pointer to the buffer.
 ///
 /// It will be aligned to the memory sharing granule size supported by the hypervisor.
-pub fn alloc_shared(size: usize) -> hyp::Result<NonNull<u8>> {
-    let layout = shared_buffer_layout(size)?;
-    let granule = layout.align();
+pub fn alloc_shared(layout: Layout) -> hyp::Result<NonNull<u8>> {
+    assert_ne!(layout.size(), 0);
+    let granule = get_hypervisor().memory_protection_granule()?;
+    let layout = layout.align_to(granule).unwrap().pad_to_align();
+    if let Some(shared_pool) = SHARED_POOL.get() {
+        // SAFETY - layout has a non-zero size.
+        let buffer = unsafe { shared_pool.alloc_zeroed(layout) };
 
-    // Safe because `shared_buffer_layout` panics if the size is 0, so the layout must have a
-    // non-zero size.
+        let Some(buffer) = NonNull::new(buffer) else {
+            handle_alloc_error(layout);
+        };
+
+        trace!("Allocated shared buffer at {buffer:?} with {layout:?}");
+        return Ok(buffer);
+    }
+
+    // SAFETY - layout has a non-zero size.
     let buffer = unsafe { alloc_zeroed(layout) };
 
     let Some(buffer) = NonNull::new(buffer) else {
@@ -322,6 +374,7 @@
     // be reused while maybe still partially shared with the host.
     share_range(&(paddr..paddr + layout.size()), granule)?;
 
+    trace!("Allocated shared memory at {buffer:?} with {layout:?}");
     Ok(buffer)
 }
 
@@ -333,9 +386,17 @@
 ///
 /// The memory must have been allocated by `alloc_shared` with the same size, and not yet
 /// deallocated.
-pub unsafe fn dealloc_shared(vaddr: NonNull<u8>, size: usize) -> hyp::Result<()> {
-    let layout = shared_buffer_layout(size)?;
-    let granule = layout.align();
+pub unsafe fn dealloc_shared(vaddr: NonNull<u8>, layout: Layout) -> hyp::Result<()> {
+    let granule = get_hypervisor().memory_protection_granule()?;
+    let layout = layout.align_to(granule).unwrap().pad_to_align();
+    if let Some(shared_pool) = SHARED_POOL.get() {
+        // 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 { shared_pool.dealloc(vaddr.as_ptr(), layout) };
+
+        trace!("Deallocated shared buffer at {vaddr:?} with {layout:?}");
+        return Ok(());
+    }
 
     let paddr = virt_to_phys(vaddr);
     unshare_range(&(paddr..paddr + layout.size()), granule)?;
@@ -343,23 +404,10 @@
     // the layout is the same as was used then.
     unsafe { dealloc(vaddr.as_ptr(), layout) };
 
+    trace!("Deallocated shared memory at {vaddr:?} with {layout:?}");
     Ok(())
 }
 
-/// Returns the layout to use for allocating a buffer of at least the given size shared with the
-/// host.
-///
-/// It will be aligned to the memory sharing granule size supported by the hypervisor.
-///
-/// Panics if `size` is 0.
-fn shared_buffer_layout(size: usize) -> hyp::Result<Layout> {
-    assert_ne!(size, 0);
-    let granule = get_hypervisor().memory_protection_granule()?;
-    let allocated_size =
-        align_up(size, granule).expect("Memory protection granule was not a power of two");
-    Ok(Layout::from_size_align(allocated_size, granule).unwrap())
-}
-
 /// Returns an iterator which yields the base address of each 4 KiB page within the given range.
 fn page_iterator(range: &MemoryRange) -> impl Iterator<Item = usize> {
     (page_4kb_of(range.start)..range.end).step_by(SIZE_4KB)
diff --git a/pvmfw/src/mmu.rs b/pvmfw/src/mmu.rs
index fa94e85..e33dcba 100644
--- a/pvmfw/src/mmu.rs
+++ b/pvmfw/src/mmu.rs
@@ -15,6 +15,7 @@
 //! Memory management.
 
 use crate::helpers;
+use crate::helpers::PVMFW_PAGE_SIZE;
 use aarch64_paging::idmap::IdMap;
 use aarch64_paging::paging::Attributes;
 use aarch64_paging::paging::MemoryRegion;
@@ -44,6 +45,13 @@
     start..end
 }
 
+/// Region allocated for the stack.
+pub fn stack_range() -> Range<usize> {
+    const STACK_PAGES: usize = 8;
+
+    layout::stack_range(STACK_PAGES * PVMFW_PAGE_SIZE)
+}
+
 impl PageTable {
     const ASID: usize = 1;
     const ROOT_LEVEL: usize = 1;
@@ -53,7 +61,8 @@
         let mut page_table = Self { idmap: IdMap::new(Self::ASID, Self::ROOT_LEVEL) };
 
         page_table.map_code(&layout::text_range())?;
-        page_table.map_data(&layout::writable_region())?;
+        page_table.map_data(&layout::scratch_range())?;
+        page_table.map_data(&stack_range())?;
         page_table.map_rodata(&layout::rodata_range())?;
         page_table.map_data(&appended_payload_range())?;
 
diff --git a/pvmfw/src/virtio/hal.rs b/pvmfw/src/virtio/hal.rs
index 5f70b33..ec0b9d8 100644
--- a/pvmfw/src/virtio/hal.rs
+++ b/pvmfw/src/virtio/hal.rs
@@ -1,93 +1,126 @@
+// 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.
+
+//! HAL for the virtio_drivers crate.
+
 use super::pci::PCI_INFO;
+use crate::helpers::RangeExt as _;
 use crate::memory::{alloc_shared, dealloc_shared, phys_to_virt, virt_to_phys};
-use core::{
-    ops::Range,
-    ptr::{copy_nonoverlapping, NonNull},
-};
-use log::debug;
+use core::alloc::Layout;
+use core::mem::size_of;
+use core::ptr::{copy_nonoverlapping, NonNull};
+use log::trace;
 use virtio_drivers::{BufferDirection, Hal, PhysAddr, PAGE_SIZE};
 
 pub struct HalImpl;
 
-impl Hal for HalImpl {
+/// Implements the `Hal` trait for `HalImpl`.
+///
+/// # Safety
+///
+/// Callers of this implementatation must follow the safety requirements documented in the `Hal`
+/// trait for the unsafe methods.
+unsafe impl Hal for HalImpl {
+    /// Allocates the given number of contiguous physical pages of DMA memory for VirtIO use.
+    ///
+    /// # Implementation Safety
+    ///
+    /// `dma_alloc` ensures the returned DMA buffer is not aliased with any other allocation or
+    ///  reference in the program until it is deallocated by `dma_dealloc` by allocating a unique
+    ///  block of memory using `alloc_shared` and returning a non-null pointer to it that is
+    ///  aligned to `PAGE_SIZE`.
     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");
+        let vaddr = alloc_shared(dma_layout(pages))
+            .expect("Failed to allocate and share VirtIO DMA range with host");
         let paddr = virt_to_phys(vaddr);
         (paddr, vaddr)
     }
 
-    fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
-        debug!("dma_dealloc: paddr={:#x}, pages={}", paddr, pages);
-        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.
-        unsafe {
-            dealloc_shared(vaddr, size).expect("Failed to unshare VirtIO DMA range with host");
-        }
+    unsafe fn dma_dealloc(_paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
+        // SAFETY - Memory was allocated by `dma_alloc` using `alloc_shared` with the same size.
+        unsafe { dealloc_shared(vaddr, dma_layout(pages)) }
+            .expect("Failed to unshare VirtIO DMA range with host");
         0
     }
 
-    fn mmio_phys_to_virt(paddr: PhysAddr, size: usize) -> NonNull<u8> {
+    /// Converts a physical address used for MMIO to a virtual address which the driver can access.
+    ///
+    /// # Implementation Safety
+    ///
+    /// `mmio_phys_to_virt` satisfies the requirement by checking that the mapped memory region
+    /// is within the PCI MMIO range.
+    unsafe fn mmio_phys_to_virt(paddr: PhysAddr, size: usize) -> NonNull<u8> {
         let pci_info = PCI_INFO.get().expect("VirtIO HAL used before PCI_INFO was initialised");
+        let bar_range = {
+            let start = pci_info.bar_range.start.try_into().unwrap();
+            let end = pci_info.bar_range.end.try_into().unwrap();
+
+            start..end
+        };
+        let mmio_range = paddr..paddr.checked_add(size).expect("PCI MMIO region end overflowed");
+
         // Check that the region is within the PCI MMIO range that we read from the device tree. If
         // not, the host is probably trying to do something malicious.
-        if !contains_range(
-            &pci_info.bar_range,
-            &(paddr.try_into().expect("PCI MMIO region start was outside of 32-bit address space")
-                ..paddr
-                    .checked_add(size)
-                    .expect("PCI MMIO region end overflowed")
-                    .try_into()
-                    .expect("PCI MMIO region end was outside of 32-bit address space")),
-        ) {
-            panic!("PCI MMIO region was outside of expected BAR range.");
-        }
+        assert!(
+            mmio_range.is_within(&bar_range),
+            "PCI MMIO region was outside of expected BAR range.",
+        );
+
         phys_to_virt(paddr)
     }
 
-    fn share(buffer: NonNull<[u8]>, direction: BufferDirection) -> PhysAddr {
+    unsafe fn share(buffer: NonNull<[u8]>, direction: BufferDirection) -> PhysAddr {
         let size = buffer.len();
 
         // TODO: Copy to a pre-shared region rather than allocating and sharing each time.
         // Allocate a range of pages, copy the buffer if necessary, and share the new range instead.
-        let copy =
-            alloc_shared(size).expect("Failed to allocate and share VirtIO buffer with host");
+        let bounce = alloc_shared(bb_layout(size))
+            .expect("Failed to allocate and share VirtIO bounce buffer with host");
+        let paddr = virt_to_phys(bounce);
         if direction == BufferDirection::DriverToDevice {
-            unsafe {
-                copy_nonoverlapping(buffer.as_ptr() as *mut u8, copy.as_ptr(), size);
-            }
+            let src = buffer.cast::<u8>().as_ptr().cast_const();
+            trace!("VirtIO bounce buffer at {bounce:?} (PA:{paddr:#x}) initialized from {src:?}");
+            // SAFETY - Both regions are valid, properly aligned, and don't overlap.
+            unsafe { copy_nonoverlapping(src, bounce.as_ptr(), size) };
         }
-        virt_to_phys(copy)
+
+        paddr
     }
 
-    fn unshare(paddr: PhysAddr, buffer: NonNull<[u8]>, direction: BufferDirection) {
-        let vaddr = phys_to_virt(paddr);
+    unsafe fn unshare(paddr: PhysAddr, buffer: NonNull<[u8]>, direction: BufferDirection) {
+        let bounce = phys_to_virt(paddr);
         let size = buffer.len();
         if direction == BufferDirection::DeviceToDriver {
-            debug!(
-                "Copying VirtIO buffer back from {:#x} to {:#x}.",
-                paddr,
-                buffer.as_ptr() as *mut u8 as usize
-            );
-            unsafe {
-                copy_nonoverlapping(vaddr.as_ptr(), buffer.as_ptr() as *mut u8, size);
-            }
+            let dest = buffer.cast::<u8>().as_ptr();
+            trace!("VirtIO bounce buffer at {bounce:?} (PA:{paddr:#x}) copied back to {dest:?}");
+            // SAFETY - Both regions are valid, properly aligned, and don't overlap.
+            unsafe { copy_nonoverlapping(bounce.as_ptr(), dest, size) };
         }
 
-        // Unshare and deallocate the shared copy of the buffer.
-        debug!("Unsharing VirtIO buffer {:#x}", paddr);
-        // Safe because the memory was allocated by `share` using `alloc_shared`, and the size is
-        // the same as was used then.
-        unsafe {
-            dealloc_shared(vaddr, size).expect("Failed to unshare VirtIO buffer with host");
-        }
+        // SAFETY - Memory was allocated by `share` using `alloc_shared` with the same size.
+        unsafe { dealloc_shared(bounce, bb_layout(size)) }
+            .expect("Failed to unshare and deallocate VirtIO bounce buffer");
     }
 }
 
-/// Returns true if `inner` is entirely contained within `outer`.
-fn contains_range(outer: &Range<u32>, inner: &Range<u32>) -> bool {
-    inner.start >= outer.start && inner.end <= outer.end
+fn dma_layout(pages: usize) -> Layout {
+    let size = pages.checked_mul(PAGE_SIZE).unwrap();
+    Layout::from_size_align(size, PAGE_SIZE).unwrap()
+}
+
+fn bb_layout(size: usize) -> Layout {
+    // In theory, it would be legal to align to 1-byte but use a larger alignment for good measure.
+    const VIRTIO_BOUNCE_BUFFER_ALIGN: usize = size_of::<u128>();
+    Layout::from_size_align(size, VIRTIO_BOUNCE_BUFFER_ALIGN).unwrap()
 }
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index 99e07b6..03fa107 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -28,11 +28,13 @@
     paging::{Attributes, MemoryRegion},
 };
 use buddy_system_allocator::LockedHeap;
+use core::ops::Range;
 use hyp::get_hypervisor;
 use log::{debug, error, info};
-use vmbase::{main, power::reboot};
+use vmbase::{layout, main, power::reboot};
 
 const SZ_1K: usize = 1024;
+const SZ_4K: usize = 4 * SZ_1K;
 const SZ_64K: usize = 64 * SZ_1K;
 const SZ_1M: usize = 1024 * SZ_1K;
 const SZ_1G: usize = 1024 * SZ_1M;
@@ -63,16 +65,8 @@
 
 static mut HEAP: [u8; SZ_64K] = [0; SZ_64K];
 
-unsafe fn kimg_ptr(sym: &u8) -> *const u8 {
-    sym as *const u8
-}
-
-unsafe fn kimg_addr(sym: &u8) -> usize {
-    kimg_ptr(sym) as usize
-}
-
-unsafe fn kimg_region(begin: &u8, end: &u8) -> MemoryRegion {
-    MemoryRegion::new(kimg_addr(begin), kimg_addr(end))
+fn into_memreg(r: &Range<usize>) -> MemoryRegion {
+    MemoryRegion::new(r.start, r.end)
 }
 
 fn init_heap() {
@@ -80,28 +74,28 @@
     unsafe {
         HEAP_ALLOCATOR.lock().init(&mut HEAP as *mut u8 as usize, HEAP.len());
     }
-    info!("Initialized heap.");
 }
 
 fn init_kernel_pgt(pgt: &mut IdMap) -> Result<()> {
     // The first 1 GiB of address space is used by crosvm for MMIO.
     let reg_dev = MemoryRegion::new(0, SZ_1G);
-    // SAFETY: Taking addresses of kernel image sections to set up page table
-    // mappings. Not taking ownerhip of the memory.
-    let reg_text = unsafe { kimg_region(&text_begin, &text_end) };
-    let reg_rodata = unsafe { kimg_region(&rodata_begin, &rodata_end) };
-    let reg_data = unsafe { kimg_region(&data_begin, &boot_stack_end) };
+    let reg_text = into_memreg(&layout::text_range());
+    let reg_rodata = into_memreg(&layout::rodata_range());
+    let reg_scratch = into_memreg(&layout::scratch_range());
+    let reg_stack = into_memreg(&layout::stack_range(40 * SZ_4K));
 
     debug!("Preparing kernel page table.");
     debug!("  dev:    {}-{}", reg_dev.start(), reg_dev.end());
     debug!("  text:   {}-{}", reg_text.start(), reg_text.end());
     debug!("  rodata: {}-{}", reg_rodata.start(), reg_rodata.end());
-    debug!("  data:   {}-{}", reg_data.start(), reg_data.end());
+    debug!("  scratch:{}-{}", reg_scratch.start(), reg_scratch.end());
+    debug!("  stack:  {}-{}", reg_stack.start(), reg_stack.end());
 
     pgt.map_range(&reg_dev, PROT_DEV)?;
     pgt.map_range(&reg_text, PROT_RX)?;
     pgt.map_range(&reg_rodata, PROT_RO)?;
-    pgt.map_range(&reg_data, PROT_RW)?;
+    pgt.map_range(&reg_scratch, PROT_RW)?;
+    pgt.map_range(&reg_stack, PROT_RW)?;
 
     pgt.activate();
     info!("Activated kernel page table.");
@@ -121,7 +115,6 @@
 
 fn try_main() -> Result<()> {
     info!("Welcome to Rialto!");
-    init_heap();
 
     let mut pgt = IdMap::new(PT_ASID, PT_ROOT_LEVEL);
     init_kernel_pgt(&mut pgt)?;
@@ -130,6 +123,7 @@
 
 /// Entry point for Rialto.
 pub fn main(_a0: u64, _a1: u64, _a2: u64, _a3: u64) {
+    init_heap();
     if try_init_logger().is_err() {
         // Don't log anything if the logger initialization fails.
         reboot();
@@ -143,13 +137,4 @@
     }
 }
 
-extern "C" {
-    static text_begin: u8;
-    static text_end: u8;
-    static rodata_begin: u8;
-    static rodata_end: u8;
-    static data_begin: u8;
-    static boot_stack_end: u8;
-}
-
 main!(main);
diff --git a/service_vm/client_apk/Android.bp b/service_vm/client_apk/Android.bp
new file mode 100644
index 0000000..e5084d4
--- /dev/null
+++ b/service_vm/client_apk/Android.bp
@@ -0,0 +1,37 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_app {
+    name: "ServiceVmClientApp",
+    installable: true,
+    jni_libs: ["libservice_vm_client"],
+    jni_uses_platform_apis: true,
+    use_embedded_native_libs: true,
+    sdk_version: "system_current",
+    compile_multilib: "first",
+    apex_available: ["com.android.virt"],
+}
+
+rust_defaults {
+    name: "service_vm_client_defaults",
+    crate_name: "service_vm_client",
+    srcs: ["src/main.rs"],
+    prefer_rlib: true,
+    rustlibs: [
+        "libandroid_logger",
+        "libanyhow",
+        "liblog_rust",
+        "libvm_payload_bindgen",
+    ],
+}
+
+rust_ffi {
+    name: "libservice_vm_client",
+    defaults: ["service_vm_client_defaults"],
+    // TODO(b/250854486): Remove the sanitize section once the bug is fixed.
+    sanitize: {
+        address: false,
+        hwaddress: false,
+    },
+}
diff --git a/service_vm/client_apk/AndroidManifest.xml b/service_vm/client_apk/AndroidManifest.xml
new file mode 100644
index 0000000..b3598fc
--- /dev/null
+++ b/service_vm/client_apk/AndroidManifest.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 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.
+-->
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+      package="com.android.virt.service_vm.client">
+     <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
+     <uses-permission android:name="android.permission.USE_CUSTOM_VIRTUAL_MACHINE" />
+
+     <application android:hasCode="false"/>
+</manifest>
diff --git a/service_vm/client_apk/assets/config.json b/service_vm/client_apk/assets/config.json
new file mode 100644
index 0000000..02749fe
--- /dev/null
+++ b/service_vm/client_apk/assets/config.json
@@ -0,0 +1,10 @@
+{
+    "os": {
+      "name": "microdroid"
+    },
+    "task": {
+      "type": "microdroid_launcher",
+      "command": "libservice_vm_client.so"
+    },
+    "export_tombstones": true
+  }
\ No newline at end of file
diff --git a/service_vm/client_apk/src/main.rs b/service_vm/client_apk/src/main.rs
new file mode 100644
index 0000000..1f8db96
--- /dev/null
+++ b/service_vm/client_apk/src/main.rs
@@ -0,0 +1,71 @@
+// 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.
+
+//! Main executable of Service VM client.
+
+use anyhow::Result;
+use log::{error, info};
+use std::{ffi::c_void, panic};
+use vm_payload_bindgen::AVmPayload_requestCertificate;
+
+/// Entry point of the Service VM client.
+#[allow(non_snake_case)]
+#[no_mangle]
+pub extern "C" fn AVmPayload_main() {
+    android_logger::init_once(
+        android_logger::Config::default()
+            .with_tag("service_vm_client")
+            .with_min_level(log::Level::Debug),
+    );
+    // Redirect panic messages to logcat.
+    panic::set_hook(Box::new(|panic_info| {
+        error!("{}", panic_info);
+    }));
+    if let Err(e) = try_main() {
+        error!("failed with {:?}", e);
+        std::process::exit(1);
+    }
+}
+
+fn try_main() -> Result<()> {
+    info!("Welcome to Service VM Client!");
+    let csr = b"Hello from Service VM";
+    let certificate = request_certificate(csr);
+    info!("Certificate: {:?}", certificate);
+    Ok(())
+}
+
+fn request_certificate(csr: &[u8]) -> Vec<u8> {
+    // SAFETY: It is safe as we only request the size of the certificate in this call.
+    let certificate_size = unsafe {
+        AVmPayload_requestCertificate(
+            csr.as_ptr() as *const c_void,
+            csr.len(),
+            [].as_mut_ptr() as *mut c_void,
+            0,
+        )
+    };
+    let mut certificate = vec![0u8; certificate_size];
+    // SAFETY: It is safe as we only write the data into the given buffer within the buffer
+    // size in this call.
+    unsafe {
+        AVmPayload_requestCertificate(
+            csr.as_ptr() as *const c_void,
+            csr.len(),
+            certificate.as_mut_ptr() as *mut c_void,
+            certificate.len(),
+        );
+    };
+    certificate
+}
diff --git a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
index c8c8660..4d043f6 100644
--- a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
@@ -18,7 +18,7 @@
 
 /** {@hide} */
 interface IBenchmarkService {
-    const int SERVICE_PORT = 5677;
+    const int PORT = 5677;
 
     /**
      * Measures the read rate for reading the given file.
diff --git a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
index 36c3aaf..a6f1c80 100644
--- a/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/ITestService.aidl
@@ -22,7 +22,7 @@
  * {@hide}
  */
 interface ITestService {
-    const long SERVICE_PORT = 5678;
+    const long PORT = 5678;
 
     const long ECHO_REVERSE_PORT = 0x80000001L; // Deliberately chosen to be > 2^31, < 2^32
 
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java
index cbb9a0a..db1d021 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java
@@ -49,7 +49,7 @@
         try {
             IBenchmarkService benchmarkService =
                     IBenchmarkService.Stub.asInterface(
-                            vm.connectToVsockServer(IBenchmarkService.SERVICE_PORT));
+                            vm.connectToVsockServer(IBenchmarkService.PORT));
             assertThat(benchmarkService).isNotNull();
 
             mListener.onPayloadReady(vm, benchmarkService);
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 6cfc71d..5d93b93 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -185,8 +185,8 @@
 Result<void> run_io_benchmark_tests() {
     auto test_service = ndk::SharedRefBase::make<IOBenchmarkService>();
     auto callback = []([[maybe_unused]] void* param) { AVmPayload_notifyPayloadReady(); };
-    AVmPayload_runVsockRpcServer(test_service->asBinder().get(), test_service->SERVICE_PORT,
-                                 callback, nullptr);
+    AVmPayload_runVsockRpcServer(test_service->asBinder().get(), test_service->PORT, callback,
+                                 nullptr);
     return {};
 }
 } // Anonymous namespace
diff --git a/tests/helper/src/java/com/android/microdroid/test/common/MetricsProcessor.java b/tests/helper/src/java/com/android/microdroid/test/common/MetricsProcessor.java
index 42eb6a1..dd68d6a 100644
--- a/tests/helper/src/java/com/android/microdroid/test/common/MetricsProcessor.java
+++ b/tests/helper/src/java/com/android/microdroid/test/common/MetricsProcessor.java
@@ -50,8 +50,8 @@
         Collections.sort(values);
 
         double sum = 0;
-        double min = Double.MAX_VALUE;
-        double max = Double.MIN_VALUE;
+        double min = Double.POSITIVE_INFINITY;
+        double max = Double.NEGATIVE_INFINITY;
         for (Double d : values) {
             sum += d;
             if (min > d) min = d;
@@ -63,7 +63,7 @@
             sqSum += (d - avg) * (d - avg);
         }
         double stdDev = Math.sqrt(sqSum / (values.size() - 1));
-        double median = Double.MIN_VALUE;
+        double median = Double.NaN;
         if (values.size() > 0) {
             int rank = values.size() / 2;
             if (values.size() % 2 == 0) {
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index 4e1d238..f58ce81 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -496,7 +496,7 @@
                         try {
                             mTestService =
                                     ITestService.Stub.asInterface(
-                                            vm.connectToVsockServer(ITestService.SERVICE_PORT));
+                                            vm.connectToVsockServer(ITestService.PORT));
                             // Make sure linkToDeath works, and include it in the log in case it's
                             // helpful.
                             mTestService
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 38bc485..bee424d 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -982,7 +982,7 @@
                         try {
                             ITestService testService =
                                     ITestService.Stub.asInterface(
-                                            vm.connectToVsockServer(ITestService.SERVICE_PORT));
+                                            vm.connectToVsockServer(ITestService.PORT));
                             vmCdis.cdiAttest = testService.insecurelyExposeAttestationCdi();
                             vmCdis.instanceSecret = testService.insecurelyExposeVmInstanceSecret();
                         } catch (Exception e) {
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index d24ddfd..7e0fc5b 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -318,7 +318,7 @@
     auto testService = ndk::SharedRefBase::make<TestService>();
 
     auto callback = []([[maybe_unused]] void* param) { AVmPayload_notifyPayloadReady(); };
-    AVmPayload_runVsockRpcServer(testService->asBinder().get(), testService->SERVICE_PORT, callback,
+    AVmPayload_runVsockRpcServer(testService->asBinder().get(), testService->PORT, callback,
                                  nullptr);
 
     return {};
diff --git a/tests/vmshareapp/src/java/com/android/microdroid/test/sharevm/VmShareServiceImpl.java b/tests/vmshareapp/src/java/com/android/microdroid/test/sharevm/VmShareServiceImpl.java
index edd6bf5..1f7ffb7 100644
--- a/tests/vmshareapp/src/java/com/android/microdroid/test/sharevm/VmShareServiceImpl.java
+++ b/tests/vmshareapp/src/java/com/android/microdroid/test/sharevm/VmShareServiceImpl.java
@@ -145,11 +145,10 @@
 
         Log.i(
                 TAG,
-                "Payload is ready, connecting to the vsock service at port "
-                        + ITestService.SERVICE_PORT);
+                "Payload is ready, connecting to the vsock service at port " + ITestService.PORT);
         ITestService testService =
                 ITestService.Stub.asInterface(
-                        mVirtualMachine.connectToVsockServer(ITestService.SERVICE_PORT));
+                        mVirtualMachine.connectToVsockServer(ITestService.PORT));
         return new RemoteTestServiceDelegate(testService);
     }
 
diff --git a/virtualizationmanager/Android.bp b/virtualizationmanager/Android.bp
index fedfd76..59e507f 100644
--- a/virtualizationmanager/Android.bp
+++ b/virtualizationmanager/Android.bp
@@ -54,7 +54,6 @@
         "libtombstoned_client_rust",
         "libvm_control",
         "libvmconfig",
-        "libvsutil",
         "libzip",
         "libvsock",
         "liblibfdt",
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 4bc25da..f57cb59 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -53,6 +53,7 @@
     self, wait_for_interface, BinderFeatures, ExceptionCode, Interface, ParcelFileDescriptor,
     Status, StatusCode, Strong,
 };
+use disk::QcowFile;
 use lazy_static::lazy_static;
 use log::{debug, error, info, warn};
 use microdroid_payload_config::{OsConfig, Task, TaskType, VmPayloadConfig};
@@ -63,7 +64,7 @@
 use std::convert::TryInto;
 use std::ffi::CStr;
 use std::fs::{read_dir, remove_file, File, OpenOptions};
-use std::io::{BufRead, BufReader, Write};
+use std::io::{BufRead, BufReader, Error, ErrorKind, Write};
 use std::num::{NonZeroU16, NonZeroU32};
 use std::os::unix::io::{FromRawFd, IntoRawFd};
 use std::os::unix::raw::pid_t;
@@ -71,7 +72,6 @@
 use std::sync::{Arc, Mutex, Weak};
 use vmconfig::VmConfig;
 use vsock::VsockStream;
-use vsutil::{clone_file, init_writable_partition};
 use zip::ZipArchive;
 
 /// The unique ID of a VM used (together with a port number) for vsock communication.
@@ -83,8 +83,19 @@
 /// Gaps in composite disk images are filled with a shared zero.img.
 const ZERO_FILLER_SIZE: u64 = 4096;
 
+/// Magic string for the instance image
+const ANDROID_VM_INSTANCE_MAGIC: &str = "Android-VM-instance";
+
+/// Version of the instance image format
+const ANDROID_VM_INSTANCE_VERSION: u16 = 1;
+
 const MICRODROID_OS_NAME: &str = "microdroid";
 
+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)
@@ -178,7 +189,45 @@
         partition_type: PartitionType,
     ) -> binder::Result<()> {
         check_manage_access()?;
-        init_writable_partition(image_fd, size_bytes, partition_type)
+        let size_bytes = size_bytes.try_into().map_err(|e| {
+            Status::new_exception_str(
+                ExceptionCode::ILLEGAL_ARGUMENT,
+                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| {
+            Status::new_service_specific_error_str(
+                -1,
+                Some(format!("Failed to reset a file: {:?}", 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)),
+            )
+        })?;
+
+        match partition_type {
+            PartitionType::RAW => Ok(()),
+            PartitionType::ANDROID_VM_INSTANCE => format_as_android_vm_instance(&mut part),
+            PartitionType::ENCRYPTEDSTORE => format_as_encryptedstore(&mut part),
+            _ => Err(Error::new(
+                ErrorKind::Unsupported,
+                format!("Unsupported partition type {:?}", partition_type),
+            )),
+        }
+        .map_err(|e| {
+            Status::new_service_specific_error_str(
+                -1,
+                Some(format!("Failed to initialize partition as {:?}: {:?}", partition_type, e)),
+            )
+        })?;
+
+        Ok(())
     }
 
     /// Creates or update the idsig file by digesting the input APK file.
@@ -435,6 +484,26 @@
     Ok(())
 }
 
+fn format_as_android_vm_instance(part: &mut dyn Write) -> std::io::Result<()> {
+    part.write_all(ANDROID_VM_INSTANCE_MAGIC.as_bytes())?;
+    part.write_all(&ANDROID_VM_INSTANCE_VERSION.to_le_bytes())?;
+    part.flush()
+}
+
+fn format_as_encryptedstore(part: &mut dyn Write) -> std::io::Result<()> {
+    part.write_all(UNFORMATTED_STORAGE_MAGIC.as_bytes())?;
+    part.flush()
+}
+
+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.
@@ -886,6 +955,16 @@
     }
 }
 
+/// Converts a `&ParcelFileDescriptor` to a `File` by cloning the file.
+pub fn clone_file(file: &ParcelFileDescriptor) -> Result<File, Status> {
+    file.as_ref().try_clone().map_err(|e| {
+        Status::new_exception_str(
+            ExceptionCode::BAD_PARCELABLE,
+            Some(format!("Failed to clone File from ParcelFileDescriptor: {:?}", e)),
+        )
+    })
+}
+
 /// Converts an `&Option<ParcelFileDescriptor>` to an `Option<File>` by cloning the file.
 fn maybe_clone_file(file: &Option<ParcelFileDescriptor>) -> Result<Option<File>, Status> {
     file.as_ref().map(clone_file).transpose()
@@ -1105,6 +1184,31 @@
             ))
         }
     }
+
+    fn requestCertificate(&self, csr: &[u8]) -> binder::Result<Vec<u8>> {
+        let cid = self.cid;
+        let Some(vm) = self.state.lock().unwrap().get_vm(cid) else {
+            error!("requestCertificate is called from an unknown CID {cid}");
+            return Err(Status::new_service_specific_error_str(
+                -1,
+                Some(format!("cannot find a VM with CID {}", cid)),
+            ))
+        };
+        let instance_img_path = vm.temporary_directory.join("rkpvm_instance.img");
+        let instance_img = OpenOptions::new()
+            .create(true)
+            .read(true)
+            .write(true)
+            .open(instance_img_path)
+            .map_err(|e| {
+                error!("Failed to create rkpvm_instance.img file: {:?}", e);
+                Status::new_service_specific_error_str(
+                    -1,
+                    Some(format!("Failed to create rkpvm_instance.img file: {:?}", e)),
+                )
+            })?;
+        GLOBAL_SERVICE.requestCertificate(csr, &ParcelFileDescriptor::new(instance_img))
+    }
 }
 
 impl VirtualMachineService {
diff --git a/virtualizationmanager/src/atom.rs b/virtualizationmanager/src/atom.rs
index 5118a8d..d6eb141 100644
--- a/virtualizationmanager/src/atom.rs
+++ b/virtualizationmanager/src/atom.rs
@@ -14,7 +14,7 @@
 
 //! Functions for creating and collecting atoms.
 
-use crate::aidl::GLOBAL_SERVICE;
+use crate::aidl::{clone_file, GLOBAL_SERVICE};
 use crate::crosvm::VmMetric;
 use crate::get_calling_uid;
 use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::DeathReason::DeathReason;
@@ -37,7 +37,6 @@
 use statslog_virtualization_rust::vm_creation_requested;
 use std::thread;
 use std::time::{Duration, SystemTime};
-use vsutil::clone_file;
 use zip::ZipArchive;
 
 const INVALID_NUM_CPUS: i32 = -1;
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index f7202da..6b39ff9 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -28,6 +28,7 @@
         "libandroid_logger",
         "libanyhow",
         "libbinder_rs",
+        "libvmclient",
         "liblibc",
         "liblog_rust",
         "libnix",
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
index 5422a48..cc59b3f 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -49,4 +49,14 @@
 
     /** Get a list of all currently running VMs. */
     VirtualMachineDebugInfo[] debugListVms();
+
+    /**
+     * Requests a certificate using the provided certificate signing request (CSR).
+     *
+     * @param csr the certificate signing request.
+     * @param instanceImgFd The file descriptor of the instance image. The file should be open for
+     *         both reading and writing.
+     * @return the X.509 encoded certificate.
+     */
+    byte[] requestCertificate(in byte[] csr, in ParcelFileDescriptor instanceImgFd);
 }
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
index 3fdb48a..7b90714 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -44,4 +44,12 @@
      * Notifies that an error has occurred inside the VM.
      */
     void notifyError(ErrorCode errorCode, in String message);
+
+    /**
+     * Requests a certificate using the provided certificate signing request (CSR).
+     *
+     * @param csr the certificate signing request.
+     * @return the X.509 encoded certificate.
+     */
+    byte[] requestCertificate(in byte[] csr);
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 3888df2..5c5a7e4 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -16,8 +16,12 @@
 
 use crate::{get_calling_pid, get_calling_uid};
 use crate::atom::{forward_vm_booted_atom, forward_vm_creation_atom, forward_vm_exited_atom};
+use crate::rkpvm::request_certificate;
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::VirtualMachineDebugInfo::VirtualMachineDebugInfo;
+use android_system_virtualizationservice::{
+    aidl::android::system::virtualizationservice::VirtualMachineDebugInfo::VirtualMachineDebugInfo,
+    binder::ParcelFileDescriptor,
+};
 use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::{
     AtomVmBooted::AtomVmBooted,
     AtomVmCreationRequested::AtomVmCreationRequested,
@@ -153,6 +157,19 @@
             .collect();
         Ok(cids)
     }
+
+    fn requestCertificate(
+        &self,
+        csr: &[u8],
+        instance_img_fd: &ParcelFileDescriptor,
+    ) -> binder::Result<Vec<u8>> {
+        check_manage_access()?;
+        info!("Received csr. Getting certificate...");
+        request_certificate(csr, instance_img_fd).map_err(|e| {
+            error!("Failed to get certificate. Error: {e:?}");
+            Status::new_exception_str(ExceptionCode::SERVICE_SPECIFIC, Some(e.to_string()))
+        })
+    }
 }
 
 #[derive(Debug, Default)]
diff --git a/virtualizationservice/src/main.rs b/virtualizationservice/src/main.rs
index 64ccb13..bf8b944 100644
--- a/virtualizationservice/src/main.rs
+++ b/virtualizationservice/src/main.rs
@@ -16,6 +16,7 @@
 
 mod aidl;
 mod atom;
+mod rkpvm;
 
 use crate::aidl::{
     remove_temporary_dir, BINDER_SERVICE_IDENTIFIER, TEMPORARY_DIRECTORY,
diff --git a/virtualizationservice/src/rkpvm.rs b/virtualizationservice/src/rkpvm.rs
new file mode 100644
index 0000000..a4649f6
--- /dev/null
+++ b/virtualizationservice/src/rkpvm.rs
@@ -0,0 +1,95 @@
+// 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.
+
+//! Handles the RKP (Remote Key Provisioning) VM and host communication.
+//! The RKP VM will be recognized and attested by the RKP server periodically and
+//! serves as a trusted platform to attest a client VM.
+
+use android_system_virtualizationservice::{
+    aidl::android::system::virtualizationservice::{
+        CpuTopology::CpuTopology, DiskImage::DiskImage, Partition::Partition,
+        PartitionType::PartitionType, VirtualMachineConfig::VirtualMachineConfig,
+        VirtualMachineRawConfig::VirtualMachineRawConfig,
+    },
+    binder::{ParcelFileDescriptor, ProcessState},
+};
+use anyhow::{anyhow, Context, Result};
+use log::info;
+use std::fs::File;
+use std::time::Duration;
+use vmclient::VmInstance;
+
+const RIALTO_PATH: &str = "/apex/com.android.virt/etc/rialto.bin";
+
+pub(crate) fn request_certificate(
+    csr: &[u8],
+    instance_img_fd: &ParcelFileDescriptor,
+) -> Result<Vec<u8>> {
+    // We need to start the thread pool for Binder to work properly, especially link_to_death.
+    ProcessState::start_thread_pool();
+
+    let virtmgr = vmclient::VirtualizationService::new().context("Failed to spawn virtmgr")?;
+    let service = virtmgr.connect().context("virtmgr failed to connect")?;
+    info!("service_vm: Connected to VirtualizationService");
+    // TODO(b/272226230): Either turn rialto into the service VM or use an empty payload here.
+    // If using an empty payload, the service code will be part of pvmfw.
+    let rialto = File::open(RIALTO_PATH).context("Failed to open Rialto kernel binary")?;
+
+    // TODO(b/272226230): Initialize the partition from virtualization manager.
+    const INSTANCE_IMG_SIZE_BYTES: i64 = 1 << 20; // 1MB
+    service
+        .initializeWritablePartition(
+            instance_img_fd,
+            INSTANCE_IMG_SIZE_BYTES,
+            PartitionType::ANDROID_VM_INSTANCE,
+        )
+        .context("Failed to initialize instange.img")?;
+    let instance_img =
+        instance_img_fd.as_ref().try_clone().context("Failed to clone instance.img")?;
+    let instance_img = ParcelFileDescriptor::new(instance_img);
+    let writable_partitions = vec![Partition {
+        label: "vm-instance".to_owned(),
+        image: Some(instance_img),
+        writable: true,
+    }];
+    info!("service_vm: Finished initializing instance.img...");
+
+    let config = VirtualMachineConfig::RawConfig(VirtualMachineRawConfig {
+        name: String::from("Service VM"),
+        kernel: None,
+        initrd: None,
+        params: None,
+        bootloader: Some(ParcelFileDescriptor::new(rialto)),
+        disks: vec![DiskImage { image: None, partitions: writable_partitions, writable: true }],
+        protectedVm: true,
+        memoryMib: 300,
+        cpuTopology: CpuTopology::ONE_CPU,
+        platformVersion: "~1.0".to_string(),
+        taskProfiles: vec![],
+        gdbPort: 0, // No gdb
+    });
+    let vm = VmInstance::create(service.as_ref(), &config, None, None, None)
+        .context("Failed to create service VM")?;
+
+    info!("service_vm: Starting Service VM...");
+    vm.start().context("Failed to start service VM")?;
+
+    // TODO(b/274441673): The host can send the CSR to the RKP VM for attestation.
+    // Wait for VM to finish.
+    vm.wait_for_death_with_timeout(Duration::from_secs(10))
+        .ok_or_else(|| anyhow!("Timed out waiting for VM exit"))?;
+
+    info!("service_vm: Finished getting the certificate");
+    Ok([b"Return: ", csr].concat())
+}
diff --git a/vm_payload/Android.bp b/vm_payload/Android.bp
index 967d1cf..77dbb6b 100644
--- a/vm_payload/Android.bp
+++ b/vm_payload/Android.bp
@@ -10,6 +10,8 @@
     srcs: ["src/*.rs"],
     include_dirs: ["include"],
     prefer_rlib: true,
+    // Require unsafe blocks for inside unsafe functions.
+    flags: ["-Dunsafe_op_in_unsafe_fn"],
     rustlibs: [
         "android.system.virtualization.payload-rust",
         "libandroid_logger",
@@ -36,7 +38,10 @@
     crate_name: "vm_payload_bindgen",
     source_stem: "bindings",
     apex_available: ["com.android.compos"],
-    visibility: ["//packages/modules/Virtualization/compos"],
+    visibility: [
+        "//packages/modules/Virtualization/compos",
+        "//packages/modules/Virtualization/service_vm/client_apk",
+    ],
     shared_libs: [
         "libvm_payload#current",
     ],
diff --git a/vm_payload/include-restricted/vm_payload_restricted.h b/vm_payload/include-restricted/vm_payload_restricted.h
index 7f17cde..1e0c3cc 100644
--- a/vm_payload/include-restricted/vm_payload_restricted.h
+++ b/vm_payload/include-restricted/vm_payload_restricted.h
@@ -22,6 +22,10 @@
 
 #include "vm_payload.h"
 
+#if !defined(__INTRODUCED_IN)
+#define __INTRODUCED_IN(__api_level) /* nothing */
+#endif
+
 // The functions declared here are restricted to VMs created with a config file;
 // they will fail if called in other VMs. The ability to create such VMs
 // requires the android.permission.USE_CUSTOM_VIRTUAL_MACHINE permission, and is
@@ -51,4 +55,18 @@
  */
 size_t AVmPayload_getDiceAttestationCdi(void* _Nullable data, size_t size);
 
+/**
+ * Requests a certificate using the provided certificate signing request (CSR).
+ *
+ * \param csr A pointer to the CSR buffer.
+ * \param csr_size The size of the CSR buffer.
+ * \param buffer A pointer to the certificate buffer.
+ * \param size number of bytes that can be written to the certificate buffer.
+ *
+ * \return the total size of the certificate
+ */
+size_t AVmPayload_requestCertificate(const void* _Nonnull csr, size_t csr_size,
+                                     void* _Nullable buffer, size_t size)
+        __INTRODUCED_IN(__ANDROID_API_V__);
+
 __END_DECLS
diff --git a/vm_payload/libvm_payload.map.txt b/vm_payload/libvm_payload.map.txt
index a2402d1..f0d867e 100644
--- a/vm_payload/libvm_payload.map.txt
+++ b/vm_payload/libvm_payload.map.txt
@@ -7,6 +7,7 @@
     AVmPayload_getDiceAttestationCdi;    # systemapi
     AVmPayload_getApkContentsPath;       # systemapi
     AVmPayload_getEncryptedStoragePath;  # systemapi
+    AVmPayload_requestCertificate;       # systemapi introduced=35
   local:
     *;
 };
diff --git a/vm_payload/src/api.rs b/vm_payload/src/api.rs
index 4b565e0..00d7299 100644
--- a/vm_payload/src/api.rs
+++ b/vm_payload/src/api.rs
@@ -14,9 +14,6 @@
 
 //! This module handles the interaction with virtual machine payload service.
 
-// We're implementing unsafe functions, but we still want warnings on unsafe usage within them.
-#![warn(unsafe_op_in_unsafe_fn)]
-
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
     ENCRYPTEDSTORE_MOUNTPOINT, IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
 use anyhow::{ensure, bail, Context, Result};
@@ -256,6 +253,52 @@
     get_vm_payload_service()?.getDiceAttestationCdi().context("Cannot get attestation CDI")
 }
 
+/// Requests a certificate using the provided certificate signing request (CSR).
+/// Panics on failure.
+///
+/// # Safety
+///
+/// Behavior is undefined if any of the following conditions are violated:
+///
+/// * `csr` must be [valid] for reads of `csr_size` bytes.
+/// * `buffer` must be [valid] for writes of `size` bytes. `buffer` can be null if `size` is 0.
+///
+/// [valid]: ptr#safety
+#[no_mangle]
+pub unsafe extern "C" fn AVmPayload_requestCertificate(
+    csr: *const u8,
+    csr_size: usize,
+    buffer: *mut u8,
+    size: usize,
+) -> usize {
+    initialize_logging();
+
+    // SAFETY: See the requirements on `csr` above.
+    let csr = unsafe { std::slice::from_raw_parts(csr, csr_size) };
+    let certificate = unwrap_or_abort(try_request_certificate(csr));
+
+    if size != 0 || buffer.is_null() {
+        // SAFETY: See the requirements on `buffer` above. The number of bytes copied doesn't exceed
+        // the length of either buffer, and `certificate` cannot overlap `buffer` because we just
+        // allocated it.
+        unsafe {
+            ptr::copy_nonoverlapping(
+                certificate.as_ptr(),
+                buffer,
+                std::cmp::min(certificate.len(), size),
+            );
+        }
+    }
+    certificate.len()
+}
+
+fn try_request_certificate(csr: &[u8]) -> Result<Vec<u8>> {
+    let certificate = get_vm_payload_service()?
+        .requestCertificate(csr)
+        .context("Failed to request certificate")?;
+    Ok(certificate)
+}
+
 /// Gets the path to the APK contents.
 #[no_mangle]
 pub extern "C" fn AVmPayload_getApkContentsPath() -> *const c_char {
diff --git a/vm_payload/src/lib.rs b/vm_payload/src/lib.rs
index 5c3ee31..4d059d1 100644
--- a/vm_payload/src/lib.rs
+++ b/vm_payload/src/lib.rs
@@ -17,6 +17,7 @@
 mod api;
 
 pub use api::{
-    AVmPayload_getDiceAttestationCdi, AVmPayload_getDiceAttestationChain,
-    AVmPayload_getVmInstanceSecret, AVmPayload_notifyPayloadReady,
+    AVmPayload_getCertificate, AVmPayload_getDiceAttestationCdi,
+    AVmPayload_getDiceAttestationChain, AVmPayload_getVmInstanceSecret,
+    AVmPayload_notifyPayloadReady,
 };
diff --git a/vmbase/entry.S b/vmbase/entry.S
index 408f5d1..9f6993a 100644
--- a/vmbase/entry.S
+++ b/vmbase/entry.S
@@ -208,8 +208,14 @@
 	stp q0, q1, [x28], #32
 	b 2b
 
-3:	/* Prepare the stack. */
-	adr_l x30, boot_stack_end
+3:	/* Prepare the exception handler stack (SP_EL1). */
+	adr_l x30, init_eh_stack_pointer
+	msr spsel, #1
+	mov sp, x30
+
+	/* Prepare the main thread stack (SP_EL0). */
+	adr_l x30, init_stack_pointer
+	msr spsel, #0
 	mov sp, x30
 
 	/* Set up exception vector. */
diff --git a/vmbase/example/src/exceptions.rs b/vmbase/example/src/exceptions.rs
index 0e637ac..0522013 100644
--- a/vmbase/example/src/exceptions.rs
+++ b/vmbase/example/src/exceptions.rs
@@ -15,56 +15,56 @@
 //! Exception handlers.
 
 use core::arch::asm;
-use vmbase::{console::emergency_write_str, eprintln, power::reboot};
+use vmbase::{eprintln, power::reboot};
 
 #[no_mangle]
 extern "C" fn sync_exception_current(_elr: u64, _spsr: u64) {
-    emergency_write_str("sync_exception_current\n");
+    eprintln!("sync_exception_current");
     print_esr();
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn irq_current(_elr: u64, _spsr: u64) {
-    emergency_write_str("irq_current\n");
+    eprintln!("irq_current");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn fiq_current(_elr: u64, _spsr: u64) {
-    emergency_write_str("fiq_current\n");
+    eprintln!("fiq_current");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn serr_current(_elr: u64, _spsr: u64) {
-    emergency_write_str("serr_current\n");
+    eprintln!("serr_current");
     print_esr();
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn sync_lower(_elr: u64, _spsr: u64) {
-    emergency_write_str("sync_lower\n");
+    eprintln!("sync_lower");
     print_esr();
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn irq_lower(_elr: u64, _spsr: u64) {
-    emergency_write_str("irq_lower\n");
+    eprintln!("irq_lower");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn fiq_lower(_elr: u64, _spsr: u64) {
-    emergency_write_str("fiq_lower\n");
+    eprintln!("fiq_lower");
     reboot();
 }
 
 #[no_mangle]
 extern "C" fn serr_lower(_elr: u64, _spsr: u64) {
-    emergency_write_str("serr_lower\n");
+    eprintln!("serr_lower");
     print_esr();
     reboot();
 }
diff --git a/vmbase/example/src/layout.rs b/vmbase/example/src/layout.rs
index e660c5f..2e9d27a 100644
--- a/vmbase/example/src/layout.rs
+++ b/vmbase/example/src/layout.rs
@@ -55,13 +55,13 @@
 
 /// Writable data region for the stack.
 pub fn boot_stack_range() -> Range<VirtualAddress> {
-    into_va_range(layout::boot_stack_range())
+    const PAGE_SIZE: usize = 4 << 10;
+    into_va_range(layout::stack_range(40 * PAGE_SIZE))
 }
 
-/// Writable data, including the stack.
-pub fn writable_region() -> MemoryRegion {
-    let r = layout::writable_region();
-    MemoryRegion::new(r.start, r.end)
+/// Writable data region for allocations.
+pub fn scratch_range() -> Range<VirtualAddress> {
+    into_va_range(layout::scratch_range())
 }
 
 fn data_load_address() -> VirtualAddress {
diff --git a/vmbase/example/src/main.rs b/vmbase/example/src/main.rs
index ed0275b..3bf850c 100644
--- a/vmbase/example/src/main.rs
+++ b/vmbase/example/src/main.rs
@@ -24,8 +24,8 @@
 extern crate alloc;
 
 use crate::layout::{
-    bionic_tls, dtb_range, print_addresses, rodata_range, stack_chk_guard, text_range,
-    writable_region, DEVICE_REGION,
+    bionic_tls, boot_stack_range, dtb_range, print_addresses, rodata_range, scratch_range,
+    stack_chk_guard, text_range, DEVICE_REGION,
 };
 use crate::pci::{check_pci, get_bar_region};
 use aarch64_paging::{idmap::IdMap, paging::Attributes};
@@ -100,7 +100,13 @@
         .unwrap();
     idmap
         .map_range(
-            &writable_region(),
+            &scratch_range().into(),
+            Attributes::NORMAL | Attributes::NON_GLOBAL | Attributes::EXECUTE_NEVER,
+        )
+        .unwrap();
+    idmap
+        .map_range(
+            &boot_stack_range().into(),
             Attributes::NORMAL | Attributes::NON_GLOBAL | Attributes::EXECUTE_NEVER,
         )
         .unwrap();
diff --git a/vmbase/example/src/pci.rs b/vmbase/example/src/pci.rs
index 117cbc8..41a3ff4 100644
--- a/vmbase/example/src/pci.rs
+++ b/vmbase/example/src/pci.rs
@@ -98,7 +98,7 @@
 
 struct HalImpl;
 
-impl Hal for HalImpl {
+unsafe impl Hal for HalImpl {
     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();
@@ -110,7 +110,7 @@
         (paddr, vaddr)
     }
 
-    fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
+    unsafe fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
         debug!("dma_dealloc: paddr={:#x}, pages={}", paddr, pages);
         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
@@ -121,17 +121,17 @@
         0
     }
 
-    fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
+    unsafe fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
         NonNull::new(paddr as _).unwrap()
     }
 
-    fn share(buffer: NonNull<[u8]>, _direction: BufferDirection) -> PhysAddr {
+    unsafe fn share(buffer: NonNull<[u8]>, _direction: BufferDirection) -> PhysAddr {
         let vaddr = buffer.cast();
         // Nothing to do, as the host already has access to all memory.
         virt_to_phys(vaddr)
     }
 
-    fn unshare(_paddr: PhysAddr, _buffer: NonNull<[u8]>, _direction: BufferDirection) {
+    unsafe fn unshare(_paddr: PhysAddr, _buffer: NonNull<[u8]>, _direction: BufferDirection) {
         // Nothing to do, as the host already has access to all memory and we didn't copy the buffer
         // anywhere else.
     }
diff --git a/vmbase/sections.ld b/vmbase/sections.ld
index 87b909d..5232d30 100644
--- a/vmbase/sections.ld
+++ b/vmbase/sections.ld
@@ -39,12 +39,10 @@
 	 * Collect together the code. This is page aligned so it can be mapped
 	 * as executable-only.
 	 */
-	.init : ALIGN(4096) {
+	.text : ALIGN(4096) {
 		text_begin = .;
 		*(.init.entry)
 		*(.init.*)
-	} >image
-	.text : {
 		*(.text.*)
 	} >image
 	text_end = .;
@@ -62,6 +60,17 @@
 	} >image
 	rodata_end = .;
 
+	.eh_stack (NOLOAD) : ALIGN(4096) {
+		/*
+		 * Get stack overflow guard from the previous page being from
+		 * .rodata and mapped read-only or left unmapped.
+		 */
+		eh_stack_limit = .;
+		. += 4096;
+		. = ALIGN(4096);
+		init_eh_stack_pointer = .;
+	} >writable_data
+
 	/*
 	 * Collect together the read-write data including .bss at the end which
 	 * will be zero'd by the entry code. This is page aligned so it can be
@@ -91,11 +100,11 @@
 		bss_end = .;
 	} >writable_data
 
+	init_stack_pointer = ORIGIN(writable_data) + LENGTH(writable_data);
 	.stack (NOLOAD) : ALIGN(4096) {
-		boot_stack_begin = .;
-		. += 40 * 4096;
-		. = ALIGN(4096);
-		boot_stack_end = .;
+		. += 4096; /* Ensure we have one guard page for overflow. */
+		stack_limit = .;
+		. = init_stack_pointer;
 	} >writable_data
 
 	/*
diff --git a/vmbase/src/layout.rs b/vmbase/src/layout.rs
index b0a5173..ead4f8e 100644
--- a/vmbase/src/layout.rs
+++ b/vmbase/src/layout.rs
@@ -14,53 +14,73 @@
 
 //! Memory layout.
 
-use crate::linker;
 use core::ops::Range;
 use core::ptr::addr_of;
 
+/// Get an address from a linker-defined symbol.
+#[macro_export]
+macro_rules! linker_addr {
+    ($symbol:ident) => {{
+        unsafe { addr_of!($crate::linker::$symbol) as usize }
+    }};
+}
+
+/// Get the address range between a pair of linker-defined symbols.
+#[macro_export]
+macro_rules! linker_region {
+    ($begin:ident,$end:ident) => {{
+        let start = linker_addr!($begin);
+        let end = linker_addr!($end);
+
+        start..end
+    }};
+}
+
 /// Memory reserved for the DTB.
 pub fn dtb_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::dtb_begin) as usize)..(addr_of!(linker::dtb_end) as usize) }
+    linker_region!(dtb_begin, dtb_end)
 }
 
 /// Executable code.
 pub fn text_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::text_begin) as usize)..(addr_of!(linker::text_end) as usize) }
+    linker_region!(text_begin, text_end)
 }
 
 /// Read-only data.
 pub fn rodata_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::rodata_begin) as usize)..(addr_of!(linker::rodata_end) as usize) }
+    linker_region!(rodata_begin, rodata_end)
 }
 
 /// Initialised writable data.
 pub fn data_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::data_begin) as usize)..(addr_of!(linker::data_end) as usize) }
+    linker_region!(data_begin, data_end)
 }
 
 /// Zero-initialised writable data.
 pub fn bss_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::bss_begin) as usize)..(addr_of!(linker::bss_end) as usize) }
+    linker_region!(bss_begin, bss_end)
 }
 
 /// Writable data region for the stack.
-pub fn boot_stack_range() -> Range<usize> {
-    unsafe {
-        (addr_of!(linker::boot_stack_begin) as usize)..(addr_of!(linker::boot_stack_end) as usize)
-    }
+pub fn stack_range(stack_size: usize) -> Range<usize> {
+    let end = linker_addr!(init_stack_pointer);
+    let start = end.checked_sub(stack_size).unwrap();
+    assert!(start >= linker_addr!(stack_limit));
+
+    start..end
 }
 
-/// Writable data, including the stack.
-pub fn writable_region() -> Range<usize> {
-    data_range().start..boot_stack_range().end
+/// All writable sections, excluding the stack.
+pub fn scratch_range() -> Range<usize> {
+    linker_region!(eh_stack_limit, bss_end)
 }
 
 /// Read-write data (original).
 pub fn data_load_address() -> usize {
-    unsafe { addr_of!(linker::data_lma) as usize }
+    linker_addr!(data_lma)
 }
 
 /// End of the binary image.
 pub fn binary_end() -> usize {
-    unsafe { addr_of!(linker::bin_end) as usize }
+    linker_addr!(bin_end)
 }
diff --git a/vmbase/src/linker.rs b/vmbase/src/linker.rs
index f4baae8..97bef3f 100644
--- a/vmbase/src/linker.rs
+++ b/vmbase/src/linker.rs
@@ -19,10 +19,6 @@
     pub static __stack_chk_guard: u64;
     /// First byte beyond the pre-loaded binary.
     pub static bin_end: u8;
-    /// First byte of the `.stack` section.
-    pub static boot_stack_begin: u8;
-    /// First byte beyond the `.stack` section.
-    pub static boot_stack_end: u8;
     /// First byte of the `.bss` section.
     pub static bss_begin: u8;
     /// First byte beyond the `.bss` section.
@@ -37,10 +33,16 @@
     pub static dtb_begin: u8;
     /// First byte beyond the `.dtb` section.
     pub static dtb_end: u8;
+    /// First byte of the region available for the exception handler stack.
+    pub static eh_stack_limit: u8;
+    /// First byte past the region available for the stack.
+    pub static init_stack_pointer: u8;
     /// First byte of the `.rodata` section.
     pub static rodata_begin: u8;
     /// First byte beyond the `.rodata` section.
     pub static rodata_end: u8;
+    /// First byte of the region available for the stack.
+    pub static stack_limit: u8;
     /// First byte of the `.text` section.
     pub static text_begin: u8;
     /// First byte beyond the `.text` section.