Add ioctrl tests

src/driver/main.cpp

@@ -1,8 +1,158 @@
 #include "std_include.hpp"
 #include "logging.hpp"
-
 #include "thread.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)
+
+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;
+}
+
+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;
+}
+
+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");
+}
+
+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;
+
+	// Assign the driver Unload routine
+	DriverObject->DriverUnload = IrpUnloadHandler;
+
+	// 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;
+}
+
 _Function_class_(DRIVER_UNLOAD)
 
 void unload(PDRIVER_OBJECT /*DriverObject*/)
@@ -25,10 +175,12 @@ extern "C" NTSTATUS DriverEntry(const PDRIVER_OBJECT DriverObject, PUNICODE_STRI
 		}
 
 		debug_log("Hello from CPU %u/%u\n", thread::get_processor_index() + 1, thread::get_processor_count());
-		++i;
+		InterlockedIncrement(&i);
 	});
 
 	debug_log("Final i = %i\n", i);
 
+	create_io_device(DriverObject);
+
 	return STATUS_SUCCESS;
 }