From 8904d88d4865f7a28ec2de9678a589eaffa58bd2 Mon Sep 17 00:00:00 2001
From: James Forcier <james.forcier@coreos.com>
Date: Thu, 8 Jun 2017 14:25:32 -0700
Subject: [PATCH] protocol: add {set,get}_current_image

Some EFI functions (namely, AllocatePool) need different arguments
depending on the currently loaded image. Exposing this API detail to
Rust consumers is unnecessary, so keep a LoadedImageProtocol struct
around representing the currently loaded image.
---
 src/base.rs         |  2 +-
 src/protocol/mod.rs | 25 +++++++++++++++++++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/base.rs b/src/base.rs
index ebb8cc2..4ce91a7 100644
--- a/src/base.rs
+++ b/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,
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 6894c42..01a18b9 100644
--- a/src/protocol/mod.rs
+++ b/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
+    }
+}
+