Got a kinda crappy post-boot services allocator working, hopefully.

Cargo.toml

@@ -3,7 +3,9 @@ name = "bootproof"
 version = "0.1.0"
 authors = ["James Martin <james@jtmar.me>"]
 edition = "2018"
+repository = "https://github.com/jamestmartin/bootproof"
 license = "GPL-3.0+"
+publish = false
 
 [dependencies]
 compiler_builtins = { git = "https://github.com/rust-lang/compiler-builtins" }

src/allocator.rs

@@ -1,28 +1,35 @@
+pub mod standard;
+pub mod uefi;
+
 use alloc::alloc::GlobalAlloc;
 use core::alloc::Layout;
 
-pub enum Allocator {
+pub enum GlobalAllocator {
     None,
+    Standard(standard::StandardAllocator),
+    Uefi(uefi::UefiAllocator),
 }
 
-unsafe impl GlobalAlloc for Allocator {
+unsafe impl GlobalAlloc for GlobalAllocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         match self {
-            Allocator::None => panic!("No allocator available!")
+            GlobalAllocator::None => panic!("No allocator available!"),
+            GlobalAllocator::Uefi(alloc) => alloc.alloc(layout),
+            GlobalAllocator::Standard(alloc) => alloc.alloc(layout),
         }
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         match self {
-            Allocator::None => {
-                panic!("No allocator available!");
-            }
+            GlobalAllocator::None => panic!("No allocator available!"),
+            GlobalAllocator::Uefi(alloc) => alloc.dealloc(ptr, layout),
+            GlobalAllocator::Standard(alloc) => alloc.dealloc(ptr, layout),
         }
     }
 }
 
 #[global_allocator]
-pub static mut ALLOCATOR: Allocator = Allocator::None;
+pub static mut ALLOCATOR: GlobalAllocator = GlobalAllocator::None;
 
 #[alloc_error_handler]
 fn handle_error(layout: Layout) -> ! {

src/allocator/standard.rs

@@ -0,0 +1,142 @@
+use alloc::alloc::GlobalAlloc;
+use alloc::vec::Vec;
+use core::alloc::Layout;
+use core::cell::UnsafeCell;
+use uefi::table::boot::MemoryMapIter;
+
+// TODO: Support granularity better than pages.
+// TODO: Use an allocation algorithm that isn't absolute garbage!!
+
+/// **This allocator only supports page-level granularity.**
+/// Be careful not to use it for small allocations.
+pub struct StandardAllocator {
+    pages: UnsafeCell<Vec<u8>>,
+}
+
+const PAGE_SIZE: usize = 4096;
+
+pub fn get_bit(bytes: &[u8], index: usize) -> bool {
+    let (byte, bit) = num_integer::div_rem(index, 8);
+    let mask = 0b10000000u8 >> bit;
+    bytes[byte] & mask > 0
+}
+
+pub fn set_bit(bytes: &mut [u8], index: usize, value: bool) {
+    let (byte, bit) = num_integer::div_rem(index, 8);
+    let mask = 0b10000000u8 >> bit;
+    if value {
+        bytes[byte] |= mask;
+    } else {
+        bytes[byte] &= !mask;
+    }
+}
+
+impl StandardAllocator {
+    /// Allocates a new allocator data structure sufficient
+    /// to use the physical memory provided in the memory map.
+    /// With what allocator does the allocator get allocated? The UEFI allocator.
+    /// This does *not* pre-populate the allocator with usage data;
+    /// by default, it will behave as though every page were allocated.
+    /// Use `populate` to fill the allocator with actual data
+    /// using the map that UEFI provides when you exit boot services.
+    pub fn new(mmap: &mut MemoryMapIter) -> StandardAllocator {
+        // Try to find the largest physical address
+        // and create a bitmap allowing the allocation of that much memory.
+        let greatest_physical_page =
+            mmap.map(|d| num_integer::div_ceil(d.phys_start as usize, PAGE_SIZE) + d.page_count as usize).max()
+                .unwrap();
+
+        let mut pages = Vec::with_capacity(num_integer::div_ceil(greatest_physical_page, 8));
+        // I can fit up to 8 pages in a byte in my bitmap.
+        pages.resize(num_integer::div_ceil(greatest_physical_page, 8), 0xFF);
+
+        StandardAllocator {
+            pages: UnsafeCell::new(pages)
+        }
+    }
+
+    pub fn populate(&mut self, mmap: &mut MemoryMapIter) {
+        let self_pages = unsafe { &mut *self.pages.get() };
+        // Mark all unsable memory as free for allocations.
+        for entry in mmap {
+            use uefi::table::boot::MemoryType;
+            if entry.ty == MemoryType::BOOT_SERVICES_CODE
+                || entry.ty == MemoryType::BOOT_SERVICES_DATA
+                || entry.ty == MemoryType::CONVENTIONAL {
+                let base = entry.phys_start as usize / PAGE_SIZE;
+                for offset in 0..entry.page_count as usize {
+                    set_bit(self_pages.as_mut_slice(), base + offset, false);
+                }
+            }
+        }
+
+        // Even if the zero address is valid memory, we *definitely* don't want to allocate it.
+        set_bit(self_pages.as_mut_slice(), 0, true);
+    }
+
+    pub fn free(&self) -> usize {
+        let self_pages = unsafe { &mut *self.pages.get() };
+        let mut free = 0;
+        for page in 0..self_pages.len() * 8 {
+            if !get_bit(self_pages.as_slice(), page) {
+                free += 1;
+            }
+        }
+        free
+    }
+
+    // Not very accurate because this includes a lot of reserved/unusable memory.
+    pub fn used(&self) -> usize {
+        let free = self.free();
+        // This line of code crashes QEMU for inexplicable reasons.
+        // I tried to figure out why and failed.
+        let self_pages = unsafe { &*self.pages.get() };
+        free - self_pages.as_slice().len() * 8
+    }
+}
+
+unsafe impl GlobalAlloc for StandardAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let min_pages = num_integer::div_ceil(layout.size(), PAGE_SIZE);
+        let mut pages = None;
+        let self_pages = &mut *self.pages.get();
+        for i in 0..(self_pages.len() * 8) {
+            if !get_bit(&self_pages.as_slice(), i) {
+                pages = match pages {
+                    Some((begin, size)) => Some((begin, size + 1)),
+                    None => Some((i, 1)),
+                };
+
+                if pages.unwrap().1 >= min_pages {
+                    break;
+                }
+            } else {
+                pages = None;
+            }
+        }
+
+        if pages.is_none() || pages.unwrap().1 < min_pages {
+            panic!("Not enough contiguous memory!");
+        }
+
+        let (begin, size) = pages.unwrap();
+
+        let address = begin * PAGE_SIZE;
+
+        for offset in 0..size {
+            set_bit(self_pages.as_mut_slice(), begin + offset, true);
+        }
+
+        address as *mut u8
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        let begin = ptr as usize / PAGE_SIZE;
+        let size = num_integer::div_ceil(layout.size(), PAGE_SIZE);
+        let self_pages = &mut *self.pages.get();
+
+        for offset in 0..size {
+            set_bit(self_pages.as_mut_slice(), begin + offset, false);
+        }
+    }
+}

