From e2d248b609ffe51e4da4331f6d6418c847338434 Mon Sep 17 00:00:00 2001
From: pjht <pterpstra@pterpstra.com>
Date: Sun, 7 Jul 2024 08:06:15 -0500
Subject: [PATCH] Fix virtual_memory::drop_table freeing each mapping frame
 multiple times

---
 src/virtual_memory.rs | 33 ++++++++++++++-------------------
 1 file changed, 14 insertions(+), 19 deletions(-)

diff --git a/src/virtual_memory.rs b/src/virtual_memory.rs
index d477a45..0f8c2d3 100644
--- a/src/virtual_memory.rs
+++ b/src/virtual_memory.rs
@@ -1,6 +1,6 @@
 mod holes;
 
-use crate::{bootinfo::BOOTINFO, dbg, physical_memory::PHYSICAL_MEMORY, println};
+use crate::{bootinfo::BOOTINFO, physical_memory::PHYSICAL_MEMORY};
 use alloc::alloc::{AllocError, Allocator, Layout};
 use core::{fmt, ops::Deref, ptr::NonNull, slice};
 use replace_with::replace_with_or_abort;
@@ -44,7 +44,7 @@ pub enum PagingError {
     PageAlreadyMapped,
     PagesAlreadyMapped,
     PageNotMapped,
-    InvalidFrameAddress(PhysAddr),
+    InvalidFrameAddress(#[allow(dead_code)] PhysAddr),
 }
 
 impl From<MapToError<Size4KiB>> for PagingError {
@@ -474,14 +474,14 @@ impl Drop for AddressSpace {
 }
 
 fn drop_table(table: &PageTable, level: u8) {
-    for (i, entry) in table.iter().enumerate() {
-        if level == 4 && i >= 256 {
-            continue;
-        }
-        if entry.flags().contains(PageTableFlags::PRESENT)
-            && !entry.flags().contains(PageTableFlags::HUGE_PAGE)
-        {
-            if level > 2 {
+    if level > 2 {
+        for (i, entry) in table.iter().enumerate() {
+            if level == 4 && i >= 256 {
+                continue;
+            }
+            if entry.flags().contains(PageTableFlags::PRESENT)
+                && !entry.flags().contains(PageTableFlags::HUGE_PAGE)
+            {
                 // SAFETY: The present flag is set on the entry, which means the child frame must
                 // contain a valid page table, so making a reference to it must be ok.
                 // Unwrap: from_start_address requires it's input to be 4KiB aligned (have none of its lower 12 bits
@@ -493,17 +493,12 @@ fn drop_table(table: &PageTable, level: u8) {
                     level - 1,
                 );
             }
-            let phys_addr = match ACTIVE_SPACE.lock().translate(VirtAddr::from_ptr(table)) {
-                TranslateResult::Mapped { frame, .. } => frame.start_address(),
-                _ => panic!("Refrence must point to mapped page!"),
-            };
-            unsafe {
-                PHYSICAL_MEMORY
-                    .lock()
-                    .deallocate_frame(PhysFrame::from_start_address(phys_addr).unwrap());
-            };
         }
     }
+    let phys_addr = ACTIVE_SPACE.lock().translate_addr(VirtAddr::from_ptr(table)).unwrap();
+    unsafe {
+        PHYSICAL_MEMORY.lock().deallocate_frame(PhysFrame::from_start_address(phys_addr).unwrap());
+    };
 }
 
 impl Translate for AddressSpace {