// 驱动代码
#include <ntddk.h>

#define MAX_FILE_LEN 4096

// 设备扩展结构体
typedef struct _FileExtend
{
	PCHAR Buffer;
	ULONG Length;
} FileExtend, *PFileExtend;

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

// 派遣函数-基本操作
NTSTATUS DispatchGeneric(IN PDEVICE_OBJECT DeviceObject, PIRP Irp);

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

// 派遣函数-写操作
NTSTATUS DispatchWrite(IN PDEVICE_OBJECT DeviceObject, PIRP Irp);

// 派遣函数-查询大小
NTSTATUS DispatchQueryInfo(IN PDEVICE_OBJECT DeviceObject, PIRP Irp);

// 入口函数
NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING deviceName = RTL_CONSTANT_STRING(L"\\Device\\FileOperation");
	UNICODE_STRING deviceSymbolicLinkName = RTL_CONSTANT_STRING(L"\\??\\FileOperation");
	PDEVICE_OBJECT DeviceObject = NULL;
	ULONG sub = 0;
	PFileExtend FileExtend = NULL;

	UNREFERENCED_PARAMETER(RegistryPath);
	KdPrint(("驱动加载\n"));

	// 设置常规派遣函数
	for (sub = 0; sub <= IRP_MJ_MAXIMUM_FUNCTION; sub++)
	{
		DriverObject->MajorFunction[sub] = DispatchGeneric;
	}

	// 设置读操作
	DriverObject->MajorFunction[IRP_MJ_READ] = DispatchRead;

	// 设置写操作
	DriverObject->MajorFunction[IRP_MJ_WRITE] = DispatchWrite;

	// 查询操作
	DriverObject->MajorFunction[IRP_MJ_QUERY_INFORMATION] = DispatchQueryInfo;

	// 创建设备
	status = IoCreateDevice(DriverObject, sizeof(FileExtend), &deviceName, FILE_DEVICE_UNKNOWN, 0, FALSE, &DeviceObject);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("创建设备失败:%x\n", status));
		return status;
	}

	// 申请空间
	FileExtend = (PFileExtend)(DeviceObject->DeviceExtension);
	FileExtend->Buffer = ExAllocatePool(PagedPool, MAX_FILE_LEN);
	if (FileExtend->Buffer == NULL)
	{
		KdPrint(("分配内存失败\n"));
		IoDeleteDevice(DeviceObject);
		DriverObject->DeviceObject = NULL;
		status = STATUS_INSUFFICIENT_RESOURCES;
		return status;
	}
	FileExtend->Length = 0;

	// 创建符号链接
	status = IoCreateSymbolicLink(&deviceSymbolicLinkName, &deviceName);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("创建符号链接失败:%x\n", status));
		IoDeleteDevice(DeviceObject);
		return status;
	}

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

	DriverObject->DriverUnload = DriverUnload;

	return status;
}

// 卸载函数
VOID DriverUnload(IN PDRIVER_OBJECT DriverObject)
{
	UNICODE_STRING deviceSymbolicLinkName = RTL_CONSTANT_STRING(L"\\??\\FileOperation");
	PFileExtend FileExtend = NULL;
	UNREFERENCED_PARAMETER(DriverObject);
	KdPrint(("驱动卸载\n"));

	if (DriverObject->DeviceObject != NULL)
	{
		FileExtend = (PFileExtend)(DriverObject->DeviceObject->DeviceExtension);
		if (FileExtend != NULL && FileExtend->Buffer != NULL)
		{
			ExFreePool(FileExtend->Buffer);
			FileExtend->Buffer = NULL;
		}
		IoDeleteSymbolicLink(&deviceSymbolicLinkName);
		IoDeleteDevice(DriverObject->DeviceObject);
	}
}

// 派遣函数-基本操作
NTSTATUS DispatchGeneric(IN PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	UNREFERENCED_PARAMETER(DeviceObject);

	return status;
}

