virtualdisk.c

Windows DDK编写的虚拟磁盘驱动程序源代码

#include <ntddk.h>
#include <ntdddisk.h>

#include "VirtualDisk.h"

NTSTATUS
DriverEntry (
	IN PDRIVER_OBJECT DriverObject,
	IN PUNICODE_STRING RegistryPath
	);

VOID
VirtualDiskUnload (
	IN PDRIVER_OBJECT DriverObject
	);

NTSTATUS
VirtualDiskCreateClose (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	);

NTSTATUS
VirtualDiskReadWrite (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	);

NTSTATUS
VirtualDiskDeviceControl (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	);

VOID
VirtualDiskThread (
	IN PVOID Context
	);

NTSTATUS
VirtualDiskOpenFile (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	);

NTSTATUS
VirtualDiskCloseFile (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	);

#ifdef ALLOC_PRAGMA
#pragma alloc_text(INIT, DriverEntry)
#pragma alloc_text(PAGE, VirtualDiskUnload)
#pragma alloc_text(PAGE, VirtualDiskOpenFile)
#pragma alloc_text(PAGE, VirtualDiskCloseFile)
#endif

NTSTATUS
DriverEntry (
	IN PDRIVER_OBJECT DriverObject,
	IN PUNICODE_STRING RegistryPath
	)
{
	NTSTATUS ntStatus;
	UNICODE_STRING ntDeviceNameString;
	UNICODE_STRING ntWin32NameString;
	PDEVICE_OBJECT deviceObject = NULL;
	PDEVICE_EXTENSION deviceExtension = NULL;
	HANDLE threadHandle = NULL;

	RtlInitUnicodeString(&ntDeviceNameString, DEVICE_NAME);

	ntStatus = IoCreateDevice(
		DriverObject,
		sizeof(DEVICE_EXTENSION),
		&ntDeviceNameString,
		FILE_DEVICE_DISK,
		0,
		FALSE,
		&deviceObject
		);

	if(!NT_SUCCESS(ntStatus))
		return ntStatus;

	deviceObject->Flags |= DO_DIRECT_IO;
	deviceExtension = (PDEVICE_EXTENSION)deviceObject->DeviceExtension;
	deviceExtension->MediaInDevice = FALSE;

	InitializeListHead(&deviceExtension->ListHead);
	KeInitializeSpinLock(&deviceExtension->ListLock);

	KeInitializeEvent(
		&deviceExtension->RequestEvent,
		SynchronizationEvent,
		FALSE
		);

	deviceExtension->TerminateThread = FALSE;

	ntStatus = PsCreateSystemThread(
		&threadHandle,
		(ACCESS_MASK)0L,
		NULL, NULL, NULL,
		VirtualDiskThread,
		deviceObject
		);

	if(!NT_SUCCESS(ntStatus))
	{
		IoDeleteDevice(deviceObject);
		return ntStatus;
	}

	ntStatus = ObReferenceObjectByHandle(
		threadHandle,
		THREAD_ALL_ACCESS,
		NULL,
		KernelMode,
		&deviceExtension->ThreadPointer,
		NULL
		);

	if(!NT_SUCCESS(ntStatus))
	{
		ZwClose(threadHandle);
		deviceExtension->TerminateThread = TRUE;

		KeSetEvent(
			&deviceExtension->RequestEvent,
			(KPRIORITY)0,
			FALSE
			);

		IoDeleteDevice(deviceObject);

		return ntStatus;
	}

	DriverObject->MajorFunction[IRP_MJ_CREATE]		= VirtualDiskCreateClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE]		= VirtualDiskCreateClose;
	DriverObject->MajorFunction[IRP_MJ_READ]		= VirtualDiskReadWrite;
	DriverObject->MajorFunction[IRP_MJ_WRITE]		= VirtualDiskReadWrite;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]	= VirtualDiskDeviceControl;
	DriverObject->DriverUnload = VirtualDiskUnload;

	RtlInitUnicodeString(&ntWin32NameString, DOS_DEVICE_NAME);

	ntStatus = IoCreateSymbolicLink(
		&ntWin32NameString, &ntDeviceNameString
		);

	if(!NT_SUCCESS(ntStatus))
	{
		ZwClose(threadHandle);
		deviceExtension->TerminateThread = TRUE;

		KeSetEvent(
			&deviceExtension->RequestEvent,
			(KPRIORITY)0,
			FALSE
			);

		KeWaitForSingleObject(
			deviceExtension->ThreadPointer,
			Executive,
			KernelMode,
			FALSE,
			NULL
			);

		ObDereferenceObject(deviceExtension->ThreadPointer);

		IoDeleteDevice(deviceObject);

		return ntStatus;
	}

	ZwClose(threadHandle);

	return STATUS_SUCCESS;
}

