Kbdclass kernel filter driver to log scan-codes

KbdclassFltrDriver.c

#include <wdm.h>
#include <ntddkbd.h>

//
// Per-device object extension
//
typedef struct _DEVICE_EXTENSTION
{
  //
  // Driver must not be deleted as long as there is a pending IRP
  //
  LONG PendingIrp;

  //
  // Reference to the attached device so we pass the IRP down to it
  //
  PDEVICE_OBJECT AttachedDevice;
} DEVICE_EXTENSTION, *PDEVICE_EXTENSTION;

//
// Linked list of scan code information
//
typedef struct _SCAN_CODE_INFORMATION
{
  LIST_ENTRY ListEntry;
  USHORT UnitId;
  USHORT MakeCode;
  USHORT Flags;
  ULONG ExtraInformation;
} SCAN_CODE_INFORMATION, *PSCAN_CODE_INFORMATION;

//
// Definition of undocumented routine ObReferenceObjectByName
//
typedef NTSTATUS (NTAPI *FPN_OB_REFERENCE_OBJECT_BY_NAME)(PUNICODE_STRING ObjectName,
                                                          ULONG Attributes,
                                                          PACCESS_STATE AccessState,
                                                          ACCESS_MASK DesiredAccess,
                                                          POBJECT_TYPE ObjectType,
                                                          KPROCESSOR_MODE AccessMode,
                                                          PVOID ParseContext,
                                                          PVOID* Object);

//
// Global variables
//

//
// Head of the list for sending scan code down to the worker thread. Use atomic methods
// for accessing the list
//
static SCAN_CODE_INFORMATION ScanCodeInformationHead;

//
// Lock to protect access to the list
//
static KSPIN_LOCK ScanCodeInformationHeadSpinLock;

//
// Event to notify the consumer about pending data
//
static KEVENT ScanCodeInformationHeadPendingDataEvent;

//
// Event to notify the consumer to exit
//
static KEVENT EndConsumerEvent;

//
// Reference to the consumer thread
//
static PETHREAD ConsumerThreadObject;

//
// Use this allocator to allocate and free entries in the list
//
static LOOKASIDE_LIST_EX ScanCodeInformationAllocator;

//
// Consumer thread to downgrade processing down to passive level
//
static
VOID
_Function_class_(KSTART_ROUTINE)
ConsumerThreadRoutine(PVOID StartContext)
{
  UNREFERENCED_PARAMETER(StartContext);

  for (;;)
  {
    PVOID WaitableObjects[2] =
    {
      &ScanCodeInformationHeadPendingDataEvent,
      &EndConsumerEvent
    };

    //
    // Wait for any event
    //
    const NTSTATUS Status = KeWaitForMultipleObjects(2,
                                                     WaitableObjects,
                                                     WaitAny,
                                                     Executive,
                                                     KernelMode,
                                                     FALSE,
                                                     NULL,
                                                     NULL);

    //
    // Consume any pending work
    //
    PLIST_ENTRY ListEntry;

    while ((ListEntry = ExInterlockedRemoveHeadList(&ScanCodeInformationHead.ListEntry,
                                                    &ScanCodeInformationHeadSpinLock)) != NULL)
    {
      const PSCAN_CODE_INFORMATION ScanCodeInformation = CONTAINING_RECORD(ListEntry, SCAN_CODE_INFORMATION, ListEntry);

      DbgPrint("[!] Consumer thread (UnitId = %hu, MakeCode = %hu, Flags = 0x%02hx, ExtraInformation = 0x%08lx)\n",
               ScanCodeInformation->UnitId,
               ScanCodeInformation->MakeCode,
               ScanCodeInformation->Flags,
               ScanCodeInformation->ExtraInformation);

      ExFreeToLookasideListEx(&ScanCodeInformationAllocator, ScanCodeInformation);
    }

    //
    // Exit if we received an exit event
    //
    if (Status == STATUS_WAIT_1)
    {
      break;
    }
  }
}

