From 6ffab3a2cb8718791c348d238054cbd8071747d9 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Tue, 17 Mar 2026 03:13:06 +0000 Subject: [PATCH 1/2] fix remap_pages bugs --- litebox/src/platform/page_mgmt.rs | 6 +++ .../src/arch/x86/mm/paging.rs | 45 ++++++++++++----- litebox_platform_linux_kernel/src/lib.rs | 2 +- .../src/arch/x86/mm/paging.rs | 50 +++++++++++++------ litebox_platform_lvbs/src/lib.rs | 2 +- 5 files changed, 75 insertions(+), 30 deletions(-) diff --git a/litebox/src/platform/page_mgmt.rs b/litebox/src/platform/page_mgmt.rs index c4fca057a..7bd612225 100644 --- a/litebox/src/platform/page_mgmt.rs +++ b/litebox/src/platform/page_mgmt.rs @@ -251,6 +251,12 @@ pub enum RemapError { AlreadyAllocated, #[error("out of memory")] OutOfMemory, + #[error("operation encountered a huge page")] + HugePage, + #[error("invalid frame address: {0:#x}")] + InvalidFrameAddress(u64), + #[error("failed to restore original mapping after remap failure (frame leaked)")] + RestoreFailed, } /// Possible errors for [`PageManagementProvider::update_permissions`] diff --git a/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs b/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs index 210b51f14..60ac10b7e 100644 --- a/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs +++ b/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs @@ -174,17 +174,36 @@ impl X64PageTable<'_, M, ALIGN> { Ok((frame, fl)) => { match unsafe { inner.map_to(new_start, frame, flags, &mut allocator) } { Ok(_) => {} - Err(e) => match e { - MapToError::PageAlreadyMapped(_) => { - return Err(page_mgmt::RemapError::AlreadyAllocated); - } - MapToError::ParentEntryHugePage => { - todo!("return Err(page_mgmt::RemapError::RemapToHugePage);") + Err(e) => { + // Restore the original mapping so the frame is not leaked. + // Safety: `start` was just unmapped so its PTE slot is free. + if let Ok(restore_flush) = + unsafe { inner.map_to(start, frame, flags, &mut allocator) } + { + if FLUSH_TLB { + restore_flush.flush(); + } + } else { + // The frame is now orphaned — we cannot recover it. + debug_assert!( + false, + "remap_pages: failed to restore mapping, frame leaked" + ); + return Err(page_mgmt::RemapError::RestoreFailed); } - MapToError::FrameAllocationFailed => { - return Err(page_mgmt::RemapError::OutOfMemory); - } - }, + + return match e { + MapToError::PageAlreadyMapped(_) => { + Err(page_mgmt::RemapError::AlreadyAllocated) + } + MapToError::ParentEntryHugePage => { + Err(page_mgmt::RemapError::HugePage) + } + MapToError::FrameAllocationFailed => { + Err(page_mgmt::RemapError::OutOfMemory) + } + }; + } } if FLUSH_TLB { fl.flush(); @@ -194,15 +213,15 @@ impl X64PageTable<'_, M, ALIGN> { unreachable!() } Err(X64UnmapError::ParentEntryHugePage) => { - todo!("return Err(page_mgmt::RemapError::RemapToHugePage);") + return Err(page_mgmt::RemapError::HugePage); } Err(X64UnmapError::InvalidFrameAddress(pa)) => { - panic!("Invalid frame address: {:#x}", pa); + return Err(page_mgmt::RemapError::InvalidFrameAddress(pa.as_u64())); } }, TranslateResult::NotMapped => {} TranslateResult::InvalidFrameAddress(pa) => { - panic!("Invalid frame address: {:#x}", pa); + return Err(page_mgmt::RemapError::InvalidFrameAddress(pa.as_u64())); } } start += 1; diff --git a/litebox_platform_linux_kernel/src/lib.rs b/litebox_platform_linux_kernel/src/lib.rs index 12f6bc2d1..3e0c2dcfd 100644 --- a/litebox_platform_linux_kernel/src/lib.rs +++ b/litebox_platform_linux_kernel/src/lib.rs @@ -466,7 +466,7 @@ impl PageManagementProvider for .ok_or(litebox::platform::page_mgmt::RemapError::Unaligned)?; let new_range = PageRange::new(new_range.start, new_range.end) .ok_or(litebox::platform::page_mgmt::RemapError::Unaligned)?; - if old_range.start.max(new_range.start) <= old_range.end.min(new_range.end) { + if old_range.start.max(new_range.start) < old_range.end.min(new_range.end) { return Err(litebox::platform::page_mgmt::RemapError::Overlapping); } unsafe { self.page_table.remap_pages(old_range, new_range) } diff --git a/litebox_platform_lvbs/src/arch/x86/mm/paging.rs b/litebox_platform_lvbs/src/arch/x86/mm/paging.rs index feab9aeee..ffb1b0cb1 100644 --- a/litebox_platform_lvbs/src/arch/x86/mm/paging.rs +++ b/litebox_platform_lvbs/src/arch/x86/mm/paging.rs @@ -258,34 +258,54 @@ impl X64PageTable<'_, M, ALIGN> { Ok((frame, _)) => { match unsafe { inner.map_to(new_start, frame, flags, &mut allocator) } { Ok(_) => {} - Err(e) => match e { - MapToError::PageAlreadyMapped(_) => { - return Err(page_mgmt::RemapError::AlreadyAllocated); - } - MapToError::ParentEntryHugePage => { - todo!("return Err(page_mgmt::RemapError::RemapToHugePage);") + Err(e) => { + // Restore the original mapping so the frame is not leaked. + // Safety: `start` was just unmapped so its PTE slot is free. + // The returned flush is intentionally ignored: the batch + // flush_tlb_range below covers this address. + if unsafe { inner.map_to(start, frame, flags, &mut allocator) } + .is_err() + { + // The frame is now orphaned — we cannot recover it. + debug_assert!( + false, + "remap_pages: failed to restore mapping, frame leaked" + ); + return Err(page_mgmt::RemapError::RestoreFailed); } - MapToError::FrameAllocationFailed => { - return Err(page_mgmt::RemapError::OutOfMemory); - } - }, + + // Flush TLB for pages already remapped in earlier iterations. + let done = (start.start_address() - flush_start.start_address()) + / Size4KiB::SIZE; + flush_tlb_range(flush_start, done.truncate()); + + return match e { + MapToError::PageAlreadyMapped(_) => { + Err(page_mgmt::RemapError::AlreadyAllocated) + } + MapToError::ParentEntryHugePage => { + Err(page_mgmt::RemapError::HugePage) + } + MapToError::FrameAllocationFailed => { + Err(page_mgmt::RemapError::OutOfMemory) + } + }; + } } } Err(X64UnmapError::PageNotMapped) => { unreachable!() } Err(X64UnmapError::ParentEntryHugePage) => { - todo!("return Err(page_mgmt::RemapError::RemapToHugePage);") + return Err(page_mgmt::RemapError::HugePage); } Err(X64UnmapError::InvalidFrameAddress(pa)) => { - // TODO: `panic!()` -> `todo!()` because user-driven interrupts or exceptions must not halt the kernel. - // We should handle this exception carefully (i.e., clean up the context and data structures belonging to an errorneous process). - todo!("Invalid frame address: {:#x}", pa); + return Err(page_mgmt::RemapError::InvalidFrameAddress(pa.as_u64())); } }, TranslateResult::NotMapped => {} TranslateResult::InvalidFrameAddress(pa) => { - todo!("Invalid frame address: {:#x}", pa); + return Err(page_mgmt::RemapError::InvalidFrameAddress(pa.as_u64())); } } start += 1; diff --git a/litebox_platform_lvbs/src/lib.rs b/litebox_platform_lvbs/src/lib.rs index 2487086d7..11e0cdd75 100644 --- a/litebox_platform_lvbs/src/lib.rs +++ b/litebox_platform_lvbs/src/lib.rs @@ -1273,7 +1273,7 @@ impl PageManagementProvider for .ok_or(litebox::platform::page_mgmt::RemapError::Unaligned)?; let new_range = PageRange::new(new_range.start, new_range.end) .ok_or(litebox::platform::page_mgmt::RemapError::Unaligned)?; - if old_range.start.max(new_range.start) <= old_range.end.min(new_range.end) { + if old_range.start.max(new_range.start) < old_range.end.min(new_range.end) { return Err(litebox::platform::page_mgmt::RemapError::Overlapping); } unsafe { From f8e6710ff0c03b120027017cbc732a2a153fd505 Mon Sep 17 00:00:00 2001 From: Sangho Lee Date: Tue, 17 Mar 2026 03:34:38 +0000 Subject: [PATCH 2/2] fix some other page flag bugs and suppress panics --- litebox/src/mm/linux.rs | 2 + litebox/src/platform/page_mgmt.rs | 6 ++ .../src/arch/x86/mm/paging.rs | 25 +++++- .../src/arch/x86/mm/paging.rs | 90 +++++++++++++------ 4 files changed, 94 insertions(+), 29 deletions(-) diff --git a/litebox/src/mm/linux.rs b/litebox/src/mm/linux.rs index f33094971..4c41f7789 100644 --- a/litebox/src/mm/linux.rs +++ b/litebox/src/mm/linux.rs @@ -1055,4 +1055,6 @@ pub enum PageFaultError { AllocationFailed, #[error("given page is part of an already mapped huge page")] HugePage, + #[error("invalid frame address: {0:#x}")] + InvalidFrameAddress(u64), } diff --git a/litebox/src/platform/page_mgmt.rs b/litebox/src/platform/page_mgmt.rs index 7bd612225..42475c3ba 100644 --- a/litebox/src/platform/page_mgmt.rs +++ b/litebox/src/platform/page_mgmt.rs @@ -235,6 +235,8 @@ pub enum DeallocationError { Unaligned, #[error("provided range contains unallocated pages")] AlreadyUnallocated, + #[error("invalid frame address: {0:#x}")] + InvalidFrameAddress(u64), } /// Possible errors for [`PageManagementProvider::remap_pages`] @@ -267,6 +269,10 @@ pub enum PermissionUpdateError { Unaligned, #[error("provided range contains unallocated pages")] Unallocated, + #[error("operation encountered a huge page")] + HugePage, + #[error("invalid frame address: {0:#x}")] + InvalidFrameAddress(u64), } /// Possible errors for [`PageManagementProvider::try_allocate_cow_pages`] diff --git a/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs b/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs index 60ac10b7e..78cad9e5d 100644 --- a/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs +++ b/litebox_platform_linux_kernel/src/arch/x86/mm/paging.rs @@ -253,7 +253,10 @@ impl X64PageTable<'_, M, ALIGN> { offset: _, flags, } => { - // If it is changed to writable, we leave it to page fault handler (COW) + // If write access is requested on a read-only mapping, leave the PTE + // read-only for now and let the write-fault path enable writes lazily. + // Keep this split explicit so a future real COW implementation can hook + // into the same fault path. let change_to_write = new_flags.contains(PageTableFlags::WRITABLE) && !flags.contains(PageTableFlags::WRITABLE); let new_flags = if change_to_write { @@ -261,7 +264,7 @@ impl X64PageTable<'_, M, ALIGN> { } else { new_flags }; - if flags != new_flags { + if flags & Self::MPROTECT_PTE_MASK != new_flags { match unsafe { inner.update_flags(page, (flags & !Self::MPROTECT_PTE_MASK) | new_flags) } { @@ -327,8 +330,22 @@ impl PageTableImpl for X64PageTabl // probably set by other threads concurrently return Ok(()); } else { - // Copy-on-Write - todo!("COW"); + // This is the deferred write-enable path used by mprotect today, not a + // real COW implementation. Keep the write-fault handling separate so a + // future COW path can replace this flag update. + match unsafe { inner.update_flags(page, flags | PageTableFlags::WRITABLE) } + { + Ok(flush) => { + if FLUSH_TLB { + flush.flush(); + } + return Ok(()); + } + Err(FlagUpdateError::PageNotMapped) => unreachable!(), + Err(FlagUpdateError::ParentEntryHugePage) => { + return Err(PageFaultError::HugePage); + } + } } } diff --git a/litebox_platform_lvbs/src/arch/x86/mm/paging.rs b/litebox_platform_lvbs/src/arch/x86/mm/paging.rs index ffb1b0cb1..929ad5340 100644 --- a/litebox_platform_lvbs/src/arch/x86/mm/paging.rs +++ b/litebox_platform_lvbs/src/arch/x86/mm/paging.rs @@ -186,7 +186,9 @@ impl X64PageTable<'_, M, ALIGN> { unreachable!("we do not support huge pages"); } Err(X64UnmapError::InvalidFrameAddress(pa)) => { - todo!("Invalid frame address: {:#x}", pa); + return Err(page_mgmt::DeallocationError::InvalidFrameAddress( + pa.as_u64(), + )); } } } @@ -343,7 +345,10 @@ impl X64PageTable<'_, M, ALIGN> { offset: _, flags, } => { - // If it is changed to writable, we leave it to page fault handler (COW) + // If write access is requested on a read-only mapping, leave the PTE + // read-only for now and let the write-fault path enable writes lazily. + // Keep this split explicit so a future real COW implementation can hook + // into the same fault path. let change_to_write = new_flags.contains(PageTableFlags::WRITABLE) && !flags.contains(PageTableFlags::WRITABLE); let new_flags = if change_to_write { @@ -351,7 +356,7 @@ impl X64PageTable<'_, M, ALIGN> { } else { new_flags }; - if flags != new_flags { + if flags & Self::MPROTECT_PTE_MASK != new_flags { match unsafe { inner.update_flags(page, (flags & !Self::MPROTECT_PTE_MASK) | new_flags) } { @@ -359,7 +364,7 @@ impl X64PageTable<'_, M, ALIGN> { Err(e) => match e { FlagUpdateError::PageNotMapped => unreachable!(), FlagUpdateError::ParentEntryHugePage => { - todo!("return Err(ProtectError::ProtectHugePage);") + return Err(page_mgmt::PermissionUpdateError::HugePage); } }, } @@ -367,7 +372,9 @@ impl X64PageTable<'_, M, ALIGN> { } TranslateResult::NotMapped => {} TranslateResult::InvalidFrameAddress(pa) => { - todo!("Invalid frame address: {:#x}", pa); + return Err(page_mgmt::PermissionUpdateError::InvalidFrameAddress( + pa.as_u64(), + )); } } } @@ -431,39 +438,59 @@ impl X64PageTable<'_, M, ALIGN> { let page: Page = Page::containing_address(pa_to_va(target_frame.start_address())); + // When exec_ranges are provided, determine per-page flags based + // on whether the frame falls within an executable region. + let page_flags = if let Some(ranges) = exec_ranges { + let frame_addr = target_frame.start_address(); + let is_exec = ranges.iter().any(|r| r.contains(&frame_addr)); + if is_exec { + // W^X: if the page is executable, it should not be writable. + (flags & !PageTableFlags::NO_EXECUTE) & !PageTableFlags::WRITABLE + } else { + flags | PageTableFlags::NO_EXECUTE + } + } else { + flags + }; + match inner.translate(page.start_address()) { TranslateResult::Mapped { frame, offset: _, - flags: _, + flags: existing_flags, } => { assert!( target_frame.start_address() == frame.start_address(), "{page:?} is already mapped to {frame:?} instead of {target_frame:?}" ); + // Update flags if the existing mapping has stale permissions. + if existing_flags != page_flags { + match unsafe { inner.update_flags(page, page_flags) } { + Ok(_) => { + // Permission changes on an existing mapping must invalidate + // stale TLB entries on local and remote cores. + flush_tlb_range(page, 1); + } + Err(FlagUpdateError::PageNotMapped) => unreachable!(), + Err(FlagUpdateError::ParentEntryHugePage) => { + unreachable!("we do not support huge pages"); + } + } + } + continue; } TranslateResult::NotMapped => {} TranslateResult::InvalidFrameAddress(pa) => { - todo!("Invalid frame address: {:#x}", pa); + debug_assert!( + false, + "map_phys_frame_range: invalid frame address: {pa:#x}" + ); + return Err(MapToError::FrameAllocationFailed); } } - // When exec_ranges are provided, determine per-page flags based - // on whether the frame falls within an executable region. - let page_flags = if let Some(ranges) = exec_ranges { - let frame_addr = target_frame.start_address(); - let is_exec = ranges.iter().any(|r| r.contains(&frame_addr)); - if is_exec { - // W^X: if the page is executable, it should not be writable. - (flags & !PageTableFlags::NO_EXECUTE) & !PageTableFlags::WRITABLE - } else { - flags | PageTableFlags::NO_EXECUTE - } - } else { - flags - }; let table_flags = page_flags - PageTableFlags::NO_EXECUTE; // parent entries should not have NO_EXECUTE match unsafe { @@ -724,8 +751,20 @@ impl PageTableImpl for X64PageTabl // probably set by other threads concurrently return Ok(()); } else { - // Copy-on-Write - todo!("COW"); + // This is the deferred write-enable path used by mprotect today, not a + // real COW implementation. Keep the write-fault handling separate so a + // future COW path can replace this flag update. + match unsafe { inner.update_flags(page, flags | PageTableFlags::WRITABLE) } + { + Ok(flush) => { + flush.flush(); + return Ok(()); + } + Err(FlagUpdateError::PageNotMapped) => unreachable!(), + Err(FlagUpdateError::ParentEntryHugePage) => { + return Err(PageFaultError::HugePage); + } + } } } @@ -739,7 +778,8 @@ impl PageTableImpl for X64PageTabl TranslateResult::NotMapped => { let mut allocator = PageTableAllocator::::new(); // TODO: if it is file-backed, we need to read the page from file - let frame = PageTableAllocator::::allocate_frame(true).unwrap(); + let frame = PageTableAllocator::::allocate_frame(true) + .ok_or(PageFaultError::AllocationFailed)?; let table_flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE | PageTableFlags::USER_ACCESSIBLE; @@ -770,7 +810,7 @@ impl PageTableImpl for X64PageTabl } } TranslateResult::InvalidFrameAddress(pa) => { - todo!("Invalid frame address: {:#x}", pa); + return Err(PageFaultError::InvalidFrameAddress(pa.as_u64())); } } Ok(())