1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
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.div_ceil(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,
)
}
}
|
