x86_64: support nested pages splitting

Author: zyuiop

src/arch/x86_64/mm/paging.rs

@@ -8,15 +8,15 @@ pub use x86_64::structures::idt::InterruptStackFrame as ExceptionStackFrame;
 use x86_64::structures::idt::PageFaultErrorCode;
 pub use x86_64::structures::paging::PageTableFlags as PageTableEntryFlags;
 use x86_64::structures::paging::frame::PhysFrameRange;
-use x86_64::structures::paging::mapper::{MapToError, MappedFrame, TranslateResult, UnmapError};
+use x86_64::structures::paging::mapper::{MapToError, MappedFrame, MapperFlush, TranslateResult, UnmapError};
 use x86_64::structures::paging::page::PageRange;
 use x86_64::structures::paging::{
-	FrameAllocator, Mapper, OffsetPageTable, Page, PageTable, PhysFrame, Size4KiB, Translate,
+	FrameAllocator, Mapper, OffsetPageTable, Page, PageTable, PhysFrame, Size1GiB, Size2MiB, Size4KiB, Translate,
 };
 
 use crate::arch::kernel::processor;
 use crate::arch::mm::{PhysAddr, VirtAddr};
-use crate::mm::{FrameAlloc, PageRangeAllocator};
+use crate::mm::{FrameAlloc, PageAlloc, PageRangeAllocator};
 use crate::scheduler;
 
 unsafe impl FrameAllocator<Size4KiB> for FrameAlloc {
@@ -191,15 +191,30 @@ pub fn map<S>(
 	where
 		M: Mapper<S>,
 		S: PageSize + fmt::Debug,
+		for<'a> OffsetPageTable<'a>: Mapper<S>,
 	{
 		let mut unmapped = false;
 		for (page, frame) in pages.zip(frames) {
 			// TODO: Require explicit unmaps
-			let unmap = mapper.unmap(page);
-			if let Ok((_frame, flush)) = unmap {
-				unmapped = true;
-				flush.flush();
-				debug!("Had to unmap page {page:?} before mapping.");
+			let unmap_result = mapper.unmap(page);
+			match unmap_result {
+				Ok((_, flush)) => {
+					unmapped = true;
+					flush.flush();
+					debug!("Had to unmap page {page:?} before mapping.");
+				}
+				Err(UnmapError::PageNotMapped) => {
+					// Expected case
+				}
+				Err(UnmapError::ParentEntryHugePage) => {
+					// Must unmap completely
+					unmapped = true;
+					unmap::<S>(page.start_address().into(), 1);
+					debug!("Had to unmap page {page:?} before mapping.");
+				}
+				Err(other) => {
+					panic!("Failed to unmap page during mapping: {other:?}");
+				}
 			}
 			let map = unsafe { mapper.map_to(page, frame, flags, &mut FrameAlloc) };
 			match map {
@@ -268,6 +283,107 @@ where
 	}
 }
 
+/// Prepare a new page table
+fn split_page<S: PageSize>(page: Page<S>) {
+	assert_ne!(S::SIZE, Size4KiB::SIZE, "cannot split small page");
+	let is_huge_page = S::SIZE == Size1GiB::SIZE;
+
+	// Allocate new page table, map it temporarily
+	let pt_frame: PhysFrame<Size4KiB> = FrameAlloc.allocate_frame().unwrap();
+	let pt_page = PageAlloc::allocate(PageLayout::from_size(Size4KiB::SIZE as usize).unwrap()).unwrap();
+	let pt_page = VirtAddr::new(pt_page.start() as u64);
+
+	let flags = PageTableEntryFlags::WRITABLE | PageTableEntryFlags::NO_EXECUTE | PageTableEntryFlags::PRESENT;
+	map::<Size4KiB>(pt_page, pt_frame.start_address().into(), 1, flags);
+
+	// Fill it with entries
+	let mut table_explorer = unsafe { identity_mapped_page_table() };
+	let (start_addr, flags) = match table_explorer.translate(page.start_address()) {
+		TranslateResult::Mapped { frame, flags, .. } => {
+			let start_addr = match frame {
+				MappedFrame::Size2MiB(frame) if S::SIZE == Size2MiB::SIZE => {
+					frame.start_address()
+				}
+				MappedFrame::Size1GiB(frame) if S::SIZE == Size1GiB::SIZE => {
+					frame.start_address()
+				}
+				MappedFrame::Size1GiB(_) if S::SIZE == Size2MiB::SIZE => {
+					// We were trying to split a large page, and we got a huge page -- we should split it
+					// Split the parent page first, then retry
+					split_page(Page::<Size1GiB>::containing_address(page.start_address()));
+					return split_page(page);
+				}
+				other => {
+					panic!("Unexpected frame mapping {other:?} when trying to split {page:?}")
+				}
+			};
+
+			(start_addr, flags)
+		}
+		TranslateResult::NotMapped => {
+			panic!("Tried to split a page that is not mapped!")
+		}
+		TranslateResult::InvalidFrameAddress(addr) => {
+			panic!("Tried to split a page that maps to invalid physical address {addr:x?}")
+		}
+	};
+
+	let flags = if is_huge_page {
+		flags // keep the HUGE flag
+	} else {
+		// Remove the large page flag, because we map to 4KiB frames
+		flags.difference(PageTableEntryFlags::HUGE_PAGE)
+	};
+
+	// Build the page table!
+	let pt = pt_page.as_mut_ptr::<PageTable>();
+	let pt = unsafe {
+		let pt = pt.as_mut().unwrap();
+		pt.zero();
+		pt
+	};
+
+	let child_page_size = if is_huge_page {
+		Size2MiB::SIZE
+	} else {
+		Size4KiB::SIZE
+	};
+
+	for (offset, entry) in pt.iter_mut().enumerate() {
+		let offset = (offset as u64) * child_page_size;
+
+		entry.set_addr(
+			start_addr + offset,
+			flags,
+		);
+	}
+
+	// We can now replace the entry in the page table
+	let offset = table_explorer.phys_offset();
+	let p4 = table_explorer.level_4_table_mut();
+	let p3 = &mut p4[page.p4_index()];
+	let p3 = offset + p3.addr().as_u64();
+	let p3 = unsafe { &mut *p3.as_mut_ptr::<PageTable>() };
+
+	if is_huge_page {
+		let flags = p3[page.p3_index()].flags() - PageTableEntryFlags::HUGE_PAGE;
+		p3[page.p3_index()].set_frame(pt_frame, flags);
+
+	} else {
+		let p2 = &mut p3[page.p3_index()];
+		let p2 = offset + p2.addr().as_u64();
+		let p2 = unsafe { &mut *p2.as_mut_ptr::<PageTable>() };
+
+		let flags = p2[page.start_address().p2_index()].flags() - PageTableEntryFlags::HUGE_PAGE;
+		p2[page.start_address().p2_index()].set_frame(pt_frame, flags);
+	}
+
+	// Unmap temporary pt mapping
+	unmap::<Size4KiB>(pt_page, 1);
+
+	MapperFlush::new(page).flush();
+}
+
 pub fn unmap<S>(virtual_address: VirtAddr, count: usize)
 where
 	S: PageSize + fmt::Debug,
@@ -279,7 +395,10 @@ where
 	let last_page = first_page + count as u64;
 	let range = Page::range(first_page, last_page);
 
-	for page in range {
+	fn unmap_page<S>(page: Page<S>)
+	where
+		S: PageSize + fmt::Debug,
+		for<'a> OffsetPageTable<'a>: Mapper<S>, {
 		let unmap_result = unsafe { identity_mapped_page_table() }.unmap(page);
 		match unmap_result {
 			Ok((_frame, flush)) => flush.flush(),
@@ -288,9 +407,23 @@ where
 			Err(UnmapError::PageNotMapped) => {
 				debug!("Tried to unmap {page:?}, which was not mapped.");
 			}
+			Err(UnmapError::ParentEntryHugePage) if S::SIZE == Size4KiB::SIZE => {
+				// Prepare new page and retry
+				split_page(Page::<Size2MiB>::containing_address(page.start_address()));
+				unmap_page(page);
+			}
+			Err(UnmapError::ParentEntryHugePage) if S::SIZE == Size2MiB::SIZE => {
+				// Prepare new page and retry
+				split_page(Page::<Size1GiB>::containing_address(page.start_address()));
+				unmap_page(page);
+			}
 			Err(err) => panic!("{err:?}"),
 		}
 	}
+
+	for page in range {
+		unmap_page(page);
+	}
 }
 
 #[cfg(not(feature = "common-os"))]