Merge changes I5d8e169d,I832b705d,Ic8af83b8,Ib3c9acb8,I2e357326, ... into main

* changes:
  vmbase: Support 16KiB MMIO_GUARD granule
  vmbase: Detect dynamic resharing of UART
  vmbase: Move MMIO_GUARD granule check out of ::hyp
  vmbase: Harden MMIO_GUARD with RAII
  vmbase: Clean up MMIO_GUARD in MemoryTracker
  vmbase: Introduce compat_android_13 feature
  vmbase: Remove now duplicate hyp::util
  vmbase: Integrate libhyp as module
  libs: Prepare libhyp for move
diff --git a/libs/hyp/Android.bp b/libs/hyp/Android.bp
deleted file mode 100644
index 404269a..0000000
--- a/libs/hyp/Android.bp
+++ /dev/null
@@ -1,27 +0,0 @@
-package {
-    default_applicable_licenses: ["Android-Apache-2.0"],
-}
-
-rust_library_rlib {
-    name: "libhyp",
-    crate_name: "hyp",
-    defaults: ["avf_build_flags_rust"],
-    srcs: ["src/lib.rs"],
-    prefer_rlib: true,
-    rustlibs: [
-        "libonce_cell_nostd",
-        "libsmccc",
-        "libuuid_nostd",
-    ],
-    no_stdlibs: true,
-    stdlibs: [
-        "libcore.rust_sysroot",
-    ],
-    enabled: false,
-    target: {
-        android_arm64: {
-            enabled: true,
-        },
-    },
-    apex_available: ["com.android.virt"],
-}
diff --git a/libs/hyp/src/util.rs b/libs/hyp/src/util.rs
deleted file mode 100644
index 56f94fd..0000000
--- a/libs/hyp/src/util.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.
-
-//! Utility functions.
-
-pub(crate) const SIZE_4KB: usize = 4 << 10;
-
-/// Computes the low memory page address of the 4KiB page containing a given address.
-pub(crate) fn page_address(addr: usize) -> u64 {
-    (addr & !(SIZE_4KB - 1)).try_into().unwrap()
-}
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index 4ee02c1..37a321d 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -19,7 +19,6 @@
         "libcstr",
         "libdiced_open_dice_nostd",
         "libfdtpci",
-        "libhyp",
         "liblibfdt",
         "liblog_rust_nostd",
         "libonce_cell_nostd",