//
// Driver unload routine
//
static
VOID
_Function_class_(DRIVER_UNLOAD)
DriverUnload(PDRIVER_OBJECT DriverObject)
{
  UNREFERENCED_PARAMETER(DriverObject);

  ULONG SuccessfulDetachments = 0;

  // Wait for 0.25 seconds
  LARGE_INTEGER WaitInterval;
  WaitInterval.QuadPart = -250000;

  //
  // Iterate over any create device, detach the filter and delete the device object
  //
  PDEVICE_OBJECT KeyboardClassFilterDevice = DriverObject->DeviceObject;

  while (KeyboardClassFilterDevice != NULL)
  {
    const PDEVICE_EXTENSTION DeviceExtenstion = KeyboardClassFilterDevice->DeviceExtension;
    const PDEVICE_OBJECT KeyboardClassFilterDeviceNext = KeyboardClassFilterDevice->NextDevice;

    //
    // Detach the filter from the stack
    //
    IoDetachDevice(DeviceExtenstion->AttachedDevice);

    //
    // By now, we no longer revive any new IRPs, but there might be some pending IRP that has not yet been completed.
    // Wait for it before deleting the device object
    //
    do
    {
      KeDelayExecutionThread(KernelMode, FALSE, &WaitInterval);
    }
    while (InterlockedCompareExchange(&DeviceExtenstion->PendingIrp, 0, 0) > 0);

    //
    // The reference counter must drop to zero
    //
    ASSERT(DeviceExtenstion->PendingIrp >= 0);

    //
    // Delete the device object
    //
    IoDeleteDevice(KeyboardClassFilterDevice);

    DbgPrint("[*] Keyboard filter device #%lu has been detached\n", SuccessfulDetachments++);

    //
    // Go to next device
    //
    KeyboardClassFilterDevice = KeyboardClassFilterDeviceNext;
  }

  //
  // Signal the consumer thread to exit
  //
  KeSetEvent(&EndConsumerEvent, 0, FALSE);

  //
  // Wait for the consumer thread exit
  //
  const NTSTATUS Status = KeWaitForSingleObject(ConsumerThreadObject, Executive, KernelMode, FALSE, NULL);

  ASSERT(Status == STATUS_WAIT_0);

  //
  // Dereference the thread object
  //
  ObDereferenceObject(ConsumerThreadObject);

  //
  // Remove the look-aside list list
  //
  ExDeleteLookasideListEx(&ScanCodeInformationAllocator);

  DbgPrint("[*] Driver has been unloaded\n");
}

//
// Pass down the current IRP to the next device in the stack
//
static
NTSTATUS
_Function_class_(DRIVER_DISPATCH)
DispatchPassThrough(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
  IoSkipCurrentIrpStackLocation(Irp);

  return IoCallDriver(((PDEVICE_EXTENSTION)DeviceObject->DeviceExtension)->AttachedDevice, Irp);
}

//
// This routine will be invoked when the IRP request is completed
//
// Reference: https://docs.microsoft.com/en-us/windows-hardware/drivers/kernel/implementing-an-iocompletion-routine
//
static
NTSTATUS
_Function_class_(IO_COMPLETION_ROUTINE)
CompletionRoutine(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context)
{
  UNREFERENCED_PARAMETER(Context);

  const PDEVICE_EXTENSTION DeviceExtenstion = DeviceObject->DeviceExtension;

  //
  // If PendingReturned is true, we must call IoMarkIrpPending on the current IRP
  //
  if (Irp->PendingReturned)
  {
    IoMarkIrpPending(Irp);
  }

  //
  // Process successful IRPs by pushing work to the worker thread
  //
  if (Irp->IoStatus.Status == STATUS_SUCCESS)
  {
    const ULONG ScanCodesCount = Irp->IoStatus.Information / sizeof(KEYBOARD_INPUT_DATA);
    const PKEYBOARD_INPUT_DATA KeyboardInputData = Irp->AssociatedIrp.SystemBuffer;


    for (ULONG I = 0; I < ScanCodesCount; ++I)
    {
      const PKEYBOARD_INPUT_DATA CurrentKeyboardInputData = &KeyboardInputData[I];


      PSCAN_CODE_INFORMATION ScanCodeInformation = ExAllocateFromLookasideListEx(&ScanCodeInformationAllocator);

      ScanCodeInformation->UnitId = CurrentKeyboardInputData->UnitId;
      ScanCodeInformation->MakeCode = CurrentKeyboardInputData->MakeCode;
      ScanCodeInformation->Flags = CurrentKeyboardInputData->Flags;
      ScanCodeInformation->ExtraInformation = CurrentKeyboardInputData->ExtraInformation;

      ExInterlockedInsertTailList(&ScanCodeInformationHead.ListEntry,
                                  &ScanCodeInformation->ListEntry,
                                  &ScanCodeInformationHeadSpinLock);

      KeSetEvent(&ScanCodeInformationHeadPendingDataEvent, 0, FALSE);
    }
  }

  //
  // Decrement the pending IRPs reference counter to allow waiting DriverUnload to continue
  //
  InterlockedDecrement(&DeviceExtenstion->PendingIrp);

  //
  // We are "hooking" into the stack and not doing actual processing so we can not determine
  // if this request needs more processing
  //
  return STATUS_SUCCESS;
}

