Add AddressSpaceProvider trait for multi-process address space management

litebox/src/platform/address_space.rs

@@ -0,0 +1,137 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+//! Address-space management types and traits for multi-process support.
+//!
+//! The [`AddressSpaceProvider`] trait is an **optional** South interface that
+//! platforms implement to manage per-process address spaces. Platforms may use
+//! separate page tables, VA-range partitioning, or other techniques to isolate
+//! address spaces.
+
+use core::ops::Range;
+use thiserror::Error;
+
+/// The result of forking an address space.
+///
+/// The variant tells the caller what kind of copy was created so it can adjust
+/// its behavior (e.g., whether to copy page contents or share them).
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ForkedAddressSpace<Id> {
+    /// Independent copy-on-write copy with the full address range. The child
+    /// has its own backing structures; CoW faults are resolved by the
+    /// platform.
+    Independent(Id),
+    /// A new VA-range partition is assigned to the child. Parent memory is
+    /// shared; the shim is responsible for copying pages as needed.
+    SharedWithParent(Id),
+}
+
+/// Errors that can occur during address-space operations.
+#[derive(Error, Debug)]
+#[non_exhaustive]
+pub enum AddressSpaceError {
+    /// No free address-space slots or VA ranges available.
+    #[error("no address space slots available")]
+    NoSpace,
+    /// The given address-space ID is not valid (already destroyed, never
+    /// created, etc.).
+    #[error("invalid address space id")]
+    InvalidId,
+    /// The platform does not support this operation.
+    #[error("operation not supported by this platform")]
+    NotSupported,
+}
+
+/// A provider for managing per-process address spaces.
+///
+/// This is an **optional** trait — platforms that do not yet support
+/// multi-process may leave all methods at the default (which returns
+/// [`AddressSpaceError::NotSupported`]).
+///
+/// # Associated Type
+///
+/// `AddressSpaceId` is an opaque, lightweight handle that identifies one
+/// address space. It must be `Copy + Eq + Send + Sync` so it can be stored
+/// inside process contexts and passed across threads.
+pub trait AddressSpaceProvider {
+    /// Opaque identifier for an address space.
+    type AddressSpaceId: Copy + Eq + Send + Sync + core::fmt::Debug;
+
+    /// Create a new, empty address space.
+    ///
+    /// The platform allocates whatever backing structures are needed for the
+    /// new address space.
+    fn create_address_space(&self) -> Result<Self::AddressSpaceId, AddressSpaceError> {
+        Err(AddressSpaceError::NotSupported)
+    }
+
+    /// Destroy an address space, releasing all associated resources.
+    ///
+    /// After this call, `id` is invalid and must not be reused.
+    fn destroy_address_space(&self, id: Self::AddressSpaceId) -> Result<(), AddressSpaceError> {
+        let _ = id;
+        Err(AddressSpaceError::NotSupported)
+    }
+
+    /// Fork an address space from `parent`.
+    ///
+    /// Returns a [`ForkedAddressSpace`] indicating what kind of fork was
+    /// performed:
+    ///
+    /// * [`Independent`](ForkedAddressSpace::Independent) — full CoW copy.
+    /// * [`SharedWithParent`](ForkedAddressSpace::SharedWithParent) — new VA
+    ///   partition, parent pages shared.
+    fn fork_address_space(
+        &self,
+        parent: Self::AddressSpaceId,
+    ) -> Result<ForkedAddressSpace<Self::AddressSpaceId>, AddressSpaceError> {
+        let _ = parent;
+        Err(AddressSpaceError::NotSupported)
+    }
+
+    /// Make `id` the active address space for the current CPU / thread.
+    fn activate_address_space(&self, id: Self::AddressSpaceId) -> Result<(), AddressSpaceError> {
+        let _ = id;
+        Err(AddressSpaceError::NotSupported)
+    }
+
+    /// Execute `f` with the given address space active, then restore the
+    /// previously active address space.
+    ///
+    /// Implementations **must** restore the prior address space even if `f`
+    /// panics (use a guard / RAII pattern).
+    ///
+    /// The default returns [`AddressSpaceError::NotSupported`]. Platforms that
+    /// implement [`activate_address_space`](Self::activate_address_space) should
+    /// also override this method with a proper save/restore sequence.
+    fn with_address_space<R>(
+        &self,
+        id: Self::AddressSpaceId,
+        f: impl FnOnce() -> R,
+    ) -> Result<R, AddressSpaceError> {
+        let _ = (id, f);
+        Err(AddressSpaceError::NotSupported)
+    }
+
+    /// Whether the platform requires eager copy-on-write snapshots during
+    /// fork instead of lazy page-fault-driven CoW.
+    ///
+    /// When `true`, the shim eagerly copies all writable guest pages before
+    /// spawning the forked child and restores them after the child execs or
+    /// exits. When `false` (the default), the shim marks writable pages
+    /// read-only and lazily snapshots individual pages on first write fault.
+    ///
+    /// Platforms where the exception/fault handler shares the guest address
+    /// space must set this to `true` because a CoW fault inside the handler
+    /// itself would be fatal.
+    const EAGER_COW_ON_FORK: bool = false;
+
+    /// Return the VA range available to the given address space.
+    fn address_space_range(
+        &self,
+        id: Self::AddressSpaceId,
+    ) -> Result<Range<usize>, AddressSpaceError> {
+        let _ = id;
+        Err(AddressSpaceError::NotSupported)
+    }
+}

