Dynamically map kernel stack, boot info, physical memory and recursive table

build.rs

@@ -1,6 +1,39 @@
 #[cfg(not(feature = "binary"))]
 fn main() {}
 
+#[cfg(feature = "binary")]
+fn address_from_env(env: &'static str) -> Option<u64> {
+    use std::env;
+    match env::var(env) {
+        Err(env::VarError::NotPresent) => None,
+        Err(env::VarError::NotUnicode(_)) => {
+            panic!("The `{}` environment variable must be valid unicode", env,)
+        }
+        Ok(s) => {
+            let addr = if s.starts_with("0x") {
+                u64::from_str_radix(&s[2..], 16)
+            } else {
+                u64::from_str_radix(&s, 10)
+            };
+
+            let addr = addr.expect(&format!(
+                "The `{}` environment variable must be an integer\
+                 (is `{}`).",
+                env, s
+            ));
+
+            if addr % 0x1000 != 0 {
+                panic!(
+                    "The `{}` environment variable must be aligned to 0x1000 (is `{:#x}`).",
+                    env, addr
+                );
+            }
+
+            Some(addr)
+        }
+    }
+}
+
 #[cfg(feature = "binary")]
 fn main() {
     use std::{
@@ -146,31 +179,29 @@ fn main() {
     // create a file with the `PHYSICAL_MEMORY_OFFSET` constant
     let file_path = out_dir.join("physical_memory_offset.rs");
     let mut file = File::create(file_path).expect("failed to create physical_memory_offset.rs");
-    let physical_memory_offset = match env::var("BOOTLOADER_PHYSICAL_MEMORY_OFFSET") {
-        Err(env::VarError::NotPresent) => 0o_177777_770_000_000_000_0000u64,
-        Err(env::VarError::NotUnicode(_)) => panic!(
-            "The `BOOTLOADER_PHYSICAL_MEMORY_OFFSET` environment variable must be valid unicode"
-        ),
-        Ok(s) => if s.starts_with("0x") {
-            u64::from_str_radix(&s[2..], 16)
-        } else {
-            u64::from_str_radix(&s, 10)
-        }
-        .expect(&format!(
-            "The `BOOTLOADER_PHYSICAL_MEMORY_OFFSET` environment variable must be an integer\
-             (is `{}`).",
-            s
-        )),
-    };
+    let physical_memory_offset = address_from_env("BOOTLOADER_PHYSICAL_MEMORY_OFFSET");
     file.write_all(
         format!(
-            "const PHYSICAL_MEMORY_OFFSET: u64 = {:#x};",
+            "const PHYSICAL_MEMORY_OFFSET: Option<u64> = {:?};",
             physical_memory_offset
         )
         .as_bytes(),
     )
     .expect("write to physical_memory_offset.rs failed");
 
+    // create a file with the `KERNEL_STACK_ADDRESS` constant
+    let file_path = out_dir.join("kernel_stack_address.rs");
+    let mut file = File::create(file_path).expect("failed to create kernel_stack_address.rs");
+    let kernel_stack_address = address_from_env("BOOTLOADER_KERNEL_STACK_ADDRESS");
+    file.write_all(
+        format!(
+            "const KERNEL_STACK_ADDRESS: Option<u64> = {:?};",
+            kernel_stack_address,
+        )
+        .as_bytes(),
+    )
+    .expect("write to kernel_stack_address.rs failed");
+
     // pass link arguments to rustc
     println!("cargo:rustc-link-search=native={}", out_dir.display());
     println!(
@@ -180,6 +211,7 @@ fn main() {
 
     println!("cargo:rerun-if-env-changed=KERNEL");
     println!("cargo:rerun-if-env-changed=BOOTLOADER_PHYSICAL_MEMORY_OFFSET");
+    println!("cargo:rerun-if-env-changed=BOOTLOADER_KERNEL_STACK_ADDRESS");
     println!("cargo:rerun-if-changed={}", kernel.display());
     println!("cargo:rerun-if-changed=build.rs");
 }

src/level4_entries.rs

@@ -0,0 +1,43 @@
+use core::convert::TryInto;
+use fixedvec::FixedVec;
+use x86_64::ux;
+use x86_64::{structures::paging::Page, VirtAddr};
+use xmas_elf::program::ProgramHeader64;
+
+pub struct UsedLevel4Entries {
+    entry_state: [bool; 512], // whether an entry is in use by the kernel
+}
+
+impl UsedLevel4Entries {
+    pub fn new(segments: &FixedVec<ProgramHeader64>) -> Self {
+        let mut used = UsedLevel4Entries {
+            entry_state: [false; 512],
+        };
+
+        used.entry_state[0] = true; // TODO: Can we do this dynamically?
+
+        for segment in segments {
+            let start_page: Page = Page::containing_address(VirtAddr::new(segment.virtual_addr));
+            let end_page: Page =
+                Page::containing_address(VirtAddr::new(segment.virtual_addr + segment.mem_size));
+
+            for p4_index in u64::from(start_page.p4_index())..=u64::from(end_page.p4_index()) {
+                used.entry_state[p4_index as usize] = true;
+            }
+        }
+
+        used
+    }
+
+    pub fn get_free_entry(&mut self) -> ux::u9 {
+        let (idx, entry) = self
+            .entry_state
+            .iter_mut()
+            .enumerate()
+            .find(|(_, &mut entry)| entry == false)
+            .expect("no usable level 4 entries found");
+
+        *entry = true;
+        ux::u9::new(idx.try_into().unwrap())
+    }
+}

src/main.rs

@@ -11,13 +11,15 @@
 compile_error!("The bootloader crate must be compiled for the `x86_64-bootloader.json` target");
 
 use bootloader::bootinfo::{BootInfo, FrameRange};
+use core::convert::TryInto;
 use core::panic::PanicInfo;
 use core::{mem, slice};
 use fixedvec::alloc_stack;
 use usize_conversions::usize_from;
+use x86_64::instructions::tlb;
 use x86_64::structures::paging::{
-    frame::PhysFrameRange, Mapper, Page, PageTableFlags, PhysFrame, RecursivePageTable, Size2MiB,
-    Size4KiB,
+    frame::PhysFrameRange, page_table::PageTableEntry, Mapper, Page, PageTable, PageTableFlags,
+    PhysFrame, RecursivePageTable, Size2MiB, Size4KiB,
 };
 use x86_64::ux::u9;
 use x86_64::{PhysAddr, VirtAddr};
@@ -26,6 +28,9 @@ use x86_64::{PhysAddr, VirtAddr};
 // the `map_physical_memory` is activated. Set by the build script.
 include!(concat!(env!("OUT_DIR"), "/physical_memory_offset.rs"));
 
+// The virtual address of the kernel stack. Set by the build script.
+include!(concat!(env!("OUT_DIR"), "/kernel_stack_address.rs"));
+
 global_asm!(include_str!("stage_1.s"));
 global_asm!(include_str!("stage_2.s"));
 global_asm!(include_str!("e820.s"));
@@ -44,6 +49,7 @@ unsafe fn context_switch(boot_info: VirtAddr, entry_point: VirtAddr, stack_point
 
 mod boot_info;
 mod frame_allocator;
+mod level4_entries;
 mod page_table;
 mod printer;
 
@@ -74,6 +80,7 @@ extern "C" {
     static __page_table_end: usize;
     static __bootloader_end: usize;
     static __bootloader_start: usize;
+    static _p4: usize;
 }
 
 #[no_mangle]
@@ -90,6 +97,7 @@ pub unsafe extern "C" fn stage_4() -> ! {
     let page_table_end = &__page_table_end as *const _ as u64;
     let bootloader_start = &__bootloader_start as *const _ as u64;
     let bootloader_end = &__bootloader_end as *const _ as u64;
+    let p4_physical = &_p4 as *const _ as u64;
 
     load_elf(
         IdentityMappedAddr(PhysAddr::new(kernel_start)),
@@ -100,6 +108,7 @@ pub unsafe extern "C" fn stage_4() -> ! {
         PhysAddr::new(page_table_end),
         PhysAddr::new(bootloader_start),
         PhysAddr::new(bootloader_end),
+        PhysAddr::new(p4_physical),
     )
 }
 
@@ -112,6 +121,7 @@ fn load_elf(
     page_table_end: PhysAddr,
     bootloader_start: PhysAddr,
     bootloader_end: PhysAddr,
+    p4_physical: PhysAddr,
 ) -> ! {
     use bootloader::bootinfo::{MemoryRegion, MemoryRegionType};
     use fixedvec::FixedVec;
@@ -149,11 +159,25 @@ fn load_elf(
         }
     }
 
+    // Mark used virtual addresses
+    let mut level4_entries = level4_entries::UsedLevel4Entries::new(&segments);
+
     // Enable support for the no-execute bit in page tables.
     enable_nxe_bit();
 
-    // Create a RecursivePageTable
-    let recursive_index = u9::new(511);
+    // Create a recursive page table entry
+    let recursive_index = u9::new(level4_entries.get_free_entry().try_into().unwrap());
+    let mut entry = PageTableEntry::new();
+    entry.set_addr(
+        p4_physical,
+        PageTableFlags::PRESENT | PageTableFlags::WRITABLE,
+    );
+
+    // Write the recursive entry into the page table
+    let page_table = unsafe { &mut *(p4_physical.as_u64() as *mut PageTable) };
+    page_table[recursive_index] = entry;
+    tlb::flush_all();
+
     let recursive_page_table_addr = Page::from_page_table_indices(
         recursive_index,
         recursive_index,
@@ -211,18 +235,14 @@ fn load_elf(
         rec_page_table.unmap(page).expect("dealloc error").1.flush();
     }
 
-    // Map kernel segments.
-    let stack_end = page_table::map_kernel(
-        kernel_start.phys(),
-        &segments,
-        &mut rec_page_table,
-        &mut frame_allocator,
-    )
-    .expect("kernel mapping failed");
-
     // Map a page for the boot info structure
     let boot_info_page = {
-        let page: Page = Page::containing_address(VirtAddr::new(0xb0071f0000));
+        let page: Page = Page::from_page_table_indices(
+            level4_entries.get_free_entry(),
+            u9::new(0),
+            u9::new(0),
+            u9::new(0),
+        );
         let frame = frame_allocator
             .allocate_frame(MemoryRegionType::BootInfo)
             .expect("frame allocation failed");
@@ -241,13 +261,38 @@ fn load_elf(
         page
     };
 
-    if cfg!(feature = "map_physical_memory") {
-        fn virt_for_phys(phys: PhysAddr) -> VirtAddr {
-            VirtAddr::new(phys.as_u64() + PHYSICAL_MEMORY_OFFSET)
-        }
+    // If no kernel stack address is provided, map the kernel stack after the boot info page
+    let kernel_stack_address = match KERNEL_STACK_ADDRESS {
+        Some(addr) => Page::containing_address(VirtAddr::new(addr)),
+        None => boot_info_page + 1,
+    };
+
+    // Map kernel segments.
+    let stack_end = page_table::map_kernel(
+        kernel_start.phys(),
+        kernel_stack_address,
+        &segments,
+        &mut rec_page_table,
+        &mut frame_allocator,
+    )
+    .expect("kernel mapping failed");
+
+    let physical_memory_offset = if cfg!(feature = "map_physical_memory") {
+        let physical_memory_offset = PHYSICAL_MEMORY_OFFSET.unwrap_or_else(|| {
+            // If offset not manually provided, find a free p4 entry and map memory here.
+            // One level 4 entry spans 2^48/512 bytes (over 500gib) so this should suffice.
+            assert!(max_phys_addr < (1 << 48) / 512);
+            Page::from_page_table_indices_1gib(level4_entries.get_free_entry(), u9::new(0))
+                .start_address()
+                .as_u64()
+        });
+
+        let virt_for_phys =
+            |phys: PhysAddr| -> VirtAddr { VirtAddr::new(phys.as_u64() + physical_memory_offset) };
 
         let start_frame = PhysFrame::<Size2MiB>::containing_address(PhysAddr::new(0));
         let end_frame = PhysFrame::<Size2MiB>::containing_address(PhysAddr::new(max_phys_addr));
+
         for frame in PhysFrame::range_inclusive(start_frame, end_frame) {
             let page = Page::containing_address(virt_for_phys(frame.start_address()));
             let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;
@@ -263,13 +308,17 @@ fn load_elf(
             .expect("Mapping of bootinfo page failed")
             .flush();
         }
-    }
+
+        physical_memory_offset
+    } else {
+        0 // Value is unused by BootInfo::new, so this doesn't matter
+    };
 
     // Construct boot info structure.
     let mut boot_info = BootInfo::new(
         memory_map,
         recursive_page_table_addr.as_u64(),
-        PHYSICAL_MEMORY_OFFSET,
+        physical_memory_offset,
     );
     boot_info.memory_map.sort();
 

src/page_table.rs

@@ -10,6 +10,7 @@ use xmas_elf::program::{self, ProgramHeader64};
 
 pub(crate) fn map_kernel(
     kernel_start: PhysAddr,
+    stack_start: Page,
     segments: &FixedVec<ProgramHeader64>,
     page_table: &mut RecursivePageTable,
     frame_allocator: &mut FrameAllocator,
@@ -18,10 +19,9 @@ pub(crate) fn map_kernel(
         map_segment(segment, kernel_start, page_table, frame_allocator)?;
     }
 
-    // create a stack
-    // TODO create a stack range dynamically (based on where the kernel is loaded)
-    let stack_start = Page::containing_address(VirtAddr::new(0x57AC_0000_0000));
+    // Create a stack
     let stack_size: u64 = 512; // in pages
+    let stack_start = stack_start + 1; // Leave the first page unmapped as a 'guard page'
     let stack_end = stack_start + stack_size;
 
     let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;

src/stage_3.s

@@ -36,9 +36,6 @@ set_up_page_tables:
     rep stosd
 
     # p4
-    lea eax, [_p4]
-    or eax, (1 | 2)
-    mov [_p4 + 511 * 8], eax # recursive mapping
     lea eax, [_p3]
     or eax, (1 | 2)
     mov [_p4], eax
@@ -63,10 +60,12 @@ set_up_page_tables:
     cmp ecx, edx
     jb map_p2_table
     # p1
-    lea eax, __bootloader_start
+    # start mapping from __page_table_start, as we need to be able to access
+    # the p4 table from rust. stop mapping at __bootloader_end
+    lea eax, __page_table_start
     and eax, 0xfffff000
     or eax, (1 | 2)
-    lea ecx, __bootloader_start
+    lea ecx, __page_table_start
     shr ecx, 12 # start page number
     lea edx, __bootloader_end
     add edx, 4096 - 1 # align up