moon-driver-utils 0.1.0

use wdk::println;
use wdk_sys::{
    ntddk::{ExAcquireFastMutex, ExReleaseFastMutex, KeInitializeEvent},
    FAST_MUTEX,
    _EVENT_TYPE::SynchronizationEvent,
};

pub struct FastMutex {
    #[allow(unused)]
    raw: Option<FAST_MUTEX>,
    ptr: *mut FAST_MUTEX,
}
unsafe impl Send for FastMutex {}
unsafe impl Sync for FastMutex {}

pub struct FastMutexGuard<'a> {
    fast_mutex: &'a FastMutex,
}

pub fn init_fast_mutex(fm: &mut FAST_MUTEX) {
    fm.Count = 1;
    fm.Owner = core::ptr::null_mut();
    fm.Contention = 0;
    unsafe { KeInitializeEvent(&mut fm.Event as *mut _, SynchronizationEvent, 0) };
}

impl FastMutex {
    pub fn new() -> Self {
        let mut fm: wdk_sys::_FAST_MUTEX = FAST_MUTEX::default();
        init_fast_mutex(&mut fm);
        let r = Self {
            raw: Some(fm),
            ptr: core::ptr::null_mut(),
        };
        r
    }

    pub fn from(ptr: *mut FAST_MUTEX) -> Self {
        Self {
            raw: Option::None,
            ptr,
        }
    }

    pub fn lock(&self) -> FastMutexGuard {
        // println!("fast lock");
        unsafe { ExAcquireFastMutex(self.as_ptr()) };
        FastMutexGuard { fast_mutex: self }
    }

    pub fn acquire(&self) {
        // println!("fast acquire:{:p}", self.as_ptr());
        unsafe { ExAcquireFastMutex(self.as_ptr()) };
    }

    pub fn release(&self) {
        // println!("fast release:{:p}", self.as_ptr());
        unsafe { ExReleaseFastMutex(self.as_ptr()) };
    }

    pub fn as_ptr(&self) -> *mut FAST_MUTEX {
        if self.ptr.is_null() {
            return self
                .raw
                .as_ref()
                .map_or(core::ptr::null_mut(), |fast_mutex| {
                    fast_mutex as *const FAST_MUTEX as *mut _
                });
        } else {
            return self.ptr;
        }
    }
}

impl Drop for FastMutexGuard<'_> {
    fn drop(&mut self) {
        println!("fast drop");
        unsafe { ExReleaseFastMutex(self.fast_mutex.ptr) };
    }
}