VOID
VirtualDiskUnload (
	IN PDRIVER_OBJECT DriverObject
	)
{
	PDEVICE_OBJECT deviceObject = NULL;
	PDEVICE_EXTENSION deviceExtension = NULL;
	UNICODE_STRING ntWin32NameString;

	deviceObject = DriverObject->DeviceObject;

	RtlInitUnicodeString(&ntWin32NameString, DOS_DEVICE_NAME);
	IoDeleteSymbolicLink(&ntWin32NameString);

	if(NULL != deviceObject)
	{
		deviceExtension = (PDEVICE_EXTENSION)deviceObject->DeviceExtension;

		deviceExtension->TerminateThread = TRUE;

		KeSetEvent(
			&deviceExtension->RequestEvent,
			(KPRIORITY)0,
			FALSE
			);

		KeWaitForSingleObject(
			deviceExtension->ThreadPointer,
			Executive,
			KernelMode,
			FALSE,
			NULL
			);

		ObDereferenceObject(deviceExtension->ThreadPointer);

		IoDeleteDevice(deviceObject);
	}
}

NTSTATUS
VirtualDiskCreateClose (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	)
{
	PAGED_CODE();

	Irp->IoStatus.Status = STATUS_SUCCESS;
	Irp->IoStatus.Information = FILE_OPENED;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return STATUS_SUCCESS;
}

NTSTATUS
VirtualDiskReadWrite (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	)
{
	PDEVICE_EXTENSION deviceExtension = NULL;
	PIO_STACK_LOCATION irpSp = NULL;

	deviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;

	if(!deviceExtension->MediaInDevice)
	{
		Irp->IoStatus.Status = STATUS_NO_MEDIA_IN_DEVICE;
		Irp->IoStatus.Information = 0;

		IoCompleteRequest(Irp, IO_NO_INCREMENT);

		return STATUS_NO_MEDIA_IN_DEVICE;
	}

	irpSp = IoGetCurrentIrpStackLocation(Irp);

	if(0 == irpSp->Parameters.Read.Length)
	{
		Irp->IoStatus.Status = STATUS_SUCCESS;
		Irp->IoStatus.Information = 0;

		return STATUS_SUCCESS;
	}

	IoMarkIrpPending(Irp);

	ExInterlockedInsertTailList(
		&deviceExtension->ListHead,
		&Irp->Tail.Overlay.ListEntry,
		&deviceExtension->ListLock
		);

	KeSetEvent(
		&deviceExtension->RequestEvent,
		(KPRIORITY)0,
		FALSE
		);

	return STATUS_PENDING;
}

