wraith/km/
ioctl.rs

1//! IOCTL (I/O Control) handling for KM<->UM communication
2
3use core::ffi::c_void;
4
5use super::error::{status, KmError, KmResult, NtStatus};
6use super::irp::Irp;
7
8/// IOCTL transfer method
9#[repr(u32)]
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum IoctlMethod {
12    /// use system buffer (small transfers, kernel copies data)
13    Buffered = 0,
14    /// use MDL for input, direct for output
15    InDirect = 1,
16    /// use direct for input, MDL for output
17    OutDirect = 2,
18    /// raw pointers (dangerous, requires manual validation)
19    Neither = 3,
20}
21
22/// IOCTL code builder
23#[derive(Debug, Clone, Copy)]
24pub struct IoctlCode(pub u32);
25
26impl IoctlCode {
27    /// create IOCTL code from components
28    pub const fn new(
29        device_type: u32,
30        function: u32,
31        method: IoctlMethod,
32        access: u32,
33    ) -> Self {
34        let code = (device_type << 16)
35            | (access << 14)
36            | (function << 2)
37            | (method as u32);
38        Self(code)
39    }
40
41    /// create IOCTL for custom device (0x8000+)
42    pub const fn custom(function: u32, method: IoctlMethod, access: IoctlAccess) -> Self {
43        Self::new(0x8000, function, method, access as u32)
44    }
45
46    /// create buffered IOCTL (most common)
47    pub const fn buffered(function: u32, access: IoctlAccess) -> Self {
48        Self::custom(function, IoctlMethod::Buffered, access)
49    }
50
51    /// get device type
52    pub const fn device_type(&self) -> u32 {
53        (self.0 >> 16) & 0xFFFF
54    }
55
56    /// get function code
57    pub const fn function(&self) -> u32 {
58        (self.0 >> 2) & 0xFFF
59    }
60
61    /// get transfer method
62    pub const fn method(&self) -> IoctlMethod {
63        match self.0 & 3 {
64            0 => IoctlMethod::Buffered,
65            1 => IoctlMethod::InDirect,
66            2 => IoctlMethod::OutDirect,
67            _ => IoctlMethod::Neither,
68        }
69    }
70
71    /// get access mode
72    pub const fn access(&self) -> u32 {
73        (self.0 >> 14) & 3
74    }
75
76    /// get raw code value
77    pub const fn code(&self) -> u32 {
78        self.0
79    }
80}
81
82/// IOCTL access modes
83#[repr(u32)]
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum IoctlAccess {
86    /// any access allowed
87    Any = 0,
88    /// read access required
89    Read = 1,
90    /// write access required
91    Write = 2,
92    /// read and write access required
93    ReadWrite = 3,
94}
95
96/// IOCTL request wrapper
97pub struct Ioctl<'a> {
98    irp: &'a mut Irp,
99    code: IoctlCode,
100}
101
102impl<'a> Ioctl<'a> {
103    /// create from IRP
104    pub fn from_irp(irp: &'a mut Irp) -> Option<Self> {
105        let code = irp.ioctl_code()?;
106        Some(Self {
107            irp,
108            code: IoctlCode(code),
109        })
110    }
111
112    /// get IOCTL code
113    pub fn code(&self) -> IoctlCode {
114        self.code
115    }
116
117    /// get function number
118    pub fn function(&self) -> u32 {
119        self.code.function()
120    }
121
122    /// get transfer method
123    pub fn method(&self) -> IoctlMethod {
124        self.code.method()
125    }
126
127    /// get input buffer length
128    pub fn input_length(&self) -> usize {
129        self.irp.input_buffer_length().unwrap_or(0) as usize
130    }
131
132    /// get output buffer length
133    pub fn output_length(&self) -> usize {
134        self.irp.output_buffer_length().unwrap_or(0) as usize
135    }
136
137    /// get input buffer as typed reference
138    pub fn input<T>(&self) -> Option<&T> {
139        if self.input_length() < core::mem::size_of::<T>() {
140            return None;
141        }
142        self.irp.input_buffer()
143    }
144
145    /// get output buffer as typed mutable reference
146    pub fn output_mut<T>(&mut self) -> Option<&mut T> {
147        if self.output_length() < core::mem::size_of::<T>() {
148            return None;
149        }
150        self.irp.output_buffer_mut()
151    }
152
153    /// get input as byte slice
154    pub fn input_bytes(&self) -> Option<&[u8]> {
155        self.irp.input_bytes()
156    }
157
158    /// get output as mutable byte slice
159    pub fn output_bytes_mut(&mut self) -> Option<&mut [u8]> {
160        self.irp.output_bytes_mut()
161    }
162
163    /// complete with success
164    pub fn complete_success(self) {
165        self.complete_with_info(status::STATUS_SUCCESS, 0);
166    }
167
168    /// complete with output size
169    pub fn complete_with_output(self, output_size: usize) {
170        self.complete_with_info(status::STATUS_SUCCESS, output_size);
171    }
172
173    /// complete with error
174    pub fn complete_error(self, error: KmError) {
175        self.complete_with_info(error.to_ntstatus(), 0);
176    }
177
178    /// complete with status and information
179    pub fn complete_with_info(mut self, status: NtStatus, information: usize) {
180        self.irp.complete_with_info(status, information);
181    }
182}
183
184/// trait for IOCTL handlers
185pub trait IoctlHandler {
186    /// handle an IOCTL request
187    fn handle(&self, ioctl: Ioctl) -> NtStatus;
188}
189
190/// IOCTL dispatcher with registered handlers
191pub struct IoctlDispatcher {
192    handlers: &'static [(u32, &'static dyn IoctlHandler)],
193    default_handler: Option<&'static dyn IoctlHandler>,
194}
195
196impl IoctlDispatcher {
197    /// create new dispatcher with static handler table
198    pub const fn new(handlers: &'static [(u32, &'static dyn IoctlHandler)]) -> Self {
199        Self {
200            handlers,
201            default_handler: None,
202        }
203    }
204
205    /// create with default handler for unknown IOCTLs
206    pub const fn with_default(
207        handlers: &'static [(u32, &'static dyn IoctlHandler)],
208        default: &'static dyn IoctlHandler,
209    ) -> Self {
210        Self {
211            handlers,
212            default_handler: Some(default),
213        }
214    }
215
216    /// dispatch IOCTL to appropriate handler
217    pub fn dispatch(&self, irp: &mut Irp) -> NtStatus {
218        let Some(ioctl) = Ioctl::from_irp(irp) else {
219            return status::STATUS_INVALID_PARAMETER;
220        };
221
222        let function = ioctl.function();
223
224        // find handler
225        for (code, handler) in self.handlers {
226            if *code == function {
227                return handler.handle(ioctl);
228            }
229        }
230
231        // try default handler
232        if let Some(handler) = self.default_handler {
233            return handler.handle(ioctl);
234        }
235
236        // no handler found
237        ioctl.complete_with_info(status::STATUS_INVALID_DEVICE_REQUEST, 0);
238        status::STATUS_INVALID_DEVICE_REQUEST
239    }
240}
241
242/// common IOCTL request structures for KM<->UM communication
243
244/// process memory read request
245#[repr(C)]
246#[derive(Debug, Clone, Copy)]
247pub struct ReadMemoryRequest {
248    pub process_id: u32,
249    pub address: u64,
250    pub size: u32,
251}
252
253/// process memory write request
254#[repr(C)]
255#[derive(Debug, Clone, Copy)]
256pub struct WriteMemoryRequest {
257    pub process_id: u32,
258    pub address: u64,
259    pub size: u32,
260    // data follows in buffer
261}
262
263/// get module base request
264#[repr(C)]
265#[derive(Debug, Clone, Copy)]
266pub struct GetModuleBaseRequest {
267    pub process_id: u32,
268    pub module_name_offset: u32, // offset in buffer where name starts
269    pub module_name_length: u32,
270}
271
272/// module base response
273#[repr(C)]
274#[derive(Debug, Clone, Copy)]
275pub struct GetModuleBaseResponse {
276    pub base_address: u64,
277    pub size: u64,
278}
279
280/// allocate memory request
281#[repr(C)]
282#[derive(Debug, Clone, Copy)]
283pub struct AllocateMemoryRequest {
284    pub process_id: u32,
285    pub size: u64,
286    pub protection: u32,
287    pub preferred_address: u64,
288}
289
290/// allocate memory response
291#[repr(C)]
292#[derive(Debug, Clone, Copy)]
293pub struct AllocateMemoryResponse {
294    pub allocated_address: u64,
295    pub actual_size: u64,
296}
297
298/// free memory request
299#[repr(C)]
300#[derive(Debug, Clone, Copy)]
301pub struct FreeMemoryRequest {
302    pub process_id: u32,
303    pub address: u64,
304}
305
306/// protect memory request
307#[repr(C)]
308#[derive(Debug, Clone, Copy)]
309pub struct ProtectMemoryRequest {
310    pub process_id: u32,
311    pub address: u64,
312    pub size: u64,
313    pub new_protection: u32,
314}
315
316/// protect memory response
317#[repr(C)]
318#[derive(Debug, Clone, Copy)]
319pub struct ProtectMemoryResponse {
320    pub old_protection: u32,
321}
322
323/// predefined IOCTL codes for memory operations
324pub mod codes {
325    use super::{IoctlAccess, IoctlCode};
326
327    /// read process memory
328    pub const READ_MEMORY: IoctlCode = IoctlCode::buffered(0x800, IoctlAccess::ReadWrite);
329
330    /// write process memory
331    pub const WRITE_MEMORY: IoctlCode = IoctlCode::buffered(0x801, IoctlAccess::ReadWrite);
332
333    /// get module base
334    pub const GET_MODULE_BASE: IoctlCode = IoctlCode::buffered(0x802, IoctlAccess::Read);
335
336    /// allocate virtual memory
337    pub const ALLOCATE_MEMORY: IoctlCode = IoctlCode::buffered(0x803, IoctlAccess::ReadWrite);
338
339    /// free virtual memory
340    pub const FREE_MEMORY: IoctlCode = IoctlCode::buffered(0x804, IoctlAccess::ReadWrite);
341
342    /// change memory protection
343    pub const PROTECT_MEMORY: IoctlCode = IoctlCode::buffered(0x805, IoctlAccess::ReadWrite);
344
345    /// query process info
346    pub const QUERY_PROCESS: IoctlCode = IoctlCode::buffered(0x806, IoctlAccess::Read);
347
348    /// copy physical memory
349    pub const COPY_PHYSICAL: IoctlCode = IoctlCode::buffered(0x807, IoctlAccess::ReadWrite);
350
351    /// map physical memory
352    pub const MAP_PHYSICAL: IoctlCode = IoctlCode::buffered(0x808, IoctlAccess::ReadWrite);
353
354    /// unmap physical memory
355    pub const UNMAP_PHYSICAL: IoctlCode = IoctlCode::buffered(0x809, IoctlAccess::ReadWrite);
356}
357
358/// macro to define IOCTL handler
359#[macro_export]
360macro_rules! define_ioctl_handler {
361    ($name:ident, |$ioctl:ident| $body:block) => {
362        struct $name;
363
364        impl $crate::km::ioctl::IoctlHandler for $name {
365            fn handle(&self, $ioctl: $crate::km::ioctl::Ioctl) -> $crate::km::error::NtStatus {
366                $body
367            }
368        }
369    };
370}
371
372/// macro to create IOCTL dispatcher
373#[macro_export]
374macro_rules! ioctl_dispatcher {
375    ($($code:expr => $handler:expr),* $(,)?) => {
376        $crate::km::ioctl::IoctlDispatcher::new(&[
377            $(($code, &$handler as &dyn $crate::km::ioctl::IoctlHandler)),*
378        ])
379    };
380}
wraith/km/ioctl.rs

wraith/km/
ioctl.rs
