Implement alignment for allocator

Author: wangrunji0408

src/alloc.rs

@@ -52,16 +52,33 @@ pub fn exit_boot_services() {
 /// Only valid for as long as the UEFI boot services are available.
 pub struct Allocator;
 
+#[allow(clippy::cast_ptr_alignment)]
 unsafe impl GlobalAlloc for Allocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         let mem_ty = MemoryType::LOADER_DATA;
         let size = layout.size();
         let align = layout.align();
 
-        // TODO: add support for other alignments.
         if align > 8 {
-            // Unsupported alignment for allocation, UEFI can only allocate 8-byte aligned addresses
-            ptr::null_mut()
+            // allocate more space for alignment
+            let ptr = if let Ok(ptr) = boot_services()
+                .as_ref()
+                .allocate_pool(mem_ty, size + align)
+                .warning_as_error()
+            {
+                ptr
+            } else {
+                return ptr::null_mut();
+            };
+            // calculate align offset
+            let mut offset = ptr.align_offset(align);
+            if offset == 0 {
+                offset = align;
+            }
+            let return_ptr = ptr.add(offset);
+            // store allocated pointer before the struct
+            (return_ptr as *mut *mut u8).sub(1).write(ptr);
+            return_ptr
         } else {
             boot_services()
                 .as_ref()
@@ -71,7 +88,10 @@ unsafe impl GlobalAlloc for Allocator {
         }
     }
 
-    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
+    unsafe fn dealloc(&self, mut ptr: *mut u8, layout: Layout) {
+        if layout.align() > 8 {
+            ptr = (ptr as *const *mut u8).sub(1).read();
+        }
         boot_services()
             .as_ref()
             .free_pool(ptr)

uefi-test-runner/src/boot/memory.rs

@@ -9,6 +9,7 @@ pub fn test(bt: &BootServices) {
 
     allocate_pages(bt);
     vec_alloc();
+    alloc_alignment();
     memmove(bt);
 
     memory_map(bt);
@@ -47,6 +48,17 @@ fn vec_alloc() {
     assert_eq!(values[..], [-5, 0, 4, 16, 23], "Failed to sort vector");
 }
 
+// Simple test to ensure our custom allocator works with correct alignment.
+fn alloc_alignment() {
+    info!("Allocating a structure with alignment to 0x100");
+
+    #[repr(align(0x100))]
+    struct Block([u8; 0x100]);
+
+    let value = vec![Block([1; 0x100])];
+    assert_eq!(value.as_ptr() as usize % 0x100, 0, "Wrong alignment");
+}
+
 // Test that the `memmove` / `memset` functions work.
 fn memmove(bt: &BootServices) {
     info!("Testing the `memmove` / `memset` functions");