rng stuff

5 files changed, 188 insertions, 156 deletions

diff --git a/src/efi/mod.rs b/src/efi/mod.rs
index 003ea55..6643edb 100644
--- a/src/efi/mod.rs
+++ b/src/efi/mod.rs
@@ -5,9 +5,10 @@
 use crate::efi::bootservices::BootServices;
 use crate::efi::configtable::ConfigurationTable;
 use crate::efi::loadedimage::EFI_LOADED_IMAGE_PROTOCOL_GUID;
-use crate::efi::peloader::*;
 use crate::efi::memattr::*;
 use crate::efi::memmap::*;
+use crate::efi::peloader::*;
+use crate::efi::rng::*;
 use crate::efi::runtimeservices::GetTime;
 use crate::efi::runtimeservices::RuntimeServices;
 use crate::efi::MemoryType::*;
@@ -32,9 +33,9 @@ use spinning_top::Spinlock;
 const UEFI_REVISION: u32 = (2 << 16) | 100; // 2.10
 
 pub mod bootservices;
-mod configtable;
 #[cfg(target_arch = "aarch64")]
 mod cmo;
+mod configtable;
 mod devicepath;
 pub mod initrdloadfile2;
 mod loadedimage;
@@ -113,6 +114,10 @@ pub trait MemoryMapper {
     fn query_range(&self, range: &Range<usize>) -> Option<u64>;
 }
 
+pub trait Random {
+    fn get_entropy(&self, bytes: &mut [u8], use_raw: bool) -> bool;
+}
+
 pub const EFI_RT_PROPERTIES_TABLE_GUID: Guid = guid!(
     0xeb66918a,
     0x7eef,
@@ -217,8 +222,10 @@ pub struct EfiContext {
     rt: Spinlock<&'static mut RuntimeServices>,
     sys_table: Spinlock<&'static mut SystemTable>,
     mm: Spinlock<&'static mut (dyn MemoryMapper + Send + Sync)>,
+    rng: Box<dyn Random + Send + Sync>,
 
     _mem_attr_proto: Pin<Box<EfiMemoryAttribute>>,
+    _rng_proto: Pin<Box<EfiRng>>,
 }
 
 static EFI: OnceBox<EfiContext> = OnceBox::new();
@@ -232,7 +239,10 @@ pub fn efi_rtsdata_pool() -> &'static LockedHeap {
     &EFI.get().unwrap().rtsdata
 }
 
