rust-kernel

//! Paging: page tables, address translation and a physical frame allocator.
//!
//! x86_64 uses 4-level page tables (L4 -> L3 -> L2 -> L1 -> 4 KiB page).
//! CR3 holds the *physical* address of the active level-4 table. But the CPU
//! (and our code) only uses *virtual* addresses, so to read or edit a page
//! table we need it mapped somewhere. The bootloader's `map_physical_memory`
//! feature maps all of physical RAM at `physical_memory_offset`, so
//!     virtual address = physical_memory_offset + physical address.
//! `OffsetPageTable` from the x86_64 crate implements exactly that.
//!
//! Address spaces (for user processes): every process gets its own level-4
//! table. Its entries are copied from the kernel's level-4 table, so the
//! kernel (code, heap, stacks, the physical-memory window) is mapped
//! identically everywhere -- just *without* the USER_ACCESSIBLE flag, so ring
//! 3 code can't touch it. User code and stack go into L4 slot 1
//! (`USER_START..USER_END`), which the kernel doesn't use. Switching CR3 then
//! swaps only the user part. (Linux does the same with the upper half; our
//! bootloader put the kernel in the *lower* half, so we use a free slot.)

use bootloader::bootinfo::{MemoryMap, MemoryRegionType};
use core::sync::atomic::{AtomicU64, Ordering};
use spin::Mutex;
use x86_64::instructions::interrupts::without_interrupts;
use x86_64::{
    PhysAddr, VirtAddr,
    registers::control::Cr3,
    structures::paging::{
        FrameAllocator, Mapper, OffsetPageTable, Page, PageTable, PageTableFlags, PhysFrame,
        Size4KiB, Translate, mapper::TranslateResult,
    },
};

/// User space: level-4 slot 1, i.e. 512 GiB starting at 512 GiB.
pub const USER_START: u64 = 0x0000_0080_0000_0000;
pub const USER_END: u64 = 0x0000_0100_0000_0000; // exclusive
const USER_L4_INDEX: usize = 1;

/// Where all physical memory is mapped (set by `init`).
static PHYS_OFFSET: AtomicU64 = AtomicU64::new(0);
/// Physical address of the kernel's (boot) level-4 table (set by `init`).
static KERNEL_L4: AtomicU64 = AtomicU64::new(0);
/// The frame allocator, once `init_frame_allocator` hands it over.
static FRAME_ALLOCATOR: Mutex<Option<BootInfoFrameAllocator>> = Mutex::new(None);

/// Create an `OffsetPageTable` for the currently active page tables.
///
/// Unsafe: the caller must guarantee that all physical memory really is
/// mapped at `physical_memory_offset`, and must call this only once (two
/// `&mut` to the same page table would be undefined behaviour).
pub unsafe fn init(physical_memory_offset: VirtAddr) -> OffsetPageTable<'static> {
    PHYS_OFFSET.store(physical_memory_offset.as_u64(), Ordering::Relaxed);
    KERNEL_L4.store(Cr3::read().0.start_address().as_u64(), Ordering::Relaxed);
    unsafe {
        let level_4_table = active_level_4_table(physical_memory_offset);
        OffsetPageTable::new(level_4_table, physical_memory_offset)
    }
}

/// Find the active level-4 table via CR3 and return a reference to it.
unsafe fn active_level_4_table(physical_memory_offset: VirtAddr) -> &'static mut PageTable {
    let (level_4_table_frame, _flags) = Cr3::read();
    let phys = level_4_table_frame.start_address();
    let virt = physical_memory_offset + phys.as_u64();
    let page_table_ptr: *mut PageTable = virt.as_mut_ptr();
    unsafe { &mut *page_table_ptr }
}

