Memory management: per-process VA range, brk frontier tracking, noreserve/LOW_2G support

litebox/src/mm/mod.rs

@@ -34,16 +34,41 @@ where
     Platform: RawSyncPrimitivesProvider + PageManagementProvider<ALIGN>,
 {
     vmem: RwLock<Platform, Vmem<Platform, ALIGN>>,
+    /// Lower bound (inclusive) of the VA range (cached for lock-free checks).
+    addr_min: usize,
+    /// Upper bound (exclusive) of the VA range (cached for lock-free checks).
+    addr_max: usize,
 }
 
 impl<Platform, const ALIGN: usize> PageManager<Platform, ALIGN>
 where
     Platform: RawSyncPrimitivesProvider + PageManagementProvider<ALIGN>,
 {
-    /// Create a new `PageManager` instance.
-    pub fn new(litebox: &LiteBox<Platform>) -> Self {
-        let vmem = RwLock::new(linux::Vmem::new(litebox.x.platform));
-        Self { vmem }
+    /// Lower bound (inclusive) of the managed VA range.
+    pub fn addr_min(&self) -> usize {
+        self.addr_min
+    }
+
+    /// Upper bound (exclusive) of the managed VA range.
+    pub fn addr_max(&self) -> usize {
+        self.addr_max
+    }
+
+    /// Create a new `PageManager` instance for the given VA `range`.
+    ///
+    /// For a single-process setup, pass the full platform range
+    /// `Platform::TASK_ADDR_MIN..Platform::TASK_ADDR_MAX`. For multi-process
+    /// setups, pass the sub-range obtained from
+    /// [`AddressSpaceProvider::address_space_range()`](crate::platform::AddressSpaceProvider::address_space_range).
+    pub fn new(litebox: &LiteBox<Platform>, range: Range<usize>) -> Self {
+        let addr_min = range.start;
+        let addr_max = range.end;
+        let vmem = RwLock::new(linux::Vmem::new(litebox.x.platform, range));
+        Self {
+            vmem,
+            addr_min,
+            addr_max,
+        }
     }
 
     /// Create a mapping with the given flags.
@@ -289,7 +314,40 @@ where
     pub fn set_initial_brk(&self, brk: usize) {
         let mut vmem = self.vmem.write();
         assert_eq!(vmem.brk, 0, "initial brk is already set");
+        vmem.brk_base = brk;
         vmem.brk = brk;
+        vmem.brk_frontier = brk.next_multiple_of(linux::PAGE_SIZE);
+    }
+
+    /// Returns the current logical program break without modifying mappings.
+    pub fn current_brk(&self) -> usize {
+        self.vmem.read().brk
+    }
+
+    /// Returns the current page-aligned heap frontier that `brk` growth uses.
+    pub fn current_brk_frontier(&self) -> usize {
+        self.vmem.read().brk_frontier
+    }
+
+    /// Advance the minimum program break, current break, and heap frontier to
+    /// at least `min_brk` without allocating new heap pages.
+    ///
+    /// This is used only when already-mapped non-heap pages, such as a
+    /// trampoline at the heap frontier, occupy space that future `brk` growth
+    /// must skip over. The skipped gap becomes a persistent floor: later
+    /// `brk` shrinks must not re-enter it, or subsequent growth would collide
+    /// with the same non-heap pages again.
+    pub fn ensure_brk_past(&self, min_brk: usize) {
+        let mut vmem = self.vmem.write();
+        if vmem.brk_base < min_brk {
+            vmem.brk_base = min_brk;
+        }
+        if vmem.brk < min_brk {
+            vmem.brk = min_brk;
+        }
+        if vmem.brk_frontier < min_brk {
+            vmem.brk_frontier = min_brk;
+        }
     }
 
     /// Set the program break to the given address.
@@ -319,10 +377,15 @@ where
             // Calling `brk` with 0 can be used to find the current location of the program break.
             return Ok(vmem.brk);
         }
+        if brk < vmem.brk_base {
+            // Linux keeps the current break unchanged when the request is
+            // below the permitted heap base.
+            return Ok(vmem.brk);
+        }
 
-        let old_brk = vmem.brk.next_multiple_of(linux::PAGE_SIZE);
+        let old_brk = vmem.brk_frontier;
         let new_brk = brk.next_multiple_of(linux::PAGE_SIZE);
-        if vmem.brk >= brk {
+        if new_brk < old_brk {
             // Shrink the memory region
             let brk = match unsafe {
                 vmem.remove_mapping(
@@ -331,6 +394,7 @@ where
             } {
                 Ok(()) => {
                     vmem.brk = brk;
+                    vmem.brk_frontier = new_brk;
                     brk
                 }
                 Err(_) => {
@@ -340,20 +404,23 @@ where
             return Ok(brk);
         }
 
-        if vmem.overlapping(old_brk..new_brk).next().is_some() {
-            return Err(MappingError::OutOfMemory);
-        }
-        if let Some(range) = PageRange::<ALIGN>::new(old_brk, new_brk) {
-            let (suggested_address, length) = range.start_and_length();
-            let perms = MemoryRegionPermissions::READ | MemoryRegionPermissions::WRITE;
-            unsafe {
-                vmem.create_pages(
-                    Some(suggested_address),
-                    length,
-                    CreatePagesFlags::FIXED_ADDR | CreatePagesFlags::POPULATE_PAGES_IMMEDIATELY,
-                    perms,
-                )
-            }?;
+        if new_brk > old_brk {
+            if vmem.overlapping(old_brk..new_brk).next().is_some() {
+                return Err(MappingError::OutOfMemory);
+            }
+            if let Some(range) = PageRange::<ALIGN>::new(old_brk, new_brk) {
+                let (suggested_address, length) = range.start_and_length();
+                let perms = MemoryRegionPermissions::READ | MemoryRegionPermissions::WRITE;
+                unsafe {
+                    vmem.create_pages(
+                        Some(suggested_address),
+                        length,
+                        CreatePagesFlags::FIXED_ADDR | CreatePagesFlags::POPULATE_PAGES_IMMEDIATELY,
+                        perms,
+                    )
+                }?;
+            }
+            vmem.brk_frontier = new_brk;
         }
         vmem.brk = brk;
         Ok(brk)
@@ -381,7 +448,9 @@ where
 
         // reset brk
         let mut vmem = self.vmem.write();
+        vmem.brk_base = 0;
         vmem.brk = 0;
+        vmem.brk_frontier = 0;
 
         Ok(())
     }
@@ -442,6 +511,24 @@ where
             Err(linux::VmemResizeError::NotExist(_)) => Err(RemapError::AlreadyUnallocated),
             Err(linux::VmemResizeError::InvalidAddr { .. }) => Err(RemapError::AlreadyAllocated),
             Err(linux::VmemResizeError::OutOfMemory) => Err(RemapError::OutOfMemory),
+            Err(linux::VmemResizeError::OutOfRange) => {
+                // Expanded range exceeds address space limits — try moving
+                if !may_move {
+                    return Err(RemapError::OutOfMemory);
+                }
+                match unsafe {
+                    vmem.move_mappings(
+                        old_range,
+                        None,
+                        NonZeroPageSize::new(new_size).ok_or(RemapError::Unaligned)?,
+                    )
+                } {
+                    Ok(new_addr) => Ok(new_addr),
+                    Err(linux::VmemMoveError::OutOfMemory) => Err(RemapError::OutOfMemory),
+                    Err(linux::VmemMoveError::UnAligned) => Err(RemapError::Unaligned),
+                    Err(linux::VmemMoveError::RemapError(err)) => Err(err),
+                }
+            }
         }
     }
 
@@ -605,16 +692,19 @@ where
     ///
     /// The `range` must be an already-mapped region with the given `permissions`.
     #[must_use]
+    #[allow(clippy::fn_params_excessive_bools)]
     pub unsafe fn register_existing_mapping(
         &self,
         range: PageRange<ALIGN>,
         permissions: MemoryRegionPermissions,
         is_file_backed: bool,
         replace: bool,
         shared: bool,
+        fd_writable: bool,
     ) -> Option<()> {
         let vma = VmArea::new(
-            VmFlags::from(permissions) | VmFlags::may_flags_for_mapping(shared, is_file_backed),
+            VmFlags::from(permissions)
+                | VmFlags::may_flags_for_mapping(shared, is_file_backed, fd_writable),
             is_file_backed,
         );
         let mut vmem = self.vmem.write();
@@ -665,22 +755,25 @@ where
     ///
     /// # Safety
     ///
-    /// This should only be called from the kernel page fault handler.
+    /// This must only be called while servicing a real page fault for the
+    /// current address space, and `fault_addr` / `error_code` must come from
+    /// that trap context. Depending on the platform, that may be a kernel fault
+    /// handler or an opted-in user-mode fault path.
     pub unsafe fn handle_page_fault(
         &self,
         fault_addr: usize,
         error_code: u64,
     ) -> Result<(), PageFaultError> {
         let fault_addr = fault_addr & !(ALIGN - 1);
-        if !(Platform::TASK_ADDR_MIN..Platform::TASK_ADDR_MAX).contains(&fault_addr) {
+        if !(self.addr_min..self.addr_max).contains(&fault_addr) {
             return Err(PageFaultError::AccessError("Invalid address"));
         }
 
         let mut vmem = self.vmem.write();
         // Find the range closest to the fault address
         let (start, vma) = {
             let (r, vma) = vmem
-                .overlapping(fault_addr..Platform::TASK_ADDR_MAX)
+                .overlapping(fault_addr..vmem.addr_max())
                 .next()
                 .ok_or(PageFaultError::AccessError("no mapping"))?;
             (r.start, *vma)
@@ -692,7 +785,7 @@ where
             }
 
             if !vmem
-                .overlapping(Platform::TASK_ADDR_MIN..fault_addr)
+                .overlapping(vmem.addr_min()..fault_addr)
                 .next_back()
                 .is_none_or(|(prev_range, prev_vma)| {
                     // Enforce gap between stack and other preceding non-stack mappings.
@@ -708,15 +801,19 @@ where
             let Some(range) = PageRange::new(fault_addr, start) else {
                 unreachable!()
             };
-            if let Err(err) = unsafe {
+            if unsafe {
                 vmem.insert_mapping(
                     range,
                     vma,
                     false,
+                    vma.noreserve(),
                     crate::platform::page_mgmt::FixedAddressBehavior::NoReplace,
+                    false,
                 )
-            } {
-                unimplemented!("failed to grow stack: {:?}", err)
+            }
+            .is_err()
+            {
+                return Err(PageFaultError::AllocationFailed);
             }
         }
 

litebox/src/mm/tests.rs

@@ -86,7 +86,10 @@ fn collect_mappings(vmm: &Vmem<DummyVmemBackend, PAGE_SIZE>) -> Vec<Range<usize>
 fn test_vmm_mapping() {
     let start_addr: usize = 0x1_0000;
     let range = PageRange::new(start_addr, start_addr + 12 * PAGE_SIZE).unwrap();
-    let mut vmm = Vmem::new(&DummyVmemBackend);
+    let mut vmm = Vmem::new(
+        &DummyVmemBackend,
+        DummyVmemBackend::TASK_ADDR_MIN..DummyVmemBackend::TASK_ADDR_MAX,
+    );
 
     // []
     unsafe {
@@ -97,7 +100,9 @@ fn test_vmm_mapping() {
                 false,
             ),
             false,
+            false,
             crate::platform::page_mgmt::FixedAddressBehavior::Replace,
+            false,
         )
     }
     .unwrap();

litebox_platform_linux_kernel/src/mm/tests.rs

@@ -218,7 +218,11 @@ fn test_vmm_page_fault() {
     let p4 = PageTableAllocator::<MockKernel>::allocate_frame(true).unwrap();
     let platform = MockKernel::new(p4.start_address());
     let litebox = LiteBox::new(platform);
-    let vmm = PageManager::<_, PAGE_SIZE>::new(&litebox);
+    let vmm = PageManager::<_, PAGE_SIZE>::new(
+        &litebox,
+        <MockKernel as litebox::platform::PageManagementProvider<PAGE_SIZE>>::TASK_ADDR_MIN
+            ..<MockKernel as litebox::platform::PageManagementProvider<PAGE_SIZE>>::TASK_ADDR_MAX,
+    );
     unsafe {
         assert_eq!(
             vmm.create_writable_pages(

litebox_platform_lvbs/src/mm/tests.rs

@@ -238,7 +238,11 @@ fn test_vmm_page_fault() {
         x86_64::PhysAddr::new(0),
     );
     let litebox = LiteBox::new(platform);
-    let vmm = PageManager::<_, PAGE_SIZE>::new(&litebox);
+    let vmm = PageManager::<_, PAGE_SIZE>::new(
+        &litebox,
+        <MockKernel as litebox::platform::PageManagementProvider<PAGE_SIZE>>::TASK_ADDR_MIN
+            ..<MockKernel as litebox::platform::PageManagementProvider<PAGE_SIZE>>::TASK_ADDR_MAX,
+    );
     unsafe {
         assert_eq!(
             vmm.create_writable_pages(

litebox_shim_linux/src/lib.rs

@@ -191,7 +191,11 @@ impl LinuxShimBuilder {
         net.set_platform_interaction(litebox::net::PlatformInteraction::Manual);
         let global = Arc::new(GlobalState {
             platform: self.platform,
-            pm: PageManager::new(&self.litebox),
+            pm: PageManager::new(
+                &self.litebox,
+                <Platform as litebox::platform::PageManagementProvider<{ PAGE_SIZE }>>::TASK_ADDR_MIN
+                    ..<Platform as litebox::platform::PageManagementProvider<{ PAGE_SIZE }>>::TASK_ADDR_MAX,
+            ),
             futex_manager: FutexManager::new(),
             pipes: Pipes::new(&self.litebox),
             net: litebox::sync::Mutex::new(net),

litebox_shim_linux/src/syscalls/mm.rs

@@ -238,6 +238,7 @@ impl<FS: ShimFS> Task<FS> {
                         true,
                         fixed_behavior == FixedAddressBehavior::Replace,
                         flags.contains(MapFlags::MAP_SHARED),
+                        false,
                     )
                 }
                 .unwrap();

litebox_shim_optee/src/lib.rs

@@ -142,7 +142,11 @@ impl OpteeShimBuilder {
     pub fn build(self) -> OpteeShim {
         let global = Arc::new(GlobalState {
             platform: self.platform,
-            pm: PageManager::new(&self.litebox),
+            pm: PageManager::new(
+                &self.litebox,
+                <Platform as litebox::platform::PageManagementProvider<PAGE_SIZE>>::TASK_ADDR_MIN
+                    ..<Platform as litebox::platform::PageManagementProvider<PAGE_SIZE>>::TASK_ADDR_MAX,
+            ),
             _litebox: self.litebox,
             ta_uuid_map: TaUuidMap::new(),
         });