From 986030732ae5b194b2046fda4e0692e649377d1b Mon Sep 17 00:00:00 2001
From: Sergio Paganoni <sergio.paganoni@gmail.com>
Date: Thu, 9 Jun 2022 20:45:27 +0200
Subject: [PATCH] Generating core ids based on the actual count of logical
 cores (#669)

* generating core ids based on the actual count of logical cores

* make clippy happy

* make fmt happy
---
 libafl/src/bolts/core_affinity.rs | 56 +++++--------------------------
 1 file changed, 9 insertions(+), 47 deletions(-)

diff --git a/libafl/src/bolts/core_affinity.rs b/libafl/src/bolts/core_affinity.rs
index ef9962b01d..09f0ef744e 100644
--- a/libafl/src/bolts/core_affinity.rs
+++ b/libafl/src/bolts/core_affinity.rs
@@ -331,29 +331,22 @@ fn set_for_current_helper(core_id: CoreId) -> Result<(), Error> {
 
 #[cfg(target_os = "windows")]
 mod windows {
-    use core::ptr::addr_of_mut;
-
     use crate::bolts::core_affinity::{CoreId, Error};
-    use alloc::{string::ToString, vec::Vec};
+    use alloc::vec::Vec;
     use windows::Win32::System::SystemInformation::GROUP_AFFINITY;
-    use windows::Win32::System::Threading::{
-        GetCurrentProcess, GetCurrentThread, GetProcessAffinityMask, SetThreadGroupAffinity,
-    };
+    use windows::Win32::System::Threading::{GetCurrentThread, SetThreadGroupAffinity};
 
     pub fn get_core_ids() -> Result<Vec<CoreId>, Error> {
-        let mask = get_affinity_mask()?;
-        // Find all active cores in the bitmask.
         let mut core_ids: Vec<CoreId> = Vec::new();
-
-        for i in 0..64 {
-            let test_mask = 1 << i;
-
-            if (mask & test_mask) == test_mask {
-                core_ids.push(CoreId { id: i });
+        match get_num_logical_cpus_ex_windows() {
+            Some(total_cores) => {
+                for i in 0..total_cores {
+                    core_ids.push(CoreId { id: i });
+                }
+                Ok(core_ids)
             }
+            None => Err(Error::unknown("Unable to get logical CPUs count!")),
         }
-
-        Ok(core_ids)
     }
 
     pub fn set_for_current(id: CoreId) -> Result<(), Error> {
@@ -389,37 +382,6 @@ mod windows {
         }
     }
 
-    fn get_affinity_mask() -> Result<u64, Error> {
-        #[cfg(target_pointer_width = "64")]
-        let mut process_mask: u64 = 0;
-        #[cfg(target_pointer_width = "32")]
-        let mut process_mask: u32 = 0;
-        #[cfg(target_pointer_width = "64")]
-        let mut system_mask: u64 = 0;
-        #[cfg(target_pointer_width = "32")]
-        let mut system_mask: u32 = 0;
-
-        let res = unsafe {
-            GetProcessAffinityMask(
-                GetCurrentProcess(),
-                addr_of_mut!(process_mask) as _,
-                addr_of_mut!(system_mask) as _,
-            )
-        };
-
-        // Successfully retrieved affinity mask
-        if res.as_bool() {
-            #[allow(trivial_numeric_casts)]
-            Ok(process_mask as _)
-        }
-        // Failed to retrieve affinity mask
-        else {
-            Err(Error::unknown(
-                "Could not get affinity mask, GetProcessAffinityMask failed.".to_string(),
-            ))
-        }
-    }
-
     #[allow(trivial_numeric_casts)]
     #[allow(clippy::cast_ptr_alignment)]
     #[allow(clippy::cast_possible_wrap)]