From 8ee3ade0d47139af9af9fa06ad6a3a8277c04cf1 Mon Sep 17 00:00:00 2001 From: momo5502 Date: Sat, 26 Mar 2022 19:09:54 +0100 Subject: [PATCH] More cleanup --- src/driver/driver_main.cpp | 238 +++++----------------------------- src/driver/irp.cpp | 142 ++++++++++++++++++++ src/driver/irp.hpp | 7 + src/driver/new.cpp | 1 - src/driver/sleep_callback.cpp | 1 - 5 files changed, 181 insertions(+), 208 deletions(-) create mode 100644 src/driver/irp.cpp create mode 100644 src/driver/irp.hpp diff --git a/src/driver/driver_main.cpp b/src/driver/driver_main.cpp index d40169b..2f3a8b0 100644 --- a/src/driver/driver_main.cpp +++ b/src/driver/driver_main.cpp @@ -1,183 +1,22 @@ #include "std_include.hpp" #include "logging.hpp" -#include "thread.hpp" #include "sleep_callback.hpp" - -#define DOS_DEV_NAME L"\\DosDevices\\HelloDev" -#define DEV_NAME L"\\Device\\HelloDev" - -#define HELLO_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_NEITHER, FILE_ANY_ACCESS) - -_Function_class_(DRIVER_DISPATCH) - -NTSTATUS IrpNotImplementedHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp) -{ - Irp->IoStatus.Information = 0; - Irp->IoStatus.Status = STATUS_NOT_SUPPORTED; - - UNREFERENCED_PARAMETER(DeviceObject); - PAGED_CODE(); - - // Complete the request - IoCompleteRequest(Irp, IO_NO_INCREMENT); - - return STATUS_NOT_SUPPORTED; -} - -_Function_class_(DRIVER_DISPATCH) - -NTSTATUS IrpCreateCloseHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp) -{ - Irp->IoStatus.Information = 0; - Irp->IoStatus.Status = STATUS_SUCCESS; - - UNREFERENCED_PARAMETER(DeviceObject); - PAGED_CODE(); - - // Complete the request - IoCompleteRequest(Irp, IO_NO_INCREMENT); - - return STATUS_SUCCESS; -} - -_Function_class_(DRIVER_DISPATCH) -VOID IrpUnloadHandler(IN PDRIVER_OBJECT DriverObject) -{ - UNICODE_STRING DosDeviceName = {0}; - - PAGED_CODE(); - - RtlInitUnicodeString(&DosDeviceName, DOS_DEV_NAME); - - // Delete the symbolic link - IoDeleteSymbolicLink(&DosDeviceName); - - // Delete the device - IoDeleteDevice(DriverObject->DeviceObject); - - debug_log("[!] Hello Driver Unloaded\n"); -} - -_Function_class_(DRIVER_DISPATCH) - -NTSTATUS IrpDeviceIoCtlHandler(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp) -{ - ULONG IoControlCode = 0; - PIO_STACK_LOCATION IrpSp = NULL; - NTSTATUS Status = STATUS_NOT_SUPPORTED; - - UNREFERENCED_PARAMETER(DeviceObject); - PAGED_CODE(); - - IrpSp = IoGetCurrentIrpStackLocation(Irp); - IoControlCode = IrpSp->Parameters.DeviceIoControl.IoControlCode; - - if (IrpSp) - { - switch (IoControlCode) - { - case HELLO_DRV_IOCTL: - debug_log("[< HelloDriver >] Hello from the Driver!\n"); - break; - default: - debug_log("[-] Invalid IOCTL Code: 0x%X\n", IoControlCode); - Status = STATUS_INVALID_DEVICE_REQUEST; - break; - } - } - - Irp->IoStatus.Status = Status; - Irp->IoStatus.Information = 0; - - // Complete the request - IoCompleteRequest(Irp, IO_NO_INCREMENT); - - return Status; -} - -NTSTATUS create_io_device(const PDRIVER_OBJECT DriverObject) -{ - UINT32 i = 0; - PDEVICE_OBJECT DeviceObject = NULL; - NTSTATUS Status = STATUS_UNSUCCESSFUL; - UNICODE_STRING DeviceName, DosDeviceName = {0}; - - PAGED_CODE(); - - RtlInitUnicodeString(&DeviceName, DEV_NAME); - RtlInitUnicodeString(&DosDeviceName, DOS_DEV_NAME); - - debug_log("[*] In DriverEntry\n"); - - // Create the device - Status = IoCreateDevice(DriverObject, - 0, - &DeviceName, - FILE_DEVICE_UNKNOWN, - FILE_DEVICE_SECURE_OPEN, - FALSE, - &DeviceObject); - - if (!NT_SUCCESS(Status)) - { - if (DeviceObject) - { - // Delete the device - IoDeleteDevice(DeviceObject); - } - - debug_log("[-] Error Initializing HelloDriver\n"); - return Status; - } - - // Assign the IRP handlers - for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++) - { - // Disable the Compiler Warning: 28169 -#pragma warning(push) -#pragma warning(disable : 28169) - DriverObject->MajorFunction[i] = IrpNotImplementedHandler; -#pragma warning(pop) - } - - // Assign the IRP handlers for Create, Close and Device Control - DriverObject->MajorFunction[IRP_MJ_CREATE] = IrpCreateCloseHandler; - DriverObject->MajorFunction[IRP_MJ_CLOSE] = IrpCreateCloseHandler; - DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = IrpDeviceIoCtlHandler; - - // Set the flags - DeviceObject->Flags |= DO_DIRECT_IO; - DeviceObject->Flags &= ~DO_DEVICE_INITIALIZING; - - // Create the symbolic link - Status = IoCreateSymbolicLink(&DosDeviceName, &DeviceName); - - // Show the banner - debug_log("[!] HelloDriver Loaded\n"); - - return Status; -} +#include "irp.hpp" +#include "exception.hpp" +#include "finally.hpp" sleep_callback* sleep_cb{nullptr}; -_Function_class_(DRIVER_UNLOAD) - -void unload(PDRIVER_OBJECT DriverObject) +void sleep_notification(const sleep_callback::type type) { - debug_log("Leaving World\n"); - IrpUnloadHandler(DriverObject); - delete sleep_cb; -} - -void throw_test() -{ - try + if (type == sleep_callback::type::sleep) { - throw 1; + debug_log("Going to sleep!"); } - catch (...) + + if (type == sleep_callback::type::wakeup) { - debug_log("Exception caught!\n"); + debug_log("Waking up!"); } } @@ -186,52 +25,39 @@ extern "C" void __cdecl __std_terminate() KeBugCheckEx(DRIVER_VIOLATION, 14, 0, 0, 0); } - -extern "C" NTSTATUS DriverEntry(const PDRIVER_OBJECT DriverObject, PUNICODE_STRING /*RegistryPath*/) +void destroy_sleep_callback() { - DriverObject->DriverUnload = unload; - debug_log("Hello World\n"); + delete sleep_cb; +} - delete(new int); +_Function_class_(DRIVER_UNLOAD) void unload(const PDRIVER_OBJECT driver_object) +{ + irp::uninitialize(driver_object); + destroy_sleep_callback(); +} - volatile long i = 0; +extern "C" NTSTATUS DriverEntry(const PDRIVER_OBJECT driver_object, PUNICODE_STRING /*registry_path*/) +{ + driver_object->DriverUnload = unload; - thread::dispatch_on_all_cores([&i]() - { - const auto index = thread::get_processor_index(); - while (i != index) - { - } - - debug_log("Hello from CPU %u/%u\n", thread::get_processor_index() + 1, thread::get_processor_count()); - InterlockedIncrement(&i); - }); - - debug_log("Final i = %i\n", i); - - throw_test(); + auto sleep_destructor = utils::finally(&destroy_sleep_callback); try { - sleep_cb = new sleep_callback([](const sleep_callback::type type) - { - if (type == sleep_callback::type::sleep) - { - debug_log("Going to sleep!"); - } - - if (type == sleep_callback::type::wakeup) - { - debug_log("Waking up!"); - } - }); + sleep_cb = new sleep_callback(sleep_notification); + irp::initialize(driver_object); + sleep_destructor.cancel(); + } + catch (std::exception& e) + { + debug_log("Error: %s\n", e.what()); + return STATUS_INTERNAL_ERROR; } catch (...) { - debug_log("Failed to register sleep callback"); + debug_log("Unknown initialization error occured"); + return STATUS_INTERNAL_ERROR; } - return create_io_device(DriverObject); - - //return STATUS_SUCCESS; + return STATUS_SUCCESS; } diff --git a/src/driver/irp.cpp b/src/driver/irp.cpp new file mode 100644 index 0000000..7b241f4 --- /dev/null +++ b/src/driver/irp.cpp @@ -0,0 +1,142 @@ +#include "irp.hpp" +#include "finally.hpp" +#include "logging.hpp" +#include "exception.hpp" + +#define DOS_DEV_NAME L"\\DosDevices\\HelloDev" +#define DEV_NAME L"\\Device\\HelloDev" + +#define HELLO_DRV_IOCTL CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_NEITHER, FILE_ANY_ACCESS) + +namespace irp +{ + namespace + { + UNICODE_STRING get_unicode_string(const wchar_t* string) + { + UNICODE_STRING unicode_string{}; + RtlInitUnicodeString(&unicode_string, string); + return unicode_string; + } + + UNICODE_STRING get_device_name() + { + return get_unicode_string(DEV_NAME); + } + + UNICODE_STRING get_dos_device_name() + { + return get_unicode_string(DOS_DEV_NAME); + } + + _Function_class_(DRIVER_DISPATCH) NTSTATUS not_supported_handler( + PDEVICE_OBJECT /*device_object*/, const PIRP irp) + { + PAGED_CODE() + + irp->IoStatus.Information = 0; + irp->IoStatus.Status = STATUS_NOT_SUPPORTED; + + IoCompleteRequest(irp, IO_NO_INCREMENT); + + return STATUS_NOT_SUPPORTED; + } + + _Function_class_(DRIVER_DISPATCH) NTSTATUS success_handler(PDEVICE_OBJECT /*device_object*/, const PIRP irp) + { + PAGED_CODE() + + irp->IoStatus.Information = 0; + irp->IoStatus.Status = STATUS_SUCCESS; + + IoCompleteRequest(irp, IO_NO_INCREMENT); + + return STATUS_SUCCESS; + } + + _Function_class_(DRIVER_DISPATCH) NTSTATUS io_ctl_handler( + PDEVICE_OBJECT /*device_object*/, const PIRP irp) + { + PAGED_CODE() + + irp->IoStatus.Information = 0; + irp->IoStatus.Status = STATUS_NOT_SUPPORTED; + + const auto irp_sp = IoGetCurrentIrpStackLocation(irp); + + if (irp_sp) + { + const auto ioctr_code = irp_sp->Parameters.DeviceIoControl.IoControlCode; + + switch (ioctr_code) + { + case HELLO_DRV_IOCTL: + debug_log("[< HelloDriver >] Hello from the Driver!\n"); + break; + default: + debug_log("[-] Invalid IOCTL Code: 0x%X\n", ioctr_code); + irp->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST; + break; + } + } + + IoCompleteRequest(irp, IO_NO_INCREMENT); + + return irp->IoStatus.Status; + } + } + + _Function_class_(DRIVER_DISPATCH) void uninitialize(const PDRIVER_OBJECT driver_object) + { + PAGED_CODE() + + auto dos_device_name = get_dos_device_name(); + + IoDeleteSymbolicLink(&dos_device_name); + IoDeleteDevice(driver_object->DeviceObject); + } + + void initialize(const PDRIVER_OBJECT driver_object) + { + PAGED_CODE() + + auto device_name = get_device_name(); + auto dos_device_name = get_dos_device_name(); + + PDEVICE_OBJECT device_object{}; + auto destructor = utils::finally([&device_object]() + { + if (device_object) + { + IoDeleteDevice(device_object); + } + }); + + auto status = IoCreateDevice(driver_object, 0, &device_name, FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN, + FALSE, &device_object); + if (!NT_SUCCESS(status)) + { + throw std::runtime_error("Unable to create device"); + } + + for (auto i = 0u; i <= IRP_MJ_MAXIMUM_FUNCTION; i++) + { + driver_object->MajorFunction[i] = not_supported_handler; + } + + driver_object->MajorFunction[IRP_MJ_CREATE] = success_handler; + driver_object->MajorFunction[IRP_MJ_CLOSE] = success_handler; + driver_object->MajorFunction[IRP_MJ_DEVICE_CONTROL] = io_ctl_handler; + + device_object->Flags |= DO_DIRECT_IO; + device_object->Flags &= ~DO_DEVICE_INITIALIZING; + + status = IoCreateSymbolicLink(&dos_device_name, &device_name); + if (!NT_SUCCESS(status)) + { + throw std::runtime_error("Unable to create symbolic link"); + } + + destructor.cancel(); + } +} diff --git a/src/driver/irp.hpp b/src/driver/irp.hpp new file mode 100644 index 0000000..2598dbc --- /dev/null +++ b/src/driver/irp.hpp @@ -0,0 +1,7 @@ +#pragma once + +namespace irp +{ + void initialize(PDRIVER_OBJECT driver_object); + void uninitialize(PDRIVER_OBJECT driver_object); +} diff --git a/src/driver/new.cpp b/src/driver/new.cpp index 494f612..bb1fcb8 100644 --- a/src/driver/new.cpp +++ b/src/driver/new.cpp @@ -1,6 +1,5 @@ #include "std_include.hpp" #include "new.hpp" -#include "logging.hpp" void* __cdecl operator new(const size_t size, const POOL_TYPE pool, const unsigned long tag) { diff --git a/src/driver/sleep_callback.cpp b/src/driver/sleep_callback.cpp index 9f650f1..f73a8ae 100644 --- a/src/driver/sleep_callback.cpp +++ b/src/driver/sleep_callback.cpp @@ -37,7 +37,6 @@ sleep_callback::~sleep_callback() } } - void sleep_callback::dispatcher(const type type) const { try