Hypervisor appears to work for all processors

driver/src/main.rs

@@ -1,105 +1,118 @@
+// This crate provides the core functionality for initializing a hypervisor environment
+// within a UEFI application. It demonstrates advanced features such as custom panic handlers,
+// early logging, and direct manipulation of loaded image properties for hypervisor initialization.
+
 #![feature(new_uninit)]
 #![feature(panic_info_message)]
 #![no_main]
 #![no_std]
 
 extern crate alloc;
 
-use core::ffi::c_void;
 use {
-    crate::{
-        processor::MpManager,
-        virtualize::{switch_stack_and_virtualize_core, zap_relocations},
-    },
-    hypervisor::{
-        global::GlobalState,
-        intel::capture::{capture_registers, GuestRegisters},
-        logger,
-    },
+    crate::processor::start_hypervisor_on_all_processors,
+    hypervisor::logger,
     log::*,
-    uefi::prelude::*,
+    uefi::{prelude::*, proto::loaded_image::LoadedImage},
 };
 
 pub mod processor;
 pub mod virtualize;
 
-// Change as you like
+/// Custom panic handler for the UEFI application.
+///
+/// # Arguments
+///
+/// * `info` - Information about the panic, including the location and optional message.
 #[cfg(not(test))]
 #[panic_handler]
 fn panic_handler(info: &core::panic::PanicInfo) -> ! {
+    // Log the file, line, and column of the panic.
     if let Some(location) = info.location() {
         error!(
             "[-] Panic in {} at ({}, {}):",
             location.file(),
             location.line(),
             location.column()
         );
+        // Log the panic message if available.
         if let Some(message) = info.message() {
             error!("[-] {}", message);
         }
     }
 
+    // Enter an infinite loop as the panic handler should not return.
     loop {}
 }
 
