Merge pull request #26 from memN0ps/development

Add CPUID Password Protection and User-Mode Hook Management

Cargo.toml

@@ -5,6 +5,7 @@ members = [
     "client",
     "hypervisor",
     "loader",
+    "shared",
     "uefi",
 ]
 

client/Cargo.toml

@@ -4,4 +4,6 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-x86 = "0.52.0" # https://crates.io/crates/x86
+x86 = "0.52.0" # https://crates.io/crates/x86
+clap = { version = "4.5.4", features = ["derive"] } # https://crates.io/crates/clap
+shared = { path = "../shared" }

client/src/hypervisor_communicator.rs

@@ -4,28 +4,7 @@
 //! using the CPUID instruction. The communication is password protected to ensure
 //! that only authorized requests are processed by the hypervisor.
 
-use std::arch::asm;
-
-/// The password used for authentication with the hypervisor.
-pub const PASSWORD: u64 = 0xDEADBEEF;
-
-/// Enumeration of possible commands that can be issued to the hypervisor.
-///
-/// This enum represents different commands that can be sent to the hypervisor for
-/// various operations such as enabling hooks or disabling page hooks.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-#[repr(u64)]
-#[allow(dead_code)]
-pub enum Commands {
-    /// Command to enable a kernel inline hook.
-    EnableKernelInlineHook = 0,
-    /// Command to enable a syscall inline hook.
-    EnableSyscallInlineHook = 1,
-    /// Command to disable a page hook.
-    DisablePageHook = 2,
-    /// Invalid command.
-    Invalid,
-}
+use {shared::PASSWORD, std::arch::asm};
 
 /// Struct to encapsulate the result of a CPUID instruction.
 #[derive(Debug)]
