Added Extended Page Tables (EPT) Page and inline hooks.

- Two sets of EPTs across all logical processors one for the default view, and the other for the hooked view.
- The MTRRs are not integrated using this way. Took the help of Matthias' code for the public version, which will help integration with the AMD hypervisor if needed. Private one might be different.
- **Note**: This is needs to be tested and is completely different from the dev branch. The dev branch is just an experiment, which did not work out well. This method is more popular but needs to be modified correctly and tested. Somethings are missing and there are crashes.

driver/src/hook.rs

@@ -0,0 +1,47 @@
+//! This module provides a safe wrapper around the system's memory validation function, typically named `MmIsAddressValid`.
+//! It allows checking the validity of addresses in a way that integrates with a system's memory management routines.
+//! The implementation uses a global atomic pointer to hold and replace the original system function with a custom hook,
+//! ensuring that any calls to check memory validity are routed through this custom implementation.
+//! Credits to Matthias: https://github.com/not-matthias/amd_hypervisor/blob/main/driver/src/hook.rs
+
+use core::{
+    mem, ptr,
+    sync::atomic::{AtomicPtr, Ordering},
+};
+
+// Extern block for interfacing with LLVM intrinsic for getting the return address.
+extern "C" {
+    // Links to the LLVM intrinsic to get the address of the return address.
+    #[link_name = "llvm.addressofreturnaddress"]
+    fn return_address() -> *const u64;
+}
+
+/// A global atomic pointer to hold the original `mm_is_address_valid` function.
+/// It's initialized to a null mutable pointer and will be set during runtime to the actual function.
+pub static ORIGINAL: AtomicPtr<u64> = AtomicPtr::new(ptr::null_mut());
+
+/// A safe wrapper around the `MmIsAddressValid` function.
+///
+/// ## Parameters
+/// - `ptr`: The pointer to check for validity.
+///
+/// ## Returns
+/// Returns `true` if the address is valid, `false` otherwise.
+///
+/// ## Safety
+/// This function assumes that the original `MmIsAddressValid` function is correctly set and points to a valid function.
+/// The caller must ensure this is the case to avoid undefined behavior.
+pub fn mm_is_address_valid(ptr: u64) -> bool {
+    // Log the address from which `MmIsAddressValid` was called.
+    log::info!("MmIsAddressValid called from {:#x}", unsafe {
+        return_address().read_volatile() // Reads the return address in a volatile manner to prevent optimizations.
+    });
+
+    // Load the original function pointer from the global atomic pointer.
+    let fn_ptr = ORIGINAL.load(Ordering::Relaxed); // Using relaxed ordering for atomic loading.
+                                                   // Transmute the function pointer to the expected function type.
+    let fn_ptr = unsafe { mem::transmute::<_, fn(u64) -> bool>(fn_ptr) };
+
+    // Call the original `MmIsAddressValid` function with the provided pointer.
+    fn_ptr(ptr)
+}

driver/src/lib.rs

@@ -5,8 +5,11 @@
 
 #![no_std]
 #![allow(unused_mut)]
+#![feature(allocator_api, new_uninit)]
+#![feature(link_llvm_intrinsics)]
 
 // Set up a panic handler for non-test configurations.
+extern crate alloc;
 #[cfg(not(test))]
 extern crate wdk_panic;
 
@@ -20,13 +23,26 @@ use wdk_alloc::WDKAllocator;
 #[global_allocator]
 static GLOBAL: hypervisor::utils::alloc::KernelAlloc = hypervisor::utils::alloc::KernelAlloc;
 
+use alloc::boxed::Box;
+use alloc::vec;
+use core::sync::atomic::Ordering;
 use {
-    hypervisor::Hypervisor,
+    hypervisor::{
+        intel::ept::{
+            access::AccessType,
+            hooks::{Hook, HookManager, HookType},
+            paging::Ept,
+        },
+        utils::alloc::PhysicalAllocator,
+        Hypervisor,
+    },
     log::LevelFilter,
     log::{self},
     wdk_sys::{DRIVER_OBJECT, NTSTATUS, PUNICODE_STRING, STATUS_SUCCESS, STATUS_UNSUCCESSFUL},
 };
 
+pub mod hook;
+
 /// The main entry point for the driver.
 ///
 /// This function is invoked by the system when the driver is loaded. It initializes
