From 842fd84cc81b87c504286c51f2dd2c71bd172cbf Mon Sep 17 00:00:00 2001
From: Mark Poliakov <mark@alnyan.me>
Date: Mon, 6 Nov 2023 10:39:45 +0200
Subject: [PATCH] Properly reserve the kernel image

---
 src/elf.rs  | 21 +++++++++++++++++----
 src/main.rs |  6 +++---
 src/mem.rs  |  3 ++-
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/src/elf.rs b/src/elf.rs
index 0ad724d1..1ad4fd41 100644
--- a/src/elf.rs
+++ b/src/elf.rs
@@ -4,8 +4,9 @@ use bytemuck::Zeroable;
 use log::{debug, error, info};
 // TODO use 'elf' crate
 use uefi::{
+    prelude::BootServices,
     proto::media::file::{File, FileAttribute, FileMode, RegularFile},
-    table::boot::MemoryMap,
+    table::boot::{AllocateType, MemoryMap, MemoryType},
     CStr16, Error, Status,
 };
 use yboot_proto::LoadProtocolV1;
@@ -152,7 +153,7 @@ impl Object {
     }
 
     // Returns the image entry point
-    pub fn load(&mut self, mmap: &MemoryMap) -> Result<LoadedObject, Error> {
+    pub fn load(&mut self, bs: &BootServices, mmap: &MemoryMap) -> Result<LoadedObject, Error> {
         // 1. Locate the protocol data structure
         let Some(loc_proto) = self.locate_protocol_data()? else {
             error!("The image does not have a valid protocol structure");
@@ -218,7 +219,19 @@ impl Object {
             }
         }
 
-        // TODO maybe reserve the kernel memory in the boot services?
+        assert!(image_end > image_start);
+        assert_eq!(image_start & 0xFFF, 0);
+        assert_eq!(image_end & 0xFFF, 0);
+
+        info!("Image start: {:#x}, end: {:#x}", image_start, image_end);
+        // Reserve the kernel memory
+        let reserved_addr = bs.allocate_pages(
+            AllocateType::Address(image_start),
+            MemoryType::RUNTIME_SERVICES_DATA,
+            (image_end - image_start) as usize / 0x1000,
+        )?;
+        assert_eq!(reserved_addr, image_start);
+
         // 3. Load the segments
         for i in 0..self.ehdr.phnum {
             let phdr = self.read_phdr(i as _)?;
@@ -248,7 +261,7 @@ impl Object {
                 self.file.read_exact(dst)?;
             }
 
-            if phdr.memsz > 0 {
+            if phdr.memsz > phdr.filesz {
                 let dst = unsafe {
                     core::slice::from_raw_parts_mut(
                         (phdr.paddr + phdr.filesz) as *mut u8,
diff --git a/src/main.rs b/src/main.rs
index ad7e271b..88d98ca5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -98,7 +98,7 @@ fn load_kernel<'a, 'b>(
     let mut root = fs.open_volume()?;
 
     let mut kernel_obj = Object::open(&mut root, cstr16!("kernel.elf"))?;
-    let loaded_obj = kernel_obj.load(&mmap)?;
+    let loaded_obj = kernel_obj.load(bs, &mmap)?;
 
     debug!("Loaded object: {:#x?}", loaded_obj);
 
@@ -157,7 +157,6 @@ unsafe fn map_and_enter_kernel(
     entry: u64,
 ) -> ! {
     let (_, mmap) = st.exit_boot_services();
-
     let mut index = 0;
 
     let mmap_data = core::slice::from_raw_parts_mut(
@@ -179,7 +178,7 @@ unsafe fn map_and_enter_kernel(
     proto_data.memory_map.len = index as _;
 
     let cr3 = mem::map_image();
-    asm!("mov {0}, %cr3", in(reg) cr3, options(att_syntax));
+    asm!("cli; wbinvd; mov {0}, %cr3", in(reg) cr3, options(att_syntax));
 
     asm!("jmp *{0}", in(reg) entry, in("eax") LOADER_MAGIC, options(noreturn, att_syntax));
 }
@@ -189,6 +188,7 @@ fn efi_main(image_handle: Handle, mut system_table: SystemTable<Boot>) -> Status
     uefi_services::init(&mut system_table).unwrap();
 
     let (entry, mmap_memory, proto_data) = load_kernel(image_handle, &system_table).unwrap();
+
     unsafe {
         map_and_enter_kernel(system_table, proto_data, mmap_memory, entry);
     }
diff --git a/src/mem.rs b/src/mem.rs
index d4cd9bdd..f26f2249 100644
--- a/src/mem.rs
+++ b/src/mem.rs
@@ -25,9 +25,10 @@ static mut MMAP_BUFFER: MmapBuffer = MmapBuffer {
     data: [0; MMAP_BUFFER_SIZE],
 };
 
+const PD_COUNT: usize = 32;
 static mut PML4: Table = Table { data: [0; 512] };
 static mut PDPT: Table = Table { data: [0; 512] };
-static mut PDS: [Table; 4] = [Table { data: [0; 512] }; 4];
+static mut PDS: [Table; PD_COUNT] = [Table { data: [0; 512] }; PD_COUNT];
 
 pub trait MemoryMapExt {
     fn is_usable(&self, page: u64) -> bool;