From ccc24527962b3f143c7b44ae3215fff9e9dcf170 Mon Sep 17 00:00:00 2001
From: memN0ps <89628341+memN0ps@users.noreply.github.com>
Date: Tue, 30 Jul 2024 19:33:58 +1200
Subject: [PATCH] Added functions to unmap guest and large pages (unhook)

---
 hypervisor/src/error.rs                      |  6 +++
 hypervisor/src/intel/hooks/hook_manager.rs   |  4 ++
 hypervisor/src/intel/hooks/memory_manager.rs | 40 ++++++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/hypervisor/src/error.rs b/hypervisor/src/error.rs
index 8948ad5..9505119 100644
--- a/hypervisor/src/error.rs
+++ b/hypervisor/src/error.rs
@@ -238,4 +238,10 @@ pub enum HypervisorError {
 
     #[error("EPT misconfiguration error")]
     EptMisconfiguration,
+
+    #[error("Large page table unmapping error")]
+    LargePageUnmapError,
+
+    #[error("Guest page table unmapping error")]
+    GuestPageUnmapError,
 }
diff --git a/hypervisor/src/intel/hooks/hook_manager.rs b/hypervisor/src/intel/hooks/hook_manager.rs
index b15c839..d8e1e0b 100644
--- a/hypervisor/src/intel/hooks/hook_manager.rs
+++ b/hypervisor/src/intel/hooks/hook_manager.rs
@@ -440,6 +440,10 @@ impl HookManager {
         vm.primary_ept
             .swap_page(guest_page_pa.as_u64(), guest_page_pa.as_u64(), AccessType::READ_WRITE_EXECUTE, pre_alloc_pt)?;
 
+        // Update the memory manager to indicate that the guest page is no longer processed (unmapped/unhooked).
+        // This will allow the page to be reprocessed/remapped/rehooked if needed.
+        self.memory_manager.unmap_guest_from_shadow_page(guest_page_pa.as_u64())?;
+
         Ok(())
     }
 
diff --git a/hypervisor/src/intel/hooks/memory_manager.rs b/hypervisor/src/intel/hooks/memory_manager.rs
index 8ec92cf..7eb29b9 100644
--- a/hypervisor/src/intel/hooks/memory_manager.rs
+++ b/hypervisor/src/intel/hooks/memory_manager.rs
@@ -144,6 +144,46 @@ impl MemoryManager {
         Ok(())
     }
 
+    /// Unmaps a shadow page from a guest physical address, removing the associated hooks.
+    ///
+    /// # Arguments
+    /// * `guest_page_pa` - The guest physical address to unmap.
+    ///
+    /// # Returns
+    /// `Ok(())` if successful, or an error if the page was not mapped.
+    pub fn unmap_guest_from_shadow_page(&mut self, guest_page_pa: u64) -> Result<(), HypervisorError> {
+        trace!("Unmapping guest page and shadow page for PA: {:#x}", guest_page_pa);
+
+        // Remove the mapping if it exists
+        if self.guest_page_mappings.remove(&guest_page_pa).is_some() {
+            trace!("Guest page unmapped from shadow page successfully");
+            Ok(())
+        } else {
+            trace!("Guest page PA: {:#x} was not mapped", guest_page_pa);
+            Err(HypervisorError::GuestPageUnmapError)
+        }
+    }
+
+    /// Unmaps a page table from a large guest physical address.
+    ///
+    /// # Arguments
+    /// * `guest_large_page_pa` - The large guest physical address to unmap.
+    ///
+    /// # Returns
+    /// `Ok(())` if successful, or an error if the page table was not mapped.
+    pub fn unmap_large_page_from_pt(&mut self, guest_large_page_pa: u64) -> Result<(), HypervisorError> {
+        trace!("Unmapping large page and page table for PA: {:#x}", guest_large_page_pa);
+
+        // Remove the mapping if it exists
+        if self.large_page_table_mappings.remove(&guest_large_page_pa).is_some() {
+            trace!("Large page unmapped from page table successfully");
+            Ok(())
+        } else {
+            trace!("Large page PA: {:#x} was not mapped", guest_large_page_pa);
+            Err(HypervisorError::LargePageUnmapError)
+        }
+    }
+
     /// Retrieves a mutable reference to the page table associated with a large guest physical address.
     ///
     /// # Arguments