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

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

@ -331,30 +331,23 @@ 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 {
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!")),
}
}
pub fn set_for_current(id: CoreId) -> Result<(), Error> {
let id: usize = id.into();
@ -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)]