SYS-OSS/FRET-LibAFL
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Generating core ids based on the actual count of logical cores (#669)

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

* make clippy happy

* make fmt happy
This commit is contained in:
Sergio Paganoni 2022-06-09 20:45:27 +02:00 committed by GitHub
parent 395b616718
commit 986030732a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
libafl/src/bolts/core_affinity.rs
View File

@ -331,29 +331,22 @@ fn set_for_current_helper(core_id: CoreId) -> Result<(), Error> {
#[cfg(target_os = "windows")] #[cfg(target_os = "windows")]
mod windows { mod windows {
use core::ptr::addr_of_mut;
use crate::bolts::core_affinity::{CoreId, Error}; 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::SystemInformation::GROUP_AFFINITY;
use windows::Win32::System::Threading::{ use windows::Win32::System::Threading::{GetCurrentThread, SetThreadGroupAffinity};
GetCurrentProcess, GetCurrentThread, GetProcessAffinityMask, SetThreadGroupAffinity,
};
pub fn get_core_ids() -> Result<Vec<CoreId>, Error> { 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(); let mut core_ids: Vec<CoreId> = Vec::new();
match get_num_logical_cpus_ex_windows() {
for i in 0..64 { Some(total_cores) => {
let test_mask = 1 << i; for i in 0..total_cores {
core_ids.push(CoreId { id: i });
if (mask & test_mask) == test_mask { }
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> { 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(trivial_numeric_casts)]
#[allow(clippy::cast_ptr_alignment)] #[allow(clippy::cast_ptr_alignment)]
#[allow(clippy::cast_possible_wrap)] #[allow(clippy::cast_possible_wrap)]