construct memory in memory mapped by other page table

This has the advantage that we don't need to map the memory map into
the bootloader's page tables. Up until now, we've assumed that UEFI
only maps memory into regions covered by the first PML4 entry, but it
turns out that some implementations map the frame buffer into memory
mapped by other PML4 entries (this is totally legal, the frame buffer
need not be identity mapped). Previously we removed all but the first
PML4 entry, so that we can map our own memory there, but it's now clear
that this can lead to problems. The solution to this is to not modify
the page tables constructed by UEFI and keep all its PML4 entries, but
this is problematic because the code constructing the boot info and
memory map assumes that it can map them into the same addresses used
for the kernel's address space, but UEFI may have already mapped some
memory in there. Instead of mapping the boot info and memory map into
the bootloader's address space, we now only map them into the kernel's
address space. When we want to write to them, we first look up the
physical addresses in the kernel's page table and write to the identity
mapping. RemoteMemoryRegion implements this.
We have some unit tests for constructing the memory map and we want
those unit tests to be able to run outside the bootloader, so we can't
use RemoteMemoryRegion. To solve this, we introduce a trait
MemoryRegionSlice that only requires implementing a method for writing
a memory region and implement that slice for MemoryRegionSlice as well
as MemoryRegionSlice.

Author: Freax13

common/src/legacy_memory_region.rs

@@ -1,13 +1,9 @@
 use bootloader_api::info::{MemoryRegion, MemoryRegionKind};
-use core::{
-    cmp,
-    iter::{empty, Empty},
-    mem::MaybeUninit,
-};
+use core::{cmp, mem::MaybeUninit};
 use x86_64::{
     align_down, align_up,
-    structures::paging::{FrameAllocator, PhysFrame, Size4KiB},
-    PhysAddr,
+    structures::paging::{FrameAllocator, OffsetPageTable, PhysFrame, Size4KiB, Translate},
+    PhysAddr, VirtAddr,
 };
 
 /// A slice of memory that is used by the bootloader and needs to be reserved
@@ -159,14 +155,14 @@ where
     /// must be at least the value returned by [`len`] plus 1.
     ///
     /// The return slice is a subslice of `regions`, shortened to the actual number of regions.
-    pub fn construct_memory_map(
+    pub(crate) fn construct_memory_map(
         self,
-        regions: &mut [MaybeUninit<MemoryRegion>],
+        regions: &mut (impl MemoryRegionSlice + ?Sized),
         kernel_slice_start: PhysAddr,
         kernel_slice_len: u64,
         ramdisk_slice_start: Option<PhysAddr>,
         ramdisk_slice_len: u64,
-    ) -> &mut [MemoryRegion] {
+    ) -> usize {
         let used_slices = [
             UsedMemorySlice {
                 start: self.min_frame.start_address().as_u64(),
@@ -211,17 +207,12 @@ where
             }
         }
 
-        let initialized = &mut regions[..next_index];
-        unsafe {
-            // inlined variant of: `MaybeUninit::slice_assume_init_mut(initialized)`
-            // TODO: undo inlining when `slice_assume_init_mut` becomes stable
-            &mut *(initialized as *mut [_] as *mut [_])
-        }
+        next_index
     }
 
     fn split_and_add_region<'a, U>(
         mut region: MemoryRegion,
-        regions: &mut [MaybeUninit<MemoryRegion>],
+        regions: &mut (impl MemoryRegionSlice + ?Sized),
         next_index: &mut usize,
         used_slices: U,
     ) where
@@ -279,24 +270,97 @@ where
 
     fn add_region(
         region: MemoryRegion,
-        regions: &mut [MaybeUninit<MemoryRegion>],
+        regions: &mut (impl MemoryRegionSlice + ?Sized),
         next_index: &mut usize,
     ) {
         if region.start == region.end {
             // skip zero sized regions
             return;
         }
-        unsafe {
-            regions
-                .get_mut(*next_index)
-                .expect("cannot add region: no more free entries in memory map")
-                .as_mut_ptr()
-                .write(region)
-        };
+        regions.set(*next_index, region);
         *next_index += 1;
     }
 }
 
