diff --git a/compos/common/compos_client.rs b/compos/common/compos_client.rs
index 9648f20..a8a176a 100644
--- a/compos/common/compos_client.rs
+++ b/compos/common/compos_client.rs
@@ -24,7 +24,10 @@
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
     CpuTopology::CpuTopology,
     IVirtualizationService::IVirtualizationService,
-    VirtualMachineAppConfig::{DebugLevel::DebugLevel, Payload::Payload, VirtualMachineAppConfig},
+    VirtualMachineAppConfig::{
+        CustomConfig::CustomConfig, DebugLevel::DebugLevel, Payload::Payload,
+        VirtualMachineAppConfig,
+    },
     VirtualMachineConfig::VirtualMachineConfig,
 };
 use anyhow::{anyhow, bail, Context, Result};
@@ -128,9 +131,10 @@
             protectedVm: protected_vm,
             memoryMib: parameters.memory_mib.unwrap_or(0), // 0 means use the default
             cpuTopology: cpu_topology,
-            taskProfiles: parameters.task_profiles.clone(),
-            gdbPort: 0, // Don't start gdb-server
-            customKernelImage: None,
+            customConfig: Some(CustomConfig {
+                taskProfiles: parameters.task_profiles.clone(),
+                ..Default::default()
+            }),
         });
 
         // Let logs go to logcat.
diff --git a/docs/getting_started/index.md b/docs/getting_started/index.md
index 0e4f2be..9dcd4fa 100644
--- a/docs/getting_started/index.md
+++ b/docs/getting_started/index.md
@@ -103,9 +103,12 @@
 on pVM. You can run a Microdroid with empty payload using the following command:
 
 ```shell
-adb shell /apex/com.android.virt/bin/vm run-microdroid --debug full
+adb shell /apex/com.android.virt/bin/vm run-microdroid
 ```
 
+which spawns a "debuggable" VM by default to allow access to guest kernel logs.
+To run a production non-debuggable VM, pass `--debug none`.
+
 ## 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/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index 2ce2b62..b400eeb 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -61,7 +61,6 @@
 @SystemApi
 public final class VirtualMachineConfig {
     private static final String TAG = "VirtualMachineConfig";
-    private static final String[] EMPTY_STRING_ARRAY = {};
 
     // These define the schema of the config file persisted on disk.
     private static final int VERSION = 6;
@@ -502,8 +501,6 @@
                 vsConfig.cpuTopology = android.system.virtualizationservice.CpuTopology.ONE_CPU;
                 break;
         }
-        // Don't allow apps to set task profiles ... at least for now.
-        vsConfig.taskProfiles = EMPTY_STRING_ARRAY;
         return vsConfig;
     }
 
