feat: paging + cleanup across memory management and other stuff

Makefile

@@ -117,26 +117,29 @@ copy-iso-files:
 partition-iso: copy-iso-files
 		# Make empty ISO of 64M in size
 		dd if=/dev/zero of=${IMAGE_PATH} bs=1M count=0 seek=${ISO_SIZE}
-		parted -s ${IMAGE_PATH} mklabel gpt
-		parted -s ${IMAGE_PATH} mkpart BIOSBOOT 1024s 2047s
-		parted -s ${IMAGE_PATH} set 1 bios_grub on
-
-		parted -s ${IMAGE_PATH} mkpart ESP fat${ESP_BITS} 2048s 262144s
+ifneq (${UEFI},)
+	parted -s ${IMAGE_PATH} mklabel gpt
+	parted -s ${IMAGE_PATH} mkpart ESP fat${ESP_BITS} 2048s 262144s
+	parted -s ${IMAGE_PATH} set 1 esp on
+else
+	parted -s ${IMAGE_PATH} mklabel msdos
+	parted -s ${IMAGE_PATH} mkpart primary fat${ESP_BITS} 2048s 262144s
+	parted -s ${IMAGE_PATH} set 1 boot on
+endif
 
 		# Make ISO with 1 partition starting at sector 2048 that is 32768 sectors, or 16MiB, in size
 		# Then a second partition spanning the rest of the disk
 		parted -s ${IMAGE_PATH} mkpart primary 262145s 100%
-		parted -s ${IMAGE_PATH} set 2 esp on
 
 build-iso: partition-iso 
-ifeq (${ARCH},x86_64)
-		# Install the Limine bootloader for bios installs
-		./limine/limine bios-install ${IMAGE_PATH}
+ifeq (${UEFI},)
+	# install limine for legacy bios
+	./limine/limine bios-install ${IMAGE_PATH}
 endif
 
 		sudo losetup -Pf --show ${IMAGE_PATH} > loopback_dev
