moon-driver-utils 0.1.0

use core::ffi::c_void;

use wdk::println;
use wdk_sys::{
    ntddk::{KeWaitForSingleObject, PsCreateSystemThread, PsTerminateSystemThread},
    LARGE_INTEGER, NT_SUCCESS, STATUS_SUCCESS, THREAD_ALL_ACCESS,
    _KWAIT_REASON::Executive,
    _MODE::KernelMode,
};

use crate::{
    memory::npp::NPP,
    wrap::{ethread::Ethread, event::Event, handle::Handle},
};

pub struct SystemThread<T> {
    wait_start_event: Event,
    stop_event: Event,
    thread_handle: Handle,
    thread_objet: Ethread,
    thread_fn: ThreadFn<T>,
    thread_args: Option<T>,
    timer: Option<u64>,
}

type ThreadFn<T> = unsafe fn(args: &mut Option<T>);

impl<T> SystemThread<T> {
    pub fn set_stop_event(&mut self) {
        self.stop_event.set()
    }

    /// # Safety
    ///
    /// call system func
    pub unsafe extern "C" fn thread_proc(thread: *mut SystemThread<T>) {
        // some wrong in this
        println!("thread start");
        let t = &mut *thread;

        t.wait_start_event.set();

        let function = &mut t.thread_fn;
        let args = &mut t.thread_args;

        if t.timer.is_none() {
            function(args);
        } else {
            let mut time = LARGE_INTEGER::default();
            time.QuadPart = -1 * 10 * 1000 * t.timer.unwrap() as i64;

            loop {
                // wait exit
                let status = KeWaitForSingleObject(
                    t.stop_event.as_ptr() as _,
                    Executive as _,
                    KernelMode as _,
                    0,
                    &mut time as _,
                );

                // drop or no timer
                if status == STATUS_SUCCESS {
                    break;
                } else {
                    function(args);
                }
            }
        }

        println!("LOG Thread Exit");
        let _ = PsTerminateSystemThread(0);
    }

    pub fn as_mut_raw(&mut self) -> *mut SystemThread<T> {
        self as *mut _
    }

    pub fn new(
        thread: ThreadFn<T>,
        args: Option<T>,
        timer: Option<u64>,
    ) -> Result<NPP<Self>, &'static str> {
        let stop_event = Event::new()?;
        let wait_start_event = Event::new()?;

        let r = NPP::new(Self {
            wait_start_event,
            stop_event,
            thread_handle: Handle::default(),
            thread_objet: Ethread::default(),
            thread_fn: thread,
            thread_args: args,
            timer,
        });

        match r {
            Ok(rr) => Ok(rr),
            Err(_) => Err("Memory Allocate Error"),
        }
    }

    pub fn start(&mut self) -> bool {
        unsafe {
            let t = core::mem::transmute::<
                unsafe extern "C" fn(*mut SystemThread<T>),
                unsafe extern "C" fn(*mut c_void),
            >(Self::thread_proc);

            let status = PsCreateSystemThread(
                &mut self.thread_handle as *mut _ as _,
                THREAD_ALL_ACCESS,
                core::ptr::null_mut(),
                core::ptr::null_mut(),
                core::ptr::null_mut(),
                Some(t),
                self.as_mut_raw() as _,
            );

            if !NT_SUCCESS(status) {
                println!("KernelThread Create Error:{}", status);
                return false;
            }

            // wait thread satrt
            let _ = KeWaitForSingleObject(
                self.wait_start_event.as_ptr() as _,
                Executive as _,
                KernelMode as _,
                0,
                core::ptr::null_mut(),
            );

            self.thread_objet = Ethread::from_handle(&mut self.thread_handle as *mut _ as _);
        }
        true
    }
}

impl<T> Drop for SystemThread<T> {
    fn drop(&mut self) {
        unsafe {
            let _ = KeWaitForSingleObject(
                self.thread_objet.as_mut_raw() as _,
                Executive as _,
                KernelMode as _,
                0,
                core::ptr::null_mut(),
            );
            println!("thread_object exit");
        }
    }
}