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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
|
#![feature(collections)]
#![no_std]
#![allow(bad_style)]
#[macro_use] extern crate km;
#[macro_use] extern crate collections;
use km::*;
use core::mem;
use core::ptr;
// Helper for converting b"string" to UNICODE_STRING
fn a2u(s: &[u8]) -> UnicodeString {
let a = AnsiString::from(s);
let mut u = UnicodeString::default();
unsafe { RtlAnsiStringToUnicodeString(&mut u, &a, true) };
return u;
}
// Our per-device settings payload
struct DEVICE_PARAMS
{
pub name: UNICODE_STRING,
pub link: UNICODE_STRING,
}
impl Drop for DEVICE_PARAMS
{
fn drop(&mut self) {
unsafe {
RtlFreeUnicodeString(&mut self.name);
RtlFreeUnicodeString(&mut self.link);
}
}
}
#[no_mangle]
pub extern "system" fn DriverEntry(driver: &mut km::DRIVER_OBJECT, _path: &km::string::UnicodeString) -> NTSTATUS
{
KdPrint!("[rs] hello, rust!\n");
KdPrint!("[rs] we are DriverObject at 0x%p, sizeof 0x%X (expected size 0xA8 or 0x150)\n",
driver as *mut km::DRIVER_OBJECT, mem::size_of::<km::DRIVER_OBJECT>());
driver.DriverUnload = Some(DriverUnload);
KdPrint!("[rs] make params\n");
let params = DEVICE_PARAMS { name: a2u(b"\\Device\\RandomUDev\0"), link: a2u(b"\\??\\urandom\0"), };
// create device
KdPrint!("[rs] create device `%ws` (%d bytes len)\n", params.name.Buffer, params.name.Length as u32);
let mut device: *mut DEVICE_OBJECT = ptr::null_mut();
check_unsafe!(IoCreateDevice(driver, mem::size_of::<DEVICE_PARAMS>() as u32, ¶ms.name, 34, 0, false, &mut device));
// store out custom params to DeviceExtension allocated memory
KdPrint!("[rs] store params at device\n");
let device = unsafe { &mut *device };
let pparams = device.DeviceExtension as *mut DEVICE_PARAMS;
let params = unsafe {
ptr::write(pparams, params);
&*pparams
};
// create symlink
KdPrint!("[rs] create symlink `%ws`\n", params.link.Buffer);
let st = unsafe { IoCreateSymbolicLink(¶ms.link, ¶ms.name) };
if st.is_err() {
DriverUnload(driver);
return st;
}
// setup I/O processing handlers
use km::irp::IRP_MJ;
driver.MajorFunction[IRP_MJ::CREATE as usize] = Some(DispatchCreateClose);
driver.MajorFunction[IRP_MJ::CLOSE as usize] = Some(DispatchCreateClose);
driver.MajorFunction[IRP_MJ::READ as usize] = Some(DispatchRead);
device.Flags |= DEVICE_FLAGS::DO_BUFFERED_IO as u32;
KdPrint!("[rs] loaded.\n");
return Status::success;
}
extern "system" fn DriverUnload(driver: &mut km::DRIVER_OBJECT)
{
KdPrint!("[rs] unload:\n");
unsafe {
// for each created device (driver.DeviceObject->NextDevice linked list)
// delete symbolic link and delete device itself
let mut pdevice = driver.DeviceObject;
while !pdevice.is_null() {
KdPrint_u!("[rs] free device\n");
let device = &mut *pdevice;
if !device.DeviceExtension.is_null() {
KdPrint_u!("[rs] drop params\n");
let params = &mut *(device.DeviceExtension as *mut DEVICE_PARAMS);
IoDeleteSymbolicLink(¶ms.link);
drop(params);
}
KdPrint_u!("[rs] delete device\n");
pdevice = device.NextDevice;
IoDeleteDevice(device);
}
}
KdPrint!("[rs] unloaded.\n");
}
extern "system" fn DispatchCreateClose(_device: &mut DEVICE_OBJECT, irp: &mut IRP) -> NTSTATUS {
KdPrint!("[rs] dispatch create/close \n");
irp.IoStatus.Information = 0;
return irp.complete_request(Status::success);
}
extern "system" fn DispatchRead(device: &mut DEVICE_OBJECT, irp: &mut IRP) -> NTSTATUS
{
KdPrint!("[rs] dispatch read\n");
if (device.Flags & DEVICE_FLAGS::DO_BUFFERED_IO as u32) == 0 {
KdPrint!("[rs] error: nonbuffered io!\n");
return irp.complete_request(Status::unsuccessful);
}
// process IRP request
let size =
{
let io = irp.get_current_stack_location();
let args = io.ParametersRead();
args.Length as usize
};
KdPrint!("[rs] read size %d\n", size as i32);
if size == 0 {
KdPrint!("[rs] error: empty buffer!\n");
return irp.complete_request(Status::unsuccessful);
}
KdPrint!("[rs] generate random\n");
let buf = irp.SystemBuffer; // AssociatedIrp.SystemBuffer union.
let r = GenerateRandom(buf, size);
let st = if let Ok(size) = r {
irp.IoStatus.Information = size;
Status::success
} else {
r.err().unwrap()
};
return irp.complete_request(st);
}
// RtlRandom: Random number generator based on MacLaren and Marsaglia.
// RtlRandomEx is twice as fast and produces better random numbers
// since the period of the random numbers generated is comparatively higher.
fn GenerateRandom(buffer: PVOID, size: usize) -> Result<usize> {
let mut seed = km::time::QueryTickCount() as u32;
let data = buffer as *mut u32;
let dwords = size / 4;
let tail = size % 4;
KdPrint!("[rs] generate random for %d bytes as %d words and %d padding\n", size as u32, dwords as u32, tail as u32);
unsafe {
let mut i = 0;
while i < dwords {
let word = km::rtl::RtlRandomEx(&mut seed);
*data.offset(i as isize) = word;
i += 1;
}
if tail != 0 {
let word = km::rtl::RtlRandomEx(&mut seed);
km::crt::memcpy(data.offset(dwords as isize) as *mut u8, &word as *const u32 as *const u8, tail);
}
}
KdPrint!("[rs] generate complete\n");
return Ok(size);
}
|