d76c4a65d32a969a29dc2406f128ddc0460f68eb
[reactos.git] / rostests / kmtests / kmtest_drv / kmtest_standalone.c
1 /*
2 * PROJECT: ReactOS kernel-mode tests
3 * LICENSE: GPLv2+ - See COPYING in the top level directory
4 * PURPOSE: Kernel-Mode Test Suite Example Test Driver
5 * PROGRAMMER: Thomas Faber <thfabba@gmx.de>
6 */
7
8 #include <ntddk.h>
9 #include <ntifs.h>
10 #include <ndk/ketypes.h>
11
12 #define KMT_DEFINE_TEST_FUNCTIONS
13 #include <kmt_test.h>
14
15 //#define NDEBUG
16 #include <debug.h>
17
18 #include <kmt_public.h>
19
20 /* types */
21 typedef struct
22 {
23 UCHAR MajorFunction;
24 PDEVICE_OBJECT DeviceObject;
25 PKMT_IRP_HANDLER IrpHandler;
26 } KMT_IRP_HANDLER_ENTRY, *PKMT_IRP_HANDLER_ENTRY;
27
28 typedef struct
29 {
30 ULONG ControlCode;
31 PDEVICE_OBJECT DeviceObject;
32 PKMT_MESSAGE_HANDLER MessageHandler;
33 } KMT_MESSAGE_HANDLER_ENTRY, *PKMT_MESSAGE_HANDLER_ENTRY;
34
35 /* Prototypes */
36 DRIVER_INITIALIZE DriverEntry;
37 static DRIVER_UNLOAD DriverUnload;
38 static DRIVER_DISPATCH DriverDispatch;
39 static KMT_IRP_HANDLER DeviceControlHandler;
40
41 /* Globals */
42 static PDEVICE_OBJECT TestDeviceObject;
43 static PDEVICE_OBJECT KmtestDeviceObject;
44
45 #define KMT_MAX_IRP_HANDLERS 256
46 static KMT_IRP_HANDLER_ENTRY IrpHandlers[KMT_MAX_IRP_HANDLERS] = { { 0 } };
47 #define KMT_MAX_MESSAGE_HANDLERS 256
48 static KMT_MESSAGE_HANDLER_ENTRY MessageHandlers[KMT_MAX_MESSAGE_HANDLERS] = { { 0 } };
49
50 /**
51 * @name DriverEntry
52 *
53 * Driver entry point.
54 *
55 * @param DriverObject
56 * Driver Object
57 * @param RegistryPath
58 * Driver Registry Path
59 *
60 * @return Status
61 */
62 NTSTATUS
63 NTAPI
64 DriverEntry(
65 IN PDRIVER_OBJECT DriverObject,
66 IN PUNICODE_STRING RegistryPath)
67 {
68 NTSTATUS Status = STATUS_SUCCESS;
69 WCHAR DeviceNameBuffer[128] = L"\\Device\\Kmtest-";
70 UNICODE_STRING KmtestDeviceName;
71 PFILE_OBJECT KmtestFileObject;
72 PKMT_DEVICE_EXTENSION KmtestDeviceExtension;
73 UNICODE_STRING DeviceName;
74 PCWSTR DeviceNameSuffix;
75 INT Flags = 0;
76 int i;
77 PKPRCB Prcb;
78
79 PAGED_CODE();
80
81 DPRINT("DriverEntry\n");
82
83 Prcb = KeGetCurrentPrcb();
84 KmtIsCheckedBuild = (Prcb->BuildType & PRCB_BUILD_DEBUG) != 0;
85 KmtIsMultiProcessorBuild = (Prcb->BuildType & PRCB_BUILD_UNIPROCESSOR) == 0;
86
87 /* get the Kmtest device, so that we get a ResultBuffer pointer */
88 RtlInitUnicodeString(&KmtestDeviceName, KMTEST_DEVICE_DRIVER_PATH);
89 Status = IoGetDeviceObjectPointer(&KmtestDeviceName, FILE_ALL_ACCESS, &KmtestFileObject, &KmtestDeviceObject);
90
91 if (!NT_SUCCESS(Status))
92 {
93 DPRINT1("Failed to get Kmtest device object pointer\n");
94 goto cleanup;
95 }
96
97 Status = ObReferenceObjectByPointer(KmtestDeviceObject, FILE_ALL_ACCESS, NULL, KernelMode);
98
99 if (!NT_SUCCESS(Status))
100 {
101 DPRINT1("Failed to reference Kmtest device object\n");
102 goto cleanup;
103 }
104
105 ObDereferenceObject(KmtestFileObject);
106 KmtestFileObject = NULL;
107 KmtestDeviceExtension = KmtestDeviceObject->DeviceExtension;
108 ResultBuffer = KmtestDeviceExtension->ResultBuffer;
109 DPRINT("KmtestDeviceObject: %p\n", (PVOID)KmtestDeviceObject);
110 DPRINT("KmtestDeviceExtension: %p\n", (PVOID)KmtestDeviceExtension);
111 DPRINT("Setting ResultBuffer: %p\n", (PVOID)ResultBuffer);
112
113 /* call TestEntry */
114 RtlInitUnicodeString(&DeviceName, DeviceNameBuffer);
115 DeviceName.MaximumLength = sizeof DeviceNameBuffer;
116 TestEntry(DriverObject, RegistryPath, &DeviceNameSuffix, &Flags);
117
118 /* create test device */
119 if (!(Flags & TESTENTRY_NO_CREATE_DEVICE))
120 {
121 RtlAppendUnicodeToString(&DeviceName, DeviceNameSuffix);
122 Status = IoCreateDevice(DriverObject, 0, &DeviceName,
123 FILE_DEVICE_UNKNOWN,
124 FILE_DEVICE_SECURE_OPEN | FILE_READ_ONLY_DEVICE,
125 TRUE, &TestDeviceObject);
126
127 if (!NT_SUCCESS(Status))
128 {
129 DPRINT1("Could not create device object %wZ\n", &DeviceName);
130 goto cleanup;
131 }
132
133 DPRINT("DriverEntry. Created DeviceObject %p\n",
134 TestDeviceObject);
135 }
136
137 /* initialize dispatch functions */
138 if (!(Flags & TESTENTRY_NO_REGISTER_UNLOAD))
139 DriverObject->DriverUnload = DriverUnload;
140 if (!(Flags & TESTENTRY_NO_REGISTER_DISPATCH))
141 for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; ++i)
142 DriverObject->MajorFunction[i] = DriverDispatch;
143
144 cleanup:
145 if (TestDeviceObject && !NT_SUCCESS(Status))
146 {
147 IoDeleteDevice(TestDeviceObject);
148 TestDeviceObject = NULL;
149 }
150
151 if (KmtestDeviceObject && !NT_SUCCESS(Status))
152 {
153 ObDereferenceObject(KmtestDeviceObject);
154 KmtestDeviceObject = NULL;
155 if (KmtestFileObject)
156 ObDereferenceObject(KmtestFileObject);
157 }
158
159 return Status;
160 }
161
162 /**
163 * @name DriverUnload
164 *
165 * Driver cleanup funtion.
166 *
167 * @param DriverObject
168 * Driver Object
169 */
170 static
171 VOID
172 NTAPI
173 DriverUnload(
174 IN PDRIVER_OBJECT DriverObject)
175 {
176 PAGED_CODE();
177
178 UNREFERENCED_PARAMETER(DriverObject);
179
180 DPRINT("DriverUnload\n");
181
182 TestUnload(DriverObject);
183
184 if (TestDeviceObject)
185 IoDeleteDevice(TestDeviceObject);
186
187 if (KmtestDeviceObject)
188 ObDereferenceObject(KmtestDeviceObject);
189 }
190
191 /**
192 * @name KmtRegisterIrpHandler
193 *
194 * Register a handler with the IRP Dispatcher.
195 * If multiple registered handlers match an IRP, it is unspecified which of
196 * them is called on IRP reception
197 *
198 * @param MajorFunction
199 * IRP major function code to be handled
200 * @param DeviceObject
201 * Device Object to handle IRPs for.
202 * Can be NULL to indicate any device object
203 * @param IrpHandler
204 * Handler function to register.
205 *
206 * @return Status
207 */
208 NTSTATUS
209 KmtRegisterIrpHandler(
210 IN UCHAR MajorFunction,
211 IN PDEVICE_OBJECT DeviceObject OPTIONAL,
212 IN PKMT_IRP_HANDLER IrpHandler)
213 {
214 NTSTATUS Status = STATUS_SUCCESS;
215 int i;
216
217 if (MajorFunction > IRP_MJ_MAXIMUM_FUNCTION)
218 {
219 Status = STATUS_INVALID_PARAMETER_1;
220 goto cleanup;
221 }
222
223 if (IrpHandler == NULL)
224 {
225 Status = STATUS_INVALID_PARAMETER_3;
226 goto cleanup;
227 }
228
229 for (i = 0; i < sizeof IrpHandlers / sizeof IrpHandlers[0]; ++i)
230 if (IrpHandlers[i].IrpHandler == NULL)
231 {
232 IrpHandlers[i].MajorFunction = MajorFunction;
233 IrpHandlers[i].DeviceObject = DeviceObject;
234 IrpHandlers[i].IrpHandler = IrpHandler;
235 goto cleanup;
236 }
237
238 Status = STATUS_ALLOTTED_SPACE_EXCEEDED;
239
240 cleanup:
241 return Status;
242 }
243
244 /**
245 * @name KmtUnregisterIrpHandler
246 *
247 * Unregister a handler with the IRP Dispatcher.
248 * Parameters must be specified exactly as in the call to
249 * KmtRegisterIrpHandler. Only the first matching entry will be removed
250 * if multiple exist
251 *
252 * @param MajorFunction
253 * IRP major function code of the handler to be removed
254 * @param DeviceObject
255 * Device Object to of the handler to be removed
256 * @param IrpHandler
257 * Handler function of the handler to be removed
258 *
259 * @return Status
260 */
261 NTSTATUS
262 KmtUnregisterIrpHandler(
263 IN UCHAR MajorFunction,
264 IN PDEVICE_OBJECT DeviceObject OPTIONAL,
265 IN PKMT_IRP_HANDLER IrpHandler)
266 {
267 NTSTATUS Status = STATUS_SUCCESS;
268 int i;
269
270 for (i = 0; i < sizeof IrpHandlers / sizeof IrpHandlers[0]; ++i)
271 if (IrpHandlers[i].MajorFunction == MajorFunction &&
272 IrpHandlers[i].DeviceObject == DeviceObject &&
273 IrpHandlers[i].IrpHandler == IrpHandler)
274 {
275 IrpHandlers[i].IrpHandler = NULL;
276 goto cleanup;
277 }
278
279 Status = STATUS_NOT_FOUND;
280
281 cleanup:
282 return Status;
283 }
284
285 /**
286 * @name DriverDispatch
287 *
288 * Driver Dispatch function
289 *
290 * @param DeviceObject
291 * Device Object
292 * @param Irp
293 * I/O request packet
294 *
295 * @return Status
296 */
297 static
298 NTSTATUS
299 NTAPI
300 DriverDispatch(
301 IN PDEVICE_OBJECT DeviceObject,
302 IN PIRP Irp)
303 {
304 NTSTATUS Status = STATUS_SUCCESS;
305 PIO_STACK_LOCATION IoStackLocation;
306 int i;
307
308 PAGED_CODE();
309
310 IoStackLocation = IoGetCurrentIrpStackLocation(Irp);
311
312 DPRINT("DriverDispatch: Function=%s, Device=%p\n",
313 KmtMajorFunctionNames[IoStackLocation->MajorFunction],
314 DeviceObject);
315
316 for (i = 0; i < sizeof IrpHandlers / sizeof IrpHandlers[0]; ++i)
317 {
318 if (IrpHandlers[i].MajorFunction == IoStackLocation->MajorFunction &&
319 (IrpHandlers[i].DeviceObject == NULL || IrpHandlers[i].DeviceObject == DeviceObject) &&
320 IrpHandlers[i].IrpHandler != NULL)
321 return IrpHandlers[i].IrpHandler(DeviceObject, Irp, IoStackLocation);
322 }
323
324 /* default handler for DeviceControl */
325 if (IoStackLocation->MajorFunction == IRP_MJ_DEVICE_CONTROL ||
326 IoStackLocation->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL)
327 return DeviceControlHandler(DeviceObject, Irp, IoStackLocation);
328
329 /* default handler */
330 Irp->IoStatus.Status = Status;
331 Irp->IoStatus.Information = 0;
332
333 IoCompleteRequest(Irp, IO_NO_INCREMENT);
334
335 return Status;
336 }
337
338 /**
339 * @name KmtRegisterMessageHandler
340 *
341 * Register a handler with the DeviceControl Dispatcher.
342 * If multiple registered handlers match a message, it is unspecified which of
343 * them is called on message reception.
344 * NOTE: message handlers registered with this function will not be called
345 * if a custom IRP handler matching the corresponding IRP is installed!
346 *
347 * @param ControlCode
348 * Control code to be handled, as passed by the application.
349 * Can be 0 to indicate any control code
350 * @param DeviceObject
351 * Device Object to handle IRPs for.
352 * Can be NULL to indicate any device object
353 * @param MessageHandler
354 * Handler function to register.
355 *
356 * @return Status
357 */
358 NTSTATUS
359 KmtRegisterMessageHandler(
360 IN ULONG ControlCode OPTIONAL,
361 IN PDEVICE_OBJECT DeviceObject OPTIONAL,
362 IN PKMT_MESSAGE_HANDLER MessageHandler)
363 {
364 NTSTATUS Status = STATUS_SUCCESS;
365 int i;
366
367 if (ControlCode >= 0x400)
368 {
369 Status = STATUS_INVALID_PARAMETER_1;
370 goto cleanup;
371 }
372
373 if (MessageHandler == NULL)
374 {
375 Status = STATUS_INVALID_PARAMETER_2;
376 goto cleanup;
377 }
378
379 for (i = 0; i < sizeof MessageHandlers / sizeof MessageHandlers[0]; ++i)
380 if (MessageHandlers[i].MessageHandler == NULL)
381 {
382 MessageHandlers[i].ControlCode = ControlCode;
383 MessageHandlers[i].DeviceObject = DeviceObject;
384 MessageHandlers[i].MessageHandler = MessageHandler;
385 goto cleanup;
386 }
387
388 Status = STATUS_ALLOTTED_SPACE_EXCEEDED;
389
390 cleanup:
391 return Status;
392 }
393
394 /**
395 * @name KmtUnregisterMessageHandler
396 *
397 * Unregister a handler with the DeviceControl Dispatcher.
398 * Parameters must be specified exactly as in the call to
399 * KmtRegisterMessageHandler. Only the first matching entry will be removed
400 * if multiple exist
401 *
402 * @param ControlCode
403 * Control code of the handler to be removed
404 * @param DeviceObject
405 * Device Object to of the handler to be removed
406 * @param MessageHandler
407 * Handler function of the handler to be removed
408 *
409 * @return Status
410 */
411 NTSTATUS
412 KmtUnregisterMessageHandler(
413 IN ULONG ControlCode OPTIONAL,
414 IN PDEVICE_OBJECT DeviceObject OPTIONAL,
415 IN PKMT_MESSAGE_HANDLER MessageHandler)
416 {
417 NTSTATUS Status = STATUS_SUCCESS;
418 int i;
419
420 for (i = 0; i < sizeof MessageHandlers / sizeof MessageHandlers[0]; ++i)
421 if (MessageHandlers[i].ControlCode == ControlCode &&
422 MessageHandlers[i].DeviceObject == DeviceObject &&
423 MessageHandlers[i].MessageHandler == MessageHandler)
424 {
425 MessageHandlers[i].MessageHandler = NULL;
426 goto cleanup;
427 }
428
429 Status = STATUS_NOT_FOUND;
430
431 cleanup:
432 return Status;
433 }
434
435 /**
436 * @name DeviceControlHandler
437 *
438 * Default IRP_MJ_DEVICE_CONTROL/IRP_MJ_INTERNAL_DEVICE_CONTROL handler
439 *
440 * @param DeviceObject
441 * Device Object.
442 * This is guaranteed not to have been touched by the dispatch function
443 * before the call to the IRP handler
444 * @param Irp
445 * Device Object.
446 * This is guaranteed not to have been touched by the dispatch function
447 * before the call to the IRP handler, except for passing it to
448 * IoGetCurrentStackLocation
449 * @param IoStackLocation
450 * Device Object.
451 * This is guaranteed not to have been touched by the dispatch function
452 * before the call to the IRP handler
453 *
454 * @return Status
455 */
456 static
457 NTSTATUS
458 DeviceControlHandler(
459 IN PDEVICE_OBJECT DeviceObject,
460 IN PIRP Irp,
461 IN PIO_STACK_LOCATION IoStackLocation)
462 {
463 NTSTATUS Status = STATUS_SUCCESS;
464 ULONG ControlCode = (IoStackLocation->Parameters.DeviceIoControl.IoControlCode & 0x00000FFC) >> 2;
465 SIZE_T OutLength = IoStackLocation->Parameters.DeviceIoControl.OutputBufferLength;
466 int i;
467
468 for (i = 0; i < sizeof MessageHandlers / sizeof MessageHandlers[0]; ++i)
469 {
470 if ((MessageHandlers[i].ControlCode == 0 ||
471 MessageHandlers[i].ControlCode == ControlCode) &&
472 (MessageHandlers[i].DeviceObject == NULL || MessageHandlers[i].DeviceObject == DeviceObject) &&
473 MessageHandlers[i].MessageHandler != NULL)
474 {
475 Status = MessageHandlers[i].MessageHandler(DeviceObject, ControlCode, Irp->AssociatedIrp.SystemBuffer,
476 IoStackLocation->Parameters.DeviceIoControl.InputBufferLength,
477 &OutLength);
478 break;
479 }
480 }
481
482 Irp->IoStatus.Status = Status;
483 Irp->IoStatus.Information = OutLength;
484
485 IoCompleteRequest(Irp, IO_NO_INCREMENT);
486
487 return Status;
488 }