Properly handle allocation failure

uefi-exts/Cargo.toml

@@ -7,4 +7,3 @@ edition = "2018"
 
 [dependencies]
 uefi = { path = ".." }
-uefi-services = { path = "../uefi-services" }

uefi-exts/src/file.rs

@@ -1,5 +1,6 @@
-use alloc::{alloc::alloc, boxed::Box};
-use core::{alloc::Layout, mem, slice};
+use super::allocate_buffer;
+use alloc::{alloc::Layout, boxed::Box};
+use core::mem;
 use uefi::prelude::*;
 use uefi::proto::media::file::{File, FileProtocolInfo};
 use uefi::Result;
@@ -19,15 +20,14 @@ pub trait FileExt: File {
             (_, Some(size)) => size,
         };
 
-        // These unsafe alloc APIs make sure our buffer is correctly aligned. We
-        // round up a size must always be a multiple of alignment. We turn the
-        // pointer into a Box<[u8]>, so it's always freed on error.
+        // We add trailing padding because the size of a rust structure must
+        // always be a multiple of alignment.
         let layout = Layout::from_size_align(size, Info::alignment())
             .unwrap()
             .pad_to_align()
             .unwrap();
-        let buffer_start = unsafe { alloc(layout) };
-        let mut buffer = unsafe { Box::from_raw(slice::from_raw_parts_mut(buffer_start, size)) };
+        let mut buffer = allocate_buffer(layout);
+        let buffer_start = buffer.as_ptr();
 
         let info = self
             .get_info(&mut buffer)

uefi-exts/src/lib.rs

@@ -4,7 +4,7 @@
 //! which add utility functions to various UEFI objects.
 
 #![no_std]
-#![feature(alloc, alloc_layout_extra)]
+#![feature(alloc, alloc_layout_extra, allocator_api)]
 
 extern crate alloc;
 
@@ -13,3 +13,24 @@ mod file;
 
 pub use self::boot::BootServicesExt;
 pub use self::file::FileExt;
+
+use alloc::{
+    alloc::{handle_alloc_error, Alloc, Global, Layout},
+    boxed::Box,
+};
+use core::slice;
+
+/// Creates a boxed byte buffer using the standard allocator
+/// # Panics
+/// Calls handle_alloc_error if the layout has a size of zero or allocation fails.
+pub fn allocate_buffer(layout: Layout) -> Box<[u8]> {
+    if layout.size() == 0 {
+        handle_alloc_error(layout);
+    }
+    unsafe {
+        match Global.alloc(layout) {
+            Ok(mem) => Box::from_raw(slice::from_raw_parts_mut(mem.as_ptr(), layout.size())),
+            Err(_) => handle_alloc_error(layout),
+        }
+    }
+}