@@ -75,7 +74,6 @@
     defaults: ["libpvmfw.test.defaults"],
     rustlibs: [
         "libdts",
-        "libhyp",
         "liblibfdt",
         "liblog_rust",
         "libpvmfw_fdt_template",
diff --git a/pvmfw/src/device_assignment.rs b/pvmfw/src/device_assignment.rs
index f826167..5edfe97 100644
--- a/pvmfw/src/device_assignment.rs
+++ b/pvmfw/src/device_assignment.rs
@@ -28,11 +28,11 @@
 use core::iter::Iterator;
 use core::mem;
 use core::ops::Range;
-// TODO(ptosi): Remove the need for this workaround.
-#[cfg(not(test))]
-use hyp::DeviceAssigningHypervisor;
 use libfdt::{Fdt, FdtError, FdtNode, FdtNodeMut, Phandle, Reg};
 use log::error;
+// TODO(b/308694211): Use vmbase::hyp::{DeviceAssigningHypervisor, Error} proper for tests.
+#[cfg(not(test))]
+use vmbase::hyp::DeviceAssigningHypervisor;
 use zerocopy::byteorder::big_endian::U32;
 use zerocopy::FromBytes as _;
 
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 72212c3..43822a5 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -23,7 +23,6 @@
 use core::num::NonZeroUsize;
 use core::ops::Range;
 use core::slice;
-use hyp::{get_mem_sharer, get_mmio_guard};
 use log::debug;
 use log::error;
 use log::info;
@@ -32,6 +31,7 @@
 use vmbase::util::RangeExt as _;
 use vmbase::{
     configure_heap, console,
+    hyp::{get_mem_sharer, get_mmio_guard},
     layout::{self, crosvm},
     main,
     memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_128KB, SIZE_4KB},
@@ -249,7 +249,7 @@
     config_entries.bcc.zeroize();
 
     info!("Expecting a bug making MMIO_GUARD_UNMAP return NOT_SUPPORTED on success");
-    MEMORY.lock().as_mut().unwrap().mmio_unmap_all().map_err(|e| {
+    MEMORY.lock().as_mut().unwrap().unshare_all_mmio().map_err(|e| {
         error!("Failed to unshare MMIO ranges: {e}");
         RebootReason::InternalError
     })?;
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index 6038f12..9206588 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -46,6 +46,7 @@
 use static_assertions::const_assert;
 use tinyvec::ArrayVec;
 use vmbase::fdt::SwiotlbInfo;
+use vmbase::hyp;
 use vmbase::layout::{crosvm::MEM_START, MAX_VIRT_ADDR};
 use vmbase::memory::SIZE_4KB;
 use vmbase::util::flatten;
diff --git a/rialto/Android.bp b/rialto/Android.bp
index d7aac35..33fe189 100644
--- a/rialto/Android.bp
+++ b/rialto/Android.bp
@@ -15,7 +15,6 @@
         "libciborium_nostd",
         "libcstr",
         "libdiced_open_dice_nostd",
-        "libhyp",
         "libfdtpci",
         "liblibfdt",
         "liblog_rust_nostd",
diff --git a/rialto/src/error.rs b/rialto/src/error.rs
index d2bdbbe..033159b 100644
--- a/rialto/src/error.rs
+++ b/rialto/src/error.rs
@@ -18,10 +18,9 @@
 use core::{fmt, result};
 use diced_open_dice::DiceError;
 use fdtpci::PciError;
-use hyp::Error as HypervisorError;
 use libfdt::FdtError;
 use service_vm_comm::RequestProcessingError;
-use vmbase::{memory::MemoryTrackerError, virtio::pci};
+use vmbase::{hyp::Error as HypervisorError, memory::MemoryTrackerError, virtio::pci};
 
 pub type Result<T> = result::Result<T, Error>;
 
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index 11e67cb..864f5e4 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -34,7 +34,6 @@
 use core::slice;
 use diced_open_dice::{bcc_handover_parse, DiceArtifacts};
 use fdtpci::PciInfo;
-use hyp::{get_mem_sharer, get_mmio_guard};
 use libfdt::FdtError;
 use log::{debug, error, info};
 use service_vm_comm::{ServiceVmRequest, VmType};
@@ -48,6 +47,7 @@
 use vmbase::{
     configure_heap,
     fdt::SwiotlbInfo,
+    hyp::{get_mem_sharer, get_mmio_guard},
     layout::{self, crosvm},
     main,
     memory::{MemoryTracker, PageTable, MEMORY, PAGE_SIZE, SIZE_128KB},
diff --git a/vmbase/Android.bp b/vmbase/Android.bp
index 07e1b4c..f01e8aa 100644
--- a/vmbase/Android.bp
+++ b/vmbase/Android.bp
@@ -79,13 +79,14 @@
         "libbuddy_system_allocator",
         "libcstr",
         "libfdtpci",
-        "libhyp",
         "liblibfdt",
         "liblog_rust_nostd",
         "libonce_cell_nostd",
         "libsmccc",
         "libspin_nostd",
+        "libstatic_assertions",
         "libtinyvec_nostd",
+        "libuuid_nostd",
         "libvirtio_drivers",
         "libzerocopy_nostd",
         "libzeroize_nostd",
@@ -93,7 +94,9 @@
     whole_static_libs: [
         "librust_baremetal",
     ],
+    // TODO(b/277859415, b/277860860): Drop "compat_android_13".
     features: [
+        "compat_android_13",
         "cpu_feat_hafdbs",
     ],
 }
diff --git a/vmbase/src/entry.rs b/vmbase/src/entry.rs
index b19efce..bb5ccef 100644
--- a/vmbase/src/entry.rs
+++ b/vmbase/src/entry.rs
@@ -15,19 +15,35 @@
 //! Rust entry point.
 
 use crate::{
-    bionic, console, heap, logger,
+    bionic, console, heap, hyp, logger,
+    memory::{page_4kb_of, SIZE_16KB, SIZE_4KB},
     power::{reboot, shutdown},
     rand,
 };
 use core::mem::size_of;
-use hyp::{self, get_mmio_guard};
+use static_assertions::const_assert_eq;
 
 fn try_console_init() -> Result<(), hyp::Error> {
     console::init();
 
-    if let Some(mmio_guard) = get_mmio_guard() {
+    if let Some(mmio_guard) = hyp::get_mmio_guard() {
         mmio_guard.enroll()?;
-        mmio_guard.validate_granule()?;
+
+        // TODO(ptosi): Use MmioSharer::share() to properly track this MMIO_GUARD_MAP.
+        //
+        // The following call shares the UART but also anything else present in 0..granule.
+        //
+        // For 4KiB, that's only the UARTs. For 16KiB, it also covers the RTC and watchdog but, as
+        // neither is used by vmbase clients (and as both are outside of the UART page), they
+        // will never have valid stage-1 mappings to those devices. As a result, this
+        // MMIO_GUARD_MAP isn't affected by the granule size in any visible way. Larger granule
+        // sizes will need to be checked separately, if needed.
+        assert!({
+            let granule = mmio_guard.granule()?;
+            granule == SIZE_4KB || granule == SIZE_16KB
+        });
+        // Validate the assumption above by ensuring that the UART is not moved to another page:
+        const_assert_eq!(page_4kb_of(console::BASE_ADDRESS), 0);
         mmio_guard.map(console::BASE_ADDRESS)?;
     }
 
diff --git a/libs/hyp/src/lib.rs b/vmbase/src/hyp.rs
similarity index 78%
rename from libs/hyp/src/lib.rs
rename to vmbase/src/hyp.rs
index 6a23585..1cc2ca7 100644
--- a/libs/hyp/src/lib.rs
+++ b/vmbase/src/hyp.rs
@@ -14,16 +14,10 @@
 
 //! This library provides wrappers around various hypervisor backends.
 
-#![no_std]
-
 mod error;
 mod hypervisor;
-mod util;
 
-pub use crate::hypervisor::DeviceAssigningHypervisor;
 pub use error::{Error, Result};
 pub use hypervisor::{
-    get_device_assigner, get_mem_sharer, get_mmio_guard, KvmError, MMIO_GUARD_GRANULE_SIZE,
+    get_device_assigner, get_mem_sharer, get_mmio_guard, DeviceAssigningHypervisor, KvmError,
 };
-
-use hypervisor::GeniezoneError;
diff --git a/libs/hyp/src/error.rs b/vmbase/src/hyp/error.rs
similarity index 86%
rename from libs/hyp/src/error.rs
rename to vmbase/src/hyp/error.rs
index 3fdad70..e9c37e1 100644
--- a/libs/hyp/src/error.rs
+++ b/vmbase/src/hyp/error.rs
@@ -14,9 +14,9 @@
 
 //! Error and Result types for hypervisor.
 
-use crate::GeniezoneError;
-use crate::KvmError;
 use core::{fmt, result};
+
+use super::hypervisor::{GeniezoneError, KvmError};
 use uuid::Uuid;
 
 /// Result type with hypervisor error.
@@ -33,8 +33,6 @@
     GeniezoneError(GeniezoneError, u32),
     /// Unsupported Hypervisor
     UnsupportedHypervisorUuid(Uuid),
-    /// The MMIO_GUARD granule used by the hypervisor is not supported.
-    UnsupportedMmioGuardGranule(usize),
 }
 
 impl fmt::Display for Error {
@@ -53,9 +51,6 @@
             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.rs b/vmbase/src/hyp/hypervisor.rs
similarity index 95%
rename from libs/hyp/src/hypervisor.rs
rename to vmbase/src/hyp/hypervisor.rs
index c53b886..1b45f38 100644
--- a/libs/hyp/src/hypervisor.rs
+++ b/vmbase/src/hyp/hypervisor.rs
@@ -14,19 +14,15 @@
 
 //! Wrappers around hypervisor back-ends.
 
-extern crate alloc;
-
 mod common;
 mod geniezone;
 mod gunyah;
 mod kvm;
 
-use crate::error::{Error, Result};
+use super::{Error, Result};
 use alloc::boxed::Box;
 use common::Hypervisor;
-pub use common::{
-    DeviceAssigningHypervisor, MemSharingHypervisor, MmioGuardedHypervisor, MMIO_GUARD_GRANULE_SIZE,
-};
+pub use common::{DeviceAssigningHypervisor, MemSharingHypervisor, MmioGuardedHypervisor};
 pub use geniezone::GeniezoneError;
 use geniezone::GeniezoneHypervisor;
 use gunyah::GunyahHypervisor;
diff --git a/libs/hyp/src/hypervisor/common.rs b/vmbase/src/hyp/hypervisor/common.rs
similarity index 84%
rename from libs/hyp/src/hypervisor/common.rs
rename to vmbase/src/hyp/hypervisor/common.rs
index eaac652..de0fe12 100644
--- a/libs/hyp/src/hypervisor/common.rs
+++ b/vmbase/src/hyp/hypervisor/common.rs
@@ -14,11 +14,7 @@
 
 //! This module regroups some common traits shared by all the hypervisors.
 
-use crate::error::{Error, Result};
-use crate::util::SIZE_4KB;
-
-/// Expected MMIO guard granule size, validated during MMIO guard initialization.
-pub const MMIO_GUARD_GRANULE_SIZE: usize = SIZE_4KB;
+use crate::hyp::Result;
 
 /// Trait for the hypervisor.
 pub trait Hypervisor {
@@ -53,15 +49,6 @@
 
     /// Returns the MMIO guard granule size in bytes.
     fn granule(&self) -> Result<usize>;
-
-    // TODO(ptosi): Fully move granule validation to client code.
-    /// Validates the MMIO guard granule size.
-    fn validate_granule(&self) -> Result<()> {
-        match self.granule()? {
-            MMIO_GUARD_GRANULE_SIZE => Ok(()),
-            granule => Err(Error::UnsupportedMmioGuardGranule(granule)),
-        }
-    }
 }
 
 pub trait MemSharingHypervisor {
diff --git a/libs/hyp/src/hypervisor/geniezone.rs b/vmbase/src/hyp/hypervisor/geniezone.rs
similarity index 94%
rename from libs/hyp/src/hypervisor/geniezone.rs
rename to vmbase/src/hyp/hypervisor/geniezone.rs
index ad18e17..fcb9b42 100644
--- a/libs/hyp/src/hypervisor/geniezone.rs
+++ b/vmbase/src/hyp/hypervisor/geniezone.rs
@@ -14,10 +14,14 @@
 
 //! Wrappers around calls to the GenieZone hypervisor.
 
-use super::common::{Hypervisor, MemSharingHypervisor, MmioGuardedHypervisor};
-use crate::error::{Error, Result};
-use crate::util::page_address;
 use core::fmt::{self, Display, Formatter};
+
+use super::{Hypervisor, MemSharingHypervisor, MmioGuardedHypervisor};
+use crate::{
+    hyp::{Error, Result},
+    memory::page_4kb_of,
+};
+
 use smccc::{
     error::{positive_or_error_64, success_or_error_64},
     hvc64,
@@ -107,14 +111,14 @@
 
     fn map(&self, addr: usize) -> Result<()> {
         let mut args = [0u64; 17];
-        args[0] = page_address(addr);
+        args[0] = page_4kb_of(addr).try_into().unwrap();
 
         checked_hvc64_expect_zero(VENDOR_HYP_GZVM_MMIO_GUARD_MAP_FUNC_ID, args)
     }
 
     fn unmap(&self, addr: usize) -> Result<()> {
         let mut args = [0u64; 17];
-        args[0] = page_address(addr);
+        args[0] = page_4kb_of(addr).try_into().unwrap();
 
         checked_hvc64_expect_zero(VENDOR_HYP_GZVM_MMIO_GUARD_UNMAP_FUNC_ID, args)
     }
diff --git a/libs/hyp/src/hypervisor/gunyah.rs b/vmbase/src/hyp/hypervisor/gunyah.rs
similarity index 100%
rename from libs/hyp/src/hypervisor/gunyah.rs
rename to vmbase/src/hyp/hypervisor/gunyah.rs
diff --git a/libs/hyp/src/hypervisor/kvm.rs b/vmbase/src/hyp/hypervisor/kvm.rs
similarity index 82%
rename from libs/hyp/src/hypervisor/kvm.rs
rename to vmbase/src/hyp/hypervisor/kvm.rs
index 720318e..8450bed 100644
--- a/libs/hyp/src/hypervisor/kvm.rs
+++ b/vmbase/src/hyp/hypervisor/kvm.rs
@@ -14,12 +14,14 @@
 
 //! Wrappers around calls to the KVM hypervisor.
 
-use super::common::{
-    DeviceAssigningHypervisor, Hypervisor, MemSharingHypervisor, MmioGuardedHypervisor,
-};
-use crate::error::{Error, Result};
-use crate::util::page_address;
 use core::fmt::{self, Display, Formatter};
+
+use super::{DeviceAssigningHypervisor, Hypervisor, MemSharingHypervisor, MmioGuardedHypervisor};
+use crate::{
+    hyp::{Error, Result},
+    memory::page_4kb_of,
+};
+
 use smccc::{
     error::{positive_or_error_64, success_or_error_32, success_or_error_64},
     hvc64,
@@ -113,24 +115,32 @@
 
     fn map(&self, addr: usize) -> Result<()> {
         let mut args = [0u64; 17];
-        args[0] = page_address(addr);
+        args[0] = page_4kb_of(addr).try_into().unwrap();
 
-        // TODO(b/277859415): pKVM returns a i32 instead of a i64 in T.
-        // Drop this hack once T reaches EoL.
-        success_or_error_32(hvc64(VENDOR_HYP_KVM_MMIO_GUARD_MAP_FUNC_ID, args)[0] as u32)
-            .map_err(|e| Error::KvmError(e, VENDOR_HYP_KVM_MMIO_GUARD_MAP_FUNC_ID))
+        if cfg!(feature = "compat_android_13") {
+            let res = hvc64(VENDOR_HYP_KVM_MMIO_GUARD_MAP_FUNC_ID, args)[0];
+            // pKVM returns i32 instead of the intended i64 in Android 13.
+            return success_or_error_32(res as u32)
+                .map_err(|e| Error::KvmError(e, VENDOR_HYP_KVM_MMIO_GUARD_MAP_FUNC_ID));
+        }
+
+        checked_hvc64_expect_zero(VENDOR_HYP_KVM_MMIO_GUARD_MAP_FUNC_ID, args)
     }
 
     fn unmap(&self, addr: usize) -> Result<()> {
         let mut args = [0u64; 17];
-        args[0] = page_address(addr);
+        args[0] = page_4kb_of(addr).try_into().unwrap();
 
-        // TODO(b/277860860): pKVM returns NOT_SUPPORTED for SUCCESS in T.
-        // Drop this hack once T reaches EoL.
-        match success_or_error_64(hvc64(VENDOR_HYP_KVM_MMIO_GUARD_UNMAP_FUNC_ID, args)[0]) {
-            Err(KvmError::NotSupported) | Ok(_) => Ok(()),
-            Err(e) => Err(Error::KvmError(e, VENDOR_HYP_KVM_MMIO_GUARD_UNMAP_FUNC_ID)),
+        if cfg!(feature = "compat_android_13") {
+            let res = hvc64(VENDOR_HYP_KVM_MMIO_GUARD_UNMAP_FUNC_ID, args)[0];
+            // pKVM returns NOT_SUPPORTED for SUCCESS in Android 13.
+            return match success_or_error_64(res) {
+                Err(KvmError::NotSupported) | Ok(_) => Ok(()),
+                Err(e) => Err(Error::KvmError(e, VENDOR_HYP_KVM_MMIO_GUARD_UNMAP_FUNC_ID)),
+            };
         }
+
+        checked_hvc64_expect_zero(VENDOR_HYP_KVM_MMIO_GUARD_UNMAP_FUNC_ID, args)
     }
 
     fn granule(&self) -> Result<usize> {
diff --git a/vmbase/src/lib.rs b/vmbase/src/lib.rs
index 431e899..630834b 100644
--- a/vmbase/src/lib.rs
+++ b/vmbase/src/lib.rs
@@ -26,6 +26,7 @@
 pub mod fdt;
 pub mod heap;
 mod hvc;
+pub mod hyp;
 pub mod layout;
 pub mod linker;
 pub mod logger;
diff --git a/vmbase/src/memory.rs b/vmbase/src/memory.rs
index 2f72fc4..299d50f 100644
--- a/vmbase/src/memory.rs
+++ b/vmbase/src/memory.rs
@@ -26,8 +26,8 @@
     handle_permission_fault, handle_translation_fault, MemoryRange, MemoryTracker, MEMORY,
 };
 pub use util::{
-    flush, flushed_zeroize, min_dcache_line_size, page_4kb_of, PAGE_SIZE, SIZE_128KB, SIZE_2MB,
-    SIZE_4KB, SIZE_4MB, SIZE_64KB,
+    flush, flushed_zeroize, min_dcache_line_size, page_4kb_of, PAGE_SIZE, SIZE_128KB, SIZE_16KB,
+    SIZE_2MB, SIZE_4KB, SIZE_4MB, SIZE_64KB,
 };
 
 pub(crate) use shared::{alloc_shared, dealloc_shared};
diff --git a/vmbase/src/memory/error.rs b/vmbase/src/memory/error.rs
index 273db56..4d08f1e 100644
--- a/vmbase/src/memory/error.rs
+++ b/vmbase/src/memory/error.rs
@@ -16,6 +16,8 @@
 
 use core::fmt;
 
+use crate::hyp;
+
 /// Errors for MemoryTracker operations.
 #[derive(Debug, Clone)]
 pub enum MemoryTrackerError {
@@ -47,6 +49,10 @@
     FlushRegionFailed,
     /// Failed to set PTE dirty state.
     SetPteDirtyFailed,
+    /// Attempting to MMIO_GUARD_MAP more than once the same region.
+    DuplicateMmioShare(usize),
+    /// The MMIO_GUARD granule used by the hypervisor is not supported.
+    UnsupportedMmioGuardGranule(usize),
 }
 
 impl fmt::Display for MemoryTrackerError {
@@ -66,6 +72,12 @@
             Self::InvalidPte => write!(f, "Page table entry is not valid"),
             Self::FlushRegionFailed => write!(f, "Failed to flush memory region"),
             Self::SetPteDirtyFailed => write!(f, "Failed to set PTE dirty state"),
+            Self::DuplicateMmioShare(addr) => {
+                write!(f, "Attempted to share the same MMIO region at {addr:#x} twice")
+            }
+            Self::UnsupportedMmioGuardGranule(g) => {
+                write!(f, "Unsupported MMIO guard granule: {g}")
+            }
         }
     }
 }
diff --git a/vmbase/src/memory/shared.rs b/vmbase/src/memory/shared.rs
index dd433d4..5a25d9f 100644
--- a/vmbase/src/memory/shared.rs
+++ b/vmbase/src/memory/shared.rs
@@ -18,14 +18,18 @@
 use super::error::MemoryTrackerError;
 use super::page_table::{PageTable, MMIO_LAZY_MAP_FLAG};
 use super::util::{page_4kb_of, virt_to_phys};
+use crate::console;
 use crate::dsb;
 use crate::exceptions::HandleExceptionError;
+use crate::hyp::{self, get_mem_sharer, get_mmio_guard};
+use crate::util::unchecked_align_down;
 use crate::util::RangeExt as _;
 use aarch64_paging::paging::{
-    Attributes, Descriptor, MemoryRegion as VaRange, VirtualAddress, BITS_PER_LEVEL, PAGE_SIZE,
+    Attributes, Descriptor, MemoryRegion as VaRange, VirtualAddress, PAGE_SIZE,
 };
 use alloc::alloc::{alloc_zeroed, dealloc, handle_alloc_error};
 use alloc::boxed::Box;
+use alloc::collections::BTreeSet;
 use alloc::vec::Vec;
 use buddy_system_allocator::{FrameAllocator, LockedFrameAllocator};
 use core::alloc::Layout;
@@ -35,7 +39,6 @@
 use core::ops::Range;
 use core::ptr::NonNull;
 use core::result;
-use hyp::{get_mem_sharer, get_mmio_guard, MMIO_GUARD_GRANULE_SIZE};
 use log::{debug, error, trace};
 use once_cell::race::OnceBox;
 use spin::mutex::SpinMutex;
@@ -77,6 +80,7 @@
     mmio_regions: ArrayVec<[MemoryRange; MemoryTracker::MMIO_CAPACITY]>,
     mmio_range: MemoryRange,
     payload_range: Option<MemoryRange>,
+    mmio_sharer: MmioSharer,
 }
 
 impl MemoryTracker {
@@ -113,6 +117,7 @@
             mmio_regions: ArrayVec::new(),
             mmio_range,
             payload_range: payload_range.map(|r| r.start.0..r.end.0),
+            mmio_sharer: MmioSharer::new().unwrap(),
         }
     }
 
@@ -248,17 +253,10 @@
         Ok(self.regions.last().unwrap().range.clone())
     }
 
-    /// Unmaps all tracked MMIO regions from the MMIO guard.
-    ///
-    /// Note that they are not unmapped from the page table.
-    pub fn mmio_unmap_all(&mut self) -> Result<()> {
-        if get_mmio_guard().is_some() {
-            for range in &self.mmio_regions {
-                self.page_table
-                    .walk_range(&get_va_range(range), &mmio_guard_unmap_page)
-                    .map_err(|_| MemoryTrackerError::FailedToUnmap)?;
-            }
-        }
+    /// Unshares any MMIO region previously shared with the MMIO guard.
+    pub fn unshare_all_mmio(&mut self) -> Result<()> {
+        self.mmio_sharer.unshare_all();
+
         Ok(())
     }
 
@@ -320,15 +318,21 @@
     /// Handles translation fault for blocks flagged for lazy MMIO mapping by enabling the page
     /// table entry and MMIO guard mapping the block. Breaks apart a block entry if required.
     fn handle_mmio_fault(&mut self, addr: VirtualAddress) -> Result<()> {
-        let page_start = VirtualAddress(page_4kb_of(addr.0));
-        assert_eq!(page_start.0 % MMIO_GUARD_GRANULE_SIZE, 0);
-        let page_range: VaRange = (page_start..page_start + MMIO_GUARD_GRANULE_SIZE).into();
-        let mmio_guard = get_mmio_guard().unwrap();
+        let shared_range = self.mmio_sharer.share(addr)?;
+        self.map_lazy_mmio_as_valid(&shared_range)?;
+
+        Ok(())
+    }
+
+    /// Modify the PTEs corresponding to a given range from (invalid) "lazy MMIO" to valid MMIO.
+    ///
+    /// Returns an error if any PTE in the range is not an invalid lazy MMIO mapping.
+    fn map_lazy_mmio_as_valid(&mut self, page_range: &VaRange) -> Result<()> {
         // This must be safe and free from break-before-make (BBM) violations, given that the
         // initial lazy mapping has the valid bit cleared, and each newly created valid descriptor
         // created inside the mapping has the same size and alignment.
         self.page_table
-            .modify_range(&page_range, &|_: &VaRange, desc: &mut Descriptor, _: usize| {
+            .modify_range(page_range, &|_: &VaRange, desc: &mut Descriptor, _: usize| {
                 let flags = desc.flags().expect("Unsupported PTE flags set");
                 if flags.contains(MMIO_LAZY_MAP_FLAG) && !flags.contains(Attributes::VALID) {
                     desc.modify_flags(Attributes::VALID, Attributes::empty());
@@ -337,8 +341,7 @@
                     Err(())
                 }
             })
-            .map_err(|_| MemoryTrackerError::InvalidPte)?;
-        Ok(mmio_guard.map(page_start.0)?)
+            .map_err(|_| MemoryTrackerError::InvalidPte)
     }
 
     /// Flush all memory regions marked as writable-dirty.
@@ -376,6 +379,71 @@
     }
 }
 
+struct MmioSharer {
+    granule: usize,
+    frames: BTreeSet<usize>,
+}
+
+impl MmioSharer {
+    fn new() -> Result<Self> {
+        let granule = Self::get_granule()?;
+        let frames = BTreeSet::new();
+
+        // Allows safely calling util::unchecked_align_down().
+        assert!(granule.is_power_of_two());
+
+        Ok(Self { granule, frames })
+    }
+
+    fn get_granule() -> Result<usize> {
+        let Some(mmio_guard) = get_mmio_guard() else {
+            return Ok(PAGE_SIZE);
+        };
+        match mmio_guard.granule()? {
+            granule if granule % PAGE_SIZE == 0 => Ok(granule), // For good measure.
+            granule => Err(MemoryTrackerError::UnsupportedMmioGuardGranule(granule)),
+        }
+    }
+
+    /// Share the MMIO region aligned to the granule size containing addr (not validated as MMIO).
+    fn share(&mut self, addr: VirtualAddress) -> Result<VaRange> {
+        // This can't use virt_to_phys() since 0x0 is a valid MMIO address and we are ID-mapped.
+        let phys = addr.0;
+        let base = unchecked_align_down(phys, self.granule);
+
+        // TODO(ptosi): Share the UART using this method and remove the hardcoded check.
+        if self.frames.contains(&base) || base == page_4kb_of(console::BASE_ADDRESS) {
+            return Err(MemoryTrackerError::DuplicateMmioShare(base));
+        }
+
+        if let Some(mmio_guard) = get_mmio_guard() {
+            mmio_guard.map(base)?;
+        }
+
+        let inserted = self.frames.insert(base);
+        assert!(inserted);
+
+        let base_va = VirtualAddress(base);
+        Ok((base_va..base_va + self.granule).into())
+    }
+
+    fn unshare_all(&mut self) {
+        let Some(mmio_guard) = get_mmio_guard() else {
+            return self.frames.clear();
+        };
+
+        while let Some(base) = self.frames.pop_first() {
+            mmio_guard.unmap(base).unwrap();
+        }
+    }
+}
+
+impl Drop for MmioSharer {
+    fn drop(&mut self) {
+        self.unshare_all();
+    }
+}
+
 /// Allocates a memory range of at least the given size and alignment that is shared with the host.
 /// Returns a pointer to the buffer.
 pub(crate) fn alloc_shared(layout: Layout) -> hyp::Result<NonNull<u8>> {
@@ -479,41 +547,6 @@
     }
 }
 
-/// MMIO guard unmaps page
-fn mmio_guard_unmap_page(
-    va_range: &VaRange,
-    desc: &Descriptor,
-    level: usize,
-) -> result::Result<(), ()> {
-    let flags = desc.flags().expect("Unsupported PTE flags set");
-    // This function will be called on an address range that corresponds to a device. Only if a
-    // page has been accessed (written to or read from), will it contain the VALID flag and be MMIO
-    // guard mapped. Therefore, we can skip unmapping invalid pages, they were never MMIO guard
-    // mapped anyway.
-    if flags.contains(Attributes::VALID) {
-        assert!(
-            flags.contains(MMIO_LAZY_MAP_FLAG),
-            "Attempting MMIO guard unmap for non-device pages"
-        );
-        const MMIO_GUARD_GRANULE_SHIFT: u32 = MMIO_GUARD_GRANULE_SIZE.ilog2() - PAGE_SIZE.ilog2();
-        const MMIO_GUARD_GRANULE_LEVEL: usize =
-            3 - (MMIO_GUARD_GRANULE_SHIFT as usize / BITS_PER_LEVEL);
-        assert_eq!(
-            level, MMIO_GUARD_GRANULE_LEVEL,
-            "Failed to break down block mapping before MMIO guard mapping"
-        );
-        let page_base = va_range.start().0;
-        assert_eq!(page_base % MMIO_GUARD_GRANULE_SIZE, 0);
-        // Since mmio_guard_map takes IPAs, if pvmfw moves non-ID address mapping, page_base
-        // should be converted to IPA. However, since 0x0 is a valid MMIO address, we don't use
-        // virt_to_phys here, and just pass page_base instead.
-        get_mmio_guard().unwrap().unmap(page_base).map_err(|e| {
-            error!("Error MMIO guard unmapping: {e}");
-        })?;
-    }
-    Ok(())
-}
-
 /// Handles a translation fault with the given fault address register (FAR).
 #[inline]
 pub fn handle_translation_fault(far: VirtualAddress) -> result::Result<(), HandleExceptionError> {
diff --git a/vmbase/src/memory/util.rs b/vmbase/src/memory/util.rs
index 2b75414..e9f867f 100644
--- a/vmbase/src/memory/util.rs
+++ b/vmbase/src/memory/util.rs
@@ -22,6 +22,8 @@
 
 /// The size of a 4KB memory in bytes.
 pub const SIZE_4KB: usize = 4 << 10;
+/// The size of a 16KB memory in bytes.
+pub const SIZE_16KB: usize = 16 << 10;
 /// The size of a 64KB memory in bytes.
 pub const SIZE_64KB: usize = 64 << 10;
 /// The size of a 128KB memory in bytes.