From df9a41065f66eeb8df8b832977cff093dc315332 Mon Sep 17 00:00:00 2001
From: Mark <alnyan@airmail.cc>
Date: Fri, 18 Sep 2020 14:18:32 +0300
Subject: [PATCH] Increase memory map buffer size

---
 src/main.rs  | 28 +++++++++++++++-------------
 src/proto.rs | 24 +++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 3b467f5..f870ba5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,20 +25,21 @@ mod video;
 mod elf;
 
 fn main() -> efi::Result<()> {
-    let mut desc_array = [0u8; 4096];
+    let mut desc_array = [0u8; 16384];
     let mut mmap = efi::MemoryMap::new(&mut desc_array);
     let bs = &system_table().boot_services;
 
+    println!("Getting memory map");
     bs.get_memory_map(&mut mmap)?;
 
+    println!("Getting RSDP");
     let rsdp = system_table()
         .config_iter()
         .find(|x| matches!(x, ConfigurationTableEntry::Acpi10Table(_)))
         .map(|x| match x {
             ConfigurationTableEntry::Acpi10Table(ptr)   => ptr,
             _                                           => panic!()
-        })
-        .unwrap();
+        });
 
     let mut root = image_handle().get_boot_path()?.open_partition()?;
 
@@ -54,17 +55,17 @@ fn main() -> efi::Result<()> {
         &obj)?;
 
     // Set video mode
-    video::set_mode(bs, data)?;
+    use proto::LoadProtocol;
+    data.set_initrd(initrd_base, initrd_size);
+    data.set_acpi_rsdp(rsdp.unwrap_or(core::ptr::null_mut()) as usize);
+    data.set_loader_magic();
 
     // Get the new memory map and terminate boot services
     bs.get_memory_map(&mut mmap)?;
-    bs.exit_boot_services(mmap.key)?;
-
-    use proto::LoadProtocol;
     data.set_mmap(&mmap);
-    data.set_initrd(initrd_base, initrd_size);
-    data.set_acpi_rsdp(rsdp as usize);
-    data.set_loader_magic();
+    video::set_mode(bs, data)?;
+
+    bs.exit_boot_services(mmap.key)?;
 
     unsafe {
         let entry_fn: unsafe fn () -> ! = core::mem::transmute(entry);
@@ -75,9 +76,10 @@ fn main() -> efi::Result<()> {
 #[no_mangle]
 extern "C" fn efi_main(ih: *mut ImageHandle, st: *mut SystemTable) -> u64 {
     efi::init(ih, st);
-    let ret = efi::Termination::to_efi(&main());
-    println!("result -> {}", ret);
-    return ret;
+    let res = &main();
+    println!("result -> {:?}", res);
+
+    efi::Termination::to_efi(res)
 }
 
 use core::panic::PanicInfo;
diff --git a/src/proto.rs b/src/proto.rs
index abeca84..12b343a 100644
--- a/src/proto.rs
+++ b/src/proto.rs
@@ -75,12 +75,30 @@ impl LoadProtocol for V1 {
     }
 
     fn set_mmap(&mut self, map: &MemoryMap) {
-        self.memory_map_data = (map.storage_ref.as_ptr() as usize).try_into().unwrap();
+        let ptr = map.storage_ref.as_ptr();
+        if ptr >= 0x100000000 as *const _ {
+            panic!("Memory map pointer crosses 4GiB");
+        }
+        if map.size > self.memory_map_size as usize {
+            panic!("Can't fit memory map");
+        }
+
+        unsafe {
+            extern "C" {
+                fn memcpy(dst: *mut u8, src: *const u8, count: usize) -> *mut u8;
+            }
+            memcpy(self.memory_map_data as *mut _, ptr, map.size);
+        }
+
+        //self.memory_map_data = (map.storage_ref.as_ptr() as usize).try_into().unwrap();
         self.memory_map_size = map.size.try_into().unwrap();
         self.memory_map_entsize = map.descriptor_size.try_into().unwrap();
     }
 
     fn set_initrd(&mut self, base: usize, size: usize) {
+        if self.initrd_base + self.initrd_size >= 0x100000000 {
+            panic!("Initrd crosses 4GiB");
+        }
         self.initrd_base = base.try_into().unwrap();
         self.initrd_size = size.try_into().unwrap();
     }
@@ -99,6 +117,10 @@ impl LoadProtocol for V1 {
     }
 
     fn set_video_info(&mut self, info: &VideoInfo) {
+        if self.video_framebuffer >= 0x100000000 {
+            panic!("Video framebuffer address is above 4GiB");
+        }
+
         self.video_width = info.width;
         self.video_height = info.height;
         self.video_format = info.format as u32;