dearxan 0.5.5

use std::{
    ptr::{NonNull, null_mut},
    sync::{
        LazyLock,
        atomic::{AtomicU64, Ordering::Relaxed},
    },
    time::{Duration, Instant},
};

use pelite::pe64::Pe;
use windows_sys::Win32::{
    Foundation::{
        CloseHandle, DUPLICATE_SAME_ACCESS, DuplicateHandle, GetLastError, HANDLE, NTSTATUS,
    },
    System::{
        Diagnostics::Debug::{CONTEXT, CONTEXT_FULL_AMD64, GetThreadContext},
        LibraryLoader::{GetModuleHandleA, GetProcAddress},
        Threading::{
            GetCurrentProcess, GetCurrentThreadId, GetThreadId, ResumeThread, SuspendThread,
            THREAD_ACCESS_RIGHTS, THREAD_ALL_ACCESS, THREAD_QUERY_INFORMATION,
            THREAD_SUSPEND_RESUME,
        },
    },
};

#[allow(
    non_camel_case_types,
    non_snake_case,
    non_upper_case_globals,
    dead_code
)]
mod ntdll {
    use super::*;

    #[cfg_attr(target_os = "windows", link(name = "ntdll", kind = "raw-dylib"))]
    unsafe extern "C" {
        pub fn NtGetNextThread(
            process_handle: HANDLE,
            thread_handle: HANDLE,
            desired_access: THREAD_ACCESS_RIGHTS,
            handle_attributes: u32,
            flags: u32,
            new_thread_handle: *mut HANDLE,
        ) -> NTSTATUS;

        pub fn NtQueryInformationThread(
            thread_handle: HANDLE,
            thread_information_class: u32,
            thread_information: *mut (),
            thread_information_length: usize,
            return_length: *mut usize,
        ) -> NTSTATUS;
    }

    pub const ThreadSuspendCount: u32 = 35;
    pub const ThreadQuerySetWin32StartAddress: u32 = 9;
}

use fxhash::FxHashMap;

use crate::{analysis::ImageView, disabler::game::game};

#[derive(Debug, PartialEq, Eq)]
pub struct OwnedHandle(NonNull<std::ffi::c_void>);

impl OwnedHandle {
    pub fn new(handle: HANDLE) -> Option<Self> {
        NonNull::new(handle).map(Self)
    }

    pub fn raw(&self) -> HANDLE {
        self.0.as_ptr()
    }
}

impl Clone for OwnedHandle {
    fn clone(&self) -> Self {
        let mut handle_out = null_mut();
        unsafe {
            let cproc = GetCurrentProcess();
            let success = DuplicateHandle(
                cproc,
                self.raw(),
                cproc,
                &mut handle_out,
                0,
                0,
                DUPLICATE_SAME_ACCESS,
            );
            if success == 0 {
                panic!("DuplicateHandle failed: {}", GetLastError());
            }
        }
        handle_out.into()
    }
}

impl Drop for OwnedHandle {
    fn drop(&mut self) {
        unsafe { CloseHandle(self.raw()) };
    }
}

impl From<HANDLE> for OwnedHandle {
    fn from(value: HANDLE) -> Self {
        Self(NonNull::new(value).unwrap())
    }
}

/// Find the virtual address of the global security cookie.
///
/// This is designed to work for MSVC binaries, but the pattern used might work on code
/// generated by other compilers.
///
/// The binary must have at least a few 10s of stack-protected functions for the analysis to be
/// successful.
fn find_gs_cookie_va(image: impl ImageView) -> Option<u64> {
    use memchr::memmem;

    const XOR_RAX_RSP: &[u8; 3] = b"\x48\x33\xc4";
    const MIN_COUNT: usize = 16;

    let mut va_counts: FxHashMap<u64, usize> = FxHashMap::default();

    image
        .sections()
        .flat_map(|(va, slice)| memmem::find_iter(slice, XOR_RAX_RSP).map(move |o| va + o as u64))
        .filter_map(|xor_va| {
            let offset = i32::from_le_bytes(image.read(xor_va - 4, 4)?[..4].try_into().unwrap());
            xor_va.checked_add_signed(offset.into())
        })
        .find(|&gs_cookie| {
            let count = va_counts.entry(gs_cookie).or_default();
            *count += 1;
            *count >= MIN_COUNT
        })
}

/// Checks if the global security cookie of the main executable has been initialized.
///
/// If the address of the GS cookie could not be found, returns [`None`].
fn is_gs_cookie_initialized() -> Option<bool> {
    static GS_COOKIE_ADDR: LazyLock<Option<u64>> = LazyLock::new(|| find_gs_cookie_va(game().pe));
    let gs_cookie_ptr = unsafe { AtomicU64::from_ptr((*GS_COOKIE_ADDR)? as *mut u64) };

    const UNITNIT_GS_COOKIE: u64 = 0x2b992ddfa232;
    Some(gs_cookie_ptr.load(Relaxed) != UNITNIT_GS_COOKIE)
}

