Windows: Implement cpu_physical_cores and cpu_cores

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]

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<usize, ReadoutError> {
-        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<usize, ReadoutError> {
-        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::<usize>();
+                            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<HashMap<String, Variant>> =
+                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::<usize>() {
+                        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<usize, ReadoutError> {