// 派遣函数-读操作
NTSTATUS DispatchRead(IN PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	ULONG Length = 0;
	ULONG Offset = 0;
	ULONG ReadLength = 0;
	PFileExtend FileExtend = NULL;
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);

	KdPrint(("读操作\n"));

	FileExtend = (PFileExtend)DeviceObject->DeviceExtension;
	__try
	{
		Length = stack->Parameters.Read.Length;
		Offset = stack->Parameters.Read.ByteOffset.LowPart;
		if (Length + Offset < MAX_FILE_LEN)
		{
			if (Length + Offset > FileExtend->Length)
			{
				ReadLength = FileExtend->Length - Offset;
			}
			else
			{
				ReadLength = Length;
			}

			RtlCopyMemory(Irp->AssociatedIrp.SystemBuffer, FileExtend->Buffer + Offset, ReadLength);
		}
		else
		{
			status = STATUS_BUFFER_TOO_SMALL;
			ReadLength = 0;
		}
	}

	__except(EXCEPTION_EXECUTE_HANDLER)
	{
		status = GetExceptionCode();
		ReadLength = 0;
	}

	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = ReadLength;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return status;
}

// 派遣函数-写操作
NTSTATUS DispatchWrite(IN PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	ULONG Length = 0;
	ULONG Offset = 0;
	ULONG WriteLength = 0;
	PFileExtend FileExtend = (PFileExtend)(DeviceObject->DeviceExtension);
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);

	KdPrint(("写操作\n"));
	__try
	{
		Length = stack->Parameters.Write.Length;
		Offset = stack->Parameters.Write.ByteOffset.LowPart;
		if (Length + Offset <= MAX_FILE_LEN)
		{
			if (Length + Offset > FileExtend->Length)
			{
				FileExtend->Length = Length + Offset;
			}

			RtlCopyMemory(FileExtend->Buffer + Offset, Irp->AssociatedIrp.SystemBuffer, Length);
			WriteLength = Length;
		}
		else
		{
			status = STATUS_BUFFER_TOO_SMALL;
			WriteLength = 0;
		}
	}

	__except (EXCEPTION_EXECUTE_HANDLER)
	{
		status = GetExceptionCode();
		WriteLength = 0;
	}

	Irp->IoStatus.Status = status;
	Irp->IoStatus.Information = WriteLength;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return status;
}

// 派遣函数-查询大小
NTSTATUS DispatchQueryInfo(IN PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	ULONG Length = 0;
	PFileExtend FileExtend = (PFileExtend)(DeviceObject->DeviceExtension);
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);

	__try
	{
		Length = stack->Parameters.QueryFile.Length;
		if (stack->Parameters.QueryFile.FileInformationClass == FileStandardInformation && Length >= sizeof(FILE_STANDARD_INFORMATION))
		{
			PFILE_STANDARD_INFORMATION pfsi = (PFILE_STANDARD_INFORMATION)Irp->AssociatedIrp.SystemBuffer;
			pfsi->EndOfFile.LowPart = FileExtend->Length;
			Length = sizeof(FILE_STANDARD_INFORMATION);
		}
		else
		{
			status = STATUS_BUFFER_TOO_SMALL;
			Length = 0;
		}
	}

	__except(EXCEPTION_EXECUTE_HANDLER)
	{
		status = GetExceptionCode();
		Length = 0;
	}

	Irp->IoStatus.Information = Length;
	Irp->IoStatus.Status = status;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return status;
}


// 应用代码

// HelloDDK.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include <Windows.h>
#include <stdio.h>
#include <winioctl.h>

#define IOTEST CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_BUFFERED, GENERIC_ALL)

int main()
{
  HANDLE hDevice = CreateFile(L"\\\\.\\FileOperation",
    GENERIC_ALL,
    FILE_SHARE_READ | FILE_SHARE_WRITE,
    NULL,
    OPEN_EXISTING,
    FILE_ATTRIBUTE_NORMAL,
    NULL);

  if (hDevice == INVALID_HANDLE_VALUE)
  {
    printf("设备打开失败%d\n", GetLastError());

    getchar();
    return -1;
  }

  printf("设备打开成功\n");

  DWORD dwRet = 0;
  WriteFile(hDevice, "这是一段来自用户层数据", strlen("这是一段来自用户层数据"), &dwRet, NULL);


  CHAR Buffer[1024] = { 0 };
  ReadFile(hDevice, Buffer, 1024, &dwRet, NULL);
  printf("读到的数据:%s\n", Buffer);
  DWORD dwSize = GetFileSize(hDevice, NULL);
  printf("文件长度%d\n", dwSize);
  //DeviceIoControl(hDevice, IOTEST, NULL, 0, NULL, 0, &dwRet, NULL);
  CloseHandle(hDevice);

  getchar();
  return 0;
}