diff --git a/libs/avb/Android.bp b/libs/avb/Android.bp
index 3a671e2..a2d9e1a 100644
--- a/libs/avb/Android.bp
+++ b/libs/avb/Android.bp
@@ -10,7 +10,6 @@
     visibility: ["//packages/modules/Virtualization:__subpackages__"],
     source_stem: "bindings",
     bindgen_flags: [
-        "--size_t-is-usize",
         "--constified-enum-module AvbDescriptorTag",
         "--default-enum-style rust",
         "--allowlist-type=AvbDescriptorTag",
diff --git a/libs/hyp/src/error.rs b/libs/hyp/src/error.rs
index 408150e..b8498ca 100644
--- a/libs/hyp/src/error.rs
+++ b/libs/hyp/src/error.rs
@@ -14,6 +14,7 @@
 
 //! Error and Result types for hypervisor.
 
+use crate::GeniezoneError;
 use crate::KvmError;
 use core::{fmt, result};
 use uuid::Uuid;
@@ -28,7 +29,9 @@
     MmioGuardNotsupported,
     /// Failed to invoke a certain KVM HVC function.
     KvmError(KvmError, u32),
-    /// Unsupported Hypervisor.
+    /// Failed to invoke GenieZone HVC function.
+    GeniezoneError(GeniezoneError, u32),
+    /// Unsupported Hypervisor
     UnsupportedHypervisorUuid(Uuid),
     /// The MMIO_GUARD granule used by the hypervisor is not supported.
     UnsupportedMmioGuardGranule(usize),
@@ -41,6 +44,12 @@
             Self::KvmError(e, function_id) => {
                 write!(f, "Failed to invoke the HVC function with function ID {function_id}: {e}")
             }
+            Self::GeniezoneError(e, function_id) => {
+                write!(
+                    f,
+                    "Failed to invoke GenieZone HVC function with function ID {function_id}: {e}"
+                )
+            }
             Self::UnsupportedHypervisorUuid(u) => {
                 write!(f, "Unsupported Hypervisor UUID {u}")
             }
diff --git a/libs/hyp/src/hypervisor/geniezone.rs b/libs/hyp/src/hypervisor/geniezone.rs
new file mode 100644
index 0000000..0741978
--- /dev/null
+++ b/libs/hyp/src/hypervisor/geniezone.rs
@@ -0,0 +1,159 @@
+// 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.
+
+//! Wrappers around calls to the GenieZone hypervisor.
+
+use super::common::{Hypervisor, HypervisorCap, MMIO_GUARD_GRANULE_SIZE};
+use crate::error::{Error, Result};
+use crate::util::page_address;
+use core::fmt::{self, Display, Formatter};
+use smccc::{
+    error::{positive_or_error_64, success_or_error_64},
+    hvc64,
+};
+use uuid::{uuid, Uuid};
+
+pub(super) struct GeniezoneHypervisor;
+
+const ARM_SMCCC_GZVM_FUNC_HYP_MEMINFO: u32 = 0xc6000002;
+const ARM_SMCCC_GZVM_FUNC_MEM_SHARE: u32 = 0xc6000003;
+const ARM_SMCCC_GZVM_FUNC_MEM_UNSHARE: u32 = 0xc6000004;
+
+const VENDOR_HYP_GZVM_MMIO_GUARD_INFO_FUNC_ID: u32 = 0xc6000005;
+const VENDOR_HYP_GZVM_MMIO_GUARD_ENROLL_FUNC_ID: u32 = 0xc6000006;
+const VENDOR_HYP_GZVM_MMIO_GUARD_MAP_FUNC_ID: u32 = 0xc6000007;
+const VENDOR_HYP_GZVM_MMIO_GUARD_UNMAP_FUNC_ID: u32 = 0xc6000008;
+
+impl GeniezoneHypervisor {
+    // We generate this uuid by ourselves to identify GenieZone hypervisor
+    // and share the same identification along with guest VMs.
+    // The previous uuid was removed due to duplication elsewhere.
+    pub const UUID: Uuid = uuid!("7e134ed0-3b82-488d-8cee-69c19211dbe7");
+    const CAPABILITIES: HypervisorCap = HypervisorCap::DYNAMIC_MEM_SHARE;
+}
+
+/// Error from a GenieZone HVC call.
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub enum GeniezoneError {
+    /// The call is not supported by the implementation.
+    NotSupported,
+    /// The call is not required to implement.
+    NotRequired,
+    /// One of the call parameters has a invalid value.
+    InvalidParameter,
+    /// There was an unexpected return value.
+    Unknown(i64),
+}
+
+impl From<i64> for GeniezoneError {
+    fn from(value: i64) -> Self {
+        match value {
+            -1 => GeniezoneError::NotSupported,
+            -2 => GeniezoneError::NotRequired,
+            -3 => GeniezoneError::InvalidParameter,
+            _ => GeniezoneError::Unknown(value),
+        }
+    }
+}
+
+impl From<i32> for GeniezoneError {
+    fn from(value: i32) -> Self {
+        i64::from(value).into()
+    }
+}
+
+impl Display for GeniezoneError {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        match self {
+            Self::NotSupported => write!(f, "GenieZone call not supported"),
+            Self::NotRequired => write!(f, "GenieZone call not required"),
+            Self::InvalidParameter => write!(f, "GenieZone call received invalid value"),
+            Self::Unknown(e) => write!(f, "Unknown return value from GenieZone {} ({0:#x})", e),
+        }
+    }
+}
+
+impl Hypervisor for GeniezoneHypervisor {
+    fn mmio_guard_init(&self) -> Result<()> {
+        mmio_guard_enroll()?;
+        let mmio_granule = mmio_guard_granule()?;
+        if mmio_granule != MMIO_GUARD_GRANULE_SIZE {
+            return Err(Error::UnsupportedMmioGuardGranule(mmio_granule));
+        }
+        Ok(())
+    }
+
+    fn mmio_guard_map(&self, addr: usize) -> Result<()> {
+        let mut args = [0u64; 17];
+        args[0] = page_address(addr);
+
+        checked_hvc64_expect_zero(VENDOR_HYP_GZVM_MMIO_GUARD_MAP_FUNC_ID, args)
+    }
+
+    fn mmio_guard_unmap(&self, addr: usize) -> Result<()> {
+        let mut args = [0u64; 17];
+        args[0] = page_address(addr);
+
+        checked_hvc64_expect_zero(VENDOR_HYP_GZVM_MMIO_GUARD_UNMAP_FUNC_ID, args)
+    }
+
+    fn mem_share(&self, base_ipa: u64) -> Result<()> {
+        let mut args = [0u64; 17];
+        args[0] = base_ipa;
+
+        checked_hvc64_expect_zero(ARM_SMCCC_GZVM_FUNC_MEM_SHARE, args)
+    }
+
+    fn mem_unshare(&self, base_ipa: u64) -> Result<()> {
+        let mut args = [0u64; 17];
+        args[0] = base_ipa;
+
+        checked_hvc64_expect_zero(ARM_SMCCC_GZVM_FUNC_MEM_UNSHARE, args)
+    }
+
+    fn memory_protection_granule(&self) -> Result<usize> {
+        let args = [0u64; 17];
+        let granule = checked_hvc64(ARM_SMCCC_GZVM_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> {
+    let args = [0u64; 17];
+
+    let granule = checked_hvc64(VENDOR_HYP_GZVM_MMIO_GUARD_INFO_FUNC_ID, args)?;
+    Ok(granule.try_into().unwrap())
+}
+
+fn mmio_guard_enroll() -> Result<()> {
+    let args = [0u64; 17];
+    match success_or_error_64(hvc64(VENDOR_HYP_GZVM_MMIO_GUARD_ENROLL_FUNC_ID, args)[0]) {
+        Ok(_) => Ok(()),
+        Err(GeniezoneError::NotSupported) => Err(Error::MmioGuardNotsupported),
+        Err(GeniezoneError::NotRequired) => Err(Error::MmioGuardNotsupported),
+        Err(e) => Err(Error::GeniezoneError(e, VENDOR_HYP_GZVM_MMIO_GUARD_ENROLL_FUNC_ID)),
+    }
+}
+
+fn checked_hvc64_expect_zero(function: u32, args: [u64; 17]) -> Result<()> {
+    success_or_error_64(hvc64(function, args)[0]).map_err(|e| Error::GeniezoneError(e, function))
+}
+
+fn checked_hvc64(function: u32, args: [u64; 17]) -> Result<u64> {
+    positive_or_error_64(hvc64(function, args)[0]).map_err(|e| Error::GeniezoneError(e, function))
+}
diff --git a/libs/hyp/src/hypervisor/mod.rs b/libs/hyp/src/hypervisor/mod.rs
index 923a21d..93d53fe 100644
--- a/libs/hyp/src/hypervisor/mod.rs
+++ b/libs/hyp/src/hypervisor/mod.rs
@@ -17,6 +17,7 @@
 extern crate alloc;
 
 mod common;
+mod geniezone;
 mod gunyah;
 mod kvm;
 
@@ -25,6 +26,8 @@
 pub use common::Hypervisor;
 pub use common::HypervisorCap;
 pub use common::MMIO_GUARD_GRANULE_SIZE;
+pub use geniezone::GeniezoneError;
+use geniezone::GeniezoneHypervisor;
 use gunyah::GunyahHypervisor;
 pub use kvm::KvmError;
 use kvm::KvmHypervisor;
@@ -35,6 +38,7 @@
 enum HypervisorBackend {
     Kvm,
     Gunyah,
+    Geniezone,
 }
 
 impl HypervisorBackend {
@@ -42,6 +46,7 @@
         match self {
             Self::Kvm => &KvmHypervisor,
             Self::Gunyah => &GunyahHypervisor,
+            Self::Geniezone => &GeniezoneHypervisor,
         }
     }
 }
@@ -51,6 +56,7 @@
 
     fn try_from(uuid: Uuid) -> Result<HypervisorBackend> {
         match uuid {
+            GeniezoneHypervisor::UUID => Ok(HypervisorBackend::Geniezone),
             GunyahHypervisor::UUID => Ok(HypervisorBackend::Gunyah),
             KvmHypervisor::UUID => Ok(HypervisorBackend::Kvm),
             u => Err(Error::UnsupportedHypervisorUuid(u)),
diff --git a/libs/hyp/src/lib.rs b/libs/hyp/src/lib.rs
index 2c2d1d6..32a59d1 100644
--- a/libs/hyp/src/lib.rs
+++ b/libs/hyp/src/lib.rs
@@ -22,3 +22,5 @@
 
 pub use error::{Error, Result};
 pub use hypervisor::{get_hypervisor, Hypervisor, HypervisorCap, KvmError, MMIO_GUARD_GRANULE_SIZE};
+
+use hypervisor::GeniezoneError;
diff --git a/libs/ignorabletest/Android.bp b/libs/ignorabletest/Android.bp
index 10aef8e..4ae89af 100644
--- a/libs/ignorabletest/Android.bp
+++ b/libs/ignorabletest/Android.bp
@@ -9,6 +9,8 @@
     rustlibs: [
         "liblibtest_mimic",
         "liblinkme",
+        "liblog_rust",
+        "liblogger",
     ],
     proc_macros: ["libpaste"],
     apex_available: [
diff --git a/libs/ignorabletest/src/runner.rs b/libs/ignorabletest/src/runner.rs
index 4ec3d79..fdac406 100644
--- a/libs/ignorabletest/src/runner.rs
+++ b/libs/ignorabletest/src/runner.rs
@@ -3,6 +3,7 @@
 use core::ops::{Deref, FnOnce};
 use libtest_mimic::{Arguments, Failed, Trial};
 use linkme::distributed_slice;
+use log::Level;
 use std::env;
 
 /// Command-line arguments to ignore, because they are not supported by libtest-mimic.
@@ -15,6 +16,7 @@
 
 /// Runs all tests.
 pub fn main() {
+    logger::init(logger::Config::default().with_min_level(Level::Debug));
     let args = Arguments::from_iter(env::args().filter(|arg| !IGNORED_ARGS.contains(&arg.deref())));
     let tests = IGNORABLETEST_TESTS.iter().map(|test| test()).collect();
     libtest_mimic::run(&args, tests).exit();
diff --git a/libs/libfdt/Android.bp b/libs/libfdt/Android.bp
index 2a6e75f..0540f26 100644
--- a/libs/libfdt/Android.bp
+++ b/libs/libfdt/Android.bp
@@ -8,7 +8,6 @@
     wrapper_src: "bindgen/fdt.h",
     source_stem: "bindings",
     bindgen_flags: [
-        "--size_t-is-usize",
         "--allowlist-type=fdt_.*",
         "--allowlist-function=fdt_.*",
         "--allowlist-var=FDT_.*",
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index c9909e6..bbe00b5 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -7,7 +7,6 @@
     crate_name: "pvmfw",
     defaults: ["vmbase_ffi_defaults"],
     srcs: ["src/main.rs"],
-    edition: "2021",
     // Require unsafe blocks for inside unsafe functions.
     flags: ["-Dunsafe_op_in_unsafe_fn"],
     features: [
@@ -16,7 +15,6 @@
     rustlibs: [
         "libaarch64_paging",
         "libbssl_ffi_nostd",
-        "libbuddy_system_allocator",
         "libciborium_nostd",
         "libciborium_io_nostd",
         "libdiced_open_dice_nostd",
@@ -83,7 +81,6 @@
     // partition image. This is just to package the unstripped file into the
     // symbols zip file for debugging purpose.
     installable: true,
-    native_coverage: false,
 }
 
 raw_binary {
@@ -135,11 +132,9 @@
 
 rust_library_rlib {
     name: "libpvmfw_embedded_key",
-    defaults: ["vmbase_ffi_defaults"],
-    prefer_rlib: true,
+    defaults: ["vmbase_rlib_defaults"],
     srcs: [":pvmfw_embedded_key_rs"],
     crate_name: "pvmfw_embedded_key",
-    apex_available: ["com.android.virt"],
 }
 
 prebuilt_etc {
@@ -193,8 +188,7 @@
 
 rust_library_rlib {
     name: "libpvmfw_fdt_template",
-    defaults: ["vmbase_ffi_defaults"],
-    prefer_rlib: true,
+    defaults: ["vmbase_rlib_defaults"],
     srcs: [":pvmfw_fdt_template_rs"],
     crate_name: "pvmfw_fdt_template",
 }
diff --git a/pvmfw/avb/Android.bp b/pvmfw/avb/Android.bp
index 5353a21..49c4717 100644
--- a/pvmfw/avb/Android.bp
+++ b/pvmfw/avb/Android.bp
@@ -2,8 +2,8 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
-rust_defaults {
-    name: "libpvmfw_avb_nostd_defaults",
+rust_library_rlib {
+    name: "libpvmfw_avb_nostd",
     crate_name: "pvmfw_avb",
     srcs: ["src/lib.rs"],
     prefer_rlib: true,
@@ -16,11 +16,6 @@
     whole_static_libs: [
         "libavb_baremetal",
     ],
-}
-
-rust_library_rlib {
-    name: "libpvmfw_avb_nostd",
-    defaults: ["libpvmfw_avb_nostd_defaults"],
     no_stdlibs: true,
     stdlibs: [
         "libcore.rust_sysroot",
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 0d2dfda..6f96fc0 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -17,9 +17,7 @@
 use crate::config;
 use crate::crypto;
 use crate::fdt;
-use crate::heap;
 use crate::memory;
-use crate::rand;
 use core::arch::asm;
 use core::mem::{drop, size_of};
 use core::num::NonZeroUsize;
@@ -33,11 +31,12 @@
 use log::LevelFilter;
 use vmbase::util::RangeExt as _;
 use vmbase::{
-    console,
+    configure_heap, console,
     layout::{self, crosvm},
     logger, main,
-    memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_4KB},
+    memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_128KB, SIZE_4KB},
     power::reboot,
+    rand,
 };
 use zeroize::Zeroize;
 
@@ -62,6 +61,7 @@
 }
 
 main!(start);
+configure_heap!(SIZE_128KB);
 
 /// Entry point for pVM firmware.
 pub fn start(fdt_address: u64, payload_start: u64, payload_size: u64, _arg3: u64) {
@@ -69,9 +69,6 @@
     // - can't access non-pvmfw memory (only statically-mapped memory)
     // - can't access MMIO (therefore, no logging)
 
-    // SAFETY - This function should and will only be called once, here.
-    unsafe { heap::init() };
-
     match main_wrapper(fdt_address as usize, payload_start as usize, payload_size as usize) {
         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.
diff --git a/pvmfw/src/instance.rs b/pvmfw/src/instance.rs
index 56468b2..1035559 100644
--- a/pvmfw/src/instance.rs
+++ b/pvmfw/src/instance.rs
@@ -21,7 +21,6 @@
 use crate::gpt;
 use crate::gpt::Partition;
 use crate::gpt::Partitions;
-use crate::rand;
 use core::fmt;
 use core::mem::size_of;
 use diced_open_dice::DiceMode;
@@ -30,6 +29,7 @@
 use log::trace;
 use uuid::Uuid;
 use virtio_drivers::transport::{pci::bus::PciRoot, DeviceType, Transport};
+use vmbase::rand;
 use vmbase::util::ceiling_div;
 use vmbase::virtio::pci::{PciTransportIterator, VirtIOBlk};
 use zerocopy::AsBytes;
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index c826cd8..61e2312 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -28,12 +28,9 @@
 mod exceptions;
 mod fdt;
 mod gpt;
-mod heap;
 mod helpers;
-mod hvc;
 mod instance;
 mod memory;
-mod rand;
 
 use crate::bcc::Bcc;
 use crate::dice::PartialInputs;
@@ -51,6 +48,7 @@
 use pvmfw_avb::Capability;
 use pvmfw_avb::DebugLevel;
 use pvmfw_embedded_key::PUBLIC_KEY;
+use vmbase::heap;
 use vmbase::memory::flush;
 use vmbase::memory::MEMORY;
 use vmbase::virtio::pci;
@@ -106,6 +104,11 @@
         error!("Failed to verify the payload: {e}");
         RebootReason::PayloadVerificationError
     })?;
+    let debuggable = verified_boot_data.debug_level != DebugLevel::None;
+    if debuggable {
+        info!("Successfully verified a debuggable payload.");
+        info!("Please disregard any previous libavb ERROR about initrd_normal.");
+    }
 
     if verified_boot_data.capabilities.contains(&Capability::RemoteAttest) {
         info!("Service VM capable of remote attestation detected");
@@ -148,7 +151,6 @@
     flush(next_bcc);
 
     let strict_boot = true;
-    let debuggable = verified_boot_data.debug_level != DebugLevel::None;
     modify_for_next_stage(fdt, next_bcc, new_instance, strict_boot, debug_policy, debuggable)
         .map_err(|e| {
             error!("Failed to configure device tree: {e}");
diff --git a/rialto/Android.bp b/rialto/Android.bp
index 59f8ba2..9aa4667 100644
--- a/rialto/Android.bp
+++ b/rialto/Android.bp
@@ -6,18 +6,15 @@
     name: "librialto",
     crate_name: "rialto",
     srcs: ["src/main.rs"],
-    edition: "2021",
     defaults: ["vmbase_ffi_defaults"],
     rustlibs: [
         "libaarch64_paging",
-        "libbuddy_system_allocator",
         "libhyp",
         "libfdtpci",
         "liblibfdt",
         "liblog_rust_nostd",
         "libvmbase",
     ],
-    apex_available: ["com.android.virt"],
 }
 
 cc_binary {
@@ -29,13 +26,11 @@
     ],
     static_libs: [
         "librialto",
-        "libvmbase_entry",
     ],
     linker_scripts: [
         "image.ld",
         ":vmbase_sections",
     ],
-    apex_available: ["com.android.virt"],
 }
 
 raw_binary {
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index 29056f1..ce83624 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -23,7 +23,6 @@
 extern crate alloc;
 
 use crate::error::{Error, Result};
-use buddy_system_allocator::LockedHeap;
 use core::num::NonZeroUsize;
 use core::result;
 use core::slice;
@@ -32,28 +31,14 @@
 use libfdt::FdtError;
 use log::{debug, error, info};
 use vmbase::{
+    configure_heap,
     fdt::SwiotlbInfo,
     layout::{self, crosvm},
     main,
-    memory::{MemoryTracker, PageTable, MEMORY, PAGE_SIZE},
+    memory::{MemoryTracker, PageTable, MEMORY, PAGE_SIZE, SIZE_64KB},
     power::reboot,
 };
 
-const SZ_1K: usize = 1024;
-const SZ_64K: usize = 64 * SZ_1K;
-
-#[global_allocator]
-static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
-
-static mut HEAP: [u8; SZ_64K] = [0; SZ_64K];
-
-fn init_heap() {
-    // SAFETY: Allocator set to otherwise unused, static memory.
-    unsafe {
-        HEAP_ALLOCATOR.lock().init(&mut HEAP as *mut u8 as usize, HEAP.len());
-    }
-}
-
 fn new_page_table() -> Result<PageTable> {
     let mut page_table = PageTable::default();
 
@@ -163,7 +148,6 @@
 
 /// Entry point for Rialto.
 pub fn main(fdt_addr: u64, _a1: u64, _a2: u64, _a3: u64) {
-    init_heap();
     let Ok(mmio_guard_supported) = try_init_logger() else {
         // Don't log anything if the logger initialization fails.
         reboot();
@@ -181,3 +165,4 @@
 }
 
 main!(main);
+configure_heap!(SIZE_64KB);
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 d30f0d0..e6d90ea 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
@@ -160,7 +160,7 @@
         private StringBuilder mLogOutput = new StringBuilder();
         private boolean mProcessedBootTimeMetrics = false;
 
-        private void processBootTimeMetrics(String log) {
+        private synchronized void processBootTimeMetrics(String log) {
             if (!mVcpuStartedNanoTime.isPresent()) {
                 mVcpuStartedNanoTime = OptionalLong.of(System.nanoTime());
             }
@@ -177,12 +177,8 @@
         }
 
         private void logVmOutputAndMonitorBootTimeMetrics(
-                String tag,
-                InputStream vmOutputStream,
-                String name,
-                StringBuilder result,
-                boolean monitorEvents) {
-            mProcessedBootTimeMetrics |= monitorEvents;
+                String tag, InputStream vmOutputStream, String name, StringBuilder result) {
+            mProcessedBootTimeMetrics = true;
             new Thread(
                             () -> {
                                 try {
@@ -192,7 +188,7 @@
                                     String line;
                                     while ((line = reader.readLine()) != null
                                             && !Thread.interrupted()) {
-                                        if (monitorEvents) processBootTimeMetrics(line);
+                                        processBootTimeMetrics(line);
                                         Log.i(tag, name + ": " + line);
                                         result.append(line + "\n");
                                     }
@@ -203,17 +199,6 @@
                     .start();
         }
 
-        private void logVmOutputAndMonitorBootTimeMetrics(
-                String tag, InputStream vmOutputStream, String name, StringBuilder result) {
-            logVmOutputAndMonitorBootTimeMetrics(tag, vmOutputStream, name, result, true);
-        }
-
-        /** Copy output from the VM to logcat. This is helpful when things go wrong. */
-        protected void logVmOutput(
-                String tag, InputStream vmOutputStream, String name, StringBuilder result) {
-            logVmOutputAndMonitorBootTimeMetrics(tag, vmOutputStream, name, result, false);
-        }
-
         public void runToFinish(String logTag, VirtualMachine vm)
                 throws VirtualMachineException, InterruptedException {
             vm.setCallback(mExecutorService, this);
@@ -221,7 +206,7 @@
             if (vm.getConfig().isVmOutputCaptured()) {
                 logVmOutputAndMonitorBootTimeMetrics(
                         logTag, vm.getConsoleOutput(), "Console", mConsoleOutput);
-                logVmOutput(logTag, vm.getLogOutput(), "Log", mLogOutput);
+                logVmOutputAndMonitorBootTimeMetrics(logTag, vm.getLogOutput(), "Log", mLogOutput);
             }
             mExecutorService.awaitTermination(300, TimeUnit.SECONDS);
         }
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 4f023ca..deaf08a 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -55,7 +55,6 @@
 import org.json.JSONObject;
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestName;
@@ -496,7 +495,7 @@
         vmInfo.mProcess.destroy();
     }
 
-    private void waitForCrosvmExit(CommandRunner android) throws Exception {
+    private void waitForCrosvmExit(CommandRunner android, String testStartTime) throws Exception {
         // TODO: improve crosvm exit check. b/258848245
         android.runWithTimeout(
                 15000,
@@ -504,10 +503,12 @@
                 "-m",
                 "1",
                 "-e",
-                "'virtualizationmanager::crosvm.*exited with status exit status:'");
+                "'virtualizationmanager::crosvm.*exited with status exit status:'",
+                "-T",
+                "'" + testStartTime + "'");
     }
 
-    private boolean isTombstoneReceivedFromHostLogcat() throws Exception {
+    private boolean isTombstoneReceivedFromHostLogcat(String testStartTime) throws Exception {
         // Note this method relies on logcat values being printed by the receiver on host
         // userspace crash log: virtualizationservice/src/aidl.rs
         // kernel ramdump log: virtualizationmanager/src/crosvm.rs
@@ -526,12 +527,17 @@
                         "-m",
                         "1",
                         "-e",
-                        ramdumpRegex);
+                        ramdumpRegex,
+                        "-T",
+                        testStartTime);
         return !result.trim().isEmpty();
     }
 
     private boolean isTombstoneGeneratedWithCmd(
             boolean protectedVm, String configPath, String... crashCommand) throws Exception {
+        CommandRunner android = new CommandRunner(getDevice());
+        String testStartTime = android.runWithTimeout(1000, "date", "'+%Y-%m-%d %H:%M:%S.%N'");
+
         mMicrodroidDevice =
                 MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
                         .debugLevel("full")
@@ -546,10 +552,9 @@
         microdroid.run(crashCommand);
 
         // check until microdroid is shut down
-        CommandRunner android = new CommandRunner(getDevice());
-        waitForCrosvmExit(android);
+        waitForCrosvmExit(android, testStartTime);
 
-        return isTombstoneReceivedFromHostLogcat();
+        return isTombstoneReceivedFromHostLogcat(testStartTime);
     }
 
     @Test
@@ -607,6 +612,7 @@
             throws Exception {
         // we can't use microdroid builder as it wants ADB connection (debuggable)
         CommandRunner android = new CommandRunner(getDevice());
+        String testStartTime = android.runWithTimeout(1000, "date", "'+%Y-%m-%d %H:%M:%S.%N'");
 
         android.run("rm", "-rf", TEST_ROOT + "*");
         android.run("mkdir", "-p", TEST_ROOT + "*");
@@ -627,7 +633,7 @@
         Collections.addAll(cmd, additionalArgs);
 
         android.run(cmd.toArray(new String[0]));
-        return isTombstoneReceivedFromHostLogcat();
+        return isTombstoneReceivedFromHostLogcat(testStartTime);
     }
 
     private boolean isTombstoneGeneratedWithCrashPayload(boolean debuggable) throws Exception {
@@ -845,44 +851,6 @@
         }
     }
 
-    @Ignore("b/288467613#comment9")
-    @Test
-    public void testCustomVirtualMachinePermission() throws Exception {
-        assumeTrue(
-                "Protected VMs are not supported",
-                getAndroidDevice().supportsMicrodroid(/*protectedVm=*/ true));
-        assumeTrue("Test requires adb unroot", getDevice().disableAdbRoot());
-        CommandRunner android = new CommandRunner(getDevice());
-
-        // Pull etc/microdroid.json
-        File virtApexDir = FileUtil.createTempDir("virt_apex");
-        File microdroidConfigFile = new File(virtApexDir, "microdroid.json");
-        assertThat(getDevice().pullFile(VIRT_APEX + "etc/microdroid.json", microdroidConfigFile))
-                .isTrue();
-        JSONObject config = new JSONObject(FileUtil.readStringFromFile(microdroidConfigFile));
-
-        // USE_CUSTOM_VIRTUAL_MACHINE is enforced only on protected mode
-        config.put("protected", true);
-
-        // Write updated config
-        final String configPath = TEST_ROOT + "raw_config.json";
-        getDevice().pushString(config.toString(), configPath);
-
-        // temporarily revoke the permission
-        android.run(
-                "pm",
-                "revoke",
-                SHELL_PACKAGE_NAME,
-                "android.permission.USE_CUSTOM_VIRTUAL_MACHINE");
-        final String ret =
-                android.runForResult(VIRT_APEX + "bin/vm run", configPath).getStderr().trim();
-
-        assertThat(ret)
-                .contains(
-                        "does not have the android.permission.USE_CUSTOM_VIRTUAL_MACHINE"
-                                + " permission");
-    }
-
     @Test
     public void testPathToBinaryIsRejected() throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
@@ -998,9 +966,6 @@
         prepareVirtualizationTestSetup(getDevice());
 
         getDevice().installPackage(findTestFile(APK_NAME), /* reinstall */ false);
-
-        // clear the log
-        getDevice().executeShellV2Command("logcat -c");
     }
 
     @After
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 06274c8..dd74d55 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -116,6 +116,20 @@
             .context("failed to create idsig")?;
 
     let mut output = clone_file(idsig_fd)?;
+
+    // Optimization. We don't have to update idsig file whenever a VM is started. Don't update it,
+    // if the idsig file already has the same APK digest.
+    if output.metadata()?.len() > 0 {
+        if let Ok(out_sig) = V4Signature::from_idsig(&mut output) {
+            if out_sig.signing_info.apk_digest == sig.signing_info.apk_digest {
+                debug!("idsig {:?} is up-to-date with apk {:?}.", output, input);
+                return Ok(());
+            }
+        }
+        // if we fail to read v4signature from output, that's fine. User can pass a random file.
+        // We will anyway overwrite the file to the v4signature generated from input_fd.
+    }
+
     output.set_len(0).context("failed to set_len on the idsig output")?;
     sig.write_into(&mut output).context("failed to write idsig")?;
     Ok(())
@@ -243,9 +257,6 @@
         input_fd: &ParcelFileDescriptor,
         idsig_fd: &ParcelFileDescriptor,
     ) -> binder::Result<()> {
-        // TODO(b/193504400): do this only when (1) idsig_fd is empty or (2) the APK digest in
-        // idsig_fd is different from APK digest in input_fd
-
         check_manage_access()?;
 
         create_or_update_idsig_file(input_fd, idsig_fd)
@@ -314,15 +325,13 @@
             VirtualMachineConfig::RawConfig(_) => true,
             VirtualMachineConfig::AppConfig(config) => {
                 // Some features are reserved for platform apps only, even when using
-                // VirtualMachineAppConfig:
+                // VirtualMachineAppConfig. Almost all of these features are grouped in the
+                // CustomConfig struct:
                 // - controlling CPUs;
-                // - specifying a config file in the APK;
+                // - specifying a config file in the APK; (this one is not part of CustomConfig)
                 // - gdbPort is set, meaning that crosvm will start a gdb server;
                 // - using anything other than the default kernel.
-                !config.taskProfiles.is_empty()
-                    || matches!(config.payload, Payload::ConfigPath(_))
-                    || config.gdbPort > 0
-                    || config.customKernelImage.as_ref().is_some()
+                config.customConfig.is_some() || matches!(config.payload, Payload::ConfigPath(_))
             }
         };
         if is_custom {
@@ -605,8 +614,12 @@
     let vm_config_file = File::open(vm_config_path)?;
     let mut vm_config = VmConfig::load(&vm_config_file)?.to_parcelable()?;
 
-    if let Some(file) = config.customKernelImage.as_ref() {
-        vm_config.kernel = Some(ParcelFileDescriptor::new(clone_file(file)?))
+    if let Some(custom_config) = &config.customConfig {
+        if let Some(file) = custom_config.customKernelImage.as_ref() {
+            vm_config.kernel = Some(ParcelFileDescriptor::new(clone_file(file)?))
+        }
+        vm_config.taskProfiles = custom_config.taskProfiles.clone();
+        vm_config.gdbPort = custom_config.gdbPort;
     }
 
     if config.memoryMib > 0 {
@@ -616,8 +629,6 @@
     vm_config.name = config.name.clone();
     vm_config.protectedVm = config.protectedVm;
     vm_config.cpuTopology = config.cpuTopology;
-    vm_config.taskProfiles = config.taskProfiles.clone();
-    vm_config.gdbPort = config.gdbPort;
 
     // Microdroid takes additional init ramdisk & (optionally) storage image
     add_microdroid_system_images(config, instance_file, storage_image, &mut vm_config)?;
@@ -1052,7 +1063,9 @@
 fn extract_gdb_port(config: &VirtualMachineConfig) -> Option<NonZeroU16> {
     match config {
         VirtualMachineConfig::RawConfig(config) => NonZeroU16::new(config.gdbPort as u16),
-        VirtualMachineConfig::AppConfig(config) => NonZeroU16::new(config.gdbPort as u16),
+        VirtualMachineConfig::AppConfig(config) => {
+            NonZeroU16::new(config.customConfig.as_ref().map(|c| c.gdbPort).unwrap_or(0) as u16)
+        }
     }
 }
 
@@ -1310,4 +1323,30 @@
         assert!(ret.is_err(), "should fail");
         Ok(())
     }
+
+    #[test]
+    fn test_create_or_update_idsig_does_not_update_if_already_valid() -> Result<()> {
+        use std::io::Seek;
+
+        // Pick any APK
+        let mut apk = File::open("/system/priv-app/Shell/Shell.apk").unwrap();
+        let mut idsig = tempfile::tempfile().unwrap();
+
+        create_or_update_idsig_file(
+            &ParcelFileDescriptor::new(apk.try_clone()?),
+            &ParcelFileDescriptor::new(idsig.try_clone()?),
+        )?;
+        let modified_orig = idsig.metadata()?.modified()?;
+        apk.rewind()?;
+        idsig.rewind()?;
+
+        // Call the function again
+        create_or_update_idsig_file(
+            &ParcelFileDescriptor::new(apk.try_clone()?),
+            &ParcelFileDescriptor::new(idsig.try_clone()?),
+        )?;
+        let modified_new = idsig.metadata()?.modified()?;
+        assert!(modified_orig == modified_new, "idsig file was updated unnecessarily");
+        Ok(())
+    }
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
index 5e05bb9..6a0bf7c 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/VirtualMachineAppConfig.aidl
@@ -45,6 +45,8 @@
     union Payload {
         /**
          * Path to a JSON file in an APK containing the configuration.
+         *
+         * <p>Setting this field requires android.permission.USE_CUSTOM_VIRTUAL_MACHINE
          */
         @utf8InCpp String configPath;
 
@@ -70,16 +72,6 @@
     /** Debug level of the VM */
     DebugLevel debugLevel = DebugLevel.NONE;
 
-    /**
-     * Port at which crosvm will start a gdb server to debug guest kernel.
-     * If set to zero, then gdb server won't be started.
-     *
-     * Note: Specifying a value here requires android.permission.USE_CUSTOM_VIRTUAL_MACHINE.
-     *
-     * TODO(b/286225150): move to a separate struct
-     */
-    int gdbPort = 0;
-
     /** Whether the VM should be a protected VM. */
     boolean protectedVm;
 
@@ -93,20 +85,28 @@
     CpuTopology cpuTopology = CpuTopology.ONE_CPU;
 
     /**
-     * List of task profile names to apply for the VM
-     *
-     * Note: Specifying a value here requires android.permission.USE_CUSTOM_VIRTUAL_MACHINE.
-     *
-     * TODO(b/286225150): move to a separate struct
+     * Encapsulates parameters that require android.permission.USE_CUSTOM_VIRTUAL_MACHINE.
      */
-    String[] taskProfiles;
+    parcelable CustomConfig {
+        /**
+         * If specified, boot Microdroid VM with the given kernel.
+         *
+         */
+        @nullable ParcelFileDescriptor customKernelImage;
 
-    /**
-     * If specified, boot Microdroid VM with the given kernel.
-     *
-     * Note: Specifying a value here requires android.permission.USE_CUSTOM_VIRTUAL_MACHINE.
-     *
-     * TODO(b/286225150): move to a separate struct
-     */
-    @nullable ParcelFileDescriptor customKernelImage;
+        /**
+         * Port at which crosvm will start a gdb server to debug guest kernel.
+         * If set to zero, then gdb server won't be started.
+         *
+         */
+        int gdbPort = 0;
+
+        /**
+         * List of task profile names to apply for the VM
+         */
+        String[] taskProfiles;
+    }
+
+    /** Configuration parameters guarded by android.permission.USE_CUSTOM_VIRTUAL_MACHINE */
+    @nullable CustomConfig customConfig;
 }
diff --git a/vm/src/run.rs b/vm/src/run.rs
index 27fd903..84072ca 100644
--- a/vm/src/run.rs
+++ b/vm/src/run.rs
@@ -19,7 +19,10 @@
     CpuTopology::CpuTopology,
     IVirtualizationService::IVirtualizationService,
     PartitionType::PartitionType,
-    VirtualMachineAppConfig::{DebugLevel::DebugLevel, Payload::Payload, VirtualMachineAppConfig},
+    VirtualMachineAppConfig::{
+        CustomConfig::CustomConfig, DebugLevel::DebugLevel, Payload::Payload,
+        VirtualMachineAppConfig,
+    },
     VirtualMachineConfig::VirtualMachineConfig,
     VirtualMachinePayloadConfig::VirtualMachinePayloadConfig,
     VirtualMachineState::VirtualMachineState,
@@ -137,6 +140,12 @@
 
     let payload_config_str = format!("{:?}!{:?}", apk, payload);
 
+    let custom_config = CustomConfig {
+        customKernelImage: kernel,
+        gdbPort: gdb.map(u16::from).unwrap_or(0) as i32, // 0 means no gdb
+        taskProfiles: task_profiles,
+    };
+
     let config = VirtualMachineConfig::AppConfig(VirtualMachineAppConfig {
         name: name.unwrap_or_else(|| String::from("VmRunApp")),
         apk: apk_fd.into(),
@@ -149,9 +158,7 @@
         protectedVm: protected,
         memoryMib: mem.unwrap_or(0) as i32, // 0 means use the VM default
         cpuTopology: cpu_topology,
-        taskProfiles: task_profiles,
-        gdbPort: gdb.map(u16::from).unwrap_or(0) as i32, // 0 means no gdb
-        customKernelImage: kernel,
+        customConfig: Some(custom_config),
     });
     run(service, &config, &payload_config_str, console_out_path, console_in_path, log_path)
 }
@@ -320,8 +327,13 @@
         .context("Failed to create VM")?;
     vm.start().context("Failed to start VM")?;
 
+    let debug_level = match config {
+        VirtualMachineConfig::AppConfig(config) => config.debugLevel,
+        _ => DebugLevel::NONE,
+    };
     println!(
-        "Created VM from {} with CID {}, state is {}.",
+        "Created {} from {} with CID {}, state is {}.",
+        if debug_level == DebugLevel::FULL { "debuggable VM" } else { "VM" },
         payload_config,
         vm.cid(),
         state_to_str(vm.state()?)
diff --git a/vmbase/Android.bp b/vmbase/Android.bp
index ac010b9..46f4937 100644
--- a/vmbase/Android.bp
+++ b/vmbase/Android.bp
@@ -2,11 +2,28 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
+// The hierarchy of Soong modules to produce a vmbase-based binary is
+//
+// 0. rlibs may be used to provide high-level code (see "vmbase_rlib_defaults");
+// 1. rust_ffi_static packages low-level Rust code and any rlib into a static
+//    library (see "vmbase_ffi_defaults") that cc_binary supports;
+// 2. cc_library_static may be used for extra C code (see "vmbase_cc_defaults");
+// 3. cc_binary produces an ELF from the (single) Rust-wrapping static library,
+//    optional extra C libraries, and linker script (see "vmbase_elf_defaults");
+// 4. raw_binary strips the ELF into an image that can be loaded to memory;
+
+// Used by intermediate rust_library_rlib for vmbase-based binaries.
 rust_defaults {
-    name: "vmbase_rust_defaults",
+    name: "vmbase_rlib_defaults",
     edition: "2021",
+    prefer_rlib: true,
     host_supported: false,
     enabled: false,
+    no_stdlibs: true,
+    stdlibs: [
+        "libcompiler_builtins.rust_sysroot",
+        "libcore.rust_sysroot",
+    ],
     target: {
         android_arm64: {
             enabled: true,
@@ -14,19 +31,17 @@
     },
 }
 
+// Used by the "top-level" rust_ffi_static of vmbase-based binaries.
 rust_defaults {
     name: "vmbase_ffi_defaults",
-    defaults: ["vmbase_rust_defaults"],
-    no_stdlibs: true,
-    stdlibs: [
-        "libcompiler_builtins.rust_sysroot",
-        "libcore.rust_sysroot",
-    ],
+    defaults: ["vmbase_rlib_defaults"],
 }
 
+// Used by extra cc_library_static linked into the final ELF.
 cc_defaults {
     name: "vmbase_cc_defaults",
     nocrt: true,
+    no_libcrt: true,
     system_shared_libs: [],
     stl: "none",
     installable: false,
@@ -39,8 +54,10 @@
     sanitize: {
         hwaddress: false,
     },
+    native_coverage: false,
 }
 
+// Used by cc_binary when producing the ELF of a vmbase-based binary.
 cc_defaults {
     name: "vmbase_elf_defaults",
     defaults: ["vmbase_cc_defaults"],
@@ -48,18 +65,11 @@
     static_libs: [
         "libvmbase_entry",
     ],
-    installable: false,
-    enabled: false,
-    target: {
-        android_arm64: {
-            enabled: true,
-        },
-    },
 }
 
 rust_library_rlib {
     name: "libvmbase",
-    defaults: ["vmbase_rust_defaults"],
+    defaults: ["vmbase_rlib_defaults"],
     crate_name: "vmbase",
     srcs: ["src/lib.rs"],
     rustlibs: [
@@ -76,14 +86,12 @@
         "libvirtio_drivers",
         "libzeroize_nostd",
     ],
-    no_stdlibs: true,
     whole_static_libs: [
         "librust_baremetal",
     ],
     features: [
         "cpu_feat_hafdbs",
     ],
-    apex_available: ["com.android.virt"],
 }
 
 cc_library_static {
@@ -94,8 +102,6 @@
         "exceptions.S",
         "exceptions_panic.S",
     ],
-    no_libcrt: true,
-    apex_available: ["com.android.virt"],
 }
 
 filegroup {
diff --git a/vmbase/example/Android.bp b/vmbase/example/Android.bp
index dc9a090..ae1a593 100644
--- a/vmbase/example/Android.bp
+++ b/vmbase/example/Android.bp
@@ -7,10 +7,8 @@
     defaults: ["vmbase_ffi_defaults"],
     crate_name: "vmbase_example",
     srcs: ["src/main.rs"],
-    edition: "2021",
     rustlibs: [
         "libaarch64_paging",
-        "libbuddy_system_allocator",
         "libdiced_open_dice_nostd",
         "libfdtpci",
         "liblibfdt",
diff --git a/vmbase/example/src/main.rs b/vmbase/example/src/main.rs
index 1dd8517..b3b5732 100644
--- a/vmbase/example/src/main.rs
+++ b/vmbase/example/src/main.rs
@@ -30,11 +30,10 @@
 use crate::pci::{check_pci, get_bar_region};
 use aarch64_paging::{idmap::IdMap, paging::Attributes};
 use alloc::{vec, vec::Vec};
-use buddy_system_allocator::LockedHeap;
 use fdtpci::PciInfo;
 use libfdt::Fdt;
 use log::{debug, error, info, trace, warn, LevelFilter};
-use vmbase::{cstr, logger, main};
+use vmbase::{configure_heap, cstr, logger, main, memory::SIZE_64KB};
 
 static INITIALISED_DATA: [u32; 4] = [1, 2, 3, 4];
 static mut ZEROED_DATA: [u32; 10] = [0; 10];
@@ -43,12 +42,8 @@
 const ASID: usize = 1;
 const ROOT_LEVEL: usize = 1;
 
-#[global_allocator]
-static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
-
-static mut HEAP: [u8; 65536] = [0; 65536];
-
 main!(main);
+configure_heap!(SIZE_64KB);
 
 /// Entry point for VM bootloader.
 pub fn main(arg0: u64, arg1: u64, arg2: u64, arg3: u64) {
@@ -74,10 +69,6 @@
 
     modify_fdt(fdt);
 
-    unsafe {
-        HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
-    }
-
     check_alloc();
 
     let mut idmap = IdMap::new(ASID, ROOT_LEVEL);
@@ -164,7 +155,6 @@
     unsafe {
         info!("ZEROED_DATA: {:?}", ZEROED_DATA.as_ptr());
         info!("MUTABLE_DATA: {:?}", MUTABLE_DATA.as_ptr());
-        info!("HEAP: {:?}", HEAP.as_ptr());
     }
 
     assert_eq!(INITIALISED_DATA[0], 1);
diff --git a/vmbase/src/entry.rs b/vmbase/src/entry.rs
index 8cdfe77..df0bb7c 100644
--- a/vmbase/src/entry.rs
+++ b/vmbase/src/entry.rs
@@ -14,11 +14,13 @@
 
 //! Rust entry point.
 
-use crate::{console, power::shutdown};
+use crate::{console, heap, power::shutdown};
 
 /// This is the entry point to the Rust code, called from the binary entry point in `entry.S`.
 #[no_mangle]
 extern "C" fn rust_entry(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
+    // SAFETY - Only called once, from here, and inaccessible to client code.
+    unsafe { heap::init() };
     console::init();
     unsafe {
         main(x0, x1, x2, x3);
diff --git a/pvmfw/src/heap.rs b/vmbase/src/heap.rs
similarity index 88%
rename from pvmfw/src/heap.rs
rename to vmbase/src/heap.rs
index 151049e..b00ca6f 100644
--- a/pvmfw/src/heap.rs
+++ b/vmbase/src/heap.rs
@@ -27,15 +27,31 @@
 
 use buddy_system_allocator::LockedHeap;
 
-/// 128 KiB
-const HEAP_SIZE: usize = 0x20000;
-static mut HEAP: [u8; HEAP_SIZE] = [0; HEAP_SIZE];
+/// Configures the size of the global allocator.
+#[macro_export]
+macro_rules! configure_heap {
+    ($len:expr) => {
+        static mut __HEAP_ARRAY: [u8; $len] = [0; $len];
+        #[export_name = "HEAP"]
+        // SAFETY - HEAP will only be accessed once as mut, from init().
+        static mut __HEAP: &'static mut [u8] = unsafe { &mut __HEAP_ARRAY };
+    };
+}
+
+extern "Rust" {
+    /// Slice used by the global allocator, configured using configure_heap!().
+    static mut HEAP: &'static mut [u8];
+}
 
 #[global_allocator]
 static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
 
-/// SAFETY: Must be called no more than once.
-pub unsafe fn init() {
+/// Initialize the global allocator.
+///
+/// # Safety
+///
+/// Must be called no more than once.
+pub(crate) unsafe fn init() {
     // 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()) };
diff --git a/pvmfw/src/hvc.rs b/vmbase/src/hvc.rs
similarity index 95%
rename from pvmfw/src/hvc.rs
rename to vmbase/src/hvc.rs
index e03c9d3..9a5e716 100644
--- a/pvmfw/src/hvc.rs
+++ b/vmbase/src/hvc.rs
@@ -21,7 +21,6 @@
     hvc64,
 };
 
-// TODO(b/272226230): Move all the trng functions to trng module
 const ARM_SMCCC_TRNG_VERSION: u32 = 0x8400_0050;
 #[allow(dead_code)]
 const ARM_SMCCC_TRNG_FEATURES: u32 = 0x8400_0051;
diff --git a/pvmfw/src/hvc/trng.rs b/vmbase/src/hvc/trng.rs
similarity index 100%
rename from pvmfw/src/hvc/trng.rs
rename to vmbase/src/hvc/trng.rs
diff --git a/vmbase/src/lib.rs b/vmbase/src/lib.rs
index 54f3384..88bad8b 100644
--- a/vmbase/src/lib.rs
+++ b/vmbase/src/lib.rs
@@ -23,11 +23,14 @@
 pub mod console;
 mod entry;
 pub mod fdt;
+pub mod heap;
+mod hvc;
 pub mod layout;
 mod linker;
 pub mod logger;
 pub mod memory;
 pub mod power;
+pub mod rand;
 pub mod uart;
 pub mod util;
 pub mod virtio;
diff --git a/vmbase/src/memory/mod.rs b/vmbase/src/memory/mod.rs
index 9f14691..5e78565 100644
--- a/vmbase/src/memory/mod.rs
+++ b/vmbase/src/memory/mod.rs
@@ -25,5 +25,5 @@
 pub use shared::{alloc_shared, dealloc_shared, MemoryRange, MemoryTracker, MEMORY};
 pub use util::{
     flush, flushed_zeroize, min_dcache_line_size, page_4kb_of, phys_to_virt, virt_to_phys,
-    PAGE_SIZE, SIZE_2MB, SIZE_4KB, SIZE_4MB,
+    PAGE_SIZE, SIZE_128KB, SIZE_2MB, SIZE_4KB, SIZE_4MB, SIZE_64KB,
 };
diff --git a/vmbase/src/memory/util.rs b/vmbase/src/memory/util.rs
index 04d42cd..b9ef5c9 100644
--- a/vmbase/src/memory/util.rs
+++ b/vmbase/src/memory/util.rs
@@ -22,6 +22,10 @@
 
 /// The size of a 4KB memory in bytes.
 pub const SIZE_4KB: usize = 4 << 10;
+/// The size of a 64KB memory in bytes.
+pub const SIZE_64KB: usize = 64 << 10;
+/// The size of a 128KB memory in bytes.
+pub const SIZE_128KB: usize = 128 << 10;
 /// The size of a 2MB memory in bytes.
 pub const SIZE_2MB: usize = 2 << 20;
 /// The size of a 4MB memory in bytes.
diff --git a/pvmfw/src/rand.rs b/vmbase/src/rand.rs
similarity index 94%
rename from pvmfw/src/rand.rs
rename to vmbase/src/rand.rs
index b45538a..00567b8 100644
--- a/pvmfw/src/rand.rs
+++ b/vmbase/src/rand.rs
@@ -12,10 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//! Functions and drivers for obtaining true entropy.
+
 use crate::hvc;
 use core::fmt;
 use core::mem::size_of;
 
+/// Error type for rand operations.
 pub enum Error {
     /// Error during SMCCC TRNG call.
     Trng(hvc::trng::Error),
@@ -29,6 +32,7 @@
     }
 }
 
+/// Result type for rand operations.
 pub type Result<T> = core::result::Result<T, Error>;
 
 impl fmt::Display for Error {
@@ -96,6 +100,7 @@
     }
 }
 
+/// Generate an array of fixed-size initialized with true-random bytes.
 pub fn random_array<const N: usize>() -> Result<[u8; N]> {
     let mut arr = [0; N];
     fill_with_entropy(&mut arr)?;