/// Demo: map the virtual page containing `page` to the VGA text buffer's
/// physical frame (0xb8000). Writing through that page then shows on screen.
pub fn create_example_mapping(
    page: Page,
    mapper: &mut OffsetPageTable,
    frame_allocator: &mut impl FrameAllocator<Size4KiB>,
) {
    let frame = PhysFrame::containing_address(PhysAddr::new(0xb8000));
    let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;

    // Unsafe because mapping the same frame twice could create aliasing
    // `&mut`s. Fine here: it's only the VGA buffer, for a demo.
    // `map_to` may need new page tables (L3/L2/L1) -> takes a frame allocator.
    let map_to_result = unsafe { mapper.map_to(page, frame, flags, frame_allocator) };
    // `flush` removes the stale entry from the TLB (the CPU's mapping cache).
    map_to_result.expect("map_to failed").flush();
}

/// Hands out unused physical 4 KiB frames, taken from the memory map the
/// bootloader got from the BIOS (E820). Very simple: it never frees, and
/// walks the map from the start on every call (fine for now).
pub struct BootInfoFrameAllocator {
    memory_map: &'static MemoryMap,
    next: usize, // how many frames we've handed out so far
}

impl BootInfoFrameAllocator {
    /// Unsafe: the caller must guarantee that every frame marked `Usable` in
    /// the memory map really is unused.
    pub unsafe fn init(memory_map: &'static MemoryMap) -> Self {
        BootInfoFrameAllocator { memory_map, next: 0 }
    }

    /// All usable frames, in order.
    fn usable_frames(&self) -> impl Iterator<Item = PhysFrame> {
        self.memory_map
            .iter()
            .filter(|r| r.region_type == MemoryRegionType::Usable) // only free RAM
            .map(|r| r.range.start_addr()..r.range.end_addr()) // -> address ranges
            .flat_map(|r| r.step_by(4096)) // -> each frame's start address
            .map(|addr| PhysFrame::containing_address(PhysAddr::new(addr)))
    }
}

unsafe impl FrameAllocator<Size4KiB> for BootInfoFrameAllocator {
    fn allocate_frame(&mut self) -> Option<PhysFrame> {
        let frame = self.usable_frames().nth(self.next);
        self.next += 1;
        frame
    }
}

/// Total usable RAM in bytes, according to the memory map (for printing).
pub fn usable_memory_bytes(memory_map: &MemoryMap) -> u64 {
    memory_map
        .iter()
        .filter(|r| r.region_type == MemoryRegionType::Usable)
        .map(|r| r.range.end_addr() - r.range.start_addr())
        .sum()
}

// ---------------------------------------------------------------------------
// Global frame allocator + per-process address spaces
// ---------------------------------------------------------------------------

/// Hand the frame allocator over to the kernel (after the heap is set up),
/// so processes can get page tables and pages later.
pub fn init_frame_allocator(frame_allocator: BootInfoFrameAllocator) {
    without_interrupts(|| *FRAME_ALLOCATOR.lock() = Some(frame_allocator));
    init_kernel_stack_area(); // must happen before any process exists
}

/// Frames given back with `free_frame`, as a linked list threaded through
/// the free frames themselves: the first 8 bytes of each free frame hold the
/// physical address of the next one (0 = end). No heap memory needed.
struct FreeList {
    head: u64,
    len: usize,
}
static FREE_FRAMES: Mutex<FreeList> = Mutex::new(FreeList { head: 0, len: 0 });
/// Frames handed out by the boot allocator so far (for statistics).
static BOOT_FRAMES_USED: AtomicU64 = AtomicU64::new(0);

/// One physical frame (contents undefined), or None if out of memory / not
/// initialized. Recycled frames first, then fresh ones from the memory map.
///
/// The frame locks are only ever taken with interrupts off (see the locking
/// rule in thread.rs): the scheduler and interrupt handlers free frames too.
pub fn alloc_frame() -> Option<PhysFrame> {
    without_interrupts(|| {
        {
            let mut free = FREE_FRAMES.lock();
            if free.head != 0 {
                let frame = PhysFrame::containing_address(PhysAddr::new(free.head));
                free.head = unsafe { *phys_to_virt(frame.start_address()).as_ptr::<u64>() };
                free.len -= 1;
                return Some(frame);
            }
        }
        fresh_frame()
    })
}