src/allocator/uefi.rs

@@ -0,0 +1,39 @@
+use alloc::alloc::GlobalAlloc;
+use core::alloc::Layout;
+use uefi::prelude::{Boot, SystemTable};
+use uefi::table::boot::{AllocateType, MemoryType};
+
+/// **This allocator only supports page-level granularity.**
+/// Be careful not to use it for small allocations.
+/// In particular, it should only be used to allocate data structures
+/// for the purpose of setting up another, better allocator.
+pub struct UefiAllocator {
+    // We must directly store an owned system table because:
+    //   1. It is impossible to take ownership of ST boot services, and
+    //   2. we cannot store a reference to *anything* here
+    //      because the global allocator has to be static even if we know it's really not.
+    st: SystemTable<Boot>,
+}
+
+impl UefiAllocator {
+    pub fn new(st: SystemTable<Boot>) -> UefiAllocator {
+        UefiAllocator { st: st }
+    }
+}
+
+const PAGE_SIZE: usize = 4096;
+
+unsafe impl GlobalAlloc for UefiAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        self.st.boot_services()
+            .allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA,
+                            num_integer::div_ceil(layout.size(), PAGE_SIZE))
+            .expect("Failed to allocate memory!").unwrap()
+            as *mut u8
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        self.st.boot_services().free_pages(ptr as u64, num_integer::div_ceil(layout.size(), PAGE_SIZE))
+            .expect("Failed to free memory!").unwrap();
+    }
+}

src/graphics/tty.rs

