// DebugOff
// Copyright (C) 2022 0xor0ne
//
// Licensed under:
// - GPL-3.0 when "obfuscate" feature is enabled;
// - MIT when "obfuscate" feature IS NOT enabled;

use std::{
    process,
    sync::atomic::{AtomicBool, Ordering},
};

use nix;

thread_local!(static TRACEME_DONE: AtomicBool = AtomicBool::new(false));

/// Sets the process as traceable, as with `ptrace(PTRACE_TRACEME, ...)`
#[inline(always)]
fn ptraceme() -> Result<(), nix::errno::Errno> {
    let res = nix::sys::ptrace::traceme();

    match res {
        Ok(_) => Ok(()),
        Err(_) => Err(nix::errno::Errno::EPERM),
    }
}

/// Call `ptrace(PTRACE_TRACEME, ...)` one time to detect the presence of a debugger.
///
/// This function can be called multiple times.
///
/// At the first invocation, the function expects a return value of 0 from `ptrace(PTRACE_TRACEME, ...)`.
/// In subsequent calls, `ptrace(PTRACE_TRACEME, ...)` should return -1.
///
/// If the above is not satisfied, the function calls `exit_group(0)`.
///
/// To be more effective, the function should be called at least once for each thread.
///
/// ## Examples
///
/// ```rust
/// // Import only on Linux and for "release builds"
/// #[cfg(target_os = "linux")]
/// #[cfg(not(debug_assertions))]
/// use debugoff;
///
/// // Call only on Linux and for "release" builds.
/// #[cfg(target_os = "linux")]
/// #[cfg(not(debug_assertions))]
/// debugoff::ptraceme_or_die();
/// ```
#[inline(always)]
pub fn ptraceme_or_die() {
    let res = ptraceme();

    TRACEME_DONE.with(|traceme_done| {
        // The first time this function is called, res should be Ok(_). Subsequent calls should
        // return Err(_)
        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() {
    // Be careful, optimizer in release mode can decide to remove the following code.
    // let p: *mut u32 = core::ptr::null_mut();
    // unsafe {
    //     *p = 0xFF;
    // }

    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() {
        // Reset the state for this test
        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 || {
                    // Each thread has its own TRACEME_DONE
                    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();
        }
    }
}
