#include <ntddk.h>

// 卸载函数
VOID DriverUnload(IN PDRIVER_OBJECT pDriverObject);

// 派遣函数-常规
NTSTATUS DispatchRoutine(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp);

// 派遣函数-读操作
NTSTATUS DispatchRead(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp);

// 取消函数-读操作
VOID OnCancelRead(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp);

// StartIO函数-读操作
VOID StartIORead(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp);

// 入口函数
NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject, IN PUNICODE_STRING pRegistryPath)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING DevName = { 0 };
	UNICODE_STRING DevSymbolicLinkName = { 0 };
	PDEVICE_OBJECT pDevObj = NULL;
	SIZE_T sub = 0;

	KdPrint(("驱动加载函数\n"));
	UNREFERENCED_PARAMETER(pRegistryPath);

	RtlInitUnicodeString(&DevName, L"\\Device\\CancelIrp");
	RtlInitUnicodeString(&DevSymbolicLinkName, L"\\??\\CancelIrp");

	pDriverObject->DriverUnload = DriverUnload;

	// 创建设备对象
	status = IoCreateDevice(pDriverObject,
		0,
		&DevName,
		FILE_DEVICE_UNKNOWN,
		FILE_DEVICE_SECURE_OPEN,
		FALSE,
		&pDevObj);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("创建设备失败\n"));
		return status;
	}

	// 创建符号链接
	status = IoCreateSymbolicLink(&DevSymbolicLinkName, &DevName);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("创建符号链接失败\n"));
		IoDeleteDevice(pDevObj);
		return status;
	}

	for (sub = 0; sub <= IRP_MJ_MAXIMUM_FUNCTION; sub++)
	{
		pDriverObject->MajorFunction[sub] = DispatchRoutine;
	}

	pDriverObject->MajorFunction[IRP_MJ_READ] = DispatchRead;
	pDriverObject->DriverStartIo = StartIORead;

	pDevObj->Flags |= DO_BUFFERED_IO;
	pDevObj->Flags &= ~DO_DEVICE_INITIALIZING;

	return status;
}

// 卸载函数
VOID DriverUnload(IN PDRIVER_OBJECT pDriverObject)
{
	UNICODE_STRING DevSymbolicLinkName = { 0 };
	RtlInitUnicodeString(&DevSymbolicLinkName, L"\\??\\CancelIrp");

	KdPrint(("驱动卸载函数\n"));
	IoDeleteSymbolicLink(&DevSymbolicLinkName);

	if (pDriverObject->DeviceObject != NULL)
	{
		IoDeleteDevice(pDriverObject->DeviceObject);
	}
}

// 派遣函数-常规
NTSTATUS DispatchRoutine(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
	UNREFERENCED_PARAMETER(pDeviceObject);

	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = 0;

	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

// 派遣函数-读操作
NTSTATUS DispatchRead(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
	KdPrint(("进入读操作\n"));
	UNREFERENCED_PARAMETER(pDeviceObject);
	
	// IoSetCancelRoutine(pIrp, OnCancelRead);
	IoMarkIrpPending(pIrp);

	IoStartPacket(pDeviceObject, pIrp, NULL, OnCancelRead);

	return STATUS_PENDING;
}

// 取消函数-读操作
VOID OnCancelRead(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
	KdPrint(("进入取消函数\n"));
	UNREFERENCED_PARAMETER(pDeviceObject);

	pIrp->IoStatus.Status = STATUS_CANCELLED;
	pIrp->IoStatus.Information = 0;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
}

// StartIO函数-读操作
VOID StartIORead(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
	KdPrint(("进入StartIORead\n"));

	LARGE_INTEGER timeout = RtlConvertLongToLargeInteger(-10 * 100 * 1000);
	KeDelayExecutionThread(KernelMode, FALSE, &timeout);

	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = 0;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);

	IoStartNextPacket(pDeviceObject, TRUE);
}