1 files changed, 356 insertions, 0 deletions

diff --git a/rsnix/src/paging.rs b/rsnix/src/paging.rs
new file mode 100644
index 0000000..766bde3
--- /dev/null
+++ b/rsnix/src/paging.rs
@@ -0,0 +1,356 @@
+// ---------------------------------------------------------------------------
+// Identity page-table helpers — 4-level paging (x86_64).
+//
+// The kernel keeps an identity-mapped 4-level tree (no higher half).  At
+// boot it runs on UEFI's tables; `adopt_tables()` rebuilds that map into
+// our own tables (UEFI maps its own table pages read-only, so the firmware
+// tree cannot be modified in place) and switches CR3.  After that the
+// helpers walk our tree to map / unmap / query individual 4 KB pages 1:1.
+// Page-table pages come from the physical page allocator.  Firmware may
+// have used 2 MB / 1 GB large pages for RAM, so the walk splits large pages
+// into 4 KB tables on demand.  Split tables stay allocated afterwards (no
+// refcounting) — fine for a kernel that owns all memory.
+// ---------------------------------------------------------------------------
+
+use crate::memory::{alloc_pages, free_pages, PAGE_SIZE};
+
+pub const PAGE_PRESENT: u64 = 1 << 0;
+pub const PAGE_WRITABLE: u64 = 1 << 1;
+pub const PAGE_USER: u64 = 1 << 2;
+pub const PAGE_PS: u64 = 1 << 7; // large page (2 MB PDE, 1 GB PDPTE)
+pub const PAGE_NX: u64 = 1 << 63;
+
+/// Physical address bits in a page-table entry.
+const PHYS_ADDR_MASK: u64 = 0x000F_FFFF_FFFF_F000;
+
+const PAGE_2M: usize = 1 << 21;
+const PAGE_1G: usize = 1 << 30;
+
+/// Current PML4 root (physical address, identity-mapped).
+pub fn cr3() -> usize {
+    let val: usize;
+    unsafe {
+        core::arch::asm!("mov {0}, cr3", out(reg) val, options(nostack, preserves_flags));
+    }
+    val
+}
+
+/// Full TLB flush by reloading CR3 (slow but bulletproof; map/unmap are
+/// rare operations in this kernel).
+fn flush_tlb_all() {
+    unsafe {
+        let cr3: usize;
+        core::arch::asm!("mov {0}, cr3", out(reg) cr3, options(nostack, preserves_flags));
+        core::arch::asm!("mov cr3, {0}", in(reg) cr3, options(nostack, preserves_flags));
+    }
+}
+
+/// Rebuild the identity map into our own page tables and switch CR3.
+///
+/// UEFI maps its own page-table pages read-only (R/W clear), so we cannot
+/// modify the firmware tree in place: map/unmap would #PF.  Instead, walk
+/// the current tables, replicate every present leaf (4 KB / 2 MB / 1 GB)
+/// into fresh tables from the page allocator, force R/W so the tree stays
+/// editable, then load the new PML4 into CR3.
+pub fn adopt_tables() -> bool {
+    unsafe {
+        let old_pml4 = cr3() as *const u64;
+        let Some(new_pml4) = alloc_pages(1) else {
+            return false;
+        };
+        core::ptr::write_bytes(new_pml4 as *mut u8, 0, PAGE_SIZE);
+        let new_pml4 = new_pml4 as *mut u64;
+
+        for l4 in 0..512 {
+            let l4e = old_pml4.add(l4).read_volatile();
+            if l4e & PAGE_PRESENT == 0 {
+                continue;
+            }
+            // Replicate this PDPT into a fresh one.
+            let old_pdpt = (l4e & PHYS_ADDR_MASK) as *const u64;
+            let Some(new_pdpt) = alloc_pages(1) else {
+                return false;
+            };
+            core::ptr::write_bytes(new_pdpt as *mut u8, 0, PAGE_SIZE);
+            let new_pdpt = new_pdpt as *mut u64;
+            *new_pml4.add(l4) =
+                (new_pdpt as u64) | PAGE_PRESENT | PAGE_WRITABLE | (l4e & PAGE_USER);
+
+            for l3 in 0..512 {
+                let pdpte = old_pdpt.add(l3).read_volatile();
+                if pdpte & PAGE_PRESENT == 0 {
+                    continue;
+                }
+                if pdpte & PAGE_PS != 0 {
+                    // 1 GB large page — copy the leaf, force writable.
+                    *new_pdpt.add(l3) = pdpte | PAGE_WRITABLE;
+                    continue;
+                }
+                let old_pd = (pdpte & PHYS_ADDR_MASK) as *const u64;
+                let Some(new_pd) = alloc_pages(1) else {
+                    return false;
+                };
+                core::ptr::write_bytes(new_pd as *mut u8, 0, PAGE_SIZE);
+                let new_pd = new_pd as *mut u64;
+                *new_pdpt.add(l3) =
+                    (new_pd as u64) | PAGE_PRESENT | PAGE_WRITABLE | (pdpte & PAGE_USER);
+
+                for l2 in 0..512 {
+                    let pde = old_pd.add(l2).read_volatile();
+                    if pde & PAGE_PRESENT == 0 {
+                        continue;
+                    }
+                    if pde & PAGE_PS != 0 {
+                        // 2 MB large page — copy the leaf, force writable.
+                        *new_pd.add(l2) = pde | PAGE_WRITABLE;
+                        continue;
+                    }
+                    let old_pt = (pde & PHYS_ADDR_MASK) as *const u64;
+                    let Some(new_pt) = alloc_pages(1) else {
+                        return false;
+                    };
+                    core::ptr::write_bytes(new_pt as *mut u8, 0, PAGE_SIZE);
+                    let new_pt = new_pt as *mut u64;
+                    *new_pd.add(l2) =
+                        (new_pt as u64) | PAGE_PRESENT | PAGE_WRITABLE | (pde & PAGE_USER);
+
+                    for l1 in 0..512 {
+                        let pte = old_pt.add(l1).read_volatile();
+                        if pte & PAGE_PRESENT != 0 {
+                            *new_pt.add(l1) = pte | PAGE_WRITABLE;
+                        }
+                    }
+                }
+            }
+        }
+
+        // Switch to the new tables (CR3 reload flushes the TLB).
+        let new_cr3 = new_pml4 as usize;
+        core::arch::asm!("mov cr3, {0}", in(reg) new_cr3, options(nostack, preserves_flags));
+        true
+    }
+}
+
+/// Ensure the PDPT entry points at a page-directory table, splitting a
+/// 1 GB large page at `base_1g` if present.  Returns the PD pointer.
+unsafe fn ensure_pd(pdpte: &mut u64, base_1g: usize) -> Option<*mut u64> {
+    unsafe {
+        if *pdpte & PAGE_PRESENT == 0 {
+            let page = alloc_pages(1)? as *mut u64;
+            core::ptr::write_bytes(page as *mut u8, 0, PAGE_SIZE);
+            *pdpte = (page as u64) | PAGE_PRESENT | PAGE_WRITABLE;
+            return Some(page);
+        }
+        if *pdpte & PAGE_PS != 0 {
+            // Split the 1 GB page into 512 × 2 MB pages.
+            let pd = alloc_pages(1)? as *mut u64;
+            core::ptr::write_bytes(pd as *mut u8, 0, PAGE_SIZE);
+            let inherit = *pdpte & (PAGE_USER | PAGE_NX);
+            for i in 0..512 {
+                let p = base_1g + i * PAGE_2M;
+                *pd.add(i) = (p as u64) | PAGE_PRESENT | PAGE_WRITABLE | PAGE_PS | inherit;
+            }
+            *pdpte = (pd as u64) | PAGE_PRESENT | PAGE_WRITABLE | (*pdpte & PAGE_USER);
+            return Some(pd);
+        }
+        Some(((*pdpte & PHYS_ADDR_MASK) as usize) as *mut u64)
+    }
+}
+
+/// Ensure the PD entry points at a page table, splitting a 2 MB large page
+/// at `base_2m` if present.  Returns the PT pointer.
+unsafe fn ensure_pt(pde: &mut u64, base_2m: usize) -> Option<*mut u64> {
+    unsafe {
+        if *pde & PAGE_PRESENT == 0 {
+            let page = alloc_pages(1)? as *mut u64;
+            core::ptr::write_bytes(page as *mut u8, 0, PAGE_SIZE);
+            *pde = (page as u64) | PAGE_PRESENT | PAGE_WRITABLE;
+            return Some(page);
+        }
+        if *pde & PAGE_PS != 0 {
+            // Split the 2 MB page into 512 × 4 KB pages.
+            let pt = alloc_pages(1)? as *mut u64;
+            core::ptr::write_bytes(pt as *mut u8, 0, PAGE_SIZE);
+            let inherit = *pde & (PAGE_USER | PAGE_NX);
+            for i in 0..512 {
+                let p = base_2m + i * PAGE_SIZE;
+                *pt.add(i) = (p as u64) | PAGE_PRESENT | PAGE_WRITABLE | inherit;
+            }
+            *pde = (pt as u64) | PAGE_PRESENT | PAGE_WRITABLE | (*pde & PAGE_USER);
+            return Some(pt);
+        }
+        Some(((*pde & PHYS_ADDR_MASK) as usize) as *mut u64)
+    }
+}
+
+/// Identity-map the physical page at `phys` (must be 4 KB aligned).
+/// `flags` adds to PRESENT (e.g. PAGE_WRITABLE, PAGE_USER).
+pub fn map_page(phys: usize, flags: u64) -> bool {
+    if !phys.is_multiple_of(PAGE_SIZE) {
+        return false;
+    }
+    unsafe {
+        let pml4 = cr3() as *mut u64;
+        let l4 = (phys >> 39) & 0x1FF;
+        let l3 = (phys >> 30) & 0x1FF;
+        let l2 = (phys >> 21) & 0x1FF;
+        let l1 = (phys >> 12) & 0x1FF;
+
+        let pml4e = &mut *pml4.add(l4);
+        if *pml4e & PAGE_PRESENT == 0 {
+            let Some(page) = alloc_pages(1) else {
+                return false;
+            };
+            core::ptr::write_bytes(page as *mut u8, 0, PAGE_SIZE);
+            *pml4e = (page as u64) | PAGE_PRESENT | PAGE_WRITABLE;
+        }
+        // PML4 entries are never large pages.
+
+        let base_1g = phys & !(PAGE_1G - 1);
+        let pdpte = &mut *(((*pml4e & PHYS_ADDR_MASK) as usize) as *mut u64).add(l3);
+        let Some(pd) = ensure_pd(pdpte, base_1g) else {
+            return false;
+        };
+
+        let base_2m = phys & !(PAGE_2M - 1);
+        let pde = &mut *pd.add(l2);
+        let Some(pt) = ensure_pt(pde, base_2m) else {
+            return false;
+        };
+
+        let pte = &mut *pt.add(l1);
+        *pte = (phys as u64) | flags | PAGE_PRESENT;
+        flush_tlb_all();
+        true
+    }
+}
+
+/// Identity-map `count` consecutive pages starting at `phys`.
+pub fn map_range(phys: usize, count: usize, flags: u64) -> bool {
+    for i in 0..count {
+        if !map_page(phys + i * PAGE_SIZE, flags) {
+            return false;
+        }
+    }
+    true
+}
+
+/// Unmap the 4 KB page at `phys`.  Splits large pages as needed so only
+/// this one page is removed.  Returns false if it wasn't mapped.
+pub fn unmap_page(phys: usize) -> bool {
+    if !phys.is_multiple_of(PAGE_SIZE) {
+        return false;
+    }
+    unsafe {
+        let pml4 = cr3() as *mut u64;
+        let l4 = (phys >> 39) & 0x1FF;
+        let l3 = (phys >> 30) & 0x1FF;
+        let l2 = (phys >> 21) & 0x1FF;
+        let l1 = (phys >> 12) & 0x1FF;
+
+        let pml4e = &mut *pml4.add(l4);
+        if *pml4e & PAGE_PRESENT == 0 {
+            return false;
+        }
+        let base_1g = phys & !(PAGE_1G - 1);
+        let pdpte = &mut *(((*pml4e & PHYS_ADDR_MASK) as usize) as *mut u64).add(l3);
+        if *pdpte & PAGE_PRESENT == 0 {
+            return false;
+        }
+        let Some(pd) = ensure_pd(pdpte, base_1g) else {
+            return false;
+        };
+        let base_2m = phys & !(PAGE_2M - 1);
+        let pde = &mut *pd.add(l2);
+        if *pde & PAGE_PRESENT == 0 {
+            return false;
+        }
+        let Some(pt) = ensure_pt(pde, base_2m) else {
+            return false;
+        };
+        let pte = &mut *pt.add(l1);
+        if *pte & PAGE_PRESENT == 0 {
+            return false;
+        }
+        *pte = 0;
+        flush_tlb_all();
+        true
+    }
+}
+
+/// Is the page at `phys` mapped (present at any level, incl. large pages)?
+pub fn is_mapped(phys: usize) -> bool {
+    unsafe {
+        let pml4 = cr3() as *const u64;
+        let l4 = (phys >> 39) & 0x1FF;
+        let l3 = (phys >> 30) & 0x1FF;
+        let l2 = (phys >> 21) & 0x1FF;
+        let l1 = (phys >> 12) & 0x1FF;
+
+        let pml4e = *pml4.add(l4);
+        if pml4e & PAGE_PRESENT == 0 {
+            return false;
+        }
+        let pdpte = *((pml4e & PHYS_ADDR_MASK) as *const u64).add(l3);
+        if pdpte & PAGE_PRESENT == 0 {
+            return false;
+        }
+        if pdpte & PAGE_PS != 0 {
+            return true; // 1 GB page covers phys
+        }
+        let pde = *((pdpte & PHYS_ADDR_MASK) as *const u64).add(l2);
+        if pde & PAGE_PRESENT == 0 {
+            return false;
+        }
+        if pde & PAGE_PS != 0 {
+            return true; // 2 MB page covers phys
+        }
+        let pte = *((pde & PHYS_ADDR_MASK) as *const u64).add(l1);
+        pte & PAGE_PRESENT != 0
+    }
+}
+
+/// Smoke test: allocate two physical pages, unmap one, verify the mapping
+/// state, remap it, verify data still works, and free both.
+pub fn self_test() -> bool {
+    let Some(p1) = alloc_pages(1) else {
+        return false;
+    };
+    let Some(p2) = alloc_pages(1) else {
+        return false;
+    };
+
+    // Both pages must already be identity-mapped (they're RAM).
+    if !is_mapped(p1) || !is_mapped(p2) {
+        return false;
+    }
+
+    // Unmap p1 — only p1 should disappear.
+    if !unmap_page(p1) {
+        return false;
+    }
+    if is_mapped(p1) {
+        return false;
+    }
+    if !is_mapped(p2) {
+        return false;
+    }
+
+    // Remap and verify we can write through the mapping.
+    if !map_page(p1, PAGE_WRITABLE) {
+        return false;
+    }
+    if !is_mapped(p1) {
+        return false;
+    }
+    unsafe {
+        core::ptr::write_volatile(p1 as *mut u8, 0xCD);
+        if core::ptr::read_volatile(p1 as *const u8) != 0xCD {
+            return false;
+        }
+    }
+
+    free_pages(p1, 1);
+    free_pages(p2, 1);
+    true
+}