@@ -52,18 +31,15 @@ impl HypervisorCommunicator {
     /// # Arguments
     ///
     /// * `command_rcx` - The value to be placed in the `rcx` register.
-    /// * `command_rdx` - The value to be placed in the `rdx` register.
-    /// * `command_r8` - The value to be placed in the `r8` register.
-    /// * `command_r9` - The value to be placed in the `r9` register.
     ///
     /// # Returns
     ///
     /// * `CpuidResult` - The result of the CPUID instruction.
-    pub fn call_hypervisor(&self, command_rcx: u64, command_rdx: u64, command_r8: u64, command_r9: u64) -> CpuidResult {
+    pub fn call_hypervisor(&self, command_rcx: u64) -> CpuidResult {
         let mut rax = PASSWORD;
         let mut rbx;
         let mut rcx = command_rcx;
-        let mut rdx = command_rdx;
+        let mut rdx;
 
         unsafe {
             asm!(
@@ -73,9 +49,7 @@ impl HypervisorCommunicator {
             out(reg) rbx,
             inout("rax") rax,
             inout("rcx") rcx,
-            inout("rdx") rdx,
-            in("r8") command_r8,
-            in("r9") command_r9,
+            lateout("rdx") rdx,
             options(nostack, preserves_flags),
             );
         }
@@ -88,59 +62,3 @@ impl HypervisorCommunicator {
         }
     }
 }
-
-/// Generate a unique hash
-///
-/// # Arguments
-///
-/// * `buffer` - The buffer to hash.
-///
-/// # Returns
-///
-/// * `u32` - The hash of the buffer.
-///
-/// # Example
-///
-/// ```
-/// let hash = djb2_hash(b"MmIsAddressValid");
-/// println!("Hash: {}", hash);
-/// ```
-#[allow(dead_code)]
-pub fn djb2_hash(buffer: &[u8]) -> u32 {
-    let mut hash: u32 = 5381;
-    for &byte in buffer {
-        let char = if byte >= b'a' { byte - 0x20 } else { byte };
-        hash = (hash << 5).wrapping_add(hash).wrapping_add(char as u32);
-    }
-    hash
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    /// Tests the `call_hypervisor` function for syscall hook.
-    ///
-    /// This test creates a new `HypervisorCommunicator` instance and sends a CPUID command
-    /// to set up a syscall hook with `NtQuerySystemInformation` (syscall number 0x36).
-    #[test]
-    fn test_call_hypervisor_syscall_hook() {
-        let communicator = HypervisorCommunicator::new();
-        let syscall_number = 0x36;
-        let result = communicator.call_hypervisor(Commands::EnableSyscallInlineHook as u64, syscall_number as u64, 1, 0);
-        assert_eq!(result.eax, 1);
-    }
-
-    /// Tests the `call_hypervisor` function for kernel hook.
-    ///
-    /// This test creates a new `HypervisorCommunicator` instance and sends a CPUID command
-    /// to set up a kernel inline hook with `MmIsAddressValid`.
-    #[test]
-    fn test_call_hypervisor_kernel_hook() {
-        let communicator = HypervisorCommunicator::new();
-        let function_name = b"MmIsAddressValid";
-        let function_hash = djb2_hash(function_name);
-        let result = communicator.call_hypervisor(Commands::EnableKernelInlineHook as u64, function_hash as u64, 0, 0);
-        assert_eq!(result.eax, 1);
-    }
-}

client/src/main.rs

@@ -5,20 +5,81 @@
 
 #![feature(asm_const)]
 
-use crate::hypervisor_communicator::{djb2_hash, Commands, HypervisorCommunicator};
+use {
+    crate::hypervisor_communicator::HypervisorCommunicator,
+    clap::{Parser, Subcommand},
+    shared::{djb2_hash, ClientData, Commands},
+};
 
 mod hypervisor_communicator;
 
-/// The main function demonstrating the usage of `HypervisorCommunicator`.
+/// Command line arguments for the Hypervisor Communicator.
+#[derive(Parser)]
+#[command(version, about, long_about = None)]
+struct Cli {
+    #[command(subcommand)]
+    command: CommandsArg,
+}
+
+#[derive(Subcommand)]
+enum CommandsArg {
+    /// Sets up a kernel inline hook
+    EnableKernelInlineHook {
+        /// The name of the function to hook
+        #[arg(short, long)]
+        function: String,
+    },
+    /// Unsets a kernel inline hook
+    DisableKernelInlineHook {
+        /// The name of the function to unhook
+        #[arg(short, long)]
+        function: String,
+    },
+}
+
 fn main() {
+    let cli = Cli::parse();
     let communicator = HypervisorCommunicator::new();
-    let function_name = b"MmIsAddressValid";
-    let function_hash = djb2_hash(function_name);
-    let result = communicator.call_hypervisor(Commands::EnableKernelInlineHook as u64, function_hash as u64, 0, 0);
-    println!("Result: {:#x} {:#x} {:#x} {:#x}", result.eax, result.ebx, result.ecx, result.edx);
+    match &cli.command {
+        CommandsArg::EnableKernelInlineHook { function: function_name } => {
+            let function_hash = djb2_hash(function_name.as_bytes());
+            println!("Function: {} Hash: {:#x}", function_name, function_hash);
 
-    let communicator = HypervisorCommunicator::new();
-    let syscall_number = 0x36;
-    let result = communicator.call_hypervisor(Commands::EnableSyscallInlineHook as u64, syscall_number as u64, 0, 0);
-    println!("Result: {:#x} {:#x} {:#x} {:#x}", result.eax, result.ebx, result.ecx, result.edx);
+            let client_data = ClientData {
+                command: Commands::EnableKernelInlineHook,
+                function_hash,
+            };
+
+            let client_data_ptr = client_data.as_ptr();
+            let result = communicator.call_hypervisor(client_data_ptr);
+
+            println!("Result: {:#x} {:#x} {:#x} {:#x}", result.eax, result.ebx, result.ecx, result.edx);
+
+            if result.eax == 0 {
+                println!("Failed to enable kernel inline hook");
+            } else {
+                println!("Successfully enabled kernel inline hook");
+            }
+        }
+        CommandsArg::DisableKernelInlineHook { function: function_name } => {
+            let function_hash = djb2_hash(function_name.as_bytes());
+            println!("Function: {} Hash: {:#x}", function_name, function_hash);
+
+            let client_data = ClientData {
+                command: Commands::DisableKernelInlineHook,
+                function_hash,
+            };
+
+            let client_data_ptr = client_data.as_ptr();
+            let result = communicator.call_hypervisor(client_data_ptr);
+
+            println!("Result: {:#x} {:#x} {:#x} {:#x}", result.eax, result.ebx, result.ecx, result.edx);
+
+            if result.eax == 0 {
+                println!("Failed to disable inline hook");
+            } else {
+                println!("Successfully disabled inline hook");
+            }
+        }
+    }
 }

hypervisor/Cargo.toml

@@ -30,4 +30,5 @@ bstr = { version = "1.9.0", default-features = false } # https://crates.io/crate
 derivative = { version = "2.2.0", features = ["use_core"]} # https://crates.io/crates/derivative
 spin = "0.9" # https://crates.io/crates/spin
 lde = "0.3.0" # https://crates.io/crates/lde
-heapless = "0.8.0" # https://crates.io/crates/heapless
+heapless = "0.8.0" # https://crates.io/crates/heapless
+shared = { path = "../shared" }

hypervisor/src/error.rs

@@ -223,4 +223,7 @@ pub enum HypervisorError {
 
     #[error("Shadow page already mapped")]
     ShadowPageAlreadyMapped,
+
+    #[error("Kernel hook missing")]
+    KernelHookMissing,
 }

hypervisor/src/intel/hooks/hook_manager.rs

@@ -32,18 +32,15 @@ pub enum EptHookType {
     Page,
 }
 
-/// The maximum number of hooks supported by the hypervisor. Change this value as needed.
-const MAX_ENTRIES: usize = 64;
-
 /// Represents hook manager structures for hypervisor operations.
 #[repr(C)]
 #[derive(Debug, Clone)]
 pub struct HookManager {
     /// The memory manager instance for the pre-allocated shadow pages and page tables.
-    pub memory_manager: Box<MemoryManager<MAX_ENTRIES>>,
+    pub memory_manager: Box<MemoryManager>,
 
     /// The hook instance for the Windows kernel, storing the VA and PA of ntoskrnl.exe. This is retrieved from the first LSTAR_MSR write operation, intercepted by the hypervisor.
-    pub kernel_hook: KernelHook,
+    pub kernel_hook: Option<Box<KernelHook>>,
 
     /// A flag indicating whether the CPUID cache information has been called. This will be used to perform hooks at boot time when SSDT has been initialized.
     /// KiSetCacheInformation -> KiSetCacheInformationIntel -> KiSetStandardizedCacheInformation -> __cpuid(4, 0)
@@ -69,12 +66,13 @@ impl HookManager {
     pub fn new() -> Result<Box<Self>, HypervisorError> {
         trace!("Initializing hook manager");
 
-        let memory_manager = Box::new(MemoryManager::<MAX_ENTRIES>::new()?);
+        let memory_manager = Box::new(MemoryManager::new()?);
+        let kernel_hook = Some(Box::new(KernelHook::new()?));
 
         Ok(Box::new(Self {
             memory_manager,
             has_cpuid_cache_info_been_called: false,
-            kernel_hook: Default::default(),
+            kernel_hook,
             old_rflags: None,
             mtf_counter: None,
         }))
@@ -169,6 +167,39 @@ impl HookManager {
         Ok(())
     }
 
+    /// Removes an EPT hook for a function.
+    ///
+    /// # Arguments
+    ///
+    /// * `vm` - The virtual machine instance of the hypervisor.
+    /// * `guest_function_va` - The virtual address of the function or page to be unhooked.
+    /// * `ept_hook_type` - The type of EPT hook to be removed.
+    ///
+    /// # Returns
+    ///
+    /// * Returns `Ok(())` if the hook was successfully removed, `Err(HypervisorError)` otherwise.
+    pub fn ept_unhook_function(vm: &mut Vm, guest_function_va: u64, _ept_hook_type: EptHookType) -> Result<(), HypervisorError> {
+        debug!("Removing EPT hook for function at VA: {:#x}", guest_function_va);
+
+        let guest_function_pa = PAddr::from(PhysicalAddress::pa_from_va(guest_function_va));
+        debug!("Guest function PA: {:#x}", guest_function_pa.as_u64());
+
+        let guest_page_pa = guest_function_pa.align_down_to_base_page();
+        debug!("Guest page PA: {:#x}", guest_page_pa.as_u64());
+
+        let pre_alloc_pt = vm
+            .hook_manager
+            .memory_manager
+            .get_page_table_as_mut(guest_page_pa.as_u64())
+            .ok_or(HypervisorError::PageTableNotFound)?;
+
+        // Swap the page back and restore the original page permissions
+        vm.primary_ept
+            .swap_page(guest_page_pa.as_u64(), guest_page_pa.as_u64(), AccessType::READ_WRITE_EXECUTE, pre_alloc_pt)?;
+
+        Ok(())
+    }
+
     /// Copies the guest page to the pre-allocated host shadow page.
     ///
     /// # Arguments

hypervisor/src/intel/hooks/memory_manager.rs

@@ -13,39 +13,42 @@ use {
     log::trace,
 };
 
+/// The maximum number of hooks supported by the hypervisor. Change this value as needed.
+const MAX_HOOK_ENTRIES: usize = 64;
+
 /// Represents a memory management system that pre-allocates and manages page tables
 /// and shadow pages for a hypervisor, using fixed-size arrays to avoid runtime allocation.
 #[derive(Debug, Clone)]
-pub struct MemoryManager<const N: usize> {
+pub struct MemoryManager {
     /// Active mappings of guest physical addresses to their respective page tables.
-    active_page_tables: LinearMap<u64, Box<Pt>, N>,
+    active_page_tables: LinearMap<u64, Box<Pt>, MAX_HOOK_ENTRIES>,
     /// Active mappings of guest physical addresses to their respective shadow pages.
-    active_shadow_pages: LinearMap<u64, Box<Page>, N>,
+    active_shadow_pages: LinearMap<u64, Box<Page>, MAX_HOOK_ENTRIES>,
 
     /// Pool of pre-allocated, free page tables available for assignment.
-    free_page_tables: Vec<Box<Pt>, N>,
+    free_page_tables: Vec<Box<Pt>, MAX_HOOK_ENTRIES>,
     /// Pool of pre-allocated, free shadow pages available for assignment.
-    free_shadow_pages: Vec<Box<Page>, N>,
+    free_shadow_pages: Vec<Box<Page>, MAX_HOOK_ENTRIES>,
 }
 
-impl<const N: usize> MemoryManager<N> {
+impl MemoryManager {
     /// Constructs a new `MemoryManager` instance, pre-allocating all necessary resources.
     ///
     /// # Returns
     /// A new instance of `MemoryManager` or an error if initial allocation fails.
     pub fn new() -> Result<Self, HypervisorError> {
         trace!("Initializing memory manager");
 
-        let active_page_tables = LinearMap::<u64, Box<Pt>, N>::new();
-        let active_shadow_pages = LinearMap::<u64, Box<Page>, N>::new();
+        let active_page_tables = LinearMap::<u64, Box<Pt>, MAX_HOOK_ENTRIES>::new();
+        let active_shadow_pages = LinearMap::<u64, Box<Page>, MAX_HOOK_ENTRIES>::new();
 
-        let mut free_page_tables = Vec::<Box<Pt>, N>::new();
-        let mut free_shadow_pages = Vec::<Box<Page>, N>::new();
+        let mut free_page_tables = Vec::<Box<Pt>, MAX_HOOK_ENTRIES>::new();
+        let mut free_shadow_pages = Vec::<Box<Page>, MAX_HOOK_ENTRIES>::new();
 
         trace!("Pre-allocating page tables and shadow pages");
 
         // Pre-allocate shadow pages and page tables for hooks.
-        for _ in 0..N {
+        for _ in 0..MAX_HOOK_ENTRIES {
             let pt = unsafe { box_zeroed::<Pt>() };
             let sp = unsafe { box_zeroed::<Page>() };