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     RtlAppendUnicodeToString(&DeviceName, DeviceNameSuffix);
 118
 119     /* create test device */
 120     if (!(Flags & TESTENTRY_NO_CREATE_DEVICE))
 121     {
 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     ULONG 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 }