litebox/src/platform/mock.rs

@@ -61,6 +61,12 @@ impl Provider for MockPlatform {}
 
 impl RawMessageProvider for MockPlatform {}
 
+impl AddressSpaceProvider for MockPlatform {
+    // All methods default to `Err(NotSupported)`, which is correct for the
+    // mock platform (single-process only).
+    type AddressSpaceId = u32;
+}
+
 pub(crate) struct MockRawMutex {
     inner: AtomicU32,
     internal_state: std::sync::RwLock<MockRawMutexInternalState>,

litebox/src/platform/mod.rs

@@ -7,6 +7,7 @@
 //! trait is merely a collection of subtraits that could be composed independently from various
 //! other crates that implement them upon various types.
 
+pub mod address_space;
 pub mod common_providers;
 pub mod page_mgmt;
 pub mod trivial_providers;
@@ -18,6 +19,7 @@ use either::Either;
 use thiserror::Error;
 use zerocopy::{FromBytes, IntoBytes};
 
+pub use address_space::*;
 pub use page_mgmt::PageManagementProvider;
 
 #[macro_export]
@@ -48,6 +50,7 @@ pub trait Provider:
     + PunchthroughProvider
     + DebugLogProvider
     + RawPointerProvider
+    + AddressSpaceProvider
 {
 }
 

litebox_platform_linux_kernel/src/lib.rs

@@ -87,6 +87,10 @@ impl<Host: HostInterface> Provider for LinuxKernel<Host> {}
 
 impl<Host: HostInterface> litebox::platform::RawMessageProvider for LinuxKernel<Host> {}
 
+impl<Host: HostInterface> litebox::platform::AddressSpaceProvider for LinuxKernel<Host> {
+    type AddressSpaceId = u32;
+}
+
 // TODO: implement pointer validation to ensure the pointers are in user space.
 type UserConstPtr<T> = litebox::platform::common_providers::userspace_pointers::UserConstPtr<
     litebox::platform::common_providers::userspace_pointers::NoValidation,

litebox_platform_linux_userland/src/lib.rs

@@ -447,6 +447,10 @@ impl litebox::platform::Provider for LinuxUserland {}
 
 impl litebox::platform::RawMessageProvider for LinuxUserland {}
 
+impl litebox::platform::AddressSpaceProvider for LinuxUserland {
+    type AddressSpaceId = u32;
+}
+
 impl litebox::platform::SignalProvider for LinuxUserland {
     type Signal = litebox_common_linux::signal::Signal;
 

litebox_platform_lvbs/src/lib.rs

@@ -1353,6 +1353,12 @@ impl<Host: HostInterface> litebox::platform::SystemInfoProvider for LinuxKernel<
 
 impl<Host: HostInterface> litebox::platform::RawMessageProvider for LinuxKernel<Host> {}
 
+impl<Host: HostInterface> litebox::platform::AddressSpaceProvider for LinuxKernel<Host> {
+    // All methods default to `Err(NotSupported)` — real implementation comes
+    // when LVBS multi-process (separate page tables) is added.
+    type AddressSpaceId = u32;
+}
+
 #[cfg(feature = "optee_syscall")]
 /// Checks whether the given physical addresses are contiguous with respect to ALIGN.
 fn is_contiguous<const ALIGN: usize>(addrs: &[PhysPageAddr<ALIGN>]) -> bool {

litebox_platform_windows_userland/src/lib.rs

@@ -333,6 +333,10 @@ impl WindowsUserland {
 
 impl litebox::platform::RawMessageProvider for WindowsUserland {}
 
+impl litebox::platform::AddressSpaceProvider for WindowsUserland {
+    type AddressSpaceId = u32;
+}
+
 impl litebox::platform::Provider for WindowsUserland {}
 
 impl litebox::platform::SignalProvider for WindowsUserland {