From f2db3d581d0d4ccf75bd1662d56cea122889d840 Mon Sep 17 00:00:00 2001 From: James Martin Date: Sat, 18 Jul 2020 17:23:14 -0700 Subject: [PATCH] Got a kinda crappy post-boot services allocator working, hopefully. --- Cargo.toml | 2 + src/allocator.rs | 21 ++++-- src/allocator/standard.rs | 142 ++++++++++++++++++++++++++++++++++++++ src/allocator/uefi.rs | 39 +++++++++++ src/graphics/tty.rs | 38 +++++----- src/main.rs | 85 +++++++++++++++-------- 6 files changed, 274 insertions(+), 53 deletions(-) create mode 100644 src/allocator/standard.rs create mode 100644 src/allocator/uefi.rs diff --git a/Cargo.toml b/Cargo.toml index 95fbee9..f10cdd2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,9 @@ name = "bootproof" version = "0.1.0" authors = ["James Martin "] 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" } diff --git a/src/allocator.rs b/src/allocator.rs index a94b2c0..d2b55aa 100644 --- a/src/allocator.rs +++ b/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) -> ! { diff --git a/src/allocator/standard.rs b/src/allocator/standard.rs new file mode 100644 index 0000000..fdc8bab --- /dev/null +++ b/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>, +} + +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); + } + } +} diff --git a/src/allocator/uefi.rs b/src/allocator/uefi.rs new file mode 100644 index 0000000..5e42bce --- /dev/null +++ b/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, +} + +impl UefiAllocator { + pub fn new(st: SystemTable) -> 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(); + } +} diff --git a/src/graphics/tty.rs b/src/graphics/tty.rs index 2ce810f..aaa0c59 100644 --- a/src/graphics/tty.rs +++ b/src/graphics/tty.rs @@ -19,27 +19,27 @@ pub static mut STDERR: Option = 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] diff --git a/src/main.rs b/src/main.rs index 26356da..e542ad1 100644 --- a/src/main.rs +++ b/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) -> 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) -> 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::(); + // 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) -> 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, _mmap: uefi::table::boot::MemoryMapIter) -> ! { +fn main(_st: SystemTable) -> ! { // Put whatever code you want for debugging/testing purposes here... arch::x86_64::breakpoint();