//
// Intercept keyboard strokes
//
static
NTSTATUS
_Function_class_(DRIVER_DISPATCH)
DispatchRead(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
  const PDEVICE_EXTENSTION DeviceExtenstion = DeviceObject->DeviceExtension;

  //
  // Increment the pending IRPs counter
  //
  InterlockedIncrement(&DeviceExtenstion->PendingIrp);

  IoCopyCurrentIrpStackLocationToNext(Irp);

  IoSetCompletionRoutine(Irp, CompletionRoutine, NULL, TRUE, TRUE, TRUE);

  return IoCallDriver(DeviceExtenstion->AttachedDevice, Irp);
}

//
// Driver entry routine
//
NTSTATUS
NTAPI
DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
{
  UNREFERENCED_PARAMETER(RegistryPath);

  ULONG SuccessfulAttachments = 0;

  //
  // Setup unload function
  //
  DriverObject->DriverUnload = DriverUnload;

  //
  // Setup dispatch function that just passes the IRP down the stack
  //
  for (ULONG I = 0; I < IRP_MJ_MAXIMUM_FUNCTION; ++I)
  {
    DriverObject->MajorFunction[I] = DispatchPassThrough;
  }

  //
  // We are only interested in read IRPs
  //
  DriverObject->MajorFunction[IRP_MJ_READ] = DispatchRead;

  //
  // Initialize the linked list head
  //
  InitializeListHead(&ScanCodeInformationHead.ListEntry);

  //
  // Lock to protect the ScanCodeInformationHead from concurrent access
  //
  KeInitializeSpinLock(&ScanCodeInformationHeadSpinLock);

  //
  // Event to tell the worker thread about pending data
  //
  KeInitializeEvent(&ScanCodeInformationHeadPendingDataEvent, SynchronizationEvent, FALSE);

  //
  // Event to tell the worker thread to exit
  //
  KeInitializeEvent(&EndConsumerEvent, NotificationEvent, FALSE);

  //
  // Create a look-aside list for allocating list entries
  //
  NTSTATUS Status = ExInitializeLookasideListEx(&ScanCodeInformationAllocator,
                                                NULL,
                                                NULL,
                                                NonPagedPool,
                                                EX_LOOKASIDE_LIST_EX_FLAGS_RAISE_ON_FAIL,
                                                sizeof(SCAN_CODE_INFORMATION),
                                                'kPt8',
                                                0);

  if (!NT_SUCCESS(Status))
  {
    return STATUS_UNSUCCESSFUL;
  }

  //
  // Create worker thread
  //
  HANDLE ConsumerThread;

  Status = PsCreateSystemThread(&ConsumerThread,
                                THREAD_ALL_ACCESS,
                                NULL,
                                NULL,
                                NULL,
                                ConsumerThreadRoutine,
                                NULL);

  if (!NT_SUCCESS(Status))
  {
    ExDeleteLookasideListEx(&ScanCodeInformationAllocator);
    return Status;
  }

  //
  // Get a reference to the thread object so we can wait for the thread to exit
  //
  Status = ObReferenceObjectByHandle(ConsumerThread,
                                     THREAD_ALL_ACCESS,
                                     *PsThreadType,
                                     KernelMode,
                                     (PVOID)&ConsumerThreadObject,
                                     NULL);

  //
  // We (should) have a direct reference to the object itself, we don't need the handle
  //
  ZwClose(ConsumerThread);

  // This should not fail
  if (!NT_SUCCESS(Status))
  {
    ASSERT(FALSE);
    ExDeleteLookasideListEx(&ScanCodeInformationAllocator);
    return Status;
  }

  //
  // Resolve IoDeviceObjectType and IoDeviceObjectType types
  //
  UNICODE_STRING Unicode;

  RtlInitUnicodeString(&Unicode, L"IoDriverObjectType");
  const POBJECT_TYPE* IoDriverObjectType = MmGetSystemRoutineAddress(&Unicode);

  RtlInitUnicodeString(&Unicode, L"IoDeviceObjectType");
  const POBJECT_TYPE* IoDeviceObjectType = MmGetSystemRoutineAddress(&Unicode);


  if (IoDriverObjectType != NULL && IoDeviceObjectType != NULL)
  {
    //
    // Resolve ObReferenceObjectByName
    //
    RtlInitUnicodeString(&Unicode, L"ObReferenceObjectByName");

#pragma warning(suppress: 4152)
    const FPN_OB_REFERENCE_OBJECT_BY_NAME ObReferenceObjectByName = MmGetSystemRoutineAddress(&Unicode);

    if (ObReferenceObjectByName != NULL)
    {
      //
      // Get a pointer to the KbdClass class driver
      //
      RtlInitUnicodeString(&Unicode, L"\\Driver\\KbdClass");

      PDRIVER_OBJECT KbdClassDriver;

      Status = ObReferenceObjectByName(&Unicode,
                                       OBJ_CASE_INSENSITIVE,
                                       0,
                                       0,
                                       *IoDriverObjectType,
                                       KernelMode,
                                       NULL,
                                       (PVOID)&KbdClassDriver);

      if (NT_SUCCESS(Status))
      {
        DbgPrint("[*] KbdClass driver been referenced\n");

        //
        // Get a reference to all keyboard devices created by KbdClass
        //

        PDEVICE_OBJECT KeyboardClassDevice = KbdClassDriver->DeviceObject;

        while (KeyboardClassDevice != NULL)
        {
          //
          // Reference the keyboard device
          //
          Status = ObReferenceObjectByPointer(KeyboardClassDevice,
                                              0,
                                              *IoDeviceObjectType,
                                              KernelMode);

          if (NT_SUCCESS(Status))
          {
            DbgPrint("[*] Referenced a keyboard class #%lu created by KbdClass\n",
                     SuccessfulAttachments);

            //
            // Attach our drivers to the keyboard device object
            //
            PDEVICE_OBJECT KeyboardClassFilterDevice;

            Status = IoCreateDevice(DriverObject,
                                    sizeof(DEVICE_EXTENSTION),
                                    NULL,
                                    FILE_DEVICE_KEYBOARD,
                                    0,
                                    TRUE,
                                    &KeyboardClassFilterDevice);

            if (NT_SUCCESS(Status))
            {
              //
              // Attach our object to the top of the chain
              //
              const PDEVICE_EXTENSTION DeviceExtenstion = KeyboardClassFilterDevice->DeviceExtension;

              DeviceExtenstion->PendingIrp = 0;

              DeviceExtenstion->AttachedDevice = IoAttachDeviceToDeviceStack(KeyboardClassFilterDevice,
                                                                             KeyboardClassDevice);

              //
              // We have failed to attach this device, delete it
              //
              if (DeviceExtenstion->AttachedDevice == NULL)
              {
                IoDeleteDevice(KeyboardClassFilterDevice);
              }
              else
              {
                DbgPrint("[*] Filter has been attached to the stack of a keyboard device #%lu\n",
                         SuccessfulAttachments);
                //
                // The keyboard class driver is using buffered IO. We need to clear the flag DO_DEVICE_INITIALIZING
                // meaning that we have done initializing and we are ready to receive IRPs
                //
                KeyboardClassFilterDevice->Flags |= DO_BUFFERED_IO;
                KeyboardClassFilterDevice->Flags &= ~DO_DEVICE_INITIALIZING;
                
                ++SuccessfulAttachments;
              }
            }

            //
            // We no longer need this reference
            //
            ObDereferenceObject(KeyboardClassDevice);
          }

          KeyboardClassDevice = KeyboardClassDevice->NextDevice;
        }

        //
        // KbdClassDriver reference is no longer needed
        //
        ObDereferenceObject(KbdClassDriver);
      }
    }
  }

  DbgPrint("[*] Driver has been loaded\n");

  //
  // Did we create at least one filter?
  //
  return SuccessfulAttachments != 0 ? STATUS_SUCCESS : STATUS_UNSUCCESSFUL;
}