/// A never-used frame straight from the boot allocator.
fn fresh_frame() -> Option<PhysFrame> {
    let frame = without_interrupts(|| FRAME_ALLOCATOR.lock().as_mut()?.allocate_frame())?;
    BOOT_FRAMES_USED.fetch_add(1, Ordering::Relaxed);
    Some(frame)
}

/// Give a frame back. The caller promises nothing uses it any more (no page
/// table maps it, no device writes to it).
pub fn free_frame(frame: PhysFrame) {
    without_interrupts(|| {
        let mut free = FREE_FRAMES.lock();
        unsafe { *phys_to_virt(frame.start_address()).as_mut_ptr::<u64>() = free.head };
        free.head = frame.start_address().as_u64();
        free.len += 1;
    })
}

/// How many frames are in use right now (handed out and not freed).
pub fn frames_in_use() -> usize {
    without_interrupts(|| BOOT_FRAMES_USED.load(Ordering::Relaxed) as usize - FREE_FRAMES.lock().len)
}

/// `FrameAllocator` adapter around the global allocator (for `map_to`, which
/// may need frames for new page tables).
struct GlobalFrameAllocator;

unsafe impl FrameAllocator<Size4KiB> for GlobalFrameAllocator {
    fn allocate_frame(&mut self) -> Option<PhysFrame> {
        alloc_frame()
    }
}

/// Virtual address through which the kernel can reach physical address `phys`.
pub fn phys_to_virt(phys: PhysAddr) -> VirtAddr {
    VirtAddr::new(PHYS_OFFSET.load(Ordering::Relaxed) + phys.as_u64())
}

/// The kernel's level-4 table frame (what kernel threads run with).
pub fn kernel_l4_frame() -> PhysFrame {
    match KERNEL_L4.load(Ordering::Relaxed) {
        0 => Cr3::read().0, // `init` not called yet: we're still on the boot tables
        addr => PhysFrame::containing_address(PhysAddr::new(addr)),
    }
}

/// Does the CPU honour the NO_EXECUTE page flag? (EFER.NXE, which our
/// bootloader turns on. With it off, setting the bit would be an error.)
pub fn nx_enabled() -> bool {
    use x86_64::registers::model_specific::{Efer, EferFlags};
    Efer::read().contains(EferFlags::NO_EXECUTE_ENABLE)
}

/// `count` physically *contiguous*, zeroed frames (device DMA buffers like
/// virtio queues must be contiguous in physical memory). Our allocator hands
/// out frames in address order, so consecutive frames are usually adjacent;
/// if a run gets broken (end of a memory region) we start over and simply
/// leak the frames we skipped. Fine for a few boot-time allocations.
pub fn alloc_contiguous_frames(count: usize) -> Option<PhysAddr> {
    // Fresh frames only: recycled ones come back in no particular order.
    let mut start = fresh_frame()?.start_address();
    let mut got = 1;
    while got < count {
        let frame = fresh_frame()?.start_address();
        if frame == start + (got as u64) * 4096 {
            got += 1;
        } else {
            start = frame; // not adjacent: restart the run here
            got = 1;
        }
    }
    let ptr: *mut u8 = phys_to_virt(start).as_mut_ptr();
    unsafe { core::ptr::write_bytes(ptr, 0, count * 4096) };
    Some(start)
}

/// A frame filled with zeros (new page tables and user pages must not leak
/// old data).
fn alloc_zeroed_frame() -> Option<PhysFrame> {
    let frame = alloc_frame()?;
    let ptr: *mut u8 = phys_to_virt(frame.start_address()).as_mut_ptr();
    unsafe { core::ptr::write_bytes(ptr, 0, 4096) };
    Some(frame)
}

