Merge pull request #8 from csssuf/allocate-pool

Add allocate_pool function

Author: csssuf

src/base.rs

@@ -153,7 +153,7 @@ fn status_str() {
 }
 
 /// Type for EFI_MEMORY_TYPE
-#[derive(PartialEq, PartialOrd, Debug)]
+#[derive(PartialEq, PartialOrd, Debug, Clone, Copy)]
 #[repr(C)]
 pub enum MemoryType {
     Reserved = 0,

src/bootservices.rs

@@ -5,7 +5,7 @@ use void::{NotYetDef, CVoid};
 use base::{Event, Handle, Handles, MemoryType, Status};
 use event::{EventType, EventNotify, TimerDelay};
 use task::TPL;
-use protocol::{DevicePathProtocol, Protocol};
+use protocol::{DevicePathProtocol, Protocol, get_current_image};
 use guid;
 use table;
 
@@ -68,6 +68,17 @@ pub struct BootServices {
 }
 
 impl BootServices {
+    /// Allocate `size` bytes of memory using type `T`.
+    pub fn allocate_pool<T>(&self, size: usize) -> Result<*mut T, Status> {
+        let mut ptr: *mut u8 = 0 as *mut u8;
+
+        let result = unsafe { (self.allocate_pool)(get_current_image().image_data_type, size, &mut ptr) };
+        if result != Status::Success {
+            return Err(result);
+        }
+        Ok(ptr as *mut T)
+    }
+
     pub fn free_pool<T>(&self, p: *const T) {
         unsafe {
             (self.free_pool)(p as *mut CVoid);

src/protocol/mod.rs

@@ -1,4 +1,4 @@
-use base::{Handle, MemoryType};
+use base::{Handle, MemoryType, Status};
 use guid::Guid;
 use void::NotYetDef;
 
@@ -13,6 +13,8 @@ pub trait Protocol {
 /// GUID for UEFI protocol for loaded images
 pub static EFI_LOADED_IMAGE_PROTOCOL_GUID: Guid = Guid(0x5B1B31A1, 0x9562, 0x11d2, [0x8E,0x3F,0x00,0xA0,0xC9,0x69,0x72,0x3B]);
 
+static mut THIS_LOADED_IMAGE: *const LoadedImageProtocol = 0 as *const LoadedImageProtocol;
+
 #[derive(Debug)]
 #[repr(C)]
 pub struct LoadedImageProtocol {
@@ -27,7 +29,7 @@ pub struct LoadedImageProtocol {
     pub image_base: usize,
     pub image_size: u64,
     image_code_type: MemoryType,
-    image_data_type: MemoryType,
+    pub image_data_type: MemoryType,
 
     //unload: unsafe extern "win64" fn(handle: ::base::Handle),
     unload: *const NotYetDef,
@@ -39,3 +41,22 @@ impl Protocol for LoadedImageProtocol {
     }
 }
 
+pub fn set_current_image(handle: Handle) -> Result<&'static LoadedImageProtocol, Status> {
+    let st = ::get_system_table();
+
+    let loaded_image_proto: Result<&'static LoadedImageProtocol, Status> = st.boot_services().handle_protocol(handle);
+    if let Ok(image) = loaded_image_proto {
+        unsafe {
+            THIS_LOADED_IMAGE = image;
+        }
+    }
+
+    loaded_image_proto
+}
+
+pub fn get_current_image() -> &'static LoadedImageProtocol {
+    unsafe {
+        &*THIS_LOADED_IMAGE
+    }
+}
+