roxy_loader/
allocation_info.rs

1use core::ptr::write_bytes;
2use uefi::boot::{AllocateType, MemoryType, allocate_pages};
3
4pub const PAGE_SIZE: u64 = 4096;
5
6#[must_use]
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub struct AllocationInfo {
9    pub page_base: u64,
10    pub page_offset: u64,
11    pub alloc_size: u64,
12    pub pages: usize,
13}
14
15impl AllocationInfo {
16    #[must_use]
17    pub fn from_region(addr: u64, size: u64) -> Self {
18        let page_base = addr & !(PAGE_SIZE - 1);
19        let page_offset = addr - page_base;
20        let alloc_size = page_offset + size;
21        let pages = alloc_size.div_ceil(PAGE_SIZE) as usize;
22
23        Self {
24            page_base,
25            page_offset,
26            alloc_size,
27            pages,
28        }
29    }
30
31    pub fn alloc(self) -> uefi::Result<*mut u8> {
32        Ok(allocate_pages(
33            AllocateType::Address(self.page_base),
34            MemoryType::LOADER_DATA,
35            self.pages,
36        )?
37        .as_ptr())
38    }
39
40    pub fn alloc_zeroed(self) -> uefi::Result<*mut u8> {
41        let ptr = self.alloc()?;
42
43        // The loader expects a fully zeroed backing region before segment contents are copied in.
44        unsafe {
45            write_bytes(ptr, 0, self.pages * PAGE_SIZE as usize);
46        }
47
48        Ok(ptr)
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::{AllocationInfo, PAGE_SIZE};
55
56    #[test]
57    fn aligned_single_page_allocation() {
58        assert_eq!(
59            AllocationInfo::from_region(0x2000, 1),
60            AllocationInfo {
61                page_base: 0x2000,
62                page_offset: 0,
63                alloc_size: 1,
64                pages: 1,
65            }
66        );
67    }
68
69    #[test]
70    fn unaligned_address_accounts_for_offset() {
71        assert_eq!(
72            AllocationInfo::from_region(0x2123, 0x100),
73            AllocationInfo {
74                page_base: 0x2000,
75                page_offset: 0x123,
76                alloc_size: 0x223,
77                pages: 1,
78            }
79        );
80    }
81
82    #[test]
83    fn crossing_page_boundary_allocates_two_pages() {
84        assert_eq!(
85            AllocationInfo::from_region(PAGE_SIZE - 1, 2),
86            AllocationInfo {
87                page_base: 0,
88                page_offset: PAGE_SIZE - 1,
89                alloc_size: PAGE_SIZE + 1,
90                pages: 2,
91            }
92        );
93    }
94
95    #[test]
96    fn exact_multiple_of_page_size_does_not_round_up_extra() {
97        assert_eq!(
98            AllocationInfo::from_region(0x4000, PAGE_SIZE * 2),
99            AllocationInfo {
100                page_base: 0x4000,
101                page_offset: 0,
102                alloc_size: PAGE_SIZE * 2,
103                pages: 2,
104            }
105        );
106    }
107
108    #[test]
109    fn zero_sized_region_allocates_zero_pages() {
110        assert_eq!(
111            AllocationInfo::from_region(0x4000, 0),
112            AllocationInfo {
113                page_base: 0x4000,
114                page_offset: 0,
115                alloc_size: 0,
116                pages: 0,
117            }
118        );
119    }
120}
roxy_loader/allocation_info.rs

roxy_loader/
allocation_info.rs
