refactor initrdmain

6 files changed, 38 insertions, 41 deletions

diff --git a/src/efi/initrdloadfile2.rs b/src/efi/initrdloadfile2.rs
index cbf53dd..a23498a 100644
--- a/src/efi/initrdloadfile2.rs
+++ b/src/efi/initrdloadfile2.rs
@@ -7,10 +7,11 @@ use crate::efi::devicepath::{DevicePath, VendorMedia};
 use crate::efi::devicepath::{DevicePathSubtype::*, DevicePathType::*};
 use crate::efi::{guid, install_protocol, new_handle, uninstall_protocol, Guid, Handle};
 use crate::efi::{status::*, Bool};
-
 use crate::efi::FileLoader;
 
+use alloc::boxed::Box;
 use core::mem::MaybeUninit;
+use core::pin::Pin;
 
 pub const EFI_LOAD_FILE2_PROTOCOL_GUID: Guid = guid!(
     0x4006c0c1,
@@ -25,34 +26,32 @@ type LoadFile =
 #[repr(C)]
 pub struct LoadFile2<'a> {
     load_file: LoadFile,
-    loader: &'a dyn FileLoader,
+    loader: Box<dyn FileLoader + 'a>,
 }
 
 pub struct InitrdLoadFile2<'a> {
-    load_file: &'a LoadFile2<'a>,
+    load_file: Pin<Box<LoadFile2<'a>>>,
     handle: Handle,
 }
 
-pub fn new<'a>(loader: &'a dyn FileLoader) -> InitrdLoadFile2<'a> {
-    let p = alloc::boxed::Box::new(LoadFile2 { load_file, loader });
+pub fn new<'a>(loader: Box<dyn FileLoader + 'a>) -> InitrdLoadFile2<'a> {
     let lf = InitrdLoadFile2 {
-        load_file: alloc::boxed::Box::leak(p),
+        load_file: Box::pin(LoadFile2 { load_file, loader }),
         handle: new_handle(),
     };
-    install_protocol(lf.handle, &EFI_LOAD_FILE2_PROTOCOL_GUID, lf.load_file);
+    install_protocol(lf.handle, &EFI_LOAD_FILE2_PROTOCOL_GUID, &*lf.load_file);
     install_protocol(lf.handle, &EFI_DEVICE_PATH_PROTOCOL_GUID, &INITRD_DEV_PATH);
     lf
 }
 
 impl<'a> Drop for InitrdLoadFile2<'a> {
     fn drop(&mut self) {
-        uninstall_protocol(self.handle, &EFI_LOAD_FILE2_PROTOCOL_GUID, self.load_file);
+        uninstall_protocol(self.handle, &EFI_LOAD_FILE2_PROTOCOL_GUID, &*self.load_file);
         uninstall_protocol(
             self.handle,
             &EFI_DEVICE_PATH_PROTOCOL_GUID,
             &INITRD_DEV_PATH,
         );
-        let _ = unsafe { alloc::boxed::Box::from_raw(self.load_file as *const _ as *mut ()) };
     }
 }
 
diff --git a/src/efi/loadedimage.rs b/src/efi/loadedimage.rs
index dec801b..bd262d4 100644
--- a/src/efi/loadedimage.rs
+++ b/src/efi/loadedimage.rs
@@ -3,13 +3,17 @@
 // Author: Ard Biesheuvel <ardb@google.com>
 
 use crate::efi::guid;
+use crate::efi::initrdloadfile2;
 use crate::efi::install_protocol;
 use crate::efi::new_handle;
 use crate::efi::uninstall_protocol;
 use crate::efi::{memorytype::*, status::*, systemtable::*, Guid, Handle};
+use crate::FileLoader;
 use crate::PeImage;
 
+use alloc::boxed::Box;
 use core::marker::PhantomData;
