use std::{
process,
sync::atomic::{AtomicBool, Ordering},
};
use nix;
thread_local!(static TRACEME_DONE: AtomicBool = AtomicBool::new(false));
#[inline(always)]
fn ptraceme() -> Result<(), nix::errno::Errno> {
let res = nix::sys::ptrace::traceme();
match res {
Ok(_) => Ok(()),
Err(_) => Err(nix::errno::Errno::EPERM),
}
}
#[inline(always)]
pub fn ptraceme_or_die() {
let res = ptraceme();
TRACEME_DONE.with(|traceme_done| {
if !traceme_done.load(Ordering::SeqCst) {
match res {
Ok(_) => traceme_done.store(true, Ordering::SeqCst),
Err(_) => the_end(),
}
} else {
if res.is_ok() {
the_end()
}
}
});
}
#[inline(always)]
fn the_end() {
process::exit(0);
}
#[cfg(target_os = "linux")]
#[cfg(test)]
mod test {
use std::{thread, time::Duration};
#[test]
fn multiple_ptraceme_or_die() {
super::TRACEME_DONE.with(|s| s.store(false, std::sync::atomic::Ordering::SeqCst));
for i in 0..10 {
super::ptraceme_or_die();
println!("{}", i);
}
super::TRACEME_DONE.with(|s| assert_eq!(true, s.load(std::sync::atomic::Ordering::SeqCst)));
}
#[test]
fn multiple_threads_ptraceme_or_die() {
super::TRACEME_DONE.with(|s| s.store(false, std::sync::atomic::Ordering::SeqCst));
super::ptraceme_or_die();
super::TRACEME_DONE.with(|s| assert_eq!(true, s.load(std::sync::atomic::Ordering::SeqCst)));
let threads: Vec<_> = (0..10)
.map(|i| {
thread::spawn(move || {
super::TRACEME_DONE
.with(|s| s.store(false, std::sync::atomic::Ordering::SeqCst));
super::ptraceme_or_die();
thread::sleep(Duration::from_millis(i * 10));
eprintln!("Thread #{}", i);
super::ptraceme_or_die();
super::TRACEME_DONE
.with(|s| assert_eq!(true, s.load(std::sync::atomic::Ordering::SeqCst)));
})
})
.collect();
for thread in threads.into_iter() {
thread.join().unwrap();
}
}
}