Ensure boot services is exited and use SetVirtualAddressMap

src/arch/aarch64/paging.rs

@@ -23,7 +23,7 @@ pub unsafe fn paging_create<
     D: Disk,
     M: Iterator<Item=OsMemoryEntry>,
     V: Iterator<Item=OsVideoMode>
->(os: &mut dyn Os<D, M, V>, kernel_phys: usize) -> Option<usize> {
+>(os: &mut dyn Os<D, M, V>, kernel_phys: usize, kernel_base: usize) -> Option<usize> {
     log::error!("paging_create not implemented for aarch64");
     None
 }

src/os/uefi/arch/aarch64/mod.rs

@@ -9,7 +9,6 @@ use crate::{
 
 use super::super::{
     OsEfi,
-    exit_boot_services,
     acpi::{
         RSDPS_AREA,
         find_acpi_table_pointers,
@@ -35,8 +34,12 @@ unsafe extern "C" fn kernel_entry(
     args: *const KernelArgs,
 ) -> ! {
     // Read memory map and exit boot services
-    let key = memory_map();
-    exit_boot_services(key);
+    {
+        let mut memory_iter = memory_map();
+        memory_iter.exit_boot_services();
+        memory_iter.set_virtual_address_map(PHYS_OFFSET);
+        mem::forget(memory_iter);
+    }
 
     // Disable interrupts
     asm!("msr daifset, #2");

src/os/uefi/arch/x86_64/mod.rs

@@ -9,7 +9,6 @@ use crate::{
 
 use super::super::{
     OsEfi,
-    exit_boot_services,
     acpi::{
         RSDPS_AREA,
         find_acpi_table_pointers,
@@ -30,8 +29,12 @@ unsafe extern "C" fn kernel_entry(
     args: *const KernelArgs,
 ) -> ! {
     // Read memory map and exit boot services
-    let key = memory_map();
-    exit_boot_services(key);
+    {
+        let mut memory_iter = memory_map();
+        memory_iter.exit_boot_services();
+        memory_iter.set_virtual_address_map(PHYS_OFFSET);
+        mem::forget(memory_iter);
+    }
 
     // Disable interrupts
     llvm_asm!("cli" : : : "memory" : "intel", "volatile");

src/os/uefi/memory_map.rs

-use alloc::boxed::Box;
+use alloc::vec::Vec;
 use core::{mem, ptr};
-use log::error;
 use uefi::memory::{MemoryDescriptor, MemoryType};
 
 use crate::os::{OsMemoryEntry, OsMemoryKind};
 
 use super::status_to_result;
 
+const EFI_MEMORY_RUNTIME: u64 = 0x8000000000000000;
+
 pub struct MemoryMapIter {
-    map: Box<[u8]>,
+    map: Vec<u8>,
     map_key: usize,
     descriptor_size: usize,
+    descriptor_version: u32,
     i: usize,
 }
 
@@ -30,63 +32,99 @@ impl MemoryMapIter {
             &mut descriptor_size,
             &mut descriptor_version
         )).expect("Failed to get UEFI memory map");
+
+        // Ensure descriptor size is usable
+        assert!(descriptor_size >= mem::size_of::<MemoryDescriptor>());
+
+        // Ensure descriptor version is supported
+        assert_eq!(descriptor_version, 1);
+
+        // Reduce map size to returned value
         map.truncate(map_size);
 
         Self {
-            map: map.into_boxed_slice(),
+            map,
             map_key,
             descriptor_size,
+            descriptor_version,
             i: 0,
         }
     }
+
+    pub fn exit_boot_services(&self) {
+        let handle = std::handle();
+        let uefi = std::system_table();
+
+        status_to_result((uefi.BootServices.ExitBootServices)(
+            handle,
+            self.map_key
+        )).expect("Failed to exit UEFI boot services");
+    }
+
+    pub fn set_virtual_address_map(&mut self, phys_offset: u64) {
+        let uefi = std::system_table();
+
+        for i in 0..self.map.len()/self.descriptor_size {
+            let descriptor_ptr = unsafe { self.map.as_mut_ptr().add(i * self.descriptor_size) };
+            let descriptor = unsafe { &mut *(descriptor_ptr as *mut MemoryDescriptor) };
+            if descriptor.Attribute & EFI_MEMORY_RUNTIME == EFI_MEMORY_RUNTIME {
+                descriptor.VirtualStart.0 = descriptor.PhysicalStart.0 + phys_offset;
+            }
+        }
+
+        status_to_result((uefi.RuntimeServices.SetVirtualAddressMap)(
+            self.map.len(),
+            self.descriptor_size,
+            self.descriptor_version,
+            self.map.as_ptr() as *const MemoryDescriptor
+        )).expect("Failed to set UEFI runtime services virtual address map");
+    }
 }
 
 impl Iterator for MemoryMapIter {
     type Item=OsMemoryEntry;
     fn next(&mut self) -> Option<Self::Item> {
-        if self.descriptor_size >= mem::size_of::<MemoryDescriptor>() {
-            if self.i < self.map.len()/self.descriptor_size {
-                let descriptor_ptr = unsafe { self.map.as_ptr().add(self.i * self.descriptor_size) };
-                self.i += 1;
-
-                let descriptor = unsafe { ptr::read(descriptor_ptr as *const MemoryDescriptor) };
-                let descriptor_type: MemoryType = unsafe { mem::transmute(descriptor.Type) };
-
-                Some(OsMemoryEntry {
-                    base: descriptor.PhysicalStart.0,
-                    //TODO: do not hard code page size
-                    size: descriptor.NumberOfPages * 4096,
-                    kind: match descriptor_type {
-                        MemoryType::EfiLoaderCode |
-                        MemoryType::EfiLoaderData |
-                        MemoryType::EfiBootServicesCode |
-                        MemoryType::EfiBootServicesData |
-                        MemoryType::EfiConventionalMemory => {
-                            OsMemoryKind::Free
-                        },
-                        //TODO: mark ACPI memory as reclaim
-                        _ => {
-                            OsMemoryKind::Reserved
-                        }
+        if self.i < self.map.len()/self.descriptor_size {
+            let descriptor_ptr = unsafe { self.map.as_ptr().add(self.i * self.descriptor_size) };
+            self.i += 1;
+
+            let descriptor = unsafe { ptr::read(descriptor_ptr as *const MemoryDescriptor) };
+            let descriptor_type: MemoryType = unsafe { mem::transmute(descriptor.Type) };
+
+            Some(OsMemoryEntry {
+                base: descriptor.PhysicalStart.0,
+                //TODO: do not hard code page size
+                size: descriptor.NumberOfPages * 4096,
+                kind: match descriptor_type {
+                    MemoryType::EfiLoaderCode |
+                    MemoryType::EfiLoaderData |
+                    MemoryType::EfiBootServicesCode |
+                    MemoryType::EfiBootServicesData |
+                    MemoryType::EfiConventionalMemory => {
+                        OsMemoryKind::Free
+                    },
+                    //TODO: mark ACPI memory as reclaim
+                    _ => {
+                        OsMemoryKind::Reserved
                     }
-                })
-            } else {
-                None
-            }
+                }
+            })
         } else {
-            error!("Unknown memory descriptor size: {}", self.descriptor_size);
             None
         }
     }
 }
 
-pub unsafe fn memory_map() -> usize {
-    let iter = MemoryMapIter::new();
-    let map_key = iter.map_key;
+pub unsafe fn memory_map() -> MemoryMapIter {
+    let mut iter = MemoryMapIter::new();
 
-    for (i, entry) in iter.enumerate() {
-        crate::AREAS[i] = entry;
+    // Using next to avoid consuming iterator
+    while let Some(entry) = iter.next() {
+        crate::AREAS[iter.i] = entry;
     }
 
-    map_key
+    // Rewind iterator
+    iter.i = 0;
+
+    iter
 }

src/os/uefi/mod.rs

@@ -202,13 +202,6 @@ impl Os<
     }
 }
 
-unsafe fn exit_boot_services(key: usize) {
-    let handle = std::handle();
-    let uefi = std::system_table();
-
-    let _ = (uefi.BootServices.ExitBootServices)(handle, key);
-}
-
 fn status_to_result(status: Status) -> Result<usize> {
     match status.branch() {
         ControlFlow::Continue(ok) => Ok(ok),