feat(mm): introduce MappedPageBox

Author: RajeshRk18

src/arch/aarch64/mm/paging.rs

@@ -79,7 +79,6 @@ bitflags! {
 }
 
 impl PageTableEntryFlags {
-	#[expect(dead_code)]
 	pub fn present(&mut self) -> &mut Self {
 		self.insert(PageTableEntryFlags::PRESENT);
 		self
@@ -102,7 +101,6 @@ impl PageTableEntryFlags {
 		self
 	}
 
-	#[expect(dead_code)]
 	pub fn read_only(&mut self) -> &mut Self {
 		self.insert(PageTableEntryFlags::READ_ONLY);
 		self

src/arch/x86_64/kernel/apic.rs

@@ -25,7 +25,7 @@ use crate::arch::mm::paging::{
 	BasePageSize, PageSize, PageTableEntryFlags, PageTableEntryFlagsExt,
 };
 use crate::arch::swapgs;
-use crate::mm::PageBox;
+use crate::mm::{MappedPageBox, PageBox};
 use crate::scheduler::CoreId;
 use crate::{arch, scheduler};
 
@@ -420,18 +420,13 @@ fn detect_from_mp() -> Result<PhysAddr, ()> {
 		info!("Virtual-Wire mode implemented");
 	}
 
-	let layout = PageLayout::from_size(BasePageSize::SIZE as usize).unwrap();
-	let page_range = PageBox::new(layout).unwrap();
-	let virtual_address = VirtAddr::from(page_range.start());
-
 	let mut flags = PageTableEntryFlags::empty();
 	flags.normal().writable();
-	paging::map::<BasePageSize>(
-		virtual_address,
-		PhysAddr::from((mp_float.mp_config as usize).align_down(BasePageSize::SIZE as usize)),
-		1,
-		flags,
-	);
+	let layout = PageLayout::from_size(BasePageSize::SIZE as usize).unwrap();
+	let phys_addr =
+		PhysAddr::from((mp_float.mp_config as usize).align_down(BasePageSize::SIZE as usize));
+	let page_range = unsafe { MappedPageBox::map_phys(phys_addr, layout, flags).unwrap() };
+	let virtual_address = VirtAddr::from(page_range.start());
 
 	let mut addr: usize =
 		(virtual_address | (u64::from(mp_float.mp_config) & (BasePageSize::SIZE - 1))) as usize;

src/arch/x86_64/mm/mod.rs

@@ -16,28 +16,26 @@ pub fn create_new_root_page_table() -> usize {
 	use free_list::PageLayout;
 	use x86_64::registers::control::Cr3;
 
-	use crate::mm::PageBox;
+	use crate::mm::MappedPageBox;
 
 	let layout = PageLayout::from_size(BasePageSize::SIZE as usize).unwrap();
 	let frame_range = FrameAlloc::allocate(layout).unwrap();
 	let physaddr = PhysAddr::from(frame_range.start());
 
-	let layout = PageLayout::from_size(2 * BasePageSize::SIZE as usize).unwrap();
-	let page_range = PageBox::new(layout).unwrap();
-	let virtaddr = VirtAddr::from(page_range.start());
 	let mut flags = PageTableEntryFlags::empty();
 	flags.normal().writable();
 
 	let entry: u64 = unsafe {
 		let (frame, _flags) = Cr3::read();
-		paging::map::<BasePageSize>(virtaddr, frame.start_address().into(), 1, flags);
-		let entry: &u64 = &*virtaddr.as_ptr();
+		let page_range =
+			MappedPageBox::map_phys(frame.start_address().into(), layout, flags).unwrap();
+		let entry: &u64 = &*VirtAddr::from(page_range.start()).as_ptr();
 
 		*entry
 	};
 
-	let slice_addr = virtaddr + BasePageSize::SIZE;
-	paging::map::<BasePageSize>(slice_addr, physaddr, 1, flags);
+	let page_range = unsafe { MappedPageBox::map_phys(physaddr, layout, flags).unwrap() };
+	let slice_addr = VirtAddr::from(page_range.start());
 
 	unsafe {
 		let pml4 = slice::from_raw_parts_mut(slice_addr.as_mut_ptr(), 512);
@@ -53,8 +51,6 @@ pub fn create_new_root_page_table() -> usize {
 		pml4[511] = physaddr.as_u64() + 0x3; // PG_PRESENT | PG_RW
 	};
 
-	paging::unmap::<BasePageSize>(virtaddr, 2);
-
 	physaddr.as_usize()
 }
 

src/mm/mapped_page_box.rs

@@ -0,0 +1,79 @@
+use core::alloc::AllocError;
+use core::ops::Deref;
+
+use free_list::{PageLayout, PageRange};
+use memory_addresses::{PhysAddr, VirtAddr};
+
+use crate::arch::mm::paging::{self, BasePageSize, PageSize, PageTableEntryFlags};
+use crate::mm::{FrameBox, PageBox};
+
+/// A range of pages that is mapped for as long as this box is alive.
+pub struct MappedPageBox {
+	frames: Option<FrameBox>,
+	pages: PageBox,
+}
+
+impl MappedPageBox {
+	/// Allocates the pages and frames described by `layout` and maps them.
+	pub fn new(layout: PageLayout, flags: PageTableEntryFlags) -> Result<Self, AllocError> {
+		let frames = FrameBox::new(layout)?;
+		let pages = PageBox::new(layout)?;
+		map(&pages, PhysAddr::from(frames.start()), flags);
+		Ok(Self {
+			frames: Some(frames),
+			pages,
+		})
+	}
+
+	/// Allocates the pages described by `layout` and maps them to `phys_addr`.
+	///
+	/// # Safety
+	///
+	/// - The frames at `phys_addr` must not be deallocated while the returned box is alive.
+	pub unsafe fn map_phys(
+		phys_addr: PhysAddr,
+		layout: PageLayout,
+		flags: PageTableEntryFlags,
+	) -> Result<Self, AllocError> {
+		let pages = PageBox::new(layout)?;
+		map(&pages, phys_addr, flags);
+		Ok(Self {
+			frames: None,
+			pages,
+		})
+	}
+}
+
+impl Drop for MappedPageBox {
+	fn drop(&mut self) {
+		paging::unmap::<BasePageSize>(VirtAddr::from(self.pages.start()), page_count(&self.pages));
+
+		// `paging::unmap` only flushes the TLB of this CPU. Request the flush on
+		// all other CPUs before the frames can be reused.
+		#[cfg(all(target_arch = "x86_64", feature = "smp"))]
+		crate::arch::kernel::apic::ipi_tlb_flush();
+
+		drop(self.frames.take());
+	}
+}
+
+impl Deref for MappedPageBox {
+	type Target = PageRange;
+
+	fn deref(&self) -> &Self::Target {
+		&self.pages
+	}
+}
+
+fn page_count(pages: &PageBox) -> usize {
+	pages.len().get() / BasePageSize::SIZE as usize
+}
+
+fn map(pages: &PageBox, phys_addr: PhysAddr, flags: PageTableEntryFlags) {
+	paging::map::<BasePageSize>(
+		VirtAddr::from(pages.start()),
+		phys_addr,
+		page_count(pages),
+		flags,
+	);
+}

src/mm/mod.rs

@@ -41,6 +41,7 @@
 //! ```
 
 pub(crate) mod device_alloc;
+mod mapped_page_box;
 mod page_range_alloc;
 mod physicalmem;
 mod virtualmem;
@@ -57,6 +58,7 @@ use talc::TalcLock;
 #[cfg(target_os = "none")]
 use talc::source::Manual;
 
+pub use self::mapped_page_box::MappedPageBox;
 pub use self::page_range_alloc::{PageRangeAllocator, PageRangeBox};
 pub use self::physicalmem::{FrameAlloc, FrameBox};
 pub use self::virtualmem::{PageAlloc, PageBox};