diff --git a/Cargo.toml b/Cargo.toml index ef3d1af9..41df3774 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,9 +47,10 @@ wmi = "0.12.0" winreg = "0.10.1" windows = { version = "0.39.0", features = [ "Win32_Foundation", + "Win32_System_Diagnostics_Debug", "Win32_System_Power", "Win32_System_SystemInformation", - "Win32_System_WindowsProgramming" + "Win32_System_WindowsProgramming", ]} [target.'cfg(not(target_os = "windows"))'.dependencies] diff --git a/src/windows/mod.rs b/src/windows/mod.rs index 74fd1686..f9f7b9af 100644 --- a/src/windows/mod.rs +++ b/src/windows/mod.rs @@ -7,12 +7,14 @@ use wmi::WMIResult; use wmi::{COMLibrary, Variant, WMIConnection}; use windows::{ - core::PSTR, Win32::System::Power::GetSystemPowerStatus, + core::PSTR, + Win32::System::Power::GetSystemPowerStatus, Win32::System::Power::SYSTEM_POWER_STATUS, - Win32::System::SystemInformation::GetComputerNameExA, - Win32::System::SystemInformation::GetTickCount64, - Win32::System::SystemInformation::GlobalMemoryStatusEx, - Win32::System::SystemInformation::MEMORYSTATUSEX, + Win32::System::SystemInformation::{ + GetComputerNameExA, GetLogicalProcessorInformationEx, GetTickCount64, GlobalMemoryStatusEx, + RelationProcessorCore, GROUP_AFFINITY, MEMORYSTATUSEX, + SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX, + }, Win32::System::WindowsProgramming::GetUserNameA, }; @@ -295,11 +297,159 @@ impl GeneralReadout for WindowsGeneralReadout { } fn cpu_physical_cores(&self) -> Result { - Err(ReadoutError::NotImplemented) + // Source: https://github.com/AFLplusplus/LibAFL/blob/main/libafl/src/bolts/core_affinity.rs#L423 + + // Get the required size of the buffer, in bytes + let mut needed_size = 0; + unsafe { + GetLogicalProcessorInformationEx( + RelationProcessorCore, + std::ptr::null_mut(), + &mut needed_size, + ); + } + + // Could be 0, or some other bogus size + if needed_size != 0 { + // Allocate memory where we will store the processor info. + let mut buffer = vec![0u8; needed_size as usize]; + + // Populate the buffer with processor information + let result = unsafe { + GetLogicalProcessorInformationEx( + RelationProcessorCore, + buffer.as_mut_ptr() as *mut _, // cast to *mut _ to avoid type mismatch + &mut needed_size, + ) + }; + if result.as_bool() { + let mut n_cores: usize = 0; + + let mut byte_offset: usize = 0; + while byte_offset < needed_size as usize { + unsafe { + // Interpret the byte-array as a SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX struct + let part_ptr_raw = buffer.as_ptr().add(byte_offset); + let part_ptr: *const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX = + part_ptr_raw as *const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX; + let part: &SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX = &*part_ptr; + + // We are only interested in RelationProcessorCore information + if part.Relationship == RelationProcessorCore { + n_cores += 1; + } + + // Set the pointer to the next part as indicated by the size of this part + byte_offset += part.Size as usize; + } + } + + Ok(n_cores) + } else { + Err(ReadoutError::Other(String::from( + "Second call to \"GetLogicalProcessorInformationEx\" failed.", + ))) + } + } else { + Err(ReadoutError::Other(String::from( + "First call to \"GetLogicalProcessorInformationEx\" failed.", + ))) + } } fn cpu_cores(&self) -> Result { - Err(ReadoutError::NotImplemented) + // Source: https://github.com/AFLplusplus/LibAFL/blob/main/libafl/src/bolts/core_affinity.rs#L423 + + // Get the required size of the buffer, in bytes + let mut needed_size = 0; + unsafe { + GetLogicalProcessorInformationEx( + RelationProcessorCore, + std::ptr::null_mut(), + &mut needed_size, + ); + } + + // Could be 0, or some other bogus size + if needed_size != 0 { + // Allocate memory where we will store the processor info. + let mut buffer = vec![0u8; needed_size as usize]; + + // Populate the buffer with processor information + let result = unsafe { + GetLogicalProcessorInformationEx( + RelationProcessorCore, + buffer.as_mut_ptr() as *mut _, // cast to *mut _ to avoid type mismatch + &mut needed_size, + ) + }; + if result.as_bool() { + let mut n_logical_procs: usize = 0; + + let mut byte_offset: usize = 0; + while byte_offset < needed_size as usize { + unsafe { + // Interpret the byte-array as a SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX struct + let part_ptr_raw = buffer.as_ptr().add(byte_offset); + let part_ptr: *const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX = + part_ptr_raw as *const SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX; + let part: &SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX = &*part_ptr; + + // We are only interested in RelationProcessorCore information + if part.Relationship == RelationProcessorCore { + // The number of GROUP_AFFINITY structs in the array will be specified in the 'groupCount' + // We tentatively use the first element to get the pointer to it and reinterpret the + // entire slice with the groupCount + let groupmasks_slice: &[GROUP_AFFINITY] = std::slice::from_raw_parts( + part.Anonymous.Processor.GroupMask.as_ptr(), + part.Anonymous.Processor.GroupCount as usize, + ); + + // Count the local logical processors of the group and accumulate + let n_local_procs: usize = groupmasks_slice + .iter() + .map(|g| g.Mask.count_ones() as usize) + .sum::(); + n_logical_procs += n_local_procs; + } + + // Set the pointer to the next part as indicated by the size of this part + byte_offset += part.Size as usize; + } + } + + return Ok(n_logical_procs); + } + } + + // Alternative Implementation 1 + // Open the registry key containing the CPUs information. + if let Ok(cpu_info) = RegKey::predef(HKEY_LOCAL_MACHINE) + .open_subkey("HARDWARE\\DESCRIPTION\\System\\CentralProcessor") + { + // Count the number of subkeys, which is the number of CPU logical processors. + return Ok(cpu_info.enum_keys().count()); + } + + // Alternative Implementation 2 + // Use WMI to get the number of logical processors. + if let Ok(wmi_con) = wmi_connection() { + let results: Vec> = + wmi_con.raw_query("SELECT NumberOfLogicalProcessors FROM Win32_Processor")?; + + if let Some(result) = results.first() { + if let Some(Variant::String(val)) = result.get("NumberOfLogicalProcessors") { + if let Ok(out) = val.clone().parse::() { + return Ok(out); + } + } + } + } + + // If all else fails, return an error. + Err(ReadoutError::Other( + "Failed to get the number of logical processors.".to_string(), + )) } fn uptime(&self) -> Result {