+/// A trait for slice-like types that allow writing a memory region to given
+/// index. Usually `RemoteMemoryRegion` is used, but we use
+/// `[MaybeUninit<MemoryRegion>]` in tests.
+pub(crate) trait MemoryRegionSlice {
+    fn set(&mut self, index: usize, region: MemoryRegion);
+}
+
+#[cfg(test)]
+impl MemoryRegionSlice for [MaybeUninit<MemoryRegion>] {
+    fn set(&mut self, index: usize, region: MemoryRegion) {
+        self.get_mut(index)
+            .expect("cannot add region: no more free entries in memory map")
+            .write(region);
+    }
+}
+
+/// This type makes it possible to write to a slice mapped in a different set
+/// of page tables. For every write access, we look up the physical frame in
+/// the page tables and directly write to the physical memory. That way we
+/// don't need to map the slice into the current page tables.
+pub(crate) struct RemoteMemoryRegion<'a> {
+    page_table: &'a OffsetPageTable<'a>,
+    base: VirtAddr,
+    len: usize,
+}
+
+impl<'a> RemoteMemoryRegion<'a> {
+    /// Construct a new `RemoteMemoryRegion`.
+    ///
+    /// # Safety
+    ///
+    /// The caller has to ensure that the memory in the starting at `base`
+    /// isn't aliasing memory in the current page tables.
+    pub unsafe fn new(page_table: &'a OffsetPageTable<'a>, base: VirtAddr, len: usize) -> Self {
+        Self {
+            page_table,
+            base,
+            len,
+        }
+    }
+}
+
+impl MemoryRegionSlice for RemoteMemoryRegion<'_> {
+    fn set(&mut self, index: usize, region: MemoryRegion) {
+        assert!(
+            index < self.len,
+            "cannot add region: no more free entries in memory map"
+        );
+
+        // Cast the memory region into a byte slice. MemoryRegion has some
+        // padding bytes, so need to use `MaybeUninit<u8>` instead of `u8`.
+        let bytes = unsafe {
+            core::slice::from_raw_parts(
+                &region as *const _ as *const MaybeUninit<u8>,
+                size_of::<MemoryRegion>(),
+            )
+        };
+
+        // An entry may cross a page boundary, so write one byte at a time.
+        let addr = self.base + index * size_of::<MemoryRegion>();
+        for (addr, byte) in (addr..).zip(bytes.iter().copied()) {
+            // Lookup the physical address in the remote page table.
+            let phys_addr = self
+                .page_table
+                .translate_addr(addr)
+                .expect("memory is mapped in the page table");
+
+            // Get a pointer to the physical memory -> All physical memory is
+            // identitiy mapped.
+            let ptr = phys_addr.as_u64() as *mut MaybeUninit<u8>;
+
+            // Write the byte.
+            unsafe {
+                ptr.write(byte);
+            }
+        }
+    }
+}
+
 unsafe impl<I, D> FrameAllocator<Size4KiB> for LegacyFrameAllocator<I, D>
 where
     I: ExactSizeIterator<Item = D> + Clone,