+use core::pin::Pin;
 use core::ptr;
 
 pub const EFI_LOADED_IMAGE_PROTOCOL_GUID: Guid = guid!(
@@ -52,7 +56,7 @@ pub struct LoadedImage {
 }
 
 pub struct LoadedImageData<'a> {
-    loaded_image: &'a mut LoadedImage,
+    loaded_image: Pin<Box<LoadedImage>>,
     image_handle: Handle,
     entrypoint: *const u8,
     randomized: bool,
@@ -68,7 +72,7 @@ impl<'a> LoadedImageData<'a> {
         randomized: bool,
     ) -> LoadedImageData<'a> {
         let handle: Handle = new_handle();
-        let p = alloc::boxed::Box::new(LoadedImage {
+        let p = LoadedImage {
             revision: EFI_LOADED_IMAGE_PROTOCOL_REVISION,
             parent_handle: 0,
             system_table: SystemTable::get(),
@@ -82,18 +86,18 @@ impl<'a> LoadedImageData<'a> {
             image_code_type: code_type,
             image_data_type: data_type,
             unload: unload,
-        });
+        };
         let li = LoadedImageData {
-            loaded_image: alloc::boxed::Box::leak(p),
+            loaded_image: Box::pin(p),
             image_handle: handle,
             entrypoint: pe_image.entry_point(),
             randomized: randomized,
             marker: PhantomData,
         };
 
-        install_protocol(handle, &EFI_LOADED_IMAGE_PROTOCOL_GUID, li.loaded_image);
+        install_protocol(handle, &EFI_LOADED_IMAGE_PROTOCOL_GUID, &*li.loaded_image);
         if randomized {
-            install_protocol(handle, &LINUX_EFI_LOADED_IMAGE_RAND_GUID, li.loaded_image);
+            install_protocol(handle, &LINUX_EFI_LOADED_IMAGE_RAND_GUID, &*li.loaded_image);
         }
         li
     }
@@ -104,21 +108,22 @@ impl Drop for LoadedImageData<'_> {
         uninstall_protocol(
             self.image_handle,
             &EFI_LOADED_IMAGE_PROTOCOL_GUID,
-            self.loaded_image,
+            &*self.loaded_image,
         );
         if self.randomized {
             uninstall_protocol(
                 self.image_handle,
                 &LINUX_EFI_LOADED_IMAGE_RAND_GUID,
-                self.loaded_image,
+                &*self.loaded_image,
             );
         }
-        let _ = unsafe { alloc::boxed::Box::from_raw(self.loaded_image as *const _ as *mut ()) };
     }
 }
 
 impl LoadedImageData<'_> {