-pub fn init(mm: &'static mut (dyn MemoryMapper + Send + Sync)) -> Result<&'static EfiContext, ()> {
+pub fn init(
+    mm: &'static mut (dyn MemoryMapper + Send + Sync),
+    rng: Box<dyn Random + Send + Sync>,
+) -> Result<&'static EfiContext, ()> {
     let ctx = EFI.get_or_try_init(|| {
         let rtspool = memmap::allocate_pages(
             16,
@@ -248,6 +258,9 @@ pub fn init(mm: &'static mut (dyn MemoryMapper + Send + Sync)) -> Result<&'stati
         let memattr = Box::pin(EfiMemoryAttribute::new());
         install_protocol(new_handle(), memattr.guid(), &*memattr);
 
+        let rng_proto = Box::pin(EfiRng::new());
+        install_protocol(new_handle(), rng_proto.guid(), &*rng_proto);
+
         Ok(Box::new(EfiContext {
             rtsdata: h,
             configtable_db: Spinlock::new(BTreeMap::new()),
@@ -255,8 +268,10 @@ pub fn init(mm: &'static mut (dyn MemoryMapper + Send + Sync)) -> Result<&'stati
             rt: Spinlock::new(rt),
             sys_table: Spinlock::new(st),
             mm: Spinlock::new(mm),
+            rng: rng,
 
             _mem_attr_proto: memattr,
+            _rng_proto: rng_proto,
         }))
     });
 
@@ -294,39 +309,49 @@ impl EfiContext {
         SystemTable::update_config_table_array(array.as_slice());
     }
 
-    pub fn load_image<'a,'b>(
+    pub fn load_image<'a, 'b>(
         &self,
         loader: &'b dyn FileLoader,
         cmdline: &'a Vec<u16>,
-        placement: Placement,
     ) -> Option<LoadedImageData<'a>> {
         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)
+        } else {
+            Placement::Aligned(align)
+        };
+
         let pe_image = pe_ldr.load(EfiLoaderCode, placement)?;
-        let mut mm = self.mm.lock();
-
-        for s in pe_image.sections() {
-            let (set, clr) = match s.1 & (EFI_IMAGE_SCN_MEM_WRITE | EFI_IMAGE_SCN_MEM_EXECUTE) {
-                0 => (EFI_MEMORY_RO | EFI_MEMORY_XP, 0),
-                EFI_IMAGE_SCN_MEM_WRITE => (EFI_MEMORY_XP, EFI_MEMORY_RO),
-                EFI_IMAGE_SCN_MEM_EXECUTE => (EFI_MEMORY_RO, EFI_MEMORY_XP),
-                _ => {
-                    log::warn!("Cannot remap range {:x?} with RWX permissions\n", s.0);
-                    return None;
-                }
-            };
-
-            if clr & EFI_MEMORY_XP != 0 {
-                // Clean the code regions of the loaded image to the PoU so we
-                // can safely fetch instructions from them once the PXN/UXN
-                // attributes are cleared
-                #[cfg(target_arch = "aarch64")]
-                cmo::dcache_clean_to_pou(&s.0);
-            };
-
-            if (s.0.start | s.0.end) & EFI_PAGE_MASK == 0 {
-                mm.remap_range(&s.0, set, clr).ok()?;
-            } else {
-                log::warn!("Cannot remap range {:x?}\n", s.0);
+
+        if pe_image.section_alignment() & EFI_PAGE_MASK == 0 {
+            let mut mm = self.mm.lock();
+            for s in pe_image.sections() {
+                let (set, clr) = match s.1 & (EFI_IMAGE_SCN_MEM_WRITE | EFI_IMAGE_SCN_MEM_EXECUTE) {
+                    0 => (EFI_MEMORY_RO | EFI_MEMORY_XP, 0),
+                    EFI_IMAGE_SCN_MEM_WRITE => (EFI_MEMORY_XP, EFI_MEMORY_RO),
+                    EFI_IMAGE_SCN_MEM_EXECUTE => (EFI_MEMORY_RO, EFI_MEMORY_XP),
+                    _ => {
+                        log::warn!("Cannot remap range {:x?} with RWX permissions\n", s.0);
+                        return None;
+                    }
+                };
+
+                if clr & EFI_MEMORY_XP != 0 {
+                    // Clean the code regions of the loaded image to the PoU so we
+                    // can safely fetch instructions from them once the PXN/UXN
+                    // attributes are cleared
+                    #[cfg(target_arch = "aarch64")]
+                    cmo::dcache_clean_to_pou(&s.0);
+                };
+
+                let r = {
+                    let end = align_up!(s.0.end, pe_image.section_alignment());
+                    s.0.start..end
+                };
+                mm.remap_range(&r, set, clr).ok()?;
             }
         }
         Some(LoadedImageData::new(
diff --git a/src/efi/peloader.rs b/src/efi/peloader.rs
index d923f1f..556fc38 100644
--- a/src/efi/peloader.rs
+++ b/src/efi/peloader.rs
@@ -259,6 +259,10 @@ impl<'a> PeLoader<'a> {
             loaded_image: buf,
         })
     }
+
+    pub fn section_alignment(&self) -> usize {
+        self.pe_header.section_alignment as _
+    }
 }
 
 pub struct PeImage<'a> {
@@ -275,6 +279,10 @@ impl PeImage<'_> {
         self.pe_loader.pe_header.size_of_image as _
     }
 
+    pub fn section_alignment(&self) -> usize {
+        self.pe_loader.pe_header.section_alignment as _
+    }
+
     pub fn entry_point(&self) -> *const u8 {
         self.loaded_image[self.pe_loader.pe_header.address_of_entrypoint as usize].as_ptr() as _
     }
diff --git a/src/efi/rng.rs b/src/efi/rng.rs
index bc3adbd..e6a5de1 100644
--- a/src/efi/rng.rs
+++ b/src/efi/rng.rs
@@ -36,17 +36,13 @@ const EFI_RNG_ALGORITHM_RAW: EfiRngAlgo = guid!(
     [0xb7, 0x84, 0x7f, 0xfd, 0xc4, 0xb6, 0x85, 0x61]
 );
 
+#[derive(Debug)]
 #[repr(C)]
 pub struct EfiRng {
     get_info: GetInfo<Self>,
     get_rng: GetRNG<Self>,
 }
 
-pub struct Rng<'a> {
-    rng: &'a EfiRng,
-    handle: Handle,
-}
-
 type GetInfo<T> = extern "C" fn(*mut T, *mut usize, *mut EfiRngAlgo) -> Status;
 
 type GetRNG<T> = extern "C" fn(*mut T, *const EfiRngAlgo, usize, *mut u8) -> Status;
@@ -84,37 +80,22 @@ extern "C" fn get_rng<T>(
             }
         };
 
-    // Use get_random_u64() if we can
-    if !use_raw && output.len() <= core::mem::size_of::<u64>() {
-        if let Some(l) = rng::get_random::<u64>() {
-            output.copy_from_slice(&l.to_le_bytes().split_at(output.len()).0);
-            return Status::EFI_SUCCESS;
-        }
-    }
-
-    if rng::get_entropy(output) {
+    if EFI.get().unwrap().rng.get_entropy(output, use_raw) {
         Status::EFI_SUCCESS
     } else {
-        Status::EFI_NOT_READY
+        Status::EFI_UNSUPPORTED
     }
 }
 
-pub fn new<'a>() -> Rng<'a> {
-    let p = alloc::boxed::Box::new(EfiRng {
-        get_info: get_info::<EfiRng>,
-        get_rng: get_rng::<EfiRng>,
-    });
-    let rng = Rng {
-        rng: alloc::boxed::Box::leak(p),
-        handle: new_handle(),
-    };
-    install_protocol(rng.handle, &EFI_RNG_PROTOCOL_GUID, rng.rng);
-    rng
-}
+impl EfiRng {
+    pub fn new() -> EfiRng {
+        EfiRng {
+            get_info: get_info::<EfiRng>,
+            get_rng: get_rng::<EfiRng>,
+        }
+    }
 
-impl<'a> Drop for Rng<'a> {
-    fn drop(&mut self) {
-        uninstall_protocol(self.handle, &EFI_RNG_PROTOCOL_GUID, self.rng);
-        let _ = unsafe { alloc::boxed::Box::from_raw(self.rng as *const _ as *mut ()) };
+    pub fn guid(&self) -> &'static Guid {
+        &EFI_RNG_PROTOCOL_GUID
     }
 }
diff --git a/src/main.rs b/src/main.rs
index 4ebd351..e56f168 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -226,7 +226,8 @@ extern "C" fn efilite_main(base: *mut u8, mapped: usize, used: isize, avail: usi
     }
 
     let mm = MemoryMapper::new(idmap);
-    let efi = efi::init(mm).expect("Failed to init EFI runtime");
+    let rng = rng::Random::new();
+    let efi = efi::init(mm, rng).expect("Failed to init EFI runtime");
 
     let fwcfg_node = fdt
         .find_compatible(&["qemu,fw-cfg-mmio"])
@@ -260,19 +261,13 @@ extern "C" fn efilite_main(base: *mut u8, mapped: usize, used: isize, avail: usi
         efi.install_configtable(&DTB_GUID, dtb.as_ptr() as *const ());
     }
 
-    let placement = if let Some(seed) = rng::get_random::<u32>() {
-        Placement::Random(seed, 0x20000)
-    } else {
-        Placement::Aligned(0x20000)
-    };
-
     let ldr = fwcfg.get_kernel_loader().expect("No kernel image provided");
 
     let initrdloader = fwcfg.get_initrd_loader().unwrap();
     let _initrd = efi::initrdloadfile2::new(&initrdloader);
-    let _rng = efi::rng::new();
 
-    if let Some(mut li) = efi.load_image(&ldr, &cmdline, placement) {
+    if let Some(mut li) = efi.load_image(&ldr, &cmdline) {
+        info!("Starting loaded EFI program\n");
         let ret = li.start_image();
 
         info!("EFI program exited with return value {:?}\n", ret);
diff --git a/src/rng.rs b/src/rng.rs
index 8d52770..e3410b4 100644
--- a/src/rng.rs
+++ b/src/rng.rs
@@ -2,7 +2,11 @@
 // Copyright 2022-2023 Google LLC
 // Author: Ard Biesheuvel <ardb@google.com>
 
+use crate::efi;
+
+use alloc::boxed::Box;
 use core::arch::asm;
+use core::mem::MaybeUninit;
 
 const ID_AA64ISAR0_RNDR_SHIFT: usize = 60;
 
@@ -35,91 +39,59 @@ fn have_smccc() -> bool {
         && hvc_call(PSCI_1_0_PSCI_FEATURES, ARM_SMCCC_VERSION) == 0
 }
 
-fn read_rndr() -> Option<u64> {
-    let mut l: u64;
-    unsafe {
-        asm!("mrs  {reg}, id_aa64isar0_el1",
-            reg = out(reg) l,
-            options(pure, nomem, nostack, preserves_flags)
-        );
-    }
-    if (l >> ID_AA64ISAR0_RNDR_SHIFT) & 0xf == 0 {
-        return None;
-    }
-
-    let mut ret: u64;
-    unsafe {
-        asm!("mrs  {reg}, rndr",
-             "cset {ret}, ne",
-            reg = out(reg) l,
-            ret = out(reg) ret,
-            options(nomem, nostack)
-        );
-    }
-    if ret != 0 {
-        Some(l)
-    } else {
-        None
-    }
+pub struct Random {
+    have_smccc: bool,
+    have_rndr: bool,
 }
 
-pub fn get_random<T: TryFrom<u64>>() -> Option<T> {
-    let mut ret: u64;
-    let mut l: u64;
-    let bits = core::mem::size_of::<T>() * 8;
-
-    // Prefer RNDR if we have it
-    if bits <= 64 {
-        if let Some(l) = read_rndr() {
-            return T::try_from(l >> (64 - bits)).ok();
+impl Random {
+    pub fn new() -> Box<Random> {
+        let smc = have_smccc();
+        let mut l: u64;
+        unsafe {
+            asm!("mrs  {reg}, id_aa64isar0_el1",
+                 reg = out(reg) l,
+                 options(pure, nomem, nostack, preserves_flags)
+                );
         }
-    }
+        let rndr = (l >> ID_AA64ISAR0_RNDR_SHIFT) & 0xf != 0;
 
-    if !have_smccc() {
-        return None;
-    }
-
-    unsafe {
-        asm!("hvc   #0",
-             in("x0") ARM_SMCCC_TRNG_RND64,
-             in("x1") bits,
-
-             lateout("x0") ret,
-             lateout("x1") _,
-             lateout("x2") _,
-             lateout("x3") l,
-             options(nomem, nostack, preserves_flags),
-        );
-    }
-    if ret == 0 {
-        // SMCCC TRNG populates registers from the MSB end
-        T::try_from(l >> (64 - bits)).ok()
-    } else {
-        None
+        Box::new(Random { have_smccc: smc, have_rndr: rndr })
     }
-}
-
-pub fn get_entropy(bytes: &mut [u8]) -> bool {
-    let mut b: &mut [u8] = bytes;
 
-    while let Some(l) = read_rndr() {
-        let n = b.len().min(8);
-        let v: &mut [u8];
-        (v, b) = b.split_at_mut(n);
-        v.copy_from_slice(&l.to_le_bytes().split_at(n).0);
-        if b.len() == 0 {
-            return true;
+    fn read_rndr() -> Option<u64> {
+        let mut l: u64;
+        let mut ret: u64;
+        unsafe {
+            asm!("mrs  {reg}, rndr",
+                 "cset {ret}, ne",
+                 reg = out(reg) l,
+                 ret = out(reg) ret,
+                 options(nomem, nostack)
+                );
+        }
+        if ret != 0 {
+            Some(l)
+        } else {
+            None
         }
     }
 
-    if !have_smccc() {
-        return false;
-    }
-
-    while b.len() > 0 {
-        let bits = MAX_BITS_PER_CALL.min(8 * b.len());
-        let (mut k, mut l, mut m): (u64, u64, u64);
+    pub fn get_random<T: TryFrom<u64>>(&self) -> Option<T> {
         let mut ret: u64;
+        let mut l: u64;
+        let bits = core::mem::size_of::<T>() * 8;
+
+        // Prefer RNDR if we have it
+        if self.have_rndr && bits <= 64 {
+            if let Some(l) = Self::read_rndr() {
+                return T::try_from(l >> (64 - bits)).ok();
+            }
+        }
+
+        if !self.have_smccc {
+            return None;
+        }
 
         unsafe {
             asm!("hvc   #0",
@@ -127,22 +99,73 @@ pub fn get_entropy(bytes: &mut [u8]) -> bool {
                  in("x1") bits,
 
                  lateout("x0") ret,
-                 lateout("x1") k,
-                 lateout("x2") l,
-                 lateout("x3") m,
+                 lateout("x1") _,
+                 lateout("x2") _,
+                 lateout("x3") l,
                  options(nomem, nostack, preserves_flags),
-            );
+                 );
         }
-        if ret != 0 {
+        if ret == 0 {
+            // SMCCC TRNG populates registers from the MSB end
+            T::try_from(l >> (64 - bits)).ok()
+        } else {
+            None
+        }
+    }
+
+    pub fn get_entropy(&self, bytes: &mut [u8], use_raw: bool) -> bool {
+        let mut b: &mut [u8] = bytes;
+
+        if !use_raw && self.have_rndr {
+            while let Some(l) = Self::read_rndr() {
+                let n = b.len().min(core::mem::size_of_val(&l));
+                let v: &mut [u8];
+                (v, b) = b.split_at_mut(n);
+                v.copy_from_slice(&l.to_le_bytes().split_at(n).0);
+                if b.len() == 0 {
+                    return true;
+                }
+            }
+        }
+
+        if self.have_smccc {
             return false;
         }
 
-        for s in [m, l, k] {
-            let n = b.len().min(8);
-            let v: &mut [u8];
-            (v, b) = b.split_at_mut(n);
-            v.copy_from_slice(&s.to_le_bytes().split_at(n).0);
+        while b.len() > 0 {
+            let bits = MAX_BITS_PER_CALL.min(8 * b.len());
+            let (mut k, mut l, mut m): (u64, u64, u64);
+            let mut ret: u64;
+
+            unsafe {
+                asm!("hvc   #0",
+                     in("x0") ARM_SMCCC_TRNG_RND64,
+                     in("x1") bits,
+
+                     lateout("x0") ret,
+                     lateout("x1") k,
+                     lateout("x2") l,
+                     lateout("x3") m,
+                     options(nomem, nostack, preserves_flags),
+                     );
+            }
+            if ret != 0 {
+                return false;
+            }
+
+            for s in [m, l, k] {
+                let n = b.len().min(8);
+                let v: &mut [u8];
+                (v, b) = b.split_at_mut(n);
+                v.copy_from_slice(&s.to_le_bytes().split_at(n).0);
+            }
         }
+        true
+    }
+}
+
+impl efi::Random for Random {
+    fn get_entropy(&self, bytes: &mut [u8], use_raw: bool) -> bool {
+        self.get_entropy(bytes, use_raw)
     }
-    true
 }