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 <thomas.faber@reactos.org>
   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                                 Flags & TESTENTRY_NO_EXCLUSIVE_DEVICE ? FALSE : TRUE,
 126                                 &TestDeviceObject);
 127
 128         if (!NT_SUCCESS(Status))
 129         {
 130             DPRINT1("Could not create device object %wZ\n", &DeviceName);
 131             goto cleanup;
 132         }
 133
 134         DPRINT("DriverEntry. Created DeviceObject %p\n",
 135                  TestDeviceObject);
 136     }
 137
 138     /* initialize dispatch functions */
 139     if (!(Flags & TESTENTRY_NO_REGISTER_UNLOAD))
 140         DriverObject->DriverUnload = DriverUnload;
 141     if (!(Flags & TESTENTRY_NO_REGISTER_DISPATCH))
 142         for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; ++i)
 143             DriverObject->MajorFunction[i] = DriverDispatch;
 144
 145 cleanup:
 146     if (TestDeviceObject && !NT_SUCCESS(Status))
 147     {
 148         IoDeleteDevice(TestDeviceObject);
 149         TestDeviceObject = NULL;
 150     }
 151
 152     if (KmtestDeviceObject && !NT_SUCCESS(Status))
 153     {
 154         ObDereferenceObject(KmtestDeviceObject);
 155         KmtestDeviceObject = NULL;
 156         if (KmtestFileObject)
 157             ObDereferenceObject(KmtestFileObject);
 158     }
 159
 160     return Status;
 161 }
 162
 163 /**
 164  * @name DriverUnload
 165  *
 166  * Driver cleanup funtion.
 167  *
 168  * @param DriverObject
 169  *        Driver Object
 170  */
 171 static
 172 VOID
 173 NTAPI
 174 DriverUnload(
 175     IN PDRIVER_OBJECT DriverObject)
 176 {
 177     PAGED_CODE();
 178
 179     UNREFERENCED_PARAMETER(DriverObject);
 180
 181     DPRINT("DriverUnload\n");
 182
 183     TestUnload(DriverObject);
 184
 185     if (TestDeviceObject)
 186         IoDeleteDevice(TestDeviceObject);
 187
 188     if (KmtestDeviceObject)
 189         ObDereferenceObject(KmtestDeviceObject);
 190 }
 191
 192 /**
 193  * @name KmtRegisterIrpHandler
 194  *
 195  * Register a handler with the IRP Dispatcher.
 196  * If multiple registered handlers match an IRP, it is unspecified which of
 197  * them is called on IRP reception
 198  *
 199  * @param MajorFunction
 200  *        IRP major function code to be handled
 201  * @param DeviceObject
 202  *        Device Object to handle IRPs for.
 203  *        Can be NULL to indicate any device object
 204  * @param IrpHandler
 205  *        Handler function to register.
 206  *
 207  * @return Status
 208  */
 209 NTSTATUS
 210 KmtRegisterIrpHandler(
 211     IN UCHAR MajorFunction,
 212     IN PDEVICE_OBJECT DeviceObject OPTIONAL,
 213     IN PKMT_IRP_HANDLER IrpHandler)
 214 {
 215     NTSTATUS Status = STATUS_SUCCESS;
 216     int i;
 217
 218     if (MajorFunction > IRP_MJ_MAXIMUM_FUNCTION)
 219     {
 220         Status = STATUS_INVALID_PARAMETER_1;
 221         goto cleanup;
 222     }
 223
 224     if (IrpHandler == NULL)
 225     {
 226         Status = STATUS_INVALID_PARAMETER_3;
 227         goto cleanup;
 228     }
 229
 230     for (i = 0; i < sizeof IrpHandlers / sizeof IrpHandlers[0]; ++i)
 231         if (IrpHandlers[i].IrpHandler == NULL)
 232         {
 233             IrpHandlers[i].MajorFunction = MajorFunction;
 234             IrpHandlers[i].DeviceObject = DeviceObject;
 235             IrpHandlers[i].IrpHandler = IrpHandler;
 236             goto cleanup;
 237         }
 238
 239     Status = STATUS_ALLOTTED_SPACE_EXCEEDED;
 240
 241 cleanup:
 242     return Status;
 243 }
 244
 245 /**
 246  * @name KmtUnregisterIrpHandler
 247  *
 248  * Unregister a handler with the IRP Dispatcher.
 249  * Parameters must be specified exactly as in the call to
 250  * KmtRegisterIrpHandler. Only the first matching entry will be removed
 251  * if multiple exist
 252  *
 253  * @param MajorFunction
 254  *        IRP major function code of the handler to be removed
 255  * @param DeviceObject
 256  *        Device Object to of the handler to be removed
 257  * @param IrpHandler
 258  *        Handler function of the handler to be removed
 259  *
 260  * @return Status
 261  */
 262 NTSTATUS
 263 KmtUnregisterIrpHandler(
 264     IN UCHAR MajorFunction,
 265     IN PDEVICE_OBJECT DeviceObject OPTIONAL,
 266     IN PKMT_IRP_HANDLER IrpHandler)
 267 {
 268     NTSTATUS Status = STATUS_SUCCESS;
 269     int i;
 270
 271     for (i = 0; i < sizeof IrpHandlers / sizeof IrpHandlers[0]; ++i)
 272         if (IrpHandlers[i].MajorFunction == MajorFunction &&
 273                 IrpHandlers[i].DeviceObject == DeviceObject &&
 274                 IrpHandlers[i].IrpHandler == IrpHandler)
 275         {
 276             IrpHandlers[i].IrpHandler = NULL;
 277             goto cleanup;
 278         }
 279
 280     Status = STATUS_NOT_FOUND;
 281
 282 cleanup:
 283     return Status;
 284 }
 285
 286 /**
 287  * @name DriverDispatch
 288  *
 289  * Driver Dispatch function
 290  *
 291  * @param DeviceObject
 292  *        Device Object
 293  * @param Irp
 294  *        I/O request packet
 295  *
 296  * @return Status
 297  */
 298 static
 299 NTSTATUS
 300 NTAPI
 301 DriverDispatch(
 302     IN PDEVICE_OBJECT DeviceObject,
 303     IN PIRP Irp)
 304 {
 305     NTSTATUS Status = STATUS_SUCCESS;
 306     PIO_STACK_LOCATION IoStackLocation;
 307     int i;
 308
 309     PAGED_CODE();
 310
 311     IoStackLocation = IoGetCurrentIrpStackLocation(Irp);
 312
 313     DPRINT("DriverDispatch: Function=%s, Device=%p\n",
 314             KmtMajorFunctionNames[IoStackLocation->MajorFunction],
 315             DeviceObject);
 316
 317     for (i = 0; i < sizeof IrpHandlers / sizeof IrpHandlers[0]; ++i)
 318     {
 319         if (IrpHandlers[i].MajorFunction == IoStackLocation->MajorFunction &&
 320                 (IrpHandlers[i].DeviceObject == NULL || IrpHandlers[i].DeviceObject == DeviceObject) &&
 321                 IrpHandlers[i].IrpHandler != NULL)
 322             return IrpHandlers[i].IrpHandler(DeviceObject, Irp, IoStackLocation);
 323     }
 324
 325     /* default handler for DeviceControl */
 326     if (IoStackLocation->MajorFunction == IRP_MJ_DEVICE_CONTROL ||
 327             IoStackLocation->MajorFunction == IRP_MJ_INTERNAL_DEVICE_CONTROL)
 328         return DeviceControlHandler(DeviceObject, Irp, IoStackLocation);
 329
 330     /* default handler */
 331     Irp->IoStatus.Status = Status;
 332     Irp->IoStatus.Information = 0;
 333
 334     IoCompleteRequest(Irp, IO_NO_INCREMENT);
 335
 336     return Status;
 337 }
 338
 339 /**
 340  * @name KmtRegisterMessageHandler
 341  *
 342  * Register a handler with the DeviceControl Dispatcher.
 343  * If multiple registered handlers match a message, it is unspecified which of
 344  * them is called on message reception.
 345  * NOTE: message handlers registered with this function will not be called
 346  *       if a custom IRP handler matching the corresponding IRP is installed!
 347  *
 348  * @param ControlCode
 349  *        Control code to be handled, as passed by the application.
 350  *        Can be 0 to indicate any control code
 351  * @param DeviceObject
 352  *        Device Object to handle IRPs for.
 353  *        Can be NULL to indicate any device object
 354  * @param MessageHandler
 355  *        Handler function to register.
 356  *
 357  * @return Status
 358  */
 359 NTSTATUS
 360 KmtRegisterMessageHandler(
 361     IN ULONG ControlCode OPTIONAL,
 362     IN PDEVICE_OBJECT DeviceObject OPTIONAL,
 363     IN PKMT_MESSAGE_HANDLER MessageHandler)
 364 {
 365     NTSTATUS Status = STATUS_SUCCESS;
 366     int i;
 367
 368     if (ControlCode >= 0x400)
 369     {
 370         Status = STATUS_INVALID_PARAMETER_1;
 371         goto cleanup;
 372     }
 373
 374     if (MessageHandler == NULL)
 375     {
 376         Status = STATUS_INVALID_PARAMETER_2;
 377         goto cleanup;
 378     }
 379
 380     for (i = 0; i < sizeof MessageHandlers / sizeof MessageHandlers[0]; ++i)
 381         if (MessageHandlers[i].MessageHandler == NULL)
 382         {
 383             MessageHandlers[i].ControlCode = ControlCode;
 384             MessageHandlers[i].DeviceObject = DeviceObject;
 385             MessageHandlers[i].MessageHandler = MessageHandler;
 386             goto cleanup;
 387         }
 388
 389     Status = STATUS_ALLOTTED_SPACE_EXCEEDED;
 390
 391 cleanup:
 392     return Status;
 393 }
 394
 395 /**
 396  * @name KmtUnregisterMessageHandler
 397  *
 398  * Unregister a handler with the DeviceControl Dispatcher.
 399  * Parameters must be specified exactly as in the call to
 400  * KmtRegisterMessageHandler. Only the first matching entry will be removed
 401  * if multiple exist
 402  *
 403  * @param ControlCode
 404  *        Control code of the handler to be removed
 405  * @param DeviceObject
 406  *        Device Object to of the handler to be removed
 407  * @param MessageHandler
 408  *        Handler function of the handler to be removed
 409  *
 410  * @return Status
 411  */
 412 NTSTATUS
 413 KmtUnregisterMessageHandler(
 414     IN ULONG ControlCode OPTIONAL,
 415     IN PDEVICE_OBJECT DeviceObject OPTIONAL,
 416     IN PKMT_MESSAGE_HANDLER MessageHandler)
 417 {
 418     NTSTATUS Status = STATUS_SUCCESS;
 419     int i;
 420
 421     for (i = 0; i < sizeof MessageHandlers / sizeof MessageHandlers[0]; ++i)
 422         if (MessageHandlers[i].ControlCode == ControlCode &&
 423                 MessageHandlers[i].DeviceObject == DeviceObject &&
 424                 MessageHandlers[i].MessageHandler == MessageHandler)
 425         {
 426             MessageHandlers[i].MessageHandler = NULL;
 427             goto cleanup;
 428         }
 429
 430     Status = STATUS_NOT_FOUND;
 431
 432 cleanup:
 433     return Status;
 434 }
 435
 436 /**
 437  * @name DeviceControlHandler
 438  *
 439  * Default IRP_MJ_DEVICE_CONTROL/IRP_MJ_INTERNAL_DEVICE_CONTROL handler
 440  *
 441  * @param DeviceObject
 442  *        Device Object.
 443  *        This is guaranteed not to have been touched by the dispatch function
 444  *        before the call to the IRP handler
 445  * @param Irp
 446  *        Device Object.
 447  *        This is guaranteed not to have been touched by the dispatch function
 448  *        before the call to the IRP handler, except for passing it to
 449  *        IoGetCurrentStackLocation
 450  * @param IoStackLocation
 451  *        Device Object.
 452  *        This is guaranteed not to have been touched by the dispatch function
 453  *        before the call to the IRP handler
 454  *
 455  * @return Status
 456  */
 457 static
 458 NTSTATUS
 459 DeviceControlHandler(
 460     IN PDEVICE_OBJECT DeviceObject,
 461     IN PIRP Irp,
 462     IN PIO_STACK_LOCATION IoStackLocation)
 463 {
 464     NTSTATUS Status = STATUS_SUCCESS;
 465     ULONG ControlCode = (IoStackLocation->Parameters.DeviceIoControl.IoControlCode & 0x00000FFC) >> 2;
 466     SIZE_T OutLength = IoStackLocation->Parameters.DeviceIoControl.OutputBufferLength;
 467     int i;
 468
 469     for (i = 0; i < sizeof MessageHandlers / sizeof MessageHandlers[0]; ++i)
 470     {
 471         if ((MessageHandlers[i].ControlCode == 0 ||
 472                 MessageHandlers[i].ControlCode == ControlCode) &&
 473                 (MessageHandlers[i].DeviceObject == NULL || MessageHandlers[i].DeviceObject == DeviceObject) &&
 474                 MessageHandlers[i].MessageHandler != NULL)
 475         {
 476             Status = MessageHandlers[i].MessageHandler(DeviceObject, ControlCode, Irp->AssociatedIrp.SystemBuffer,
 477                                                         IoStackLocation->Parameters.DeviceIoControl.InputBufferLength,
 478                                                         &OutLength);
 479             break;
 480         }
 481     }
 482
 483     Irp->IoStatus.Status = Status;
 484     Irp->IoStatus.Information = OutLength;
 485
 486     IoCompleteRequest(Irp, IO_NO_INCREMENT);
 487
 488     return Status;
 489 }