@@ -83,10 +99,12 @@ pub unsafe extern "system" fn driver_entry(
 pub extern "C" fn driver_unload(_driver: *mut DRIVER_OBJECT) {
     log::info!("Driver unloaded successfully!");
     if let Some(mut hypervisor) = unsafe { HYPERVISOR.take() } {
-        core::mem::drop(hypervisor);
+        drop(hypervisor);
     }
 }
 
+static mut HOOK_MANAGER: Option<HookManager> = None;
+
 /// The main hypervisor object.
 ///
 /// This static mutable option holds the global instance of the hypervisor used by this driver.
@@ -102,7 +120,56 @@ static mut HYPERVISOR: Option<Hypervisor> = None;
 /// * `Some(())` if the system was successfully virtualized.
 /// * `None` if there was an error during virtualization.
 fn virtualize() -> Option<()> {
-    let mut hypervisor = match Hypervisor::new() {
+    // Initialize the hook and hook manager
+    //
+    let hook = Hook::hook_function("MmIsAddressValid", hook::mm_is_address_valid as *const ())?;
+    if let HookType::Function { ref inline_hook } = hook.hook_type {
+        hook::ORIGINAL.store(inline_hook.trampoline_address(), Ordering::Relaxed);
+    }
+    let hook_manager = HookManager::new(vec![hook]);
+
+    // Setup the extended page tables. Because we also have hooks, we need to change
+    // the permissions of the page tables accordingly. This will be done by the
+    // `HookManager`.
+    //
+    let mut primary_ept: Box<Ept, PhysicalAllocator> = unsafe {
+        Box::try_new_zeroed_in(PhysicalAllocator)
+            .expect("failed to allocate primary ept")
+            .assume_init()
+    };
+    let mut secondary_ept: Box<Ept, PhysicalAllocator> = unsafe {
+        Box::try_new_zeroed_in(PhysicalAllocator)
+            .expect("failed to allocate secondary ept")
+            .assume_init()
+    };
+
+    log::info!("Setting Primary EPTs");
+    primary_ept.identity_4kb(AccessType::ReadWriteExecute);
+
+    log::info!("Setting Secondary EPTs");
+    secondary_ept.identity_4kb(AccessType::ReadWrite);
+
+    let primary_eptp = match primary_ept.create_eptp_with_wb_and_4lvl_walk() {
+        Ok(eptp) => eptp,
+        Err(err) => {
+            log::info!("Failed to create primary EPTP: {}", err);
+            return None;
+        }
+    };
+
+    let secondary_eptp = match secondary_ept.create_eptp_with_wb_and_4lvl_walk() {
+        Ok(eptp) => eptp,
+        Err(err) => {
+            log::info!("Failed to create secondary EPTP: {}", err);
+            return None;
+        }
+    };
+
+    hook_manager.enable_hooks(&mut primary_ept, &mut secondary_ept);
+
+    unsafe { HOOK_MANAGER = Some(hook_manager) };
+
+    let mut hypervisor = match Hypervisor::new(primary_eptp, secondary_eptp) {
         Ok(hypervisor) => hypervisor,
         Err(err) => {
             log::info!("Failed to initialize hypervisor: {}", err);

hypervisor/Cargo.toml

@@ -23,6 +23,7 @@ static_assertions = "1.1.0" # https://crates.io/crates/static_assertions
 log = "0.4.20" # https://crates.io/crates/log
 kernel-log = "0.1.2" # https://crates.io/crates/kernel-log
 com_logger = "0.1.1" # https://crates.io/crates/com_logger
+iced-x86 = { version = "1.16.0", default-features = false, features = ["no_std", "decoder", "block_encoder", "instr_info", "no_d3now", "no_evex", "no_vex", "no_xop"] } # https://crates.io/crates/iced-x86
 
 [build-dependencies]
 wdk-build = "0.1.0"

hypervisor/src/error.rs

@@ -73,4 +73,22 @@ pub enum HypervisorError {
 
     #[error("Invalid CR3 base address")]
     InvalidCr3BaseAddress,
+
+    #[error("Failed to parse bytes of original function")]
+    InvalidBytes,
+
+    #[error("Couldn't find enough space for the jump shellcode")]
+    NotEnoughBytes,
+
+    #[error("Failed to find original instructions")]
+    NoInstructions,
+
+    #[error("Failed to encode trampoline")]
+    EncodingFailed,
+
+    #[error("Found rip-relative instruction which is not supported")]
+    RelativeInstruction,
+
+    #[error("Found unsupported instruction")]
+    UnsupportedInstruction,
 }

hypervisor/src/intel/ept/access.rs

@@ -0,0 +1,91 @@
+//! This crate provides utilities for setting up and managing the memory paging mechanism
+//! for x86 architecture within a hypervisor context. Special attention is given to the
+//! manipulation of page table flags corresponding to different access types and page sizes.
+//! Credits to Matthias: https://github.com/not-matthias/amd_hypervisor/blob/main/hypervisor/src/svm/utils/paging.rs
+
+#![allow(dead_code)] // Allows for definitions that may not be used in all configurations.
+
+use x86::bits64::paging::{PDFlags, PDPTFlags, PML4Flags, PTFlags};
+
+// Constants defining sizes for various page table entries.
+pub const _512GB: u64 = 512 * 1024 * 1024 * 1024;
+pub const _1GB: u64 = 1024 * 1024 * 1024;
+pub const _2MB: usize = 2 * 1024 * 1024;
+pub const _4KB: usize = 4 * 1024;
+
+/// `AccessType` defines the types of access that can be granted to a page in the paging structure.
+#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Copy, Clone)]
+pub enum AccessType {
+    ReadWrite,
+    ReadWriteExecute,
+}
+
+impl AccessType {
+    /// Generates the appropriate PML4 flags for the access type.
+    pub fn pml4_flags(self) -> PML4Flags {
+        match self {
+            AccessType::ReadWrite => PML4Flags::P | PML4Flags::RW | PML4Flags::US | PML4Flags::XD,
+            AccessType::ReadWriteExecute => PML4Flags::P | PML4Flags::RW | PML4Flags::US,
+        }
+    }
+
+    /// Generates the appropriate PDPT flags for the access type.
+    pub fn pdpt_flags(self) -> PDPTFlags {
+        match self {
+            AccessType::ReadWrite => PDPTFlags::P | PDPTFlags::RW | PDPTFlags::US | PDPTFlags::XD,
+            AccessType::ReadWriteExecute => PDPTFlags::P | PDPTFlags::RW | PDPTFlags::US,
+        }
+    }
+
+    /// Generates the appropriate PD flags for the access type.
+    pub fn pd_flags(self) -> PDFlags {
+        match self {
+            AccessType::ReadWrite => PDFlags::P | PDFlags::RW | PDFlags::US | PDFlags::XD,
+            AccessType::ReadWriteExecute => PDFlags::P | PDFlags::RW | PDFlags::US,
+        }
+    }
+
+    /// Generates the appropriate PT flags for the access type.
+    pub fn pt_flags(self) -> PTFlags {
+        match self {
+            AccessType::ReadWrite => {
+                PTFlags::from_iter([PTFlags::P, PTFlags::RW, PTFlags::US, PTFlags::XD])
+            }
+            AccessType::ReadWriteExecute => {
+                PTFlags::from_iter([PTFlags::P, PTFlags::RW, PTFlags::US])
+            }
+        }
+    }
+
+    /// Modifies the PDFlags for a 2MB page based on the access type.
+    pub fn modify_2mb(&self, mut flags: PDFlags) -> PDFlags {
+        match self {
+            AccessType::ReadWrite => {
+                flags.insert(PDFlags::RW); // Set the ReadWrite flag.
+                flags.insert(PDFlags::XD); // Set the Execute Disable flag.
+            }
+            AccessType::ReadWriteExecute => {
+                flags.insert(PDFlags::RW); // Set the ReadWrite flag.
+                flags.remove(PDFlags::XD); // Remove the Execute Disable flag to allow execution.
+            }
+        }
+
+        flags // Return the modified flags.
+    }
+
+    /// Modifies the PTFlags for a 4KB page based on the access type.
+    pub fn modify_4kb(&self, mut flags: PTFlags) -> PTFlags {
+        match self {
+            AccessType::ReadWrite => {
+                flags.insert(PTFlags::RW); // Set the ReadWrite flag.
+                flags.insert(PTFlags::XD); // Set the Execute Disable flag.
+            }
+            AccessType::ReadWriteExecute => {
+                flags.insert(PTFlags::RW); // Set the ReadWrite flag.
+                flags.remove(PTFlags::XD); // Remove the Execute Disable flag to allow execution.
+            }
+        }
+
+        flags // Return the modified flags.
+    }
+}