/// Page-table view of the address space whose level-4 table is `l4`.
///
/// Unsafe: the caller must not create two of these for the same table at the
/// same time (both would hand out `&mut` to the same memory).
unsafe fn mapper_for(l4: PhysFrame) -> OffsetPageTable<'static> {
    let table: *mut PageTable = phys_to_virt(l4.start_address()).as_mut_ptr();
    let offset = VirtAddr::new(PHYS_OFFSET.load(Ordering::Relaxed));
    unsafe { OffsetPageTable::new(&mut *table, offset) }
}

/// Create a new address space: a fresh level-4 table that shares all kernel
/// mappings (entries copied from the kernel's table) and has an empty user part.
pub fn new_address_space() -> Option<PhysFrame> {
    let frame = alloc_zeroed_frame()?;
    unsafe {
        let new: &mut PageTable = &mut *phys_to_virt(frame.start_address()).as_mut_ptr();
        let kernel: &PageTable = &*phys_to_virt(kernel_l4_frame().start_address()).as_ptr();
        // The kernel must not use our user slot, or user mappings would end up
        // in page tables shared with the kernel (and every other process).
        assert!(kernel[USER_L4_INDEX].is_unused(), "kernel uses the user L4 slot");
        for i in 0..512 {
            new[i] = kernel[i].clone();
        }
    }
    Some(frame)
}

/// Free everything a process's address space owns: every user page's frame,
/// the page tables of the user slot (levels 3, 2, 1) and the level-4 table.
/// The kernel entries are shared with every other address space, so they're
/// left alone.
///
/// Unsafe: `l4` must not be in use by any thread (not loaded in CR3, and the
/// thread that ran in it must have finished).
pub unsafe fn free_address_space(l4: PhysFrame) {
    // A page table, reached through the physical-memory map.
    fn table(frame: PhysFrame) -> &'static PageTable {
        unsafe { &*phys_to_virt(frame.start_address()).as_ptr() }
    }
    // The used entries' frames (user space only has 4 KiB pages, no huge ones).
    fn used(t: &'static PageTable) -> impl Iterator<Item = PhysFrame> {
        t.iter().filter(|e| !e.is_unused()).filter_map(|e| e.frame().ok())
    }
    let user_l4_entry = &table(l4)[USER_L4_INDEX];
    if let Ok(l3) = user_l4_entry.frame() {
        for l2 in used(table(l3)) {
            for l1 in used(table(l2)) {
                for page in used(table(l1)) {
                    free_frame(page);
                }
                free_frame(l1);
            }
            free_frame(l2);
        }
        free_frame(l3);
    }
    free_frame(l4);
}

/// Make sure `page` (in user space) is mapped in address space `l4` with at
/// least `flags` (USER_ACCESSIBLE + PRESENT are added). A new page gets a
/// zeroed frame; an already mapped one keeps its frame and gets the union
/// of old and new flags. Returns the backing frame.
pub fn map_user_page(l4: PhysFrame, page: Page, flags: PageTableFlags) -> Option<PhysFrame> {
    let addr = page.start_address().as_u64();
    assert!((USER_START..USER_END).contains(&addr), "not a user address");
    let flags = flags | PageTableFlags::PRESENT | PageTableFlags::USER_ACCESSIBLE;

    let mut mapper = unsafe { mapper_for(l4) };
    if let TranslateResult::Mapped { frame, flags: old, .. } = mapper.translate(page.start_address()) {
        // Already mapped (two ELF segments sharing a page): widen the flags.
        // Keep NO_EXECUTE only if both want it.
        let mut new = old | flags;
        if !(old & flags).contains(PageTableFlags::NO_EXECUTE) {
            new.remove(PageTableFlags::NO_EXECUTE);
        }
        unsafe { mapper.update_flags(page, new).ok()?.ignore() };
        return Some(PhysFrame::containing_address(frame.start_address()));
    }
    let frame = alloc_zeroed_frame()?;
    // map_to also creates missing L3/L2/L1 tables; since `flags` contains
    // USER_ACCESSIBLE, those parent entries get it too (the CPU checks the
    // user bit at every level). The TLB doesn't matter: this address space
    // may not even be active, and the mapping is new. Hence `ignore`.
    unsafe { mapper.map_to(page, frame, flags, &mut GlobalFrameAllocator).ok()?.ignore() };
    Some(frame)
}