@@ -384,14 +448,21 @@ mod tests {
         let ramdisk_slice_start = None;
         let ramdisk_slice_len = 0;
 
-        let kernel_regions = allocator.construct_memory_map(
-            &mut regions,
+        let len = allocator.construct_memory_map(
+            regions.as_mut_slice(),
             kernel_slice_start,
             kernel_slice_len,
             ramdisk_slice_start,
             ramdisk_slice_len,
         );
 
+        let initialized = &mut regions[..len];
+        let kernel_regions = unsafe {
+            // inlined variant of: `MaybeUninit::slice_assume_init_mut(initialized)`
+            // TODO: undo inlining when `slice_assume_init_mut` becomes stable
+            &mut *(initialized as *mut [_] as *mut [MemoryRegion])
+        };
+
         for region in kernel_regions.iter() {
             assert!(region.start % 0x1000 == 0);
             assert!(region.end % 0x1000 == 0);
@@ -411,13 +482,21 @@ mod tests {
         let ramdisk_slice_start = Some(PhysAddr::new(0x60000));
         let ramdisk_slice_len = 0x2000;
 
-        let kernel_regions = allocator.construct_memory_map(
-            &mut regions,
+        let len = allocator.construct_memory_map(
+            regions.as_mut_slice(),
             kernel_slice_start,
             kernel_slice_len,
             ramdisk_slice_start,
             ramdisk_slice_len,
         );
+
+        let initialized = &mut regions[..len];
+        let kernel_regions = unsafe {
+            // inlined variant of: `MaybeUninit::slice_assume_init_mut(initialized)`
+            // TODO: undo inlining when `slice_assume_init_mut` becomes stable
+            &mut *(initialized as *mut [_] as *mut [MemoryRegion])
+        };
+
         let mut kernel_regions = kernel_regions.iter();
         // usable memory before the kernel
         assert_eq!(
@@ -514,13 +593,21 @@ mod tests {
         let ramdisk_slice_start = Some(PhysAddr::new(0x60000));
         let ramdisk_slice_len = 0x2000;
 
-        let kernel_regions = allocator.construct_memory_map(
-            &mut regions,
+        let len = allocator.construct_memory_map(
+            regions.as_mut_slice(),
             kernel_slice_start,
             kernel_slice_len,
             ramdisk_slice_start,
             ramdisk_slice_len,
         );
+
+        let initialized = &mut regions[..len];
+        let kernel_regions = unsafe {
+            // inlined variant of: `MaybeUninit::slice_assume_init_mut(initialized)`
+            // TODO: undo inlining when `slice_assume_init_mut` becomes stable
+            &mut *(initialized as *mut [_] as *mut [MemoryRegion])
+        };
+
         let mut kernel_regions = kernel_regions.iter();
 
         // usable memory before the kernel

common/src/lib.rs

@@ -10,6 +10,7 @@ use bootloader_api::{
 };
 use bootloader_boot_config::{BootConfig, LevelFilter};
 use core::{alloc::Layout, arch::asm, mem::MaybeUninit, slice};
+use legacy_memory_region::RemoteMemoryRegion;
 use level_4_entries::UsedLevel4Entries;
 use usize_conversions::FromUsize;
 use x86_64::{
@@ -481,69 +482,67 @@ where
     log::info!("Allocate bootinfo");
 
     // allocate and map space for the boot info
-    let (boot_info, memory_regions) = {
-        let boot_info_layout = Layout::new::<BootInfo>();
-        let regions = frame_allocator.memory_map_max_region_count();
-        let memory_regions_layout = Layout::array::<MemoryRegion>(regions).unwrap();
-        let (combined, memory_regions_offset) =
-            boot_info_layout.extend(memory_regions_layout).unwrap();
-
-        let boot_info_addr = mapping_addr(
-            config.mappings.boot_info,
-            u64::from_usize(combined.size()),
-            u64::from_usize(combined.align()),
-            &mut mappings.used_entries,
-        )
-        .expect("boot info addr is not properly aligned");
+    let boot_info_layout = Layout::new::<BootInfo>();
+    let regions = frame_allocator.memory_map_max_region_count();
+    let memory_regions_layout = Layout::array::<MemoryRegion>(regions).unwrap();
+    let (combined, memory_regions_offset) = boot_info_layout.extend(memory_regions_layout).unwrap();
+
+    let boot_info_addr = mapping_addr(
+        config.mappings.boot_info,
+        u64::from_usize(combined.size()),
+        u64::from_usize(combined.align()),
+        &mut mappings.used_entries,
+    )
+    .expect("boot info addr is not properly aligned");
 
-        let memory_map_regions_addr = boot_info_addr + memory_regions_offset;
-        let memory_map_regions_end = boot_info_addr + combined.size();
+    let memory_map_regions_addr = boot_info_addr + memory_regions_offset;
+    let memory_map_regions_end = boot_info_addr + combined.size();
 
-        let start_page = Page::containing_address(boot_info_addr);
-        let end_page = Page::containing_address(memory_map_regions_end - 1u64);
-        for page in Page::range_inclusive(start_page, end_page) {
-            let flags =
-                PageTableFlags::PRESENT | PageTableFlags::WRITABLE | PageTableFlags::NO_EXECUTE;
-            let frame = frame_allocator
-                .allocate_frame()
-                .expect("frame allocation for boot info failed");
-            match unsafe {
-                page_tables
-                    .kernel
-                    .map_to(page, frame, flags, &mut frame_allocator)
-            } {
-                Ok(tlb) => tlb.flush(),
-                Err(err) => panic!("failed to map page {:?}: {:?}", page, err),
-            }
-            // we need to be able to access it too
-            match unsafe {
-                page_tables
-                    .bootloader
-                    .map_to(page, frame, flags, &mut frame_allocator)
-            } {
-                Ok(tlb) => tlb.flush(),
-                Err(err) => panic!("failed to map page {:?}: {:?}", page, err),
-            }
+    let start_page = Page::containing_address(boot_info_addr);
+    let end_page = Page::containing_address(memory_map_regions_end - 1u64);
+    for page in Page::range_inclusive(start_page, end_page) {
+        let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE | PageTableFlags::NO_EXECUTE;
+        let frame = frame_allocator
+            .allocate_frame()
+            .expect("frame allocation for boot info failed");
+        match unsafe {
+            page_tables
+                .kernel
+                .map_to(page, frame, flags, &mut frame_allocator)
+        } {
+            Ok(tlb) => tlb.flush(),
+            Err(err) => panic!("failed to map page {:?}: {:?}", page, err),
         }
+        // we need to be able to access it too
+        match unsafe {
+            page_tables
+                .bootloader
+                .map_to(page, frame, flags, &mut frame_allocator)
+        } {
+            Ok(tlb) => tlb.flush(),
+            Err(err) => panic!("failed to map page {:?}: {:?}", page, err),
+        }
+    }
 
-        let boot_info: &'static mut MaybeUninit<BootInfo> =
-            unsafe { &mut *boot_info_addr.as_mut_ptr() };
-        let memory_regions: &'static mut [MaybeUninit<MemoryRegion>] =
-            unsafe { slice::from_raw_parts_mut(memory_map_regions_addr.as_mut_ptr(), regions) };
-        (boot_info, memory_regions)
-    };
+    let boot_info: &'static mut MaybeUninit<BootInfo> =
+        unsafe { &mut *boot_info_addr.as_mut_ptr() };
 
     log::info!("Create Memory Map");
 
     // build memory map
-    let memory_regions = frame_allocator.construct_memory_map(
-        memory_regions,
+    let mut slice =
+        unsafe { RemoteMemoryRegion::new(&page_tables.kernel, memory_map_regions_addr, regions) };
+    let len = frame_allocator.construct_memory_map(
+        &mut slice,
         mappings.kernel_slice_start,
         mappings.kernel_slice_len,
         mappings.ramdisk_slice_phys_start,
         mappings.ramdisk_slice_len,
     );
 
+    let memory_regions =
+        unsafe { core::slice::from_raw_parts_mut(memory_map_regions_addr.as_mut_ptr(), len) };
+
     log::info!("Create bootinfo");
 
     // create boot info