From f2bc2d33f03d90083b670af0fd0c53d99f67229f Mon Sep 17 00:00:00 2001
From: Philipp Oppermann <dev@phil-opp.com>
Date: Sat, 26 Jan 2019 13:50:39 +0100
Subject: [PATCH] Rewrite translation function on top of RecursivePageTable

---
 src/main.rs   | 10 +++++---
 src/memory.rs | 67 +++++++++++++++++++++------------------------------
 2 files changed, 34 insertions(+), 43 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index c55a797f..d763429b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,7 +9,7 @@ use core::panic::PanicInfo;
 #[no_mangle]
 pub extern "C" fn _start() -> ! {
     use blog_os::interrupts::PICS;
-    use blog_os::memory::translate_addr;
+    use blog_os::memory::{self, translate_addr};
 
     println!("Hello World{}", "!");
 
@@ -19,14 +19,16 @@ pub extern "C" fn _start() -> ! {
     x86_64::instructions::interrupts::enable();
 
     const LEVEL_4_TABLE_ADDR: usize = 0o_177777_777_777_777_777_0000;
+    let recursive_page_table = unsafe { memory::init(LEVEL_4_TABLE_ADDR) };
 
     // the identity-mapped vga buffer page
-    println!("0xb8000 -> {:?}", translate_addr(0xb8000, LEVEL_4_TABLE_ADDR));
+    println!("0xb8000 -> {:?}", translate_addr(0xb8000, &recursive_page_table));
     // some code page
-    println!("0x20010a -> {:?}", translate_addr(0x20010a, LEVEL_4_TABLE_ADDR));
+    println!("0x20010a -> {:?}", translate_addr(0x20010a, &recursive_page_table));
     // some stack page
     println!("0x57ac001ffe48 -> {:?}", translate_addr(0x57ac001ffe48,
-        LEVEL_4_TABLE_ADDR));
+        &recursive_page_table));
+
 
     println!("It did not crash!");
     blog_os::hlt_loop();
diff --git a/src/memory.rs b/src/memory.rs
index 76837ef7..e9f0db2f 100644
--- a/src/memory.rs
+++ b/src/memory.rs
@@ -1,43 +1,32 @@
-use x86_64::structures::paging::PageTable;
-use x86_64::PhysAddr;
+use x86_64::structures::paging::{Mapper, Page, PageTable, RecursivePageTable};
+use x86_64::{VirtAddr, PhysAddr};
 
-/// Returns the physical address for the given virtual address, or `None` if the
-/// virtual address is not mapped.
-pub fn translate_addr(addr: usize, level_4_table_addr: usize) -> Option<PhysAddr> {
-    // retrieve the page table indices of the address that we want to translate
-    let level_4_index = (addr >> 39) & 0o777;
-    let level_3_index = (addr >> 30) & 0o777;
-    let level_2_index = (addr >> 21) & 0o777;
-    let level_1_index = (addr >> 12) & 0o777;
-    let page_offset = addr & 0o7777;
-
-    // check that level 4 entry is mapped
-    let level_4_table = unsafe { &*(level_4_table_addr as *const PageTable) };
-    if level_4_table[level_4_index].addr().is_null() {
-        return None;
-    }
-    let level_3_table_addr = (level_4_table_addr << 9) | (level_4_index << 12);
-
-    // check that level 3 entry is mapped
-    let level_3_table = unsafe { &*(level_3_table_addr as *const PageTable) };
-    if level_3_table[level_3_index].addr().is_null() {
-        return None;
-    }
-    let level_2_table_addr = (level_3_table_addr << 9) | (level_3_index << 12);
-
-    // check that level 2 entry is mapped
-    let level_2_table = unsafe { &*(level_2_table_addr as *const PageTable) };
-    if level_2_table[level_2_index].addr().is_null() {
-        return None;
-    }
-    let level_1_table_addr = (level_2_table_addr << 9) | (level_2_index << 12);
-
-    // check that level 1 entry is mapped and retrieve physical address from it
-    let level_1_table = unsafe { &*(level_1_table_addr as *const PageTable) };
-    let phys_addr = level_1_table[level_1_index].addr();
-    if phys_addr.is_null() {
-        return None;
+/// Creates a RecursivePageTable instance from the level 4 address.
+///
+/// This function is unsafe because it can break memory safety if an invalid
+/// address is passed.
+pub unsafe fn init(level_4_table_addr: usize) -> RecursivePageTable<'static> {
+    /// Rust currently treats the whole body of unsafe functions as an unsafe
+    /// block, which makes it difficult to see which operations are unsafe. To
+    /// limit the scope of unsafe we use a safe inner function.
+    fn init_inner(level_4_table_addr: usize) -> RecursivePageTable<'static> {
+        let level_4_table_ptr = level_4_table_addr as *mut PageTable;
+        let level_4_table = unsafe { &mut *level_4_table_ptr };
+        RecursivePageTable::new(level_4_table).unwrap()
     }
 
-    Some(phys_addr + page_offset)
+    init_inner(level_4_table_addr)
+}
+
+/// Returns the physical address for the given virtual address, or `None` if
+/// the virtual address is not mapped.
+pub fn translate_addr(addr: u64, recursive_page_table: &RecursivePageTable)
+    -> Option<PhysAddr>
+{
+    let addr = VirtAddr::new(addr);
+    let page: Page = Page::containing_address(addr);
+
+    // perform the translation
+    let frame = recursive_page_table.translate_page(page);
+    frame.map(|frame| frame.start_address() + u64::from(addr.page_offset()))
 }