-    pub fn start_image(&mut self) -> Status {
+    pub fn start_image<'a>(&mut self, initrd: Option<impl FileLoader + 'a>) -> Status {
+        let _initrd = initrd.map(|i| initrdloadfile2::new(Box::new(i) as Box<dyn FileLoader + 'a>));
+
         unsafe {
             start_image(
                 self.image_handle,
diff --git a/src/efi/mod.rs b/src/efi/mod.rs
index 6643edb..e9deab3 100644
--- a/src/efi/mod.rs
+++ b/src/efi/mod.rs
@@ -221,7 +221,7 @@ pub struct EfiContext {
     _bs: Pin<Box<BootServices>>,
     rt: Spinlock<&'static mut RuntimeServices>,
     sys_table: Spinlock<&'static mut SystemTable>,
-    mm: Spinlock<&'static mut (dyn MemoryMapper + Send + Sync)>,
+    mm: Spinlock<Box<dyn MemoryMapper + Send + Sync>>,
     rng: Box<dyn Random + Send + Sync>,
 
     _mem_attr_proto: Pin<Box<EfiMemoryAttribute>>,
@@ -240,7 +240,7 @@ pub fn efi_rtsdata_pool() -> &'static LockedHeap {
 }
 
 pub fn init(
-    mm: &'static mut (dyn MemoryMapper + Send + Sync),
+    mm: Box<dyn MemoryMapper + Send + Sync>,
     rng: Box<dyn Random + Send + Sync>,
 ) -> Result<&'static EfiContext, ()> {
     let ctx = EFI.get_or_try_init(|| {
@@ -317,11 +317,11 @@ impl EfiContext {
         let pe_ldr = PeLoader::new(loader)?;
 
         let align = pe_ldr.section_alignment() as u64;
-        let placement = if let Some(seed) = Some(0) {
-            //rng::new().get_random::<u32>() {
-            Placement::Random(seed, align)
+        let mut seed: [u8; 4] = [0; 4];
+        let (placement, randomized) = if self.rng.get_entropy(&mut seed, false) {
+            (Placement::Random(u32::from_le_bytes(seed), align), true)
         } else {
-            Placement::Aligned(align)
+            (Placement::Aligned(align), false)
         };
 
         let pe_image = pe_ldr.load(EfiLoaderCode, placement)?;
@@ -359,7 +359,7 @@ impl EfiContext {
             cmdline.as_slice(),
             EfiLoaderCode,
             EfiLoaderData,
-            false,
+            randomized,
         ))
     }
 
diff --git a/src/efi/rng.rs b/src/efi/rng.rs
index e6a5de1..7de8c2f 100644
--- a/src/efi/rng.rs
+++ b/src/efi/rng.rs
@@ -5,8 +5,7 @@
 
 use crate::efi::guid;
 use crate::efi::*;
-use crate::efi::{new_handle, status::*, Guid, Handle};
-use crate::rng;
+use crate::efi::{status::*, Guid};
 
 use core::slice;
 
diff --git a/src/main.rs b/src/main.rs
index e56f168..7c7f10f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -28,8 +28,8 @@ mod pl031;
 mod psci;
 mod rng;
 
+use alloc::boxed::Box;
 use core::{arch::global_asm, panic::PanicInfo, slice};
-use core::mem::MaybeUninit;
 use core::ops::Range;
 use linked_list_allocator::LockedHeap;
 use log::{debug, error, info};
@@ -86,12 +86,8 @@ struct MemoryMapper {
 }
 
 impl MemoryMapper {
-    fn new(idmap: idmap::IdMap) -> &'static mut MemoryMapper {
-        static mut _MM: MaybeUninit<MemoryMapper> = MaybeUninit::uninit();
-
-        unsafe {
-            _MM.write(MemoryMapper { idmap })
-        }
+    fn new(idmap: idmap::IdMap) -> Box<MemoryMapper> {
+        Box::new(MemoryMapper { idmap })
     }
 
     fn match_attributes(attributes: u64) -> Option<Attributes> {
@@ -263,13 +259,11 @@ extern "C" fn efilite_main(base: *mut u8, mapped: usize, used: isize, avail: usi
 
     let ldr = fwcfg.get_kernel_loader().expect("No kernel image provided");
 
-    let initrdloader = fwcfg.get_initrd_loader().unwrap();
-    let _initrd = efi::initrdloadfile2::new(&initrdloader);
-
     if let Some(mut li) = efi.load_image(&ldr, &cmdline) {
-        info!("Starting loaded EFI program\n");
-        let ret = li.start_image();
+        let initrd = fwcfg.get_initrd_loader(); 
 
+        info!("Starting loaded EFI program\n");
+        let ret = li.start_image(initrd);
         info!("EFI program exited with return value {:?}\n", ret);
     };
 }
diff --git a/src/rng.rs b/src/rng.rs
index e3410b4..4f6f5d8 100644
--- a/src/rng.rs
+++ b/src/rng.rs
@@ -6,7 +6,6 @@ use crate::efi;
 
 use alloc::boxed::Box;
 use core::arch::asm;
-use core::mem::MaybeUninit;
 
 const ID_AA64ISAR0_RNDR_SHIFT: usize = 60;
 
@@ -77,6 +76,7 @@ impl Random {
         }
     }
 
+    #[allow(dead_code)]
     pub fn get_random<T: TryFrom<u64>>(&self) -> Option<T> {
         let mut ret: u64;
         let mut l: u64;