NTSTATUS
VirtualDiskDeviceControl (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	)
{
	PDEVICE_EXTENSION deviceExtension = NULL;
	PIO_STACK_LOCATION irpSp = NULL;
	NTSTATUS ntStatus;

	deviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
	irpSp = IoGetCurrentIrpStackLocation(Irp);

	if(!deviceExtension->MediaInDevice &&
		irpSp->Parameters.DeviceIoControl.IoControlCode != IOCTL_VIRTUAL_DISK_CREATE_DISK)
	{
		Irp->IoStatus.Status = STATUS_NO_MEDIA_IN_DEVICE;
		Irp->IoStatus.Information = 0;

		IoCompleteRequest(Irp, IO_NO_INCREMENT);

		return STATUS_NO_MEDIA_IN_DEVICE;
	}

	switch(irpSp->Parameters.DeviceIoControl.IoControlCode)
	{
	case IOCTL_VIRTUAL_DISK_CREATE_DISK:
		{
			if(deviceExtension->MediaInDevice)
			{
				DbgPrint("VirtualDisk: IOCTL_VIRTUAL_DISK_CREATE_DISK: Media already opened\n");

				ntStatus = STATUS_INVALID_DEVICE_REQUEST;
				Irp->IoStatus.Information = 0;
				break;
			}

			if(irpSp->Parameters.DeviceIoControl.InputBufferLength < 
				sizeof(OPEN_FILE_INFORMATION))
			{
				ntStatus = STATUS_INVALID_PARAMETER;
				Irp->IoStatus.Information = 0;
				break;
			}

			if(irpSp->Parameters.DeviceIoControl.InputBufferLength <
				sizeof(OPEN_FILE_INFORMATION) +
				((POPEN_FILE_INFORMATION)Irp->AssociatedIrp.SystemBuffer)->FileNameLength - sizeof(UCHAR))
			{
				ntStatus = STATUS_INVALID_PARAMETER;
				Irp->IoStatus.Information = 0;
				break;
			}

			IoMarkIrpPending(Irp);

			ExInterlockedInsertTailList(
				&deviceExtension->ListHead,
				&Irp->Tail.Overlay.ListEntry,
				&deviceExtension->ListLock
				);

			KeSetEvent(
				&deviceExtension->RequestEvent,
				(KPRIORITY)0,
				FALSE
				);

			ntStatus = STATUS_PENDING;

			break;
		}

	case IOCTL_VIRTUAL_DISK_CLOSE_DISK:
		{
			IoMarkIrpPending(Irp);

			ExInterlockedInsertTailList(
				&deviceExtension->ListHead,
				&Irp->Tail.Overlay.ListEntry,
				&deviceExtension->ListLock
				);

			KeSetEvent(
				&deviceExtension->RequestEvent,
				(KPRIORITY)0,
				FALSE
				);

			ntStatus = STATUS_PENDING;

			break;
		}

	case IOCTL_VIRTUAL_DISK_QUERY_DISK:
		{
			POPEN_FILE_INFORMATION openFileInformation = NULL;

			if(irpSp->Parameters.DeviceIoControl.OutputBufferLength <
				sizeof(OPEN_FILE_INFORMATION) + deviceExtension->FileName.Length - sizeof(UCHAR))
			{
				ntStatus = STATUS_BUFFER_TOO_SMALL;
				Irp->IoStatus.Information = 0;
				break;
			}

			openFileInformation = (POPEN_FILE_INFORMATION)Irp->AssociatedIrp.SystemBuffer;

			openFileInformation->FileSize.QuadPart = deviceExtension->FileSize.QuadPart;
			openFileInformation->FileNameLength = deviceExtension->FileName.Length;

			RtlCopyMemory(
				openFileInformation->FileName,
				deviceExtension->FileName.Buffer,
				deviceExtension->FileName.Length
				);

			ntStatus = STATUS_SUCCESS;

			Irp->IoStatus.Information = sizeof(OPEN_FILE_INFORMATION) +
				openFileInformation->FileNameLength - sizeof(UCHAR);

			break;
		}

	case IOCTL_DISK_CHECK_VERIFY:
	case IOCTL_STORAGE_CHECK_VERIFY:
	case IOCTL_STORAGE_CHECK_VERIFY2:
		{
			ntStatus = STATUS_SUCCESS;
			Irp->IoStatus.Information = 0;
			break;
		}

	/* Returns information about the physical disk's geometry (media type, number of cylinders,
	 * tracks per cylinder, sectors per track, and bytes per sector).
	 */
	case IOCTL_DISK_GET_DRIVE_GEOMETRY:
		{
			PDISK_GEOMETRY diskGeometry = NULL;
			ULONGLONG length;

			if(irpSp->Parameters.DeviceIoControl.OutputBufferLength <
				sizeof(DISK_GEOMETRY))
			{
				ntStatus = STATUS_BUFFER_TOO_SMALL;
				Irp->IoStatus.Information = 0;
				break;
			}

			diskGeometry = (PDISK_GEOMETRY)Irp->AssociatedIrp.SystemBuffer;

			length = deviceExtension->FileSize.QuadPart;
			diskGeometry->Cylinders.QuadPart = length / SECTOR_SIZE / 32 / 2;
			diskGeometry->MediaType = FixedMedia;
			diskGeometry->TracksPerCylinder = 2;
			diskGeometry->SectorsPerTrack = 32;
			diskGeometry->BytesPerSector = SECTOR_SIZE;

			ntStatus = STATUS_SUCCESS;
			Irp->IoStatus.Information = sizeof(DISK_GEOMETRY);
			break;
		}

	/* Returns the length, in bytes, of the disk, partition, 
	 * or volume associated with the device object that is the target of the request
	 */
	case IOCTL_DISK_GET_LENGTH_INFO:
		{
			PGET_LENGTH_INFORMATION getLengthInformation = NULL;

			if(irpSp->Parameters.DeviceIoControl.OutputBufferLength <
				sizeof(GET_LENGTH_INFORMATION))
			{
				ntStatus = STATUS_BUFFER_TOO_SMALL;
				Irp->IoStatus.Information = 0;
				break;
			}

			getLengthInformation = (PGET_LENGTH_INFORMATION)Irp->AssociatedIrp.SystemBuffer;
			getLengthInformation->Length.QuadPart = deviceExtension->FileSize.QuadPart;

			ntStatus = STATUS_SUCCESS;
			Irp->IoStatus.Information = sizeof(GET_LENGTH_INFORMATION);
			break;
		}

	/* Returns information about the type, size, and nature of a disk partition.
	 * The PARTITION_INFORMATION structure contains partition information for a partition
	 * with a traditional AT-style Master Boot Record (MBR). 
	 */