pub unsafe fn wait_for_gs_cookie(timeout: Option<Duration>) -> Result<(), &'static str> {
    // Poll GS cookie every 10ms until it is no longer equal to the uninitialized value
    let ts = Instant::now();
    while timeout.is_none_or(|timeout| ts.elapsed() < timeout) {
        if is_gs_cookie_initialized().ok_or("global security cookie not found")? {
            return Ok(());
        }
        std::thread::sleep(Duration::from_millis(10));
    }
    Err("timed out waiting for __security_init_cookie")
}

/// Iterate over all threads in the current process, getting handles with the requested access
/// rights.
///
/// If `access` is `None`, the default rights of [`THREAD_ALL_ACCESS`] are used.
pub fn iter_threads(access: Option<THREAD_ACCESS_RIGHTS>) -> impl Iterator<Item = OwnedHandle> {
    let access = access.unwrap_or(THREAD_ALL_ACCESS);
    let proc = unsafe { GetCurrentProcess() };
    let mut thread: Option<OwnedHandle> = None;

    std::iter::from_fn(move || unsafe {
        let mut raw_thread = thread.as_ref().map(|t| t.raw()).unwrap_or_default();
        let status = ntdll::NtGetNextThread(proc, raw_thread, access, 0, 0, &mut raw_thread);
        thread = OwnedHandle::new(raw_thread);
        thread.as_ref().filter(|_| status >= 0).cloned()
    })
}

pub struct SuspendGuard {
    suspended: Vec<OwnedHandle>,
}

impl SuspendGuard {
    pub unsafe fn suspend_all_threads() -> Self {
        log::debug!("suspending all threads");

        unsafe {
            let current_thread = GetCurrentThreadId();

            let suspended = iter_threads(Some(THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION))
                .filter(|h| {
                    GetThreadId(h.raw()) != current_thread && SuspendThread(h.raw()) != u32::MAX
                })
                .collect();

            Self { suspended }
        }
    }
}

impl Drop for SuspendGuard {
    fn drop(&mut self) {
        log::debug!("resuming threads");
        for thread in std::mem::take(&mut self.suspended) {
            unsafe { ResumeThread(thread.raw()) };
        }
    }
}

pub fn process_main_thread() -> Option<OwnedHandle> {
    let pe = game().pe;
    let pe_ep = pe.optional_header().ImageBase + pe.optional_header().AddressOfEntryPoint as u64;

    iter_threads(None).find(|thread| unsafe {
        let mut thread_ep: u64 = 0;
        let status = ntdll::NtQueryInformationThread(
            thread.raw(),
            ntdll::ThreadQuerySetWin32StartAddress,
            (&raw mut thread_ep).cast(),
            size_of_val(&thread_ep),
            null_mut(),
        );
        status >= 0 && thread_ep == pe_ep
    })
}

pub fn is_created_suspended(thread: HANDLE) -> bool {
    static RTL_USER_THREAD_START: LazyLock<usize> = LazyLock::new(|| unsafe {
        let ntdll = GetModuleHandleA(c"ntdll.dll".as_ptr().cast());
        GetProcAddress(ntdll, c"RtlUserThreadStart".as_ptr().cast())
            .expect("RtlUserThreadStart not found") as usize
    });

    // Check if the thread is suspended
    let mut suspend_count: std::ffi::c_ulong = 0;
    let info_status = unsafe {
        ntdll::NtQueryInformationThread(
            thread,
            ntdll::ThreadSuspendCount,
            (&raw mut suspend_count).cast(),
            size_of_val(&suspend_count),
            null_mut(),
        )
    };
    if info_status < 0 {
        log::error!(
            "NtQueryInformationThread failed: {:x}",
            info_status.cast_unsigned()
        );
        return false;
    }
    if suspend_count == 0 {
        return false;
    }

    // Check the context to verify that it hasn't started executing its entry point yet
    let mut context = CONTEXT {
        ContextFlags: CONTEXT_FULL_AMD64,
        ..Default::default()
    };
    if unsafe { GetThreadContext(thread, &mut context) } == 0 {
        log::error!("GetThreadContext failed");
        return false;
    }

    let pe = game().pe;
    let pe_ep = pe.optional_header().ImageBase + pe.optional_header().AddressOfEntryPoint as u64;

    // We check if either are true instead of both to account for thread hijacking techniques
    // In particular, the Steam game overlay initializes itself using RIP thread hijacking
    // But another thread hijacking technique for suspended threads is overwriting RCX
    context.Rip == *RTL_USER_THREAD_START as u64 || context.Rcx == pe_ep
}

/// Check if the code calling this function is running before the game's entry point is executed.
pub fn is_pre_entry_point() -> bool {
    let Some(main_thread) = process_main_thread()
    else {
        log::warn!("process main thread not found, assuming dearxan is running pre entry point");
        return true;
    };

    let main_thread_id = unsafe { GetThreadId(main_thread.raw()) };
    let current_thread_id = unsafe { GetCurrentThreadId() };
    if main_thread_id == current_thread_id {
        !is_gs_cookie_initialized().unwrap_or_else(|| {
            log::warn!("gs cookie not found, assuming dearxan is running post entry point");
            true
        })
    }
    else {
        is_created_suspended(main_thread.raw())
    }
}