-		sudo mkfs.fat -F ${ESP_BITS} `cat loopback_dev`p2
-		sudo mount `cat loopback_dev`p2 ${ARTIFACTS_PATH}/mnt
+		sudo mkfs.fat -F ${ESP_BITS} `cat loopback_dev`p1
+		sudo mount `cat loopback_dev`p1 ${ARTIFACTS_PATH}/mnt
 		sudo cp -r ${ISO_PATH}/* ${ARTIFACTS_PATH}/mnt
 		sync
 		sudo umount ${ARTIFACTS_PATH}/mnt
@@ -158,13 +161,13 @@ compile-bootloader:
 compile-binaries:
 		cargo build ${CARGO_OPTS}
 
-ovmf-x86_64: ovmf
+ovmf-x86_64:
 	mkdir -p ovmf/ovmf-x86_64
 	@if [ ! -d "ovmf/ovmf-x86_64/OVMF.fd" ]; then \
 		cd ovmf/ovmf-x86_64 && curl -Lo OVMF.fd https://retrage.github.io/edk2-nightly/bin/RELEASEX64_OVMF.fd; \
 	fi
 
-ovmf-aarch64: ovmf
+ovmf-aarch64:
 	mkdir -p ovmf/ovmf-aarch64
 	@if [ ! -d "ovmf/ovmf-aarch64/OVMF.fd" ]; then \
 		cd ovmf/ovmf-aarch64 && curl -o OVMF.fd https://retrage.github.io/edk2-nightly/bin/RELEASEAARCH64_QEMU_EFI.fd; \
@@ -174,10 +177,14 @@ ovmf-aarch64: ovmf
 # gdb target/x86_64-unknown-none/debug/CappuccinOS.elf -ex "target remote :1234"
 
 run: build ${RUN_OPTS} run-${ARCH}
+run-serial: build ${RUN_OPTS} run-${ARCH}-serial
 
 run-x86_64:
 	tmux new-session -d -s qemu 'qemu-system-x86_64 ${QEMU_OPTS}'
 
+run-x86_64-serial:
+	qemu-system-x86_64 ${QEMU_OPTS} -boot d -display none -serial stdio -monitor none -no-reboot -no-shutdown
+
 line-count:
 		cloc --quiet --exclude-dir=bin --include-lang=Rust --csv src/ | tail -n 1 | awk -F, '{print $$5}'
 clean:

src/arch/x86_64/gdt.rs

@@ -57,10 +57,10 @@ static mut GDT: Gdt = Gdt::new();
 static mut TSS: TaskStateSegment = TaskStateSegment::new();
 static mut DOUBLE_FAULT_STACK: ExceptionStack = ExceptionStack([0; DOUBLE_FAULT_STACK_SIZE]);
 
-pub fn init() {
-    const _: () = assert!(core::mem::size_of::<TaskStateSegment>() == 104);
-    const _: () = assert!(core::mem::size_of::<GdtPointer>() == 10);
+const _: () = assert!(core::mem::size_of::<TaskStateSegment>() == 104);
+const _: () = assert!(core::mem::size_of::<GdtPointer>() == 10);
 
+pub fn init() {
     unsafe {
         let stack_bottom = core::ptr::addr_of_mut!(DOUBLE_FAULT_STACK) as u64;
         let stack_top = stack_bottom + DOUBLE_FAULT_STACK_SIZE as u64;

src/arch/x86_64/interrupts/idt.rs

@@ -89,11 +89,11 @@ impl Idt {
 
 static mut IDT: Idt = Idt::new();
 
-pub fn idt_init() {
-    const _: () = assert!(core::mem::size_of::<IdtEntry>() == 16);
-    const _: () = assert!(core::mem::size_of::<IdtPointer>() == 10);
-    const _: () = assert!(core::mem::size_of::<InterruptStackFrame>() == 40);
+const _: () = assert!(core::mem::size_of::<IdtEntry>() == 16);
+const _: () = assert!(core::mem::size_of::<IdtPointer>() == 10);
+const _: () = assert!(core::mem::size_of::<InterruptStackFrame>() == 40);
 
+pub fn idt_init() {
     let mut idt = Idt::new();
 
     exceptions::install(&mut idt);

src/arch/x86_64/mod.rs

@@ -1,5 +1,6 @@
 mod gdt;
 mod interrupts;
+pub mod paging;
 pub mod port;
 
 use core::arch::asm;

src/arch/x86_64/paging.rs

@@ -0,0 +1,643 @@
+use core::arch::asm;
+
+use crate::{
+    memory::{
+        DirectMap, FrameAllocator, KernelImage, MemoryRegion, MemoryRegionKind, PhysicalAddr,
+        PhysicalFrame, VirtualAddr,
+    },
+    println,
+};
+
+pub const PAGE_SIZE: usize = 4096;
+pub const PAGE_TABLE_ENTRIES: usize = 512;
+
+const ADDRESS_MASK: usize = 0x000F_FFFF_FFFF_F000;
+
+#[repr(transparent)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+struct PageTableEntry(u64);
+
+impl PageTableEntry {
+    const PRESENT: u64 = 1 << 0;
+    const WRITABLE: u64 = 1 << 1;
+    const USER_ACCESSIBLE: u64 = 1 << 2;
+    const LARGE_PAGE: u64 = 1 << 7;
+    const NX: u64 = 1 << 63;
+
+    // TODO: validate that page is withing the CPU's mappable address space
+    const fn new(physical_address: PhysicalAddr, permissions: PagePermissions) -> Self {
+        let mut value = physical_address.as_usize() as u64 | Self::PRESENT;
+
+        if permissions.writable {
+            value |= Self::WRITABLE;
+        }
+        if permissions.user_accessible {
+            value |= Self::USER_ACCESSIBLE;
+        }
+        if !permissions.executable {
+            value |= Self::NX;
+        }
+
+        Self(value)
+    }
+
+    fn new_table(frame: PhysicalFrame, user_accessible: bool) -> Self {
+        let mut value = frame.start_address().as_usize() as u64 | Self::PRESENT | Self::WRITABLE;
+        if user_accessible {
+            value |= Self::USER_ACCESSIBLE;
+        }
+
+        Self(value)
+    }
+
+    const fn null() -> Self {
+        Self(0)
+    }
+
+    fn physical_address(&self) -> PhysicalAddr {
+        PhysicalAddr::new(self.0 as usize & ADDRESS_MASK)
+    }
+
+    fn is_present(&self) -> bool {
+        self.0 & Self::PRESENT != 0
+    }
+
+    fn is_user_accessible(&self) -> bool {
+        self.0 & Self::USER_ACCESSIBLE != 0
+    }
+
+    fn is_huge(&self) -> bool {
+        self.0 & Self::LARGE_PAGE != 0
+    }
+
+    fn frame(&self) -> Option<PhysicalFrame> {
+        if !self.is_present() || self.is_huge() {
+            return None;
+        }
+
+        PhysicalFrame::from_start_address(self.physical_address())
+    }
+}
+
+#[repr(C, align(4096))]
+struct PageTable {
+    entries: [PageTableEntry; PAGE_TABLE_ENTRIES],
+}
+
+impl PageTable {
+    pub fn is_empty(&self) -> bool {
+        self.entries.iter().all(|&entry| !entry.is_present())
+    }
+}
+
+const _: () = assert!(core::mem::size_of::<PageTable>() == 4096);
+
+fn translate_huge_page(
+    entry: PageTableEntry,
+    virtual_addr: usize,
+    page_size: usize,
+) -> Option<PhysicalAddr> {
+    let physical_base = entry.physical_address().as_usize() & !(page_size - 1);
+
+    let offset = virtual_addr & (page_size - 1);
+
+    physical_base.checked_add(offset).map(PhysicalAddr::new)
+}
+
+#[derive(Debug)]
+pub enum PageError {
+    NotCanonical,
+    Unaligned,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct PagePermissions {
+    pub writable: bool,
+    pub executable: bool,
+    pub user_accessible: bool,
+}
+
+impl PagePermissions {
+    // TODO: give more fine grained permissions
+    pub const HHDM: Self = Self::new(true, false, false);
+    pub const KERNEL: Self = Self::new(true, true, false);
+    pub const KERNEL_DATA: Self = Self::new(true, false, false);
+
+    pub const fn new(writable: bool, executable: bool, user_accessible: bool) -> Self {
+        Self {
+            writable,
+            executable,
+            user_accessible,
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct Page {
+    start: VirtualAddr,
+}
+
+impl Page {
+    pub fn from_start_address(start: VirtualAddr) -> Result<Self, PageError> {
+        if !is_canonical_48_bit(start.as_usize()) {
+            return Err(PageError::NotCanonical);
+        }
+
+        if start.as_usize() & (4096 - 1) != 0 {
+            return Err(PageError::Unaligned);
+        }
+
+        Ok(Self { start })
+    }
+}
+
+fn p4_index(addr: usize) -> usize {
+    (addr >> 39) & 0x1FF
+}
+
+fn p3_index(addr: usize) -> usize {
+    (addr >> 30) & 0x1FF
+}
+
+fn p2_index(addr: usize) -> usize {
+    (addr >> 21) & 0x1FF
+}
+
+fn p1_index(addr: usize) -> usize {
+    (addr >> 12) & 0x1FF
+}
+
+fn page_offset(addr: usize) -> usize {
+    addr & 0xFFF
+}
+
+unsafe fn read_cr3() -> PhysicalFrame {
+    let value: usize;
+    unsafe {
+        asm!(
+            "mov {}, cr3",
+            out(reg) value,
+            options(nomem, nostack, preserves_flags),
+        );
+    }
+
+    PhysicalFrame::from_start_address(PhysicalAddr::new(value & ADDRESS_MASK))
+        .expect("CR3 contains an unaligned page-table address")
+}
+
+#[derive(Debug)]
+pub enum MapError {
+    MisalignedPage,
+    PageAlreadyMapped,
+    ParentIsHuge,
+    FrameAllocatorError,
+    AddressOverflow,
+    UserPageInKernelHalf,
+}
+
+#[derive(Debug)]
+pub enum UnmapError {
+    PageNotWithinHHDM,
+    PageNotPresent,
+    ParentIsHuge,
+    ParentNotPresent,
+    FrameAllocatorError,
+}
+
+type PageTableRoot = PhysicalFrame;
+
+#[derive(Clone, Copy)]
+struct NewTable {
+    parent: PhysicalFrame,
+    index: usize,
+    child: PhysicalFrame,
+}
+
+#[derive(Debug)]
+pub enum AddressSpaceError {
+    MalformedMemmap,
+    Map(MapError),
+    FrameAllocatorError,
+}
+
+pub struct AddressSpace {
+    root: PageTableRoot,
+    direct_map: DirectMap,
+}
+
+impl AddressSpace {
+    pub fn new<I: Iterator<Item = MemoryRegion> + Clone>(
+        direct_map: DirectMap,
+        memmap: I,
+        kernel_address: KernelImage,
+        allocator: &mut FrameAllocator,
+    ) -> Result<Self, AddressSpaceError> {
+        let frame = allocator
+            .alloc()
+            .ok_or(AddressSpaceError::FrameAllocatorError)?;
+        let mut space = AddressSpace {
+            root: frame,
+            direct_map,
+        };
+
+        // TODO: cleanup page tables if any allocations fail
+        // map hhdm
+        for region in memmap.clone() {
+            if region.kind == MemoryRegionKind::Reserved
+                || region.kind == MemoryRegionKind::BadMemory
+            {
+                continue;
+            }
+
+            println!(
+                "Mapping: {:?} at {:#X} with length {:#X}",
+                region.kind,
+                region.start.as_usize(),
+                region.length
+            );
+
+            space
+                .map_range(
+                    direct_map
+                        .translate(region.start)
+                        .ok_or(AddressSpaceError::MalformedMemmap)?,
+                    region.start,
+                    region.length,
+                    PagePermissions::HHDM,
+                    allocator,
+                )
+                .map_err(|err| AddressSpaceError::Map(err))?;
+        }
+
+        // map kernel
+        println!(
+            "Mapping kernel at {:#X} with length {:#X}",
+            kernel_address.virtual_base.as_usize(),
+            kernel_address.length
+        );
+
+        space
+            .map_range(
+                kernel_address.virtual_base,
+                kernel_address.physical_base,
+                kernel_address.length,
+                PagePermissions::KERNEL,
+                allocator,
+            )
+            .map_err(|err| AddressSpaceError::Map(err))?;
+
+        Ok(space)
+    }
+
+    fn is_active(&self) -> bool {
+        let cr3 = unsafe { read_cr3() };
+
+        cr3.start_address() == self.root.start_address()
+    }
+
+    fn table(&self, frame: PhysicalFrame) -> Option<&PageTable> {
+        let virtual_addr = self.direct_map.translate(frame.start_address())?;
+
+        Some(unsafe { &*(virtual_addr.as_ptr::<PageTable>()) })
+    }
+
+    fn table_mut(&mut self, frame: PhysicalFrame) -> Option<&mut PageTable> {
+        let virtual_addr = self.direct_map.translate(frame.start_address())?;
+
+        Some(unsafe { &mut *(virtual_addr.as_mut_ptr::<PageTable>()) })
+    }
+
+    pub fn translate(&self, addr: VirtualAddr) -> Option<PhysicalAddr> {
+        let addr = addr.as_usize();
+
+        if !is_canonical_48_bit(addr) {
+            return None;
+        }
+
+        let p4 = self.table(self.root)?;
+        let p4_entry = p4.entries[p4_index(addr)];
+
+        if !p4_entry.is_present() {
+            return None;
+        }
+
+        let p3 = self.table(p4_entry.frame()?)?;
+        let p3_entry = p3.entries[p3_index(addr)];
+
+        if !p3_entry.is_present() {
+            return None;
+        }
+
+        if p3_entry.is_huge() {
+            return translate_huge_page(p3_entry, addr, 1 << 30);
+        }
+
+        let p2 = self.table(p3_entry.frame()?)?;
+        let p2_entry = p2.entries[p2_index(addr)];
+
+        if !p2_entry.is_present() {
+            return None;
+        }
+
+        if p2_entry.is_huge() {
+            return translate_huge_page(p2_entry, addr, 1 << 21);
+        }
+
+        let p1 = self.table(p2_entry.frame()?)?;
+        let p1_entry = p1.entries[p1_index(addr)];
+
+        if !p1_entry.is_present() {
+            return None;
+        }
+
+        let physical_base = p1_entry.physical_address().as_usize();
+
+        physical_base
+            .checked_add(page_offset(addr))
+            .map(PhysicalAddr::new)
+    }
+
+    fn get_next_level(
+        &self,
+        parent: PhysicalFrame,
+        index: usize,
+    ) -> Result<PhysicalFrame, UnmapError> {
+        let parent_table = self.table(parent).ok_or(UnmapError::PageNotWithinHHDM)?;
+
+        // TODO: encode some level so we dont spuriously think a page is huge
+        if parent_table.entries[index].is_huge() {
+            return Err(UnmapError::ParentIsHuge);
+        }
+
+        Ok(parent_table.entries[index]
+            .frame()
+            .ok_or(UnmapError::ParentNotPresent)?)
+    }
+
+    fn get_next_level_or_allocate(
+        &mut self,
+        parent: PhysicalFrame,
+        index: usize,
+        user_accessible: bool,
+        allocator: &mut FrameAllocator,
+    ) -> Result<(PhysicalFrame, bool), MapError> {
+        let parent_table = self
+            .table_mut(parent)
+            .ok_or(MapError::FrameAllocatorError)?;
+
+        let entry = &mut parent_table.entries[index];
+
+        if !entry.is_present() {
+            let frame = allocator.alloc().ok_or(MapError::FrameAllocatorError)?;
+            *entry = PageTableEntry::new_table(frame, user_accessible);
+            return Ok((frame, true));
+        }
+
+        if entry.is_huge() {
+            return Err(MapError::ParentIsHuge);
+        }
+
+        // TODO: we upgrade intermediate entries, and dont carefully rollback if we fail
+        if user_accessible && !entry.is_user_accessible() {
+            entry.0 |= PageTableEntry::USER_ACCESSIBLE;
+        }
+
+        entry
+            .frame()
+            .map(|frame| (frame, false))
+            .ok_or(MapError::FrameAllocatorError)
+    }
+
+    fn rollback_tables(
+        &mut self,
+        new_tables: &[Option<NewTable>],
+        count: usize,
+        allocator: &mut FrameAllocator,
+    ) {
+        for table in new_tables[..count].iter().rev().flatten() {
+            self.table_mut(table.parent).unwrap().entries[table.index] = PageTableEntry::null();
+
+            unsafe { allocator.dealloc(table.child) };
+        }
+    }
+
+    pub fn map(
+        &mut self,
+        page: Page,
+        frame: PhysicalFrame,
+        permissions: PagePermissions,
+        allocator: &mut FrameAllocator,
+    ) -> Result<(), MapError> {
+        if permissions.user_accessible && p4_index(page.start.as_usize()) >= 256 {
+            return Err(MapError::UserPageInKernelHalf);
+        }
+
+        let mut new_tables: [Option<NewTable>; 3] = [None; 3];
+        let mut new_table_count = 0;
+
+        let result = (|| {
+            let p4_entry_index = p4_index(page.start.as_usize());
+            let (pdpt_frame, allocated) = self.get_next_level_or_allocate(
+                self.root,
+                p4_entry_index,
+                permissions.user_accessible,
+                allocator,
+            )?;
+            if allocated {
+                new_tables[new_table_count] = Some(NewTable {
+                    parent: self.root,
+                    index: p4_entry_index,
+                    child: pdpt_frame,
+                });
+                new_table_count += 1;
+            }
+
+            let p3_entry_index = p3_index(page.start.as_usize());
+            let (pd_frame, allocated) = self.get_next_level_or_allocate(
+                pdpt_frame,
+                p3_entry_index,
+                permissions.user_accessible,
+                allocator,
+            )?;
+            if allocated {
+                new_tables[new_table_count] = Some(NewTable {
+                    parent: pdpt_frame,
+                    index: p3_entry_index,
+                    child: pd_frame,
+                });
+                new_table_count += 1;
+            }
+
+            let p2_entry_index = p2_index(page.start.as_usize());
+            let (pt_frame, allocated) = self.get_next_level_or_allocate(
+                pd_frame,
+                p2_entry_index,
+                permissions.user_accessible,
+                allocator,
+            )?;
+            if allocated {
+                new_tables[new_table_count] = Some(NewTable {
+                    parent: pd_frame,
+                    index: p2_entry_index,
+                    child: pt_frame,
+                });
+                new_table_count += 1;
+            }
+
+            let pt_table = self
+                .table_mut(pt_frame)
+                .ok_or(MapError::FrameAllocatorError)?;
+            let entry = &mut pt_table.entries[p1_index(page.start.as_usize())];
+            if entry.is_present() {
+                return Err(MapError::PageAlreadyMapped);
+            }
+
+            *entry = PageTableEntry::new(frame.start_address(), permissions);
+            Ok(())
+        })();
+
+        if result.is_err() {
+            self.rollback_tables(&new_tables, new_table_count, allocator);
+            return result;
+        }
+
+        self.flush_tlb_if_active(page);
+
+        Ok(())
+    }
+
+    pub fn map_range(
+        &mut self,
+        virtual_start: VirtualAddr,
+        physical_start: PhysicalAddr,
+        length: usize,
+        permissions: PagePermissions,
+        allocator: &mut FrameAllocator,
+    ) -> Result<(), MapError> {
+        if length % PAGE_SIZE != 0 {
+            return Err(MapError::MisalignedPage);
+        }
+
+        let virtual_start =
+            Page::from_start_address(virtual_start).map_err(|_| MapError::MisalignedPage)?;
+        let physical_start =
+            PhysicalFrame::from_start_address(physical_start).ok_or(MapError::MisalignedPage)?;
+        virtual_start
+            .start
+            .as_usize()
+            .checked_add(length)
+            .ok_or(MapError::AddressOverflow)?;
+        physical_start
+            .start_address()
+            .as_usize()
+            .checked_add(length)
+            .ok_or(MapError::AddressOverflow)?;
+
+        for page_idx in 0..length / PAGE_SIZE {
+            let offset = page_idx * PAGE_SIZE;
+            let page =
+                Page::from_start_address(VirtualAddr::new(virtual_start.start.as_usize() + offset))
+                    .unwrap();
+            let frame = PhysicalFrame::from_start_address(PhysicalAddr::new(
+                physical_start.start_address().as_usize() + offset,
+            ))
+            .unwrap();
+
+            if let Err(error) = self.map(page, frame, permissions, allocator) {
+                for rollback_idx in (0..page_idx).rev() {
+                    let rollback_page = Page::from_start_address(VirtualAddr::new(
+                        virtual_start.start.as_usize() + rollback_idx * PAGE_SIZE,
+                    ))
+                    .unwrap();
+                    self.unmap(rollback_page, allocator)
+                        .expect("failed to roll back a mapped page");
+                }
+
+                return Err(error);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn unmap(
+        &mut self,
+        page: Page,
+        allocator: &mut FrameAllocator,
+    ) -> Result<PhysicalFrame, UnmapError> {
+        let pdpt_frame = self.get_next_level(self.root, p4_index(page.start.as_usize()))?;
+
+        let pd_frame = self.get_next_level(pdpt_frame, p3_index(page.start.as_usize()))?;
+
+        let pt_frame = self.get_next_level(pd_frame, p2_index(page.start.as_usize()))?;
+        let pt_table = self
+            .table_mut(pt_frame)
+            .ok_or(UnmapError::FrameAllocatorError)?;
+
+        let entry = pt_table.entries[p1_index(page.start.as_usize())];
+
+        if !entry.is_present() {
+            return Err(UnmapError::PageNotPresent);
+        }
+
+        let frame = entry.frame().ok_or(UnmapError::FrameAllocatorError)?;
+        pt_table.entries[p1_index(page.start.as_usize())] = PageTableEntry::null();
+
+        if pt_table.is_empty() {
+            self.table_mut(pd_frame).unwrap().entries[p2_index(page.start.as_usize())] =
+                PageTableEntry::null();
+            unsafe { allocator.dealloc(pt_frame) };
+
+            if self.table(pd_frame).unwrap().is_empty() {
+                self.table_mut(pdpt_frame).unwrap().entries[p3_index(page.start.as_usize())] =
+                    PageTableEntry::null();
+                unsafe { allocator.dealloc(pd_frame) };
+
+                if self.table(pdpt_frame).unwrap().is_empty() {
+                    self.table_mut(self.root).unwrap().entries[p4_index(page.start.as_usize())] =
+                        PageTableEntry::null();
+                    unsafe { allocator.dealloc(pdpt_frame) };
+                }
+            }
+        }
+
+        self.flush_tlb_if_active(page);
+
+        Ok(frame)
+    }
+
+    // TODO: unmap range
+
+    fn flush_tlb_if_active(&self, page: Page) {
+        if self.is_active() {
+            println!("Flushing TLB");
+            unsafe {
+                asm!("invlpg [{}]", in(reg) page.start.as_usize(), options(nostack, preserves_flags));
+            }
+        }
+    }
+
+    pub unsafe fn activate(&self) -> PageTableRoot {
+        println!("Activating");
+        unsafe {
+            asm!(
+                "mov cr3, {}",
+                in(reg) self.root.start_address().as_usize(),
+                options(nostack, preserves_flags)
+            );
+        }
+
+        self.root
+    }
+}
+
+fn is_canonical_48_bit(addr: usize) -> bool {
+    let upper = addr >> 48;
+    let sign_bit = (addr >> 47) & 1;
+
+    if sign_bit == 0 {
+        upper == 0
+    } else {
+        upper == 0xFFFF
+    }
+}

src/boot/limine.rs

@@ -3,6 +3,8 @@ use ::limine as limine_api;
 use limine_api::request::{ExecutableAddressRequest, HhdmRequest, MemmapRequest};
 use limine_api::{BaseRevision, RequestsEndMarker, RequestsStartMarker};
 
+use crate::memory::{KernelImage, PhysicalAddr, VirtualAddr};
+
 /// Sets the base revision to the latest revision supported by the crate.
 /// See specification for further info.
 /// Be sure to mark all limine requests with #[used], otherwise they may be removed by the compiler.
@@ -32,7 +34,7 @@ static _START_MARKER: RequestsStartMarker = RequestsStartMarker::new();
 static _END_MARKER: RequestsEndMarker = RequestsEndMarker::new();
 
 pub struct BootInfo {
-    pub kernel_address: crate::memory::PhysicalAddr,
+    pub kernel_address: KernelImage,
     pub hhdm_offset: usize,
     entries: &'static [&'static limine_api::memmap::Entry],
 }
@@ -69,6 +71,7 @@ pub enum BootError {
     FailedToGetKernelAddress,
     FailedToGetHHDMAddress,
     FailedToGetMemmap,
+    FailedToLocateKernel,
 }
 
 pub fn load_boot_info() -> Result<BootInfo, BootError> {
@@ -78,8 +81,7 @@ pub fn load_boot_info() -> Result<BootInfo, BootError> {
 
     let kernel_address = KERNEL_ADDRESS_REQUEST
         .response()
-        .ok_or(BootError::FailedToGetKernelAddress)?
-        .physical_base;
+        .ok_or(BootError::FailedToGetKernelAddress)?;
     let hhdm_offset = HHDM_REQUEST
         .response()
         .ok_or(BootError::FailedToGetHHDMAddress)?
@@ -90,8 +92,28 @@ pub fn load_boot_info() -> Result<BootInfo, BootError> {
         .ok_or(BootError::FailedToGetMemmap)?
         .entries();
 
+    let mut kernel_length = None;
+    for &entry in memmap.iter() {
+        if entry.type_ != limine_api::memmap::MEMMAP_EXECUTABLE_AND_MODULES {
+            continue;
+        }
+
+        if entry.base != kernel_address.physical_base {
+            continue;
+        }
+
+        kernel_length = Some(entry.length as usize);
+        break;
+    }
+
+    let kernel_length = kernel_length.ok_or(BootError::FailedToLocateKernel)?;
+
     Ok(BootInfo {
-        kernel_address: crate::memory::PhysicalAddr::new(kernel_address as usize),
+        kernel_address: KernelImage {
+            physical_base: PhysicalAddr::new(kernel_address.physical_base as usize),
+            virtual_base: VirtualAddr::new(kernel_address.virtual_base as usize),
+            length: kernel_length,
+        },
         hhdm_offset: hhdm_offset as usize,
         entries: memmap,
     })

src/main.rs

@@ -8,7 +8,11 @@ mod boot;
 mod debug;
 mod memory;
 
-use crate::debug::serial;
+use crate::{
+    arch::paging,
+    debug::serial,
+    memory::{PhysicalAddr, VirtualAddr},
+};
 
 #[unsafe(no_mangle)]
 pub extern "C" fn _start() -> ! {
@@ -16,38 +20,72 @@ pub extern "C" fn _start() -> ! {
 
     arch::init();
     let boot_info = boot::load_boot_info().unwrap();
+    let direct_map = memory::DirectMap::new(boot_info.hhdm_offset);
 
-    let mut allocator = memory::FrameAllocator::new(
+    let mut allocator = memory::FrameAllocator::new(boot_info.memory_regions(), direct_map)
+        .expect("failed to create frame allocator");
+
+    let mut page_table = paging::AddressSpace::new(
+        direct_map,
         boot_info.memory_regions(),
-        memory::DirectMap::new(boot_info.hhdm_offset),
+        boot_info.kernel_address,
+        &mut allocator,
     )
-    .expect("failed to create frame allocator");
+    .expect("failed to create page table");
 
-    let initial = allocator.free_frames();
+    // safety: trust me bro
+    unsafe { page_table.activate() };
 
-    let first = allocator.alloc().unwrap();
-    let second = allocator.alloc().unwrap();
-    let third = allocator.alloc().unwrap();
+    let frame = allocator.alloc().unwrap();
 
-    assert_ne!(first, second);
-    assert_ne!(second, third);
-    assert_eq!(allocator.free_frames(), initial - 3);
+    let direct_mapped = direct_map.translate(frame.start_address()).unwrap();
+    let translated = page_table.translate(direct_mapped).unwrap();
 
+    println!("{:?}", translated);
+
+    let new_virtual = VirtualAddr::new(0x8000_0000);
+
+    assert!(page_table.translate(new_virtual).is_none());
+
+    let page = paging::Page::from_start_address(new_virtual).unwrap();
+
+    page_table
+        .map(
+            page,
+            frame,
+            paging::PagePermissions::KERNEL_DATA,
+            &mut allocator,
+        )
+        .unwrap();
+
+    assert_eq!(
+        page_table.translate(new_virtual),
+        Some(frame.start_address())
+    );
+
+    assert_eq!(
+        page_table.translate(VirtualAddr::new(new_virtual.as_usize() + 123)),
+        Some(PhysicalAddr::new(frame.start_address().as_usize() + 123))
+    );
+
+    // write to the page and read it back via HHDM
     unsafe {
-        allocator.dealloc(second);
-        allocator.dealloc(first);
-        allocator.dealloc(third);
+        core::ptr::write_bytes(new_virtual.as_mut_ptr::<u8>(), 0xFF, 0x1000);
     }
 
-    assert_eq!(allocator.free_frames(), initial);
+    let slice = unsafe { core::slice::from_raw_parts(direct_mapped.as_ptr::<u8>(), 0x1000) };
 
-    let a = allocator.alloc().unwrap();
-    let b = allocator.alloc().unwrap();
-    let c = allocator.alloc().unwrap();
+    assert!(slice.iter().all(|&byte| byte == 0xFF));
 
-    assert_ne!(a, b);
-    assert_ne!(b, c);
-    assert_ne!(a, c);
+    println!("{:#X}", slice[0]);
+
+    let unmapped_frame = page_table.unmap(page, &mut allocator).unwrap();
+
+    assert_eq!(unmapped_frame, frame);
+
+    unsafe { allocator.dealloc(unmapped_frame) };
+
+    assert!(page_table.translate(new_virtual).is_none());
 
     hcf();
 }

src/memory/frame.rs

@@ -1,5 +1,3 @@
-use core::ops::Range;
-
 use crate::memory::{DirectMap, MemoryRegion, MemoryRegionKind, PhysicalAddr, VirtualAddr};
 
 pub const FRAME_SIZE: usize = 4096;
@@ -13,102 +11,53 @@ pub fn align_down_to_frame(addr: usize) -> usize {
     addr & !(FRAME_SIZE - 1)
 }
 
+#[repr(u8)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum FrameState {
+    Reserved = 0b00,
+    Free = 0b01,
+    Allocated = 0b10,
+}
+
 #[derive(Debug)]
 struct Bitmap {
     start: VirtualAddr,
-    bit_count: usize,
+    frame_count: usize,
 }
 
-// 1 = available, 0 = unavailable
 impl Bitmap {
-    fn new(start: VirtualAddr, bit_count: usize) -> Self {
-        Self { start, bit_count }
+    fn new(start: VirtualAddr, frame_count: usize) -> Self {
+        Self { start, frame_count }
     }
 
-    fn is_available(&self, idx: usize) -> bool {
-        assert!(idx < self.bit_count, "index out of bounds");
+    fn state(&self, frame_idx: usize) -> FrameState {
+        assert!(frame_idx < self.frame_count, "frame index out of bounds");
 
-        let byte = idx / 8;
-        let bit = idx % 8;
+        let byte_idx = frame_idx / 4;
+        let shift = (frame_idx % 4) * 2;
+        let byte = unsafe { self.start.as_ptr::<u8>().add(byte_idx).read() };
 
-        unsafe { self.start.as_mut_ptr::<u8>().add(byte).read() & (1 << bit) != 0 }
-    }
-
-    fn set_available(&mut self, idx: usize, available: bool) {
-        assert!(idx < self.bit_count, "index out of bounds");
-
-        let byte = idx / 8;
-        let bit = idx % 8;
-
-        let val = unsafe { self.start.as_mut_ptr::<u8>().add(byte).read() };
-        unsafe {
-            self.start.as_mut_ptr::<u8>().add(byte).write(if available {
-                val | (1 << bit)
-            } else {
-                val & !(1 << bit)
-            });
+        match (byte >> shift) & 0b11 {
+            0b00 => FrameState::Reserved,
+            0b01 => FrameState::Free,
+            0b10 => FrameState::Allocated,
+            _ => panic!("invalid frame state"),
         }
     }
 
-    // fn fill_available(&mut self, range: Range<usize>, available: bool) -> bool {
-    //     if range.start > range.end {
-    //         return false;
-    //     }
+    fn set_state(&mut self, frame_idx: usize, state: FrameState) {
+        assert!(frame_idx < self.frame_count, "frame index out of bounds");
 
-    //     if range.end > self.bit_count {
-    //         return false;
-    //     }
+        let byte_idx = frame_idx / 4;
+        let shift = (frame_idx % 4) * 2;
+        let ptr = unsafe { self.start.as_mut_ptr::<u8>().add(byte_idx) };
+        let byte = unsafe { ptr.read() };
+        let mask = 0b11 << shift;
 
-    //     if range.is_empty() {
-    //         return true;
-    //     }
-
-    //     let ptr = unsafe { self.start.as_mut_ptr::<u8>() };
-    //     let first_byte = range.start / 8;
-    //     let last_byte = (range.end - 1) / 8;
-    //     let start_bit = range.start % 8;
-    //     let end_bit = range.end % 8;
-
-    //     unsafe {
-    //         if first_byte == last_byte {
-    //             let width = range.end - range.start;
-    //             let mask = (((1u16 << width) - 1) << start_bit) as u8;
-    //             let byte = ptr.add(first_byte).read();
-
-    //             ptr.add(first_byte)
-    //                 .write(if available { byte | mask } else { byte & !mask });
-
-    //             return true;
-    //         }
-
-    //         let mut full_start = first_byte;
-
-    //         if start_bit != 0 {
-    //             let mask = u8::MAX << start_bit;
-    //             let byte = ptr.add(first_byte).read();
-
-    //             ptr.add(first_byte)
-    //                 .write(if available { byte | mask } else { byte & !mask });
-
-    //             full_start += 1;
-    //         }
-
-    //         let full_end = range.end / 8;
-
-    //         ptr.add(full_start)
-    //             .write_bytes(if available { u8::MAX } else { 0 }, full_end - full_start);
-
-    //         if end_bit != 0 {
-    //             let mask = (1u8 << end_bit) - 1;
-    //             let byte = ptr.add(full_end).read();
-
-    //             ptr.add(full_end)
-    //                 .write(if available { byte | mask } else { byte & !mask });
-    //         }
-    //     }
-
-    //     true
-    // }
+        unsafe {
+            ptr.write((byte & !mask) | ((state as u8) << shift));
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -123,8 +72,6 @@ pub enum FrameAllocatorInitError {
 #[derive(Debug)]
 pub struct FrameAllocator {
     bitmap: Bitmap,
-    bitmap_start_frame: usize,
-    bitmap_frame_count: usize,
     next_search: usize,
     allocatable_frames: usize,
     free_frames: usize,
@@ -153,7 +100,7 @@ impl FrameAllocator {
 
         let highest_frame = highest_frame.ok_or(FrameAllocatorInitError::NoUsableFrames)?;
 
-        let bitmap_bytes = highest_frame.div_ceil(8);
+        let bitmap_bytes = highest_frame.div_ceil(4);
         let bitmap_frame_count = bitmap_bytes.div_ceil(FRAME_SIZE);
         let bitmap_storage_bytes = bitmap_frame_count
             .checked_mul(FRAME_SIZE)
@@ -215,7 +162,7 @@ impl FrameAllocator {
                     continue;
                 }
 
-                bitmap.set_available(frame_idx, true);
+                bitmap.set_state(frame_idx, FrameState::Free);
                 allocatable_frames += 1;
                 free_frames += 1;
             }
@@ -223,8 +170,6 @@ impl FrameAllocator {
 
         Ok(Self {
             bitmap,
-            bitmap_start_frame: bitmap_start_frame,
-            bitmap_frame_count,
             next_search: bitmap_start_frame + bitmap_frame_count,
             allocatable_frames,
             free_frames,
@@ -233,11 +178,11 @@ impl FrameAllocator {
     }
 
     fn find_free_in(&self, start: usize, end: usize) -> Option<usize> {
-        (start..end).find(|&index| self.bitmap.is_available(index))
+        (start..end).find(|&index| self.bitmap.state(index) == FrameState::Free)
     }
 
     fn find_free_frame(&self) -> Option<usize> {
-        self.find_free_in(self.next_search, self.bitmap.bit_count)
+        self.find_free_in(self.next_search, self.bitmap.frame_count)
             .or_else(|| self.find_free_in(0, self.next_search))
     }
 
@@ -248,7 +193,7 @@ impl FrameAllocator {
 
         let frame_idx = self.find_free_frame()?;
 
-        self.bitmap.set_available(frame_idx, false);
+        self.bitmap.set_state(frame_idx, FrameState::Allocated);
         self.free_frames -= 1;
         self.next_search = frame_idx.saturating_add(1);
 
@@ -279,22 +224,17 @@ impl FrameAllocator {
     pub unsafe fn dealloc(&mut self, frame: PhysicalFrame) {
         let frame_idx = frame.index();
 
-        assert!(
-            frame_idx < self.bitmap.bit_count,
-            "frame index out of bounds"
-        );
-        assert!(
-            !self.is_bitmap_storage(frame_idx),
-            "attempted to free frame allocator bitmap"
-        );
-        assert!(
-            !self.bitmap.is_available(frame_idx),
-            "frame is already free"
-        );
-
-        self.bitmap.set_available(frame_idx, true);
-        self.free_frames += 1;
-        self.next_search = self.next_search.min(frame_idx);
+        match self.bitmap.state(frame_idx) {
+            FrameState::Allocated => {
+                self.bitmap.set_state(frame_idx, FrameState::Free);
+                self.free_frames += 1;
+                self.next_search = self.next_search.min(frame_idx);
+            }
+            FrameState::Free => panic!("attempted to free free frame"),
+            FrameState::Reserved => {
+                panic!("attempted to free reserved frame");
+            }
+        };
     }
 
     pub const fn free_frames(&self) -> usize {
@@ -305,11 +245,6 @@ impl FrameAllocator {
         self.allocatable_frames
     }
 
-    fn is_bitmap_storage(&self, index: usize) -> bool {
-        index >= self.bitmap_start_frame
-            && index < (self.bitmap_start_frame + self.bitmap_frame_count)
-    }
-
     fn usable_frame_range(
         region: MemoryRegion,
     ) -> Result<core::ops::Range<usize>, FrameAllocatorInitError> {

src/memory/mod.rs

@@ -2,6 +2,13 @@ mod frame;
 
 pub use frame::{FRAME_SIZE, FrameAllocator, PhysicalFrame};
 
+#[derive(Debug, Clone, Copy)]
+pub struct KernelImage {
+    pub physical_base: PhysicalAddr,
+    pub virtual_base: VirtualAddr,
+    pub length: usize,
+}
+
 #[repr(transparent)]
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct PhysicalAddr(usize);
@@ -32,6 +39,10 @@ impl VirtualAddr {
     pub unsafe fn as_mut_ptr<T>(self) -> *mut T {
         self.as_usize() as *mut T
     }
+
+    pub unsafe fn as_ptr<T>(self) -> *const T {
+        self.as_usize() as *const T
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]