@@ -19,27 +19,27 @@ pub static mut STDERR: Option<SerialTty> = None;
 macro_rules! print {
     // These additional single-argument cases are necessary because `format!` requires allocation,
     // which I don't necessarily *have* (not to mention it's inefficient).
-    ($s:expr) => {
+    ($s:expr) => {{
         let mut tty;
         unsafe {
             tty = crate::graphics::tty::STDOUT.clone().unwrap();
         }
         tty.puts($s);
         tty.flush();
-    };
-    ($($arg:expr),*) => {
+    }};
+    ($($arg:expr),*) => {{
         let mut tty;
         unsafe {
             tty = crate::graphics::tty::STDOUT.clone().unwrap();
         }
         tty.puts(&alloc::format!($($arg),*));
         tty.flush();
-    }
+    }}
 }
 
 #[macro_export]
 macro_rules! println {
-    ($s:expr) => {
+    ($s:expr) => {{
         let mut tty;
         unsafe {
             tty = crate::graphics::tty::STDOUT.clone().unwrap();
@@ -47,8 +47,8 @@ macro_rules! println {
         tty.puts($s);
         tty.putc('\n');
         tty.flush();
-    };
-    ($($arg:expr),*) => {
+    }};
+    ($($arg:expr),*) => {{
         let mut tty;
         unsafe {
             tty = crate::graphics::tty::STDOUT.clone().unwrap();
@@ -56,32 +56,32 @@ macro_rules! println {
         tty.puts(&alloc::format!($($arg),*));
         tty.putc('\n');
         tty.flush();
-    }
+    }}
 }
 
 #[macro_export]
 macro_rules! eprint {
-    ($s:expr) => {
+    ($s:expr) => {{
         let mut tty;
         unsafe {
             tty = crate::graphics::tty::STDERR.clone().unwrap();
         }
         tty.puts($s);
         tty.flush();
-    };
-    ($($arg:expr),*) => {
+    }};
+    ($($arg:expr),*) => {{
         let mut tty;
         unsafe {
             tty = crate::graphics::tty::STDERR.clone().unwrap();
         }
         tty.puts(&alloc::format!($($arg),*));
         tty.flush();
-    }
+    }}
 }
 
 #[macro_export]
 macro_rules! eprintln {
-    ($s:expr) => {
+    ($s:expr) => {{
         let mut tty;
         unsafe {
             tty = crate::graphics::tty::STDERR.clone().unwrap();
@@ -89,8 +89,8 @@ macro_rules! eprintln {
         tty.puts($s);
         tty.putc('\n');
         tty.flush();
-    };
-    ($($arg:expr),*) => {
+    }};
+    ($($arg:expr),*) => {{
         let mut tty;
         unsafe {
             tty = crate::graphics::tty::STDERR.clone().unwrap();
@@ -98,15 +98,15 @@ macro_rules! eprintln {
         tty.puts(&alloc::format!($($arg),*));
         tty.putc('\n');
         tty.flush();
-    }
+    }}
 }
 
 #[macro_export]
 macro_rules! panic {
-    ($($arg:expr),*) => {
+    ($($arg:expr),*) => {{
         crate::eprintln!($($arg),*);
-        crate::arch::x86_64::halt();
-    }
+        crate::arch::x86_64::halt()
+    }}
 }
 
 #[panic_handler]

src/main.rs

@@ -8,27 +8,48 @@
 extern crate alloc;
 
 mod allocator;
+mod arch;
 #[macro_use]
 mod graphics;
-mod arch;
 
-use core::mem;
-use core::slice;
-use crate::allocator::{Allocator, ALLOCATOR};
+use alloc::vec::Vec;
 use crate::graphics::tty::{Tty, STDOUT, STDERR};
 use crate::graphics::tty::serial::SerialTty;
 use uefi::prelude::*;
-use uefi::table::boot::{AllocateType, MemoryDescriptor, MemoryType};
 
 #[entry]
 fn efi_main(handle: Handle, st_boot: SystemTable<Boot>) -> Status {
-    // Although serial ports don't physically exist on modern devices,
-    // they're still supposed by emulators (for QEMU you can set `-serial stdio`),
-    // and they're extremely useful for debugging
-    // because they don't require any setup and are trivial to use.
+    use crate::allocator::{ALLOCATOR, GlobalAllocator};
     unsafe {
-        STDOUT = Some(SerialTty::new(0x3F8));
-        STDERR = Some(SerialTty::new(0x3F8));
+        // Generally speaking, we want to depend on UEFI as little as possible,
+        // so the need for a UEFI allocator may seem a bit strange.
+        // However, there's this awkward time during booting when we need
+        // to allocate space for our "real" allocator's data structures and the UEFI memory maps,
+        // the result being that we need an allocator for our allocator.
+        // In theory there are probably ways to get around it, but why bother?
+        // Just taking advantage of the UEFI allocator briefly is a lot easier.
+        // (This also lets us use `println!` prior to our main allocator being set up.)
+        use crate::allocator::uefi::UefiAllocator;
+        // ABSOLUTELY DO NOT FORGET TO DISABLE THIS AFTER LEAVING UEFI BOOT SERVICES.
+        // ALL ALLOCATIONS MUST BE STATIC OR BE FREED BEFORE BOOT SERVICES EXITS.
+        // If the're not, Rust still try to free UEFI-allocated data using the new allocator,
+        // which is undefined behavior.
+        ALLOCATOR = GlobalAllocator::Uefi(UefiAllocator::new(st_boot.unsafe_clone()));
+
+        // Although serial ports don't physically exist on modern devices,
+        // they're still supposed by emulators (for QEMU you can set `-serial stdio`),
+        // and they're extremely useful for debugging
+        // because they don't require any setup and are trivial to use.
+        STDOUT = Some({
+            let mut stdout = SerialTty::new(0x3F8);
+            stdout.clear();
+            stdout
+        });
+        STDERR = Some({
+            let mut stderr = SerialTty::new(0x3F8);
+            stderr.clear();
+            stderr
+        });
     }
 
     // Our first task is to exit the UEFI boot services.
@@ -40,25 +61,35 @@ fn efi_main(handle: Handle, st_boot: SystemTable<Boot>) -> Status {
     // which describes where a bunch of important stuff lies in memory
     // (e.g. memory-mapped devices and the ACPI tables),
     // and what memory is available for us to use safely.
-    let (st, mmap) = {
-        // We must provide a buffer (mmap_buf) for UEFI to write the memory map to.
+
+    // We can't let the memory map be de-allocated because it is allocated using the UEFI allocator,
+    // but would end up being freed using the standard allocator, which is undefined behavior.
+    // We must provide a buffer (mmap_buf) for UEFI to write the memory map to.
+    let mut mmap_buf = Vec::new();
+    let (_mmap, st) = {
         let bs = st_boot.boot_services();
-
         // More allocations can happen between the allocation of the buffer and the buffer being filled,
-        // so add space for 8 more memory descriptors (an arbitrary number) just to make sure there's enough.
-        let mmap_buf_size = bs.memory_map_size() + 8 * mem::size_of::<MemoryDescriptor>();
+        // so add space for 32 more memory descriptors (an arbitrary number) just to make sure there's enough.
+        mmap_buf.resize(bs.memory_map_size() + 1024, 0);
 
-        // In the memory map, the OS's own code and data is included as LOADER_CODE and LOADER_DATA.
-        // This is good to know so we don't accidentally write over it!
-        // We allocate the mmap_buf as LOADER_DATA for this purpose.
-        let mmap_buf = bs.allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, mmap_buf_size)
-            .expect_success("Failed to allocate memory for UEFI memory map buffer.");
-
-        // allocate_pages returns a u64 address, but we need a `&mut [u8]`.
-        let mmap_buf_slice = unsafe { slice::from_raw_parts_mut(mmap_buf as *mut u8, mmap_buf_size) };
+        // HACK: I hate having to use the UEFI allocator just to set up another allocator!
+        //   There's got to be a better way.
+        use crate::allocator::standard::StandardAllocator;
+        let mut allocator;
+        {
+            let mut mmap = bs.memory_map(mmap_buf.as_mut_slice())
+                .expect_success("Failed to exit the UEFI boot services.").1;
+            allocator = StandardAllocator::new(&mut mmap);
+        }
 
         // Finally, we actually exit the UEFI boot services.
-        st_boot.exit_boot_services(handle, mmap_buf_slice).expect_success("Failed to exit the UEFI boot services.")
+        let (st, mut mmap) = st_boot.exit_boot_services(handle, mmap_buf.as_mut_slice())
+            .expect_success("Failed to exit the UEFI boot services.");
+
+        allocator.populate(&mut mmap);
+        unsafe { ALLOCATOR = GlobalAllocator::Standard(allocator); }
+
+        (mmap, st)
     };
 
     // Set up the stuff I need to handle interrupts, which is necessary to write drivers for most devices.
@@ -76,10 +107,10 @@ fn efi_main(handle: Handle, st_boot: SystemTable<Boot>) -> Status {
     // Everything up to this point has been setting up the CPU state, drivers, etc.
     // Now we begin running actual programs
     // (or in this case, since we don't support actual programs yet, whatever debug stuff I want to run).
-    main(st, mmap)
+    main(st)
 }
 
-fn main(_st: SystemTable<uefi::table::Runtime>, _mmap: uefi::table::boot::MemoryMapIter) -> ! {
+fn main(_st: SystemTable<uefi::table::Runtime>) -> ! {
     // Put whatever code you want for debugging/testing purposes here...
     arch::x86_64::breakpoint();