From 7760697579ffbe4531f284ba1f550a53dc4dceee Mon Sep 17 00:00:00 2001
From: David CARLIER <devnexen@gmail.com>
Date: Sun, 28 Aug 2022 09:19:09 +0100
Subject: [PATCH] Netopenbsd build fix (#746)

* core affinity netbsd implementation.

* openbsd build fix
---
 libafl/src/bolts/core_affinity.rs   | 95 +++++++++++++++++++++++++++++
 libafl/src/bolts/os/unix_signals.rs | 32 +++++-----
 2 files changed, 113 insertions(+), 14 deletions(-)

diff --git a/libafl/src/bolts/core_affinity.rs b/libafl/src/bolts/core_affinity.rs
index 84629eac9d..efe6e1f1d9 100644
--- a/libafl/src/bolts/core_affinity.rs
+++ b/libafl/src/bolts/core_affinity.rs
@@ -697,6 +697,101 @@ mod freebsd {
     }
 }
 
+// NetBSD Section
+
+#[cfg(target_os = "netbsd")]
+#[inline]
+fn get_core_ids_helper() -> Result<Vec<CoreId>, Error> {
+    netbsd::get_core_ids()
+}
+
+#[cfg(target_os = "netbsd")]
+#[inline]
+fn set_for_current_helper(core_id: CoreId) -> Result<(), Error> {
+    netbsd::set_for_current(core_id)
+}
+
+#[cfg(target_os = "netbsd")]
+mod netbsd {
+    use alloc::vec::Vec;
+    use std::thread::available_parallelism;
+
+    use libc::{_cpuset, pthread_self, pthread_setaffinity_np};
+
+    use super::CoreId;
+    use crate::Error;
+
+    extern "C" {
+        fn _cpuset_create() -> *mut _cpuset;
+        fn _cpuset_destroy(c: *mut _cpuset);
+        fn _cpuset_set(i: usize, c: *mut _cpuset) -> i32;
+        fn _cpuset_size(c: *const _cpuset) -> usize;
+    }
+
+    #[allow(trivial_numeric_casts)]
+    pub fn get_core_ids() -> Result<Vec<CoreId>, Error> {
+        Ok((0..(usize::from(available_parallelism()?)))
+            .into_iter()
+            .map(|n| CoreId { id: n })
+            .collect::<Vec<_>>())
+    }
+
+    pub fn set_for_current(core_id: CoreId) -> Result<(), Error> {
+        let set = new_cpuset();
+
+        unsafe { _cpuset_set(core_id.id, set) };
+        // Set the current thread's core affinity.
+        let result = unsafe {
+            pthread_setaffinity_np(
+                pthread_self(), // Defaults to current thread
+                _cpuset_size(set),
+                set,
+            )
+        };
+
+        unsafe { _cpuset_destroy(set) };
+
+        if result < 0 {
+            Err(Error::unknown("Failed to set_for_current"))
+        } else {
+            Ok(())
+        }
+    }
+
+    fn new_cpuset() -> *mut _cpuset {
+        unsafe { _cpuset_create() }
+    }
+}
+
+#[cfg(target_os = "openbsd")]
+#[inline]
+fn get_core_ids_helper() -> Result<Vec<CoreId>, Error> {
+    openbsd::get_core_ids()
+}
+
+#[cfg(target_os = "openbsd")]
+#[inline]
+fn set_for_current_helper(_: CoreId) -> Result<(), Error> {
+    Ok(()) // There is no notion of cpu affinity on this platform
+}
+
+#[cfg(target_os = "openbsd")]
+mod openbsd {
+    use alloc::vec::Vec;
+    use std::thread::available_parallelism;
+
+    use super::CoreId;
+    use crate::Error;
+
+    #[allow(trivial_numeric_casts)]
+    pub fn get_core_ids() -> Result<Vec<CoreId>, Error> {
+        Ok((0..(usize::from(available_parallelism()?)))
+            .into_iter()
+            .map(|n| CoreId { id: n })
+            .collect::<Vec<_>>())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::thread::available_parallelism;
diff --git a/libafl/src/bolts/os/unix_signals.rs b/libafl/src/bolts/os/unix_signals.rs
index 1bc172c692..32428cb984 100644
--- a/libafl/src/bolts/os/unix_signals.rs
+++ b/libafl/src/bolts/os/unix_signals.rs
@@ -439,20 +439,24 @@ pub unsafe fn setup_signal_handler<T: 'static + Handler>(handler: &mut T) -> Res
 #[inline(always)]
 pub fn ucontext() -> Result<ucontext_t, Error> {
     let mut ucontext = unsafe { mem::zeroed() };
-    if unsafe { getcontext(&mut ucontext) } == 0 {
-        Ok(ucontext)
-    } else {
-        #[cfg(not(feature = "std"))]
-        unsafe {
-            libc::perror(b"Failed to get ucontext\n".as_ptr() as _);
-        };
-        #[cfg(not(feature = "std"))]
-        return Err(Error::unknown("Failed to get ucontex"));
+    if cfg!(not(target_os = "openbsd")) {
+        if unsafe { getcontext(&mut ucontext) } == 0 {
+            Ok(ucontext)
+        } else {
+            #[cfg(not(feature = "std"))]
+            unsafe {
+                libc::perror(b"Failed to get ucontext\n".as_ptr() as _);
+            };
+            #[cfg(not(feature = "std"))]
+            return Err(Error::unknown("Failed to get ucontex"));
 
-        #[cfg(feature = "std")]
-        Err(Error::unknown(format!(
-            "Failed to get ucontext: {:?}",
-            Errno::from_i32(errno())
-        )))
+            #[cfg(feature = "std")]
+            Err(Error::unknown(format!(
+                "Failed to get ucontext: {:?}",
+                Errno::from_i32(errno())
+            )))
+        }
+    } else {
+        Ok(ucontext)
     }
 }