wraith/km/
driver.rs

1//! Driver object management
2
3use core::ffi::c_void;
4use core::ptr::NonNull;
5
6use super::device::Device;
7use super::error::{status, KmError, KmResult, NtStatus};
8use super::irp::{Irp, IrpMajorFunction};
9use super::string::UnicodeString;
10
11/// driver object wrapper
12#[repr(C)]
13pub struct DriverObjectRaw {
14    pub type_: i16,
15    pub size: i16,
16    pub device_object: *mut c_void,
17    pub flags: u32,
18    pub driver_start: *mut c_void,
19    pub driver_size: u32,
20    pub driver_section: *mut c_void,
21    pub driver_extension: *mut DriverExtensionRaw,
22    pub driver_name: UnicodeStringRaw,
23    pub hardware_database: *mut UnicodeStringRaw,
24    pub fast_io_dispatch: *mut c_void,
25    pub driver_init: *mut c_void,
26    pub driver_start_io: *mut c_void,
27    pub driver_unload: Option<DriverUnload>,
28    pub major_function: [Option<DriverDispatch>; 28],
29}
30
31/// driver extension
32#[repr(C)]
33pub struct DriverExtensionRaw {
34    pub driver_object: *mut DriverObjectRaw,
35    pub add_device: *mut c_void,
36    pub count: u32,
37    pub service_key_name: UnicodeStringRaw,
38}
39
40/// raw unicode string (for FFI)
41#[repr(C)]
42#[derive(Clone, Copy)]
43pub struct UnicodeStringRaw {
44    pub length: u16,
45    pub maximum_length: u16,
46    pub buffer: *mut u16,
47}
48
49/// driver entry function signature
50pub type DriverEntry = unsafe extern "system" fn(
51    driver_object: *mut DriverObjectRaw,
52    registry_path: *const UnicodeStringRaw,
53) -> NtStatus;
54
55/// driver unload function signature
56pub type DriverUnload = unsafe extern "system" fn(driver_object: *mut DriverObjectRaw);
57
58/// IRP dispatch function signature
59pub type DriverDispatch = unsafe extern "system" fn(
60    device_object: *mut c_void,
61    irp: *mut c_void,
62) -> NtStatus;
63
64/// safe driver object wrapper
65pub struct Driver {
66    raw: NonNull<DriverObjectRaw>,
67}
68
69impl Driver {
70    /// wrap raw driver object pointer
71    ///
72    /// # Safety
73    /// ptr must be a valid DRIVER_OBJECT pointer
74    pub unsafe fn from_raw(ptr: *mut DriverObjectRaw) -> Option<Self> {
75        NonNull::new(ptr).map(|raw| Self { raw })
76    }
77
78    /// get raw pointer
79    pub fn as_raw(&self) -> *mut DriverObjectRaw {
80        self.raw.as_ptr()
81    }
82
83    /// set driver unload routine
84    pub fn set_unload(&mut self, unload: DriverUnload) {
85        // SAFETY: we have exclusive access through mut reference
86        unsafe {
87            (*self.raw.as_ptr()).driver_unload = Some(unload);
88        }
89    }
90
91    /// set major function handler
92    pub fn set_major_function(&mut self, function: IrpMajorFunction, handler: DriverDispatch) {
93        let index = function as usize;
94        if index < 28 {
95            // SAFETY: valid index
96            unsafe {
97                (*self.raw.as_ptr()).major_function[index] = Some(handler);
98            }
99        }
100    }
101
102    /// set all major functions to same handler
103    pub fn set_all_major_functions(&mut self, handler: DriverDispatch) {
104        // SAFETY: we have exclusive access
105        unsafe {
106            for i in 0..28 {
107                (*self.raw.as_ptr()).major_function[i] = Some(handler);
108            }
109        }
110    }
111
112    /// get driver name
113    pub fn name(&self) -> &[u16] {
114        // SAFETY: driver object is valid
115        unsafe {
116            let name = &(*self.raw.as_ptr()).driver_name;
117            if name.buffer.is_null() || name.length == 0 {
118                return &[];
119            }
120            core::slice::from_raw_parts(name.buffer, (name.length / 2) as usize)
121        }
122    }
123
124    /// get driver start address
125    pub fn start_address(&self) -> *mut c_void {
126        // SAFETY: driver object is valid
127        unsafe { (*self.raw.as_ptr()).driver_start }
128    }
129
130    /// get driver size
131    pub fn size(&self) -> u32 {
132        // SAFETY: driver object is valid
133        unsafe { (*self.raw.as_ptr()).driver_size }
134    }
135
136    /// create a device
137    pub fn create_device(
138        &mut self,
139        name: &UnicodeString,
140        device_type: u32,
141        characteristics: u32,
142        exclusive: bool,
143    ) -> KmResult<Device> {
144        let mut device_object: *mut c_void = core::ptr::null_mut();
145
146        // SAFETY: valid driver object
147        let status = unsafe {
148            IoCreateDevice(
149                self.raw.as_ptr(),
150                0, // device extension size
151                name.as_ptr() as *const _,
152                device_type,
153                characteristics,
154                if exclusive { 1 } else { 0 },
155                &mut device_object,
156            )
157        };
158
159        if !status::nt_success(status) {
160            return Err(KmError::DeviceCreationFailed {
161                reason: "IoCreateDevice failed",
162            });
163        }
164
165        // SAFETY: IoCreateDevice succeeded
166        unsafe { Device::from_raw(device_object) }
167            .ok_or(KmError::DeviceCreationFailed {
168                reason: "device object is null",
169            })
170    }
171}
172
173/// driver builder for setting up dispatch routines
174pub struct DriverBuilder {
175    driver: Driver,
176}
177
178impl DriverBuilder {
179    /// create builder from raw driver object
180    ///
181    /// # Safety
182    /// ptr must be valid DRIVER_OBJECT
183    pub unsafe fn new(ptr: *mut DriverObjectRaw) -> Option<Self> {
184        // SAFETY: caller ensures ptr is valid
185        unsafe { Driver::from_raw(ptr) }.map(|driver| Self { driver })
186    }
187
188    /// set unload routine
189    pub fn unload(mut self, handler: DriverUnload) -> Self {
190        self.driver.set_unload(handler);
191        self
192    }
193
194    /// set create handler (IRP_MJ_CREATE)
195    pub fn create(mut self, handler: DriverDispatch) -> Self {
196        self.driver.set_major_function(IrpMajorFunction::Create, handler);
197        self
198    }
199
200    /// set close handler (IRP_MJ_CLOSE)
201    pub fn close(mut self, handler: DriverDispatch) -> Self {
202        self.driver.set_major_function(IrpMajorFunction::Close, handler);
203        self
204    }
205
206    /// set device control handler (IRP_MJ_DEVICE_CONTROL)
207    pub fn device_control(mut self, handler: DriverDispatch) -> Self {
208        self.driver.set_major_function(IrpMajorFunction::DeviceControl, handler);
209        self
210    }
211
212    /// set read handler (IRP_MJ_READ)
213    pub fn read(mut self, handler: DriverDispatch) -> Self {
214        self.driver.set_major_function(IrpMajorFunction::Read, handler);
215        self
216    }
217
218    /// set write handler (IRP_MJ_WRITE)
219    pub fn write(mut self, handler: DriverDispatch) -> Self {
220        self.driver.set_major_function(IrpMajorFunction::Write, handler);
221        self
222    }
223
224    /// set custom major function handler
225    pub fn major_function(mut self, function: IrpMajorFunction, handler: DriverDispatch) -> Self {
226        self.driver.set_major_function(function, handler);
227        self
228    }
229
230    /// build and return the configured driver
231    pub fn build(self) -> Driver {
232        self.driver
233    }
234}
235
236/// helper trait for implementing driver entry
237pub trait DriverImpl {
238    /// called during driver initialization
239    fn init(driver: &mut Driver, registry_path: &UnicodeString) -> KmResult<()>;
240
241    /// called during driver unload
242    fn unload(driver: &Driver);
243
244    /// handle IRP_MJ_CREATE
245    fn create(_device: *mut c_void, _irp: &mut Irp) -> NtStatus {
246        status::STATUS_SUCCESS
247    }
248
249    /// handle IRP_MJ_CLOSE
250    fn close(_device: *mut c_void, _irp: &mut Irp) -> NtStatus {
251        status::STATUS_SUCCESS
252    }
253
254    /// handle IRP_MJ_DEVICE_CONTROL
255    fn device_control(_device: *mut c_void, _irp: &mut Irp) -> NtStatus {
256        status::STATUS_NOT_IMPLEMENTED
257    }
258
259    /// handle IRP_MJ_READ
260    fn read(_device: *mut c_void, _irp: &mut Irp) -> NtStatus {
261        status::STATUS_NOT_IMPLEMENTED
262    }
263
264    /// handle IRP_MJ_WRITE
265    fn write(_device: *mut c_void, _irp: &mut Irp) -> NtStatus {
266        status::STATUS_NOT_IMPLEMENTED
267    }
268}
269
270/// macro to generate driver entry boilerplate
271#[macro_export]
272macro_rules! driver_entry {
273    ($impl_type:ty) => {
274        #[no_mangle]
275        pub unsafe extern "system" fn DriverEntry(
276            driver_object: *mut $crate::km::driver::DriverObjectRaw,
277            registry_path: *const $crate::km::driver::UnicodeStringRaw,
278        ) -> $crate::km::error::NtStatus {
279            // SAFETY: called from kernel with valid parameters
280            unsafe { __driver_entry_impl::<$impl_type>(driver_object, registry_path) }
281        }
282
283        unsafe fn __driver_entry_impl<T: $crate::km::driver::DriverImpl>(
284            driver_object: *mut $crate::km::driver::DriverObjectRaw,
285            registry_path: *const $crate::km::driver::UnicodeStringRaw,
286        ) -> $crate::km::error::NtStatus {
287            use $crate::km::driver::DriverImpl;
288            use $crate::km::error::status;
289
290            // SAFETY: driver_object is provided by the kernel
291            let Some(mut driver) = (unsafe { $crate::km::Driver::from_raw(driver_object) }) else {
292                return status::STATUS_INVALID_PARAMETER;
293            };
294
295            // set up dispatch routines
296            driver.set_unload(__driver_unload::<T>);
297            driver.set_major_function($crate::km::IrpMajorFunction::Create, __dispatch_create::<T>);
298            driver.set_major_function($crate::km::IrpMajorFunction::Close, __dispatch_close::<T>);
299            driver.set_major_function($crate::km::IrpMajorFunction::DeviceControl, __dispatch_device_control::<T>);
300            driver.set_major_function($crate::km::IrpMajorFunction::Read, __dispatch_read::<T>);
301            driver.set_major_function($crate::km::IrpMajorFunction::Write, __dispatch_write::<T>);
302
303            // create unicode string wrapper for registry path
304            if registry_path.is_null() {
305                return status::STATUS_INVALID_PARAMETER;
306            }
307
308            let reg_string = $crate::km::UnicodeString::empty(); // simplified
309
310            match T::init(&mut driver, &reg_string) {
311                Ok(()) => status::STATUS_SUCCESS,
312                Err(e) => e.to_ntstatus(),
313            }
314        }
315
316        unsafe extern "system" fn __driver_unload<T: $crate::km::driver::DriverImpl>(
317            driver_object: *mut $crate::km::driver::DriverObjectRaw
318        ) {
319            // SAFETY: driver_object is valid from kernel
320            if let Some(driver) = unsafe { $crate::km::Driver::from_raw(driver_object) } {
321                T::unload(&driver);
322            }
323        }
324
325        unsafe extern "system" fn __dispatch_create<T: $crate::km::driver::DriverImpl>(
326            device: *mut core::ffi::c_void,
327            irp: *mut core::ffi::c_void,
328        ) -> $crate::km::error::NtStatus {
329            // SAFETY: irp is valid from kernel
330            if let Some(mut irp_wrapper) = unsafe { $crate::km::Irp::from_raw(irp) } {
331                let status = T::create(device, &mut irp_wrapper);
332                irp_wrapper.complete(status);
333                status
334            } else {
335                $crate::km::error::status::STATUS_INVALID_PARAMETER
336            }
337        }
338
339        unsafe extern "system" fn __dispatch_close<T: $crate::km::driver::DriverImpl>(
340            device: *mut core::ffi::c_void,
341            irp: *mut core::ffi::c_void,
342        ) -> $crate::km::error::NtStatus {
343            // SAFETY: irp is valid from kernel
344            if let Some(mut irp_wrapper) = unsafe { $crate::km::Irp::from_raw(irp) } {
345                let status = T::close(device, &mut irp_wrapper);
346                irp_wrapper.complete(status);
347                status
348            } else {
349                $crate::km::error::status::STATUS_INVALID_PARAMETER
350            }
351        }
352
353        unsafe extern "system" fn __dispatch_device_control<T: $crate::km::driver::DriverImpl>(
354            device: *mut core::ffi::c_void,
355            irp: *mut core::ffi::c_void,
356        ) -> $crate::km::error::NtStatus {
357            // SAFETY: irp is valid from kernel
358            if let Some(mut irp_wrapper) = unsafe { $crate::km::Irp::from_raw(irp) } {
359                let status = T::device_control(device, &mut irp_wrapper);
360                irp_wrapper.complete(status);
361                status
362            } else {
363                $crate::km::error::status::STATUS_INVALID_PARAMETER
364            }
365        }
366
367        unsafe extern "system" fn __dispatch_read<T: $crate::km::driver::DriverImpl>(
368            device: *mut core::ffi::c_void,
369            irp: *mut core::ffi::c_void,
370        ) -> $crate::km::error::NtStatus {
371            // SAFETY: irp is valid from kernel
372            if let Some(mut irp_wrapper) = unsafe { $crate::km::Irp::from_raw(irp) } {
373                let status = T::read(device, &mut irp_wrapper);
374                irp_wrapper.complete(status);
375                status
376            } else {
377                $crate::km::error::status::STATUS_INVALID_PARAMETER
378            }
379        }
380
381        unsafe extern "system" fn __dispatch_write<T: $crate::km::driver::DriverImpl>(
382            device: *mut core::ffi::c_void,
383            irp: *mut core::ffi::c_void,
384        ) -> $crate::km::error::NtStatus {
385            // SAFETY: irp is valid from kernel
386            if let Some(mut irp_wrapper) = unsafe { $crate::km::Irp::from_raw(irp) } {
387                let status = T::write(device, &mut irp_wrapper);
388                irp_wrapper.complete(status);
389                status
390            } else {
391                $crate::km::error::status::STATUS_INVALID_PARAMETER
392            }
393        }
394    };
395}
396
397// kernel driver functions
398extern "system" {
399    fn IoCreateDevice(
400        DriverObject: *mut DriverObjectRaw,
401        DeviceExtensionSize: u32,
402        DeviceName: *const c_void,
403        DeviceType: u32,
404        DeviceCharacteristics: u32,
405        Exclusive: u8,
406        DeviceObject: *mut *mut c_void,
407    ) -> NtStatus;
408}
wraith/km/driver.rs

wraith/km/
driver.rs