/// Can user code legitimately access `len` bytes at `addr` in the *current*
/// address space? Checks the range is inside user space and every page is
/// mapped USER_ACCESSIBLE (and WRITABLE if `write`). Syscalls call this
/// before touching any pointer a user program hands them; otherwise a program
/// could make the kernel read or overwrite kernel memory for it.
pub fn user_range_ok(addr: u64, len: u64, write: bool) -> bool {
    let Some(end) = addr.checked_add(len) else { return false };
    if addr < USER_START || end > USER_END {
        return false;
    }
    if len == 0 {
        return true;
    }
    let mapper = unsafe { mapper_for(Cr3::read().0) };
    let first = Page::<Size4KiB>::containing_address(VirtAddr::new(addr));
    let last = Page::<Size4KiB>::containing_address(VirtAddr::new(end - 1));
    Page::range_inclusive(first, last).all(|page| match mapper.translate(page.start_address()) {
        TranslateResult::Mapped { flags, .. } => {
            flags.contains(PageTableFlags::USER_ACCESSIBLE)
                && (!write || flags.contains(PageTableFlags::WRITABLE))
        }
        _ => false,
    })
}

// ---------------------------------------------------------------------------
// Kernel-space mappings for kernel stacks (see stack.rs)
// ---------------------------------------------------------------------------

/// The L4 slot that holds every kernel stack (0x0000_6400_0000_0000).
pub const STACKS_L4_INDEX: usize = 200;

/// Protects the kernel's own page tables. Taken with interrupts off.
static KERNEL_TABLES: Mutex<()> = Mutex::new(());

/// Give the kernel-stack slot its level-3 table *now*, before any process
/// address space is created. New address spaces copy the kernel's L4
/// entries, so they all point at this one L3 table, and every stack mapped
/// later shows up in every address space automatically.
pub fn init_kernel_stack_area() {
    without_interrupts(|| {
        let _lock = KERNEL_TABLES.lock();
        let l4: &mut PageTable = unsafe { &mut *phys_to_virt(kernel_l4_frame().start_address()).as_mut_ptr() };
        if l4[STACKS_L4_INDEX].is_unused() {
            let l3 = alloc_zeroed_frame().expect("no memory for the kernel stack area");
            l4[STACKS_L4_INDEX].set_frame(l3, PageTableFlags::PRESENT | PageTableFlags::WRITABLE);
        }
    })
}

/// Map `page` (inside the kernel stack area) to a fresh zeroed frame:
/// present, writable, not executable, kernel-only.
pub fn map_kernel_page(page: Page) -> Option<()> {
    let frame = alloc_zeroed_frame()?;
    let mut flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;
    if nx_enabled() {
        flags |= PageTableFlags::NO_EXECUTE;
    }
    let ok = without_interrupts(|| {
        let _lock = KERNEL_TABLES.lock();
        let mut mapper = unsafe { mapper_for(kernel_l4_frame()) };
        match unsafe { mapper.map_to(page, frame, flags, &mut GlobalFrameAllocator) } {
            Ok(flush) => {
                flush.flush(); // drop any stale TLB entry on *this* CPU
                true
            }
            Err(_) => false,
        }
    });
    if !ok {
        free_frame(frame);
    }
    ok.then_some(())
}

/// Unmap a page mapped by `map_kernel_page` and free its frame. (The page
/// tables stay: the next stack in the same place reuses them.)
pub fn unmap_kernel_page(page: Page) {
    without_interrupts(|| {
        let _lock = KERNEL_TABLES.lock();
        let mut mapper = unsafe { mapper_for(kernel_l4_frame()) };
        if let Ok((frame, flush)) = mapper.unmap(page) {
            flush.flush();
            free_frame(frame);
        }
    })
}