1 files changed, 196 insertions, 0 deletions

diff --git a/src/memory.rs b/src/memory.rs
new file mode 100644
index 0000000..d223f69
--- /dev/null
+++ b/src/memory.rs
@@ -0,0 +1,196 @@
+use crate::uefi::*;
+use core::alloc::{GlobalAlloc, Layout};
+use core::mem;
+use core::ptr;
+
+const MAX_MAP_ENTRIES: usize = 512;
+
+static mut MEMORY_MAP: [MemoryDescriptor; MAX_MAP_ENTRIES] = unsafe { mem::zeroed() };
+static mut MAP_ENTRIES: usize = 0;
+
+/// The map key returned by GetMemoryMap, needed for ExitBootServices.
+static mut MAP_KEY: usize = 0;
+
+/// Static working buffer for GetMemoryMap (64 KB, way more than enough).
+/// Lives in .bss; no pool allocation, so no map-key staleness.
+static mut MAP_WORK_BUF: [u8; 65536] = [0u8; 65536];
+
+// ---------------------------------------------------------------------------
+// Bump allocator — no deallocation, just pointer advancement.
+// Used as Rust's #[global_allocator] once the heap region is set up.
+// ---------------------------------------------------------------------------
+struct BumpAllocator;
+
+unsafe impl Sync for BumpAllocator {}
+
+unsafe impl GlobalAlloc for BumpAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let align = layout.align().max(16);
+        let size = layout.size();
+        unsafe {
+            let a = (HEAP_CURSOR + align - 1) & !(align - 1);
+            let e = a + size;
+            if e > HEAP_END || e < a {
+                return ptr::null_mut();
+            }
+            HEAP_CURSOR = e;
+            a as *mut u8
+        }
+    }
+
+    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
+        // Bump allocator — cannot free individual allocations.
+        // Memory is reclaimed on reboot.  Good enough for now.
+        unsafe {
+            core::ptr::read_volatile(&_ptr);
+        }
+        let _ = _layout;
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: BumpAllocator = BumpAllocator;
+
+/// Heap bump allocator bounds (`static mut` accessed only inside
+/// `GlobalAlloc` impl, which is already `unsafe`).
+static mut HEAP_START: usize = 0;
+static mut HEAP_END: usize = 0;
+static mut HEAP_CURSOR: usize = 0;
+
+/// Capture the UEFI memory map into the static buffer.  Uses a pre-allocated
+/// .bss working buffer so no pool allocation (and thus no memory-map change)
+/// occurs between GetMemoryMap and ExitBootServices.
+pub fn capture_memory_map(bs: &BootServices) -> Result<usize, Status> {
+    let mut map_key: usize = 0;
+    let mut desc_size: usize = 0;
+    let mut desc_ver: u32 = 0;
+
+    // Probe with null to get required size + descriptor size
+    let mut buf_size: usize = 0;
+    let status = unsafe {
+        (bs.get_memory_map)(
+            &mut buf_size,
+            ptr::null_mut(),
+            &mut map_key,
+            &mut desc_size,
+            &mut desc_ver,
+        )
+    };
+    if status != EFI_BUFFER_TOO_SMALL {
+        // Fallback: allocate a sensible size
+        buf_size = 64 * 1024;
+    }
+    if desc_size == 0 {
+        desc_size = mem::size_of::<MemoryDescriptor>();
+    }
+
+    // Round up to desc_size alignment
+    let buf_size = (buf_size + desc_size - 1) / desc_size * desc_size;
+
+    // If our static buffer is too small, fall back to pool (rare)
+    if buf_size > unsafe { MAP_WORK_BUF.len() } {
+        return Err(EFI_BUFFER_TOO_SMALL);
+    }
+
+    let mut actual_size = buf_size;
+    let status = unsafe {
+        (bs.get_memory_map)(
+            &mut actual_size,
+            &raw mut MAP_WORK_BUF[0] as *mut MemoryDescriptor,
+            &mut map_key,
+            &mut desc_size,
+            &mut desc_ver,
+        )
+    };
+
+    if status == EFI_BUFFER_TOO_SMALL {
+        // The map grew between the probe and the real call.  Try a larger
+        // static buffer (can't grow here without pool, so return error to
+        // let the caller retry if they can allocate pool).
+        // For our 64 KB buffer this should never happen on OVMF.
+        return Err(EFI_BUFFER_TOO_SMALL);
+    }
+    if status != EFI_SUCCESS {
+        return Err(status);
+    }
+
+    // Copy descriptors into our indexed array
+    let count = actual_size / desc_size;
+    let count = count.min(MAX_MAP_ENTRIES);
+    unsafe {
+        for i in 0..count {
+            let p = (&raw const MAP_WORK_BUF[..]) as *const u8;
+            let off = i * desc_size;
+            let src = p.add(off) as *const MemoryDescriptor;
+            MEMORY_MAP[i] = ptr::read(src);
+        }
+        MAP_ENTRIES = count;
+        MAP_KEY = map_key;
+    }
+
+    // Do NOT free anything — the static buffer stays.  Map key is still
+    // valid because no pool allocation/free happened after GetMemoryMap.
+
+    init_bump_allocator();
+    Ok(map_key)
+}
+
+/// Retry ExitBootServices with a fresh memory map.
+pub fn refresh_and_exit(bs: &BootServices, image_handle: Handle) -> Result<usize, Status> {
+    let key = capture_memory_map(bs)?;
+    let status = unsafe { (bs.exit_boot_services)(image_handle, key) };
+    if status == EFI_SUCCESS {
+        Ok(key)
+    } else {
+        Err(status)
+    }
+}
+
+fn init_bump_allocator() {
+    unsafe {
+        let mut best_start = 0usize;
+        let mut best_size = 0usize;
+        for i in 0..MAP_ENTRIES {
+            let d = &MEMORY_MAP[i];
+            if d.kind == MEMORY_TYPE_CONVENTIONAL {
+                let sz = d.number_of_pages as usize * 4096;
+                if sz > best_size {
+                    best_start = d.physical_start as usize;
+                    best_size = sz;
+                }
+            }
+        }
+        if best_size > 2 * 1024 * 1024 {
+            let s = best_start + 2 * 1024 * 1024;
+            let e = best_start + best_size;
+            if s < e {
+                HEAP_START = s;
+                HEAP_END = e;
+                HEAP_CURSOR = s;
+            }
+        }
+    }
+}
+
+#[allow(dead_code)]
+pub fn allocate(size: usize) -> *mut u8 {
+    unsafe {
+        let a = (HEAP_CURSOR + 15) & !15;
+        let e = a + size;
+        if e > HEAP_END || e < a {
+            return ptr::null_mut();
+        }
+        HEAP_CURSOR = e;
+        a as *mut u8
+    }
+}
+
+#[allow(dead_code)]
+pub fn get_map() -> (&'static [MemoryDescriptor], usize) {
+    unsafe {
+        (
+            core::slice::from_raw_parts(&MEMORY_MAP[0], MAP_ENTRIES),
+            MAP_KEY,
+        )
+    }
+}