diff --git a/driver/src/main.rs b/driver/src/main.rs index a616fcb..952995d 100644 --- a/driver/src/main.rs +++ b/driver/src/main.rs @@ -1,3 +1,7 @@ +// 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] @@ -5,28 +9,25 @@ 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 ({}, {}):", @@ -34,72 +35,84 @@ fn panic_handler(info: &core::panic::PanicInfo) -> ! { 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) -> 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) { + let boot_service = system_table.boot_services(); + + // Obtain the current loaded image protocol. + let loaded_image = boot_service + .open_protocol_exclusive::(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. + } +} diff --git a/driver/src/processor.rs b/driver/src/processor.rs index 2a00de4..98ce5b0 100644 --- a/driver/src/processor.rs +++ b/driver/src/processor.rs @@ -1,10 +1,14 @@ -//! 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}, @@ -12,58 +16,150 @@ use { }, }; -/// 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 { let handle = bt.get_handle_for_protocol::()?; let mp_services = bt.open_protocol_exclusive::(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 { 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 { 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) -> 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); + } +} diff --git a/driver/src/virtualize.rs b/driver/src/virtualize.rs index 4881456..671bc8f 100644 --- a/driver/src/virtualize.rs +++ b/driver/src/virtualize.rs @@ -1,79 +1,36 @@ -use core::ffi::c_void; use { alloc::alloc::{alloc_zeroed, handle_alloc_error}, core::{alloc::Layout, arch::global_asm}, - hypervisor::{global::GlobalState, intel::page::Page, vmm::start_hypervisor}, - log::debug, - uefi::{ - proto::loaded_image::LoadedImage, - table::{Boot, SystemTable}, + hypervisor::{ + intel::{capture::GuestRegisters, page::Page}, + vmm::start_hypervisor, }, + log::debug, }; -pub fn zap_relocations(system_table: &SystemTable) { - let boot_service = system_table.boot_services(); - - // Open the loaded image protocol to get the current image base and image size. - let loaded_image = boot_service - .open_protocol_exclusive::(boot_service.image_handle()) - .unwrap(); - - // Get the current image base and image 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; - debug!("Image base: {:#x?}", image_range); - - // Prevent relocation by zapping the Relocation Table in the PE header. UEFI - // keeps the list of runtime drivers and applies patches into their code and - // data according to relocation information, as address translation switches - // from physical-mode to virtual-mode when the OS starts. This causes a problem - // with us because the host part keeps running under physical-mode, as the - // host has its own page tables. Relocation ends up breaking the host code. - // The easiest way is prevented this from happening is to nullify the relocation - // table. - unsafe { - *((image_base + 0x128) as *mut u32) = 0; - *((image_base + 0x12c) as *mut u32) = 0; - } -} - -pub fn allocate_stack() -> u64 { +/// Installs the hypervisor on the current processor. +pub fn virtualize_system(guest_registers: &GuestRegisters) -> ! { // Allocate separate stack space. This is never freed. let layout = Layout::array::(0x10).unwrap(); - let stack = unsafe { alloc_zeroed(layout) }; - if stack.is_null() { handle_alloc_error(layout); } - let stack_base = stack as u64 + layout.size() as u64 - 0x10; - debug!("Stack range: {:#x?}", stack_base..stack as u64); - - stack_base -} - -pub extern "efiapi" fn switch_stack_and_virtualize_core(procedure_argument: *mut c_void) { - let global_state = unsafe { &mut *(procedure_argument as *mut GlobalState) }; - let stack_base = allocate_stack(); + debug!("Stack range: {:#x?}", (stack as u64..stack_base)); - unsafe { switch_stack(global_state, start_hypervisor as usize, stack_base) }; + unsafe { switch_stack(guest_registers, start_hypervisor as usize, stack_base) }; } extern "efiapi" { /// Jumps to the landing code with the new stack pointer. - fn switch_stack(global_state: &mut GlobalState, landing_code: usize, stack_base: u64) -> !; + fn switch_stack(guest_registers: &GuestRegisters, landing_code: usize, stack_base: u64) -> !; } global_asm!( r#" -// The module containing the `switch_stack` function. -// Jumps to the landing code with the new stack pointer. -// -// fn switch_stack(global_state: &mut GlobalState, landing_code: usize, stack_base: u64) -> ! +// The module containing the `switch_stack` function. Jumps to the landing code with the new stack pointer. +// fn switch_stack(guest_registers: &GuestRegisters, landing_code: usize, stack_base: u64) -> ! .global switch_stack switch_stack: xchg bx, bx diff --git a/hypervisor/src/global.rs b/hypervisor/src/global.rs deleted file mode 100644 index 1ed56b0..0000000 --- a/hypervisor/src/global.rs +++ /dev/null @@ -1,11 +0,0 @@ -use crate::intel::capture::GuestRegisters; - -pub struct GlobalState { - pub guest_registers: GuestRegisters, -} - -impl GlobalState { - pub fn new(guest_registers: GuestRegisters) -> Self { - Self { guest_registers } - } -} diff --git a/hypervisor/src/lib.rs b/hypervisor/src/lib.rs index 91ad532..1dceba8 100644 --- a/hypervisor/src/lib.rs +++ b/hypervisor/src/lib.rs @@ -14,7 +14,6 @@ extern crate alloc; extern crate static_assertions; pub mod error; -pub mod global; pub mod intel; pub mod logger; pub mod vmm; diff --git a/hypervisor/src/vmm.rs b/hypervisor/src/vmm.rs index 9ae1d74..084fb93 100644 --- a/hypervisor/src/vmm.rs +++ b/hypervisor/src/vmm.rs @@ -1,7 +1,6 @@ use { crate::{ error::HypervisorError, - global::GlobalState, intel::{ capture::GuestRegisters, support::{vmread, vmwrite}, @@ -27,14 +26,9 @@ use { }; /// Starts the hypervisor. -pub fn start_hypervisor(global_state: &mut GlobalState) -> ! { +pub fn start_hypervisor(guest_registers: &GuestRegisters) -> ! { debug!("Starting hypervisor"); - let guest_registers = global_state.guest_registers; - - debug!("Guest registers: {:#x?}", guest_registers); - //debug!("Shared data: {:#p}", shared_data); - match check_supported_cpu() { Ok(_) => debug!("CPU is supported"), Err(e) => panic!("CPU is not supported: {:?}", e),