More cleanup

This commit is contained in:
momo5502 2022-03-26 19:09:54 +01:00
parent a5dab1e9ba
commit 8ee3ade0d4
5 changed files with 181 additions and 208 deletions

View File

@ -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;
}

142
src/driver/irp.cpp Normal file
View File

@ -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();
}
}

7
src/driver/irp.hpp Normal file
View File

@ -0,0 +1,7 @@
#pragma once
namespace irp
{
void initialize(PDRIVER_OBJECT driver_object);
void uninitialize(PDRIVER_OBJECT driver_object);
}

View File

@ -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)
{

View File

@ -37,7 +37,6 @@ sleep_callback::~sleep_callback()
}
}
void sleep_callback::dispatcher(const type type) const
{
try