+/// Entry point for the UEFI application.
+///
+/// Initializes logging, UEFI services, and attempts to start the hypervisor on all processors.
+///
+/// # Arguments
+///
+/// * `_image_handle` - Handle to the loaded image of the application.
+/// * `system_table` - Reference to the UEFI System Table.
+///
+/// # Returns
+///
+/// The status of the application execution. Returns `Status::SUCCESS` on successful execution,
+/// or `Status::ABORTED` if the hypervisor fails to install.
 #[entry]
 fn main(_image_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
-    // Initialize the COM2 port logger with level filter set to Info.
+    // Initialize logging with the COM2 port and set the level filter to Trace.
     logger::init(LevelFilter::Trace);
 
-    //com_logger::builder().base(0x2f8).filter(LevelFilter::Trace).setup();
-
+    // Initialize UEFI services.
     uefi_services::init(&mut system_table).unwrap();
 
     info!("The Matrix is an illusion");
 
-    // Get the MP Services (MultiProcessor Services) Protocol
-    let mp_manager = match MpManager::new(system_table.boot_services()) {
-        Ok(mp_manager) => mp_manager,
-        Err(e) => panic!("Failed to get MP Services: {:?}", e),
-    };
-
-    // Get the processor count
-    let processor_count = match mp_manager.processor_count() {
-        Ok(processor_count) => processor_count,
-        Err(e) => panic!("Failed to get processor count: {:?}", e),
-    };
-
-    info!("Total processors: {}", processor_count.total);
-    info!("Enabled processors: {}", processor_count.enabled);
-
     zap_relocations(&system_table);
 
-    // Capture the register values to be used as an initial state of the VM.
-    let mut guest_registers = GuestRegisters::default();
-    unsafe { capture_registers(&mut guest_registers) }
-
-    // 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 !mp_manager.is_virtualized() {
-        debug!("Virtualizing the system");
-        mp_manager.set_virtualized();
-
-        let mut global_state = GlobalState::new(guest_registers);
-
-        if processor_count.enabled == 1 {
-            info!("Found only one processor, virtualizing it");
-            switch_stack_and_virtualize_core(&mut global_state as *mut _ as *mut c_void);
-        } else {
-            info!("Found multiple processors, virtualizing all of them");
-            match mp_manager.start_virtualization_on_all_processors(
-                switch_stack_and_virtualize_core,
-                &mut global_state as *mut _ as *mut c_void,
-            ) {
-                Ok(_) => debug!("Virtualization started on all processors"),
-                Err(e) => panic!("Failed to start virtualization on all processors: {:?}", e),
-            }
+    // Attempt to start the hypervisor on all processors.
+    match start_hypervisor_on_all_processors(&system_table) {
+        Ok(_) => debug!("Hypervisor installed successfully"),
+        Err(e) => {
+            error!("Failed to install hypervisor: {:?}", e);
+            return Status::ABORTED;
         }
     }
 
-    info!("The hypervisor has been installed successfully!");
-
-    system_table.boot_services().stall(20_000_000);
-
+    // Return success status to UEFI environment.
     Status::SUCCESS
 }
+
+/// Nullifies the relocation table of the loaded UEFI image to prevent relocation.
+///
+/// This function manipulates the loaded image's PE header to zero out the relocation table,
+/// preventing UEFI from applying patches to the hypervisor code during the transition
+/// from physical-mode to virtual-mode addressing by the operating system.
+///
+/// # Arguments
+///
+/// * `system_table` - Reference to the UEFI System Table.
+pub fn zap_relocations(system_table: &SystemTable<Boot>) {
+    let boot_service = system_table.boot_services();
+
+    // Obtain the current loaded image protocol.
+    let loaded_image = boot_service
+        .open_protocol_exclusive::<LoadedImage>(boot_service.image_handle())
+        .unwrap();
+
+    // Extract the image base address and size.
+    let (image_base, image_size) = loaded_image.info();
+    let image_base = image_base as usize;
+    let image_range = image_base..image_base + image_size as usize;
+
+    // Log the image base address range for debugging purposes.
+    debug!("Image base: {:#x?}", image_range);
+
+    // Unsafe block to directly modify the PE header of the loaded image.
+    // This operation nullifies the relocation table to prevent UEFI from
+    // applying relocations to the hypervisor code.
+    unsafe {
+        *((image_base + 0x128) as *mut u32) = 0; // Zero out the relocation table offset.
+        *((image_base + 0x12c) as *mut u32) = 0; // Zero out the relocation table size.
+    }
+}

driver/src/processor.rs

@@ -1,69 +1,165 @@
-//! This module provides utility functions for processor-related operations in UEFI.
+//! This module provides utility functions for processor-related operations in UEFI,
+//! facilitating the initialization of virtualization across multiple processors.
 
 use {
+    crate::virtualize::virtualize_system,
     core::{
         ffi::c_void,
         sync::atomic::{AtomicU64, Ordering},
     },
+    hypervisor::intel::capture::{capture_registers, GuestRegisters},
+    log::*,
     uefi::{
         prelude::*,
         proto::pi::mp::{MpServices, Procedure, ProcessorCount},
         table::boot::ScopedProtocol,
     },
 };
 
-/// Atomic bitset used to track which processors have been virtualized.
-static VIRTUALIZED_BITSET: AtomicU64 = AtomicU64::new(0);
+/// Tracks the virtualization status of processors.
+///
+/// Each bit in this `AtomicU64` represents the virtualization status of a processor:
+/// a set bit indicates that the processor has been virtualized.
+pub static VIRTUALIZED_BITSET: AtomicU64 = AtomicU64::new(0);
 
 pub struct MpManager<'a> {
+    /// UEFI MP Services Protocol instance.
     mp_services: ScopedProtocol<'a, MpServices>,
 }
 
 impl<'a> MpManager<'a> {
-    /// Creates a new instance of MpManager, acquiring the MP Services Protocol.
+    /// Creates a new `MpManager` instance by acquiring the MP Services Protocol.
+    ///
+    /// # Arguments
+    ///
+    /// * `bt` - A reference to the UEFI Boot Services.
+    ///
+    /// # Returns
+    ///
+    /// A result containing the new `MpManager` instance or an error.
     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 })
     }
 
+    /// Initiates virtualization on all processors by executing the provided procedure.
+    ///
+    /// # Arguments
+    ///
+    /// * `procedure` - The function to execute on all application processors.
+    /// * `procedure_argument` - A pointer to the argument to pass to the procedure.
+    ///
+    /// # Returns
+    ///
+    /// A result indicating success or failure of the operation.
     pub fn start_virtualization_on_all_processors(
         &self,
         procedure: Procedure,
         procedure_argument: *mut c_void,
     ) -> uefi::Result<()> {
-        // The `procedure` is an `extern "efiapi" fn(_: *mut c_void)` compatible with `Procedure`
-        // and performs the necessary actions to initialize virtualization per-processor.
-        self.mp_services.startup_all_aps(
-            false, // Run on all processors simultaneously.
-            procedure,
-            procedure_argument,
-            None, // No associated event.
-            None, // No timeout.
-        )
+        self.mp_services
+            .startup_all_aps(false, procedure, procedure_argument, None, None)
     }
 
     /// Determines if the current processor is already virtualized.
+    ///
+    /// # Returns
+    ///
+    /// True if the current processor is virtualized, false otherwise.
     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
+        VIRTUALIZED_BITSET.load(Ordering::SeqCst) & 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);
+        VIRTUALIZED_BITSET.fetch_or(bit, Ordering::SeqCst);
     }
 
-    /// Returns the number of active logical processors.
+    /// Retrieves the number of active logical processors.
+    ///
+    /// # Returns
+    ///
+    /// A result containing the processor count or an error.
     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.
+    /// Identifies the index of the logical processor that is calling this method.
+    ///
+    /// # Returns
+    ///
+    /// A result containing the processor index or an error.
     pub fn current_processor_index(&self) -> uefi::Result<usize> {
         self.mp_services.who_am_i()
     }
 }
+
+/// Starts the hypervisor on all processors.
+///
+/// # Arguments
+///
+/// * `system_table` - A reference to the UEFI System Table.
+///
+/// # Returns
+///
+/// A result indicating the success or failure of starting the hypervisor.
+pub fn start_hypervisor_on_all_processors(system_table: &SystemTable<Boot>) -> uefi::Result<()> {
+    let mp_manager = MpManager::new(system_table.boot_services())?;
+    let processor_count = mp_manager.processor_count()?;
+
+    info!(
+        "Total processors: {}, Enabled processors: {}",
+        processor_count.total, processor_count.enabled
+    );
+
+    if processor_count.enabled == 1 {
+        info!("Found only one processor, virtualizing it");
+        start_hypervisor(&mp_manager);
+    } else {
+        info!("Found multiple processors, virtualizing all of them");
+        mp_manager.start_virtualization_on_all_processors(
+            start_hypervisor_on_ap,
+            &mp_manager as *const _ as *mut _,
+        )?;
+    }
+
+    info!("The hypervisor has been installed successfully!");
+
+    Ok(())
+}
+
+/// Hypervisor initialization procedure for Application Processors (APs).
+///
+/// # Arguments
+///
+/// * `procedure_argument` - A pointer to the `MpManager` instance.
+extern "efiapi" fn start_hypervisor_on_ap(procedure_argument: *mut c_void) {
+    let mp_manager = unsafe { &*(procedure_argument as *const MpManager) };
+    start_hypervisor(mp_manager);
+}
+
+/// Initiates the virtualization process.
+///
+/// # Arguments
+///
+/// * `mp_manager` - A reference to the `MpManager` to check and set virtualization status.
+fn start_hypervisor(mp_manager: &MpManager) {
+    let mut guest_registers = GuestRegisters::default();
+    // Unsafe block to capture the current CPU's register state.
+    unsafe { capture_registers(&mut guest_registers) };
+
+    // After capturing RIP, Guest execution will begin here. We then check for an existing hypervisor:
+    // if absent, proceed with installation; otherwise, no further action is needed.
+
+    // Proceed with virtualization only if the current processor is not yet virtualized.
+    if !mp_manager.is_virtualized() {
+        debug!("Virtualizing the system");
+        mp_manager.set_virtualized();
+        virtualize_system(&guest_registers);
+    }
+}