Hypervisor works without any crashes. CPUID vmexits mostly.

Author: memN0ps

driver/src/main.rs

@@ -4,6 +4,7 @@
 
 extern crate alloc;
 
+use crate::processor::MpManager;
 use {
     crate::virtualize::virtualize_system,
     alloc::boxed::Box,
@@ -14,12 +15,12 @@ use {
             ept::paging::{AccessType, Ept},
             shared_data::SharedData,
         },
-        vmm::is_hypervisor_present,
     },
     log::*,
     uefi::prelude::*,
 };
 
+pub mod processor;
 pub mod virtualize;
 
 #[entry]
@@ -34,6 +35,17 @@ fn main(_image_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
 
     info!("The Matrix is an illusion");
 
+    let mp_manager =
+        MpManager::new(system_table.boot_services()).expect("Failed to create MpManager");
+
+    // Get the processor count
+    if let Ok(processor_count) = mp_manager.processor_count() {
+        let total = processor_count.total;
+        let enabled = processor_count.enabled;
+        info!("Total processors: {}", total);
+        info!("Enabled processors: {}", enabled);
+    }
+
     let mut shared_data = match setup_ept() {
         Ok(shared_data) => shared_data,
         Err(e) => panic!("Failed to setup EPT: {:?}", e),
@@ -46,8 +58,9 @@ fn main(_image_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
     // Since we captured RIP just above, the VM will start running from here.
     // Check if our hypervisor is already loaded. If so, done, otherwise, continue
     // installing the hypervisor.
-    if !is_hypervisor_present() {
+    if !mp_manager.is_virtualized() {
         debug!("Virtualizing the system");
+        mp_manager.set_virtualized();
         virtualize_system(&regs, &mut shared_data, &system_table);
     }
 

driver/src/processor.rs

@@ -0,0 +1,46 @@
+//! This module provides utility functions for processor-related operations in UEFI.
+
+use core::sync::atomic::{AtomicU64, Ordering};
+use uefi::prelude::*;
+use uefi::proto::pi::mp::{MpServices, ProcessorCount};
+use uefi::table::boot::ScopedProtocol;
+
+/// Atomic bitset used to track which processors have been virtualized.
+static VIRTUALIZED_BITSET: AtomicU64 = AtomicU64::new(0);
+
+pub struct MpManager<'a> {
+    mp_services: ScopedProtocol<'a, MpServices>,
+}
+
+impl<'a> MpManager<'a> {
+    /// Creates a new instance of MpManager, acquiring the MP Services Protocol.
+    pub fn new(bt: &'a BootServices) -> uefi::Result<Self> {
+        let handle = bt.get_handle_for_protocol::<MpServices>()?;
+        let mp_services = bt.open_protocol_exclusive::<MpServices>(handle)?;
+        Ok(Self { mp_services })
+    }
+
+    /// Determines if the current processor is already virtualized.
+    pub fn is_virtualized(&self) -> bool {
+        let current_processor_index = self.current_processor_index().unwrap_or(0);
+        let bit = 1 << current_processor_index;
+        VIRTUALIZED_BITSET.load(Ordering::Relaxed) & bit != 0
+    }
+
+    /// Marks the current processor as virtualized.
+    pub fn set_virtualized(&self) {
+        let current_processor_index = self.current_processor_index().unwrap_or(0);
+        let bit = 1 << current_processor_index;
+        VIRTUALIZED_BITSET.fetch_or(bit, Ordering::Relaxed);
+    }
+
+    /// Returns the number of active logical processors.
+    pub fn processor_count(&self) -> uefi::Result<ProcessorCount> {
+        self.mp_services.get_number_of_processors()
+    }
+
+    /// Gets the processor number of the logical processor that the caller is running on.
+    pub fn current_processor_index(&self) -> uefi::Result<usize> {
+        self.mp_services.who_am_i()
+    }
+}

hypervisor/src/vmm.rs

@@ -26,15 +26,6 @@ use {
     },
 };
 
-pub const CPUID_VENDOR_AND_MAX_FUNCTIONS: u32 = 0x4000_0000;
-pub const VENDOR_NAME: u32 = 0x5441_4c48; // "HLAT"
-
-/// Checks if this hypervisor is already installed.
-pub fn is_hypervisor_present() -> bool {
-    let regs = cpuid!(CPUID_VENDOR_AND_MAX_FUNCTIONS);
-    (regs.ebx == regs.ecx) && (regs.ecx == regs.edx) && (regs.edx == VENDOR_NAME)
-}
-
 // pass shared data to the hypervisor soon too
 pub fn start_hypervisor(guest_registers: &GuestRegisters, shared_data: &mut SharedData) -> ! {
     debug!("Starting hypervisor");
@@ -57,7 +48,7 @@ pub fn start_hypervisor(guest_registers: &GuestRegisters, shared_data: &mut Shar
     loop {
         if let Ok(basic_exit_reason) = vm.run() {
             let exit_type = match basic_exit_reason {
-                //VmxBasicExitReason::ExceptionOrNmi => handle_exception(guest_registers, vmx),
+                //VmxBasicExitReason::ExceptionOrNmi => handle_exception(vmx),
                 VmxBasicExitReason::Cpuid => handle_cpuid(&mut vm.guest_registers),
 
                 // Grouping multiple exit reasons that are handled by the same function