moon-driver-core 0.1.3

Windows WDK driver core helpers: device, IOCTL, IRP, symlink
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

use core::ffi::c_void;

use alloc::boxed::Box;
use wdk_sys::{
    ntddk::IoDeleteDevice, DEVICE_OBJECT, IRP, IRP_MJ_CLEANUP, IRP_MJ_CLOSE, IRP_MJ_CREATE,
    NTSTATUS, STATUS_SUCCESS,
};

use crate::io_request::IoRequest;

pub struct Device {
    pub raw: *mut DEVICE_OBJECT,
}

impl Drop for Device {
    fn drop(&mut self) {
        if self.raw.is_null() {
            return;
        }
        unsafe {
            let extension = (*self.raw).DeviceExtension as *mut DeviceExtension;
            let vtable = (*extension).vtable;
            if let Some(release) = (*vtable).release {
                release(self.raw);
            }
            IoDeleteDevice(self.raw);
        }
    }
}

impl Device {
    pub fn from_raw(raw: *mut DEVICE_OBJECT) -> Self {
        Self { raw }
    }
    pub fn as_raw(&self) -> *const DEVICE_OBJECT {
        self.raw as *const _
    }
    pub fn as_raw_mut(&self) -> *mut DEVICE_OBJECT {
        self.raw
    }
    pub fn into_raw(mut self) -> *mut DEVICE_OBJECT {
        core::mem::replace(&mut self.raw, core::ptr::null_mut())
    }
    pub(crate) fn extension(&self) -> &DeviceExtension {
        unsafe { &*((*self.raw).DeviceExtension as *const DeviceExtension) }
    }
    pub(crate) fn extension_mut(&mut self) -> &mut DeviceExtension {
        unsafe { &mut *((*self.raw).DeviceExtension as *mut DeviceExtension) }
    }
    pub(crate) fn vtable(&self) -> &DeviceOperationsImpl {
        unsafe { &*(self.extension().vtable as *const _) }
    }
    pub fn data<T: DeviceOperations>(&self) -> &T {
        unsafe { &*(self.extension().data as *const T) }
    }
}

pub unsafe extern "C" fn dispatch_device(device: *mut DEVICE_OBJECT, irp: *mut IRP) -> NTSTATUS {
    let stack_location =
        unsafe { &*moon_driver_utils::kernel_fucntion::io_get_current_irp_stack_location(irp) };
    let device = Device::from_raw(device);
    let vtable = device.vtable();
    match vtable.dispatch {
        Some(dispatch) => dispatch(device.into_raw(), irp, stack_location.MajorFunction),
        _ => {
            device.into_raw();
            STATUS_SUCCESS
        }
    }
}

pub trait DeviceOperations {
    fn create(&self, _device: &Device, request: &mut IoRequest) -> Result<(), NTSTATUS> {
        request.complete(Ok(0));
        Ok(())
    }
    fn close(&self, _device: &Device, request: &mut IoRequest) -> Result<(), NTSTATUS> {
        request.complete(Ok(0));
        Ok(())
    }
    fn cleanup(&self, _device: &Device, request: &mut IoRequest) -> Result<(), NTSTATUS> {
        request.complete(Ok(0));
        Ok(())
    }
    fn others(&self, _device: &Device, request: &mut IoRequest) -> Result<(), NTSTATUS> {
        request.complete(Ok(0));
        Ok(())
    }
}

extern "C" fn dispatch_callback<T: DeviceOperations>(
    device: *mut DEVICE_OBJECT,
    irp: *mut IRP,
    major: u8,
) -> NTSTATUS {
    let device = Device::from_raw(device);
    let data: &T = device.data();
    let mut request = IoRequest::from_raw(irp);
    let device_ref = &device;
    let status = match major as _ {
        IRP_MJ_CREATE => data.create(device_ref, &mut request),
        IRP_MJ_CLOSE => data.close(device_ref, &mut request),
        IRP_MJ_CLEANUP => data.cleanup(device_ref, &mut request),
        _ => data.others(device_ref, &mut request),
    };
    device.into_raw();
    match status {
        Ok(()) => unsafe { (*request.as_raw()).IoStatus.__bindgen_anon_1.Status },
        Err(e) => {
            request.complete_status(e, 0);
            e
        }
    }
}

extern "C" fn release_callback<T: DeviceOperations>(device: *mut DEVICE_OBJECT) {
    unsafe {
        let extension = (*device).DeviceExtension as *mut DeviceExtension;
        let ptr = core::mem::replace(&mut (*extension).data, core::ptr::null_mut());
        let _ = Box::from_raw(ptr as *mut T);
    }
}

#[repr(C)]
pub struct DeviceExtension {
    pub(crate) vtable: *const DeviceOperationsImpl,
    pub(crate) data: *mut c_void,
}

#[repr(C)]
pub struct DeviceOperationsImpl {
    pub(crate) dispatch: Option<extern "C" fn(*mut DEVICE_OBJECT, *mut IRP, u8) -> NTSTATUS>,
    pub(crate) release: Option<extern "C" fn(*mut DEVICE_OBJECT)>,
}

pub struct DeviceOperationsVtable<T>(core::marker::PhantomData<T>);
impl<T: DeviceOperations> DeviceOperationsVtable<T> {
    pub const VTABLE: DeviceOperationsImpl = DeviceOperationsImpl {
        dispatch: Some(dispatch_callback::<T>),
        release: Some(release_callback::<T>),
    };
}