virtualdisk.c

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

	case IOCTL_DISK_GET_PARTITION_INFO:
		{
			PPARTITION_INFORMATION partitionInformation = NULL;
			ULONGLONG length;

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

			partitionInformation = (PPARTITION_INFORMATION)Irp->AssociatedIrp.SystemBuffer;

			length = deviceExtension->FileSize.QuadPart;
			
			partitionInformation->StartingOffset.QuadPart = 0;
			partitionInformation->PartitionLength.QuadPart = length;
			partitionInformation->HiddenSectors = 1;
			partitionInformation->PartitionNumber = 0;
			partitionInformation->PartitionType = 0;
			partitionInformation->BootIndicator = FALSE;
			partitionInformation->RecognizedPartition = FALSE;
			partitionInformation->RewritePartition = FALSE;

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

	/* Returns information about the type, size, and nature of a disk partition.
	 * PARTITION_INFORMATION_EX is the extended version of the PARTITION_INFORMATION structure.
	 * It holds information both for partitions with a Master Boot Record and for partitions with a GUID Partition Table.
	 */
	case IOCTL_DISK_GET_PARTITION_INFO_EX:
		{
			PPARTITION_INFORMATION_EX partitionInformationEx = NULL;
			ULONGLONG length;

			if(irpSp->Parameters.DeviceIoControl.OutputBufferLength <
				sizeof(PARTITION_INFORMATION_EX))
			{
				ntStatus = STATUS_BUFFER_TOO_SMALL;
				Irp->IoStatus.Information = 0;
				break;
			}
			
			partitionInformationEx = (PPARTITION_INFORMATION_EX)Irp->AssociatedIrp.SystemBuffer;

			length = deviceExtension->FileSize.QuadPart;

			partitionInformationEx->PartitionStyle = PARTITION_STYLE_MBR;
			partitionInformationEx->StartingOffset.QuadPart = 0;
			partitionInformationEx->PartitionLength.QuadPart = length;
			partitionInformationEx->PartitionNumber = 0;
			partitionInformationEx->RewritePartition = FALSE;
			partitionInformationEx->Mbr.PartitionType = 0;
			partitionInformationEx->Mbr.BootIndicator = FALSE;
			partitionInformationEx->Mbr.RecognizedPartition = FALSE;
			partitionInformationEx->Mbr.HiddenSectors = 1;

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

	/* Determines whether a disk is writable */
	case IOCTL_DISK_IS_WRITABLE:
		{
			ntStatus = STATUS_SUCCESS;

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

	case IOCTL_DISK_MEDIA_REMOVAL:
	case IOCTL_STORAGE_MEDIA_REMOVAL:
		{
			ntStatus = STATUS_SUCCESS;
			Irp->IoStatus.Information = 0;
		}

	/* Changes the partition type of the specified disk partition.
	 * SET_PARTITION_INFORMATION is used with IOCTL_DISK_SET_PARTITION_INFO to 
	 * change the partition type of a specified Master Boot Record (MBR) disk partition.
	 */ 
	case IOCTL_DISK_SET_PARTITION_INFO:
		{
			if(irpSp->Parameters.DeviceIoControl.InputBufferLength <
				sizeof(SET_PARTITION_INFORMATION))
			{
				ntStatus = STATUS_INVALID_PARAMETER;
				Irp->IoStatus.Information = 0;
				break;
			}

			ntStatus = STATUS_SUCCESS;
			Irp->IoStatus.Information = 0;

			break;
		}

	/* Performs a logical format of a specified extent on a disk.
	 */
	case IOCTL_DISK_VERIFY:
		{
			PVERIFY_INFORMATION verifyInformation = NULL;

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

			verifyInformation = (PVERIFY_INFORMATION)Irp->AssociatedIrp.SystemBuffer;

			ntStatus = STATUS_SUCCESS;
			Irp->IoStatus.Information = verifyInformation->Length;
			break;
		}

	default:
		{
			DbgPrint("VirtualDisk: Unknown IoControlCode: %#x\n",
				irpSp->Parameters.DeviceIoControl.IoControlCode);

			ntStatus = STATUS_INVALID_DEVICE_REQUEST;
			Irp->IoStatus.Information = 0;
		}
	}	/* end of switch */

	if(STATUS_PENDING != ntStatus)
	{
		Irp->IoStatus.Status = ntStatus;

		IoCompleteRequest(Irp, IO_NO_INCREMENT);
	}

	return ntStatus;
}	/* end of VirtualDiskDeviceControl */

VOID
VirtualDiskThread (
	IN PVOID Context
	)
{
	PDEVICE_OBJECT deviceObject = NULL;
	PDEVICE_EXTENSION deviceExtension = NULL;
	PLIST_ENTRY request = NULL;
	PIRP irp;
	PIO_STACK_LOCATION irpSp = NULL;
	PUCHAR systemBuffer;
	PUCHAR buffer;

	ASSERT(NULL != Context);

	deviceObject = (PDEVICE_OBJECT)Context;
	deviceExtension = (PDEVICE_EXTENSION)deviceObject->DeviceExtension;

	KeSetPriorityThread(KeGetCurrentThread(), LOW_REALTIME_PRIORITY);

	while(TRUE)
	{
		KeWaitForSingleObject(
			&deviceExtension->RequestEvent,
			Executive,
			KernelMode,
			FALSE,
			NULL
			);

		/* Terminates the current system thread */
		if(deviceExtension->TerminateThread)
		{
			PsTerminateSystemThread(STATUS_SUCCESS);
		}

		while(request = ExInterlockedRemoveHeadList(
			&deviceExtension->ListHead,
			&deviceExtension->ListLock))
		{
			irp = CONTAINING_RECORD(request, IRP, Tail.Overlay.ListEntry);
			irpSp = IoGetCurrentIrpStackLocation(irp);

			switch(irpSp->MajorFunction)
			{
			case IRP_MJ_READ:
				systemBuffer = (PUCHAR)MmGetSystemAddressForMdlSafe(irp->MdlAddress, NormalPagePriority);
				if(NULL == systemBuffer)
				{
					irp->IoStatus.Status = STATUS_INSUFFICIENT_RESOURCES;
					irp->IoStatus.Information = 0;
					break;
				}
				buffer = (PUCHAR)ExAllocatePool(PagedPool, irpSp->Parameters.Read.Length);
				if(NULL == buffer)
				{
					irp->IoStatus.Status = STATUS_INSUFFICIENT_RESOURCES;
					irp->IoStatus.Information = 0;
					break;
				}

				/* Redirect to our file handle */
				ZwReadFile(
					deviceExtension->FileHandle,
					NULL, NULL, NULL,
					&irp->IoStatus,
					buffer,
					irpSp->Parameters.Read.Length,
					&irpSp->Parameters.Read.ByteOffset,
					NULL
					);
				RtlCopyMemory(systemBuffer, buffer, irpSp->Parameters.Read.Length);
				ExFreePool(buffer);
				break;
			case IRP_MJ_WRITE:
				if(irpSp->Parameters.Write.ByteOffset.QuadPart +
					irpSp->Parameters.Write.Length > deviceExtension->FileSize.QuadPart)
				{
					irp->IoStatus.Status = STATUS_INVALID_PARAMETER;
					irp->IoStatus.Information = 0;
					break;
				}
				ZwWriteFile(
					deviceExtension->FileHandle,
					NULL, NULL, NULL,
					&irp->IoStatus,
					MmGetSystemAddressForMdlSafe(irp->MdlAddress, NormalPagePriority),
					irpSp->Parameters.Write.Length,
					&irpSp->Parameters.Write.ByteOffset,
					NULL
					);
				break;
			case IRP_MJ_DEVICE_CONTROL:
				switch(irpSp->Parameters.DeviceIoControl.IoControlCode)
				{
				case IOCTL_VIRTUAL_DISK_CREATE_DISK:
					irp->IoStatus.Status = VirtualDiskOpenFile(deviceObject, irp);
					break;
				case IOCTL_VIRTUAL_DISK_CLOSE_DISK:
					irp->IoStatus.Status = VirtualDiskCloseFile(deviceObject, irp);
					break;
				default:
					irp->IoStatus.Status = STATUS_DRIVER_INTERNAL_ERROR;
					break;
				}
				break;
			default:
				irp->IoStatus.Status = STATUS_DRIVER_INTERNAL_ERROR;
				break;
			}	/* end of switch */

			IoCompleteRequest(
				irp,
				(CCHAR)(NT_SUCCESS(irp->IoStatus.Status) ? IO_DISK_INCREMENT : IO_NO_INCREMENT)
				);
		}	/* end of while */
	}	/* end of while(TRUE) */
}

NTSTATUS
VirtualDiskOpenFile (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	)
{
	PDEVICE_EXTENSION deviceExtension = NULL;
	POPEN_FILE_INFORMATION openFileInformation = NULL;
	UNICODE_STRING uFileName;
	NTSTATUS ntStatus;
	OBJECT_ATTRIBUTES objectAttributes;
	FILE_END_OF_FILE_INFORMATION fileEof;
	FILE_BASIC_INFORMATION fileBasic;
	FILE_STANDARD_INFORMATION fileStandard;
	FILE_ALIGNMENT_INFORMATION fileAlignment;

	PAGED_CODE();

	ASSERT(NULL != DeviceObject);
	ASSERT(NULL != Irp);

	deviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
	openFileInformation = (POPEN_FILE_INFORMATION)Irp->AssociatedIrp.SystemBuffer;

	deviceExtension->FileName.Length = openFileInformation->FileNameLength;
	deviceExtension->FileName.MaximumLength = openFileInformation->FileNameLength;
	deviceExtension->FileName.Buffer = ExAllocatePool(NonPagedPool, openFileInformation->FileNameLength);

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

	/* Must deallocate the buffer by calling RtlFreeUnicodeString */
	ntStatus = RtlAnsiStringToUnicodeString(
		&uFileName,
		&deviceExtension->FileName,
		TRUE
		);
	if(!NT_SUCCESS(ntStatus))
	{
		ExFreePool(deviceExtension->FileName.Buffer);
		Irp->IoStatus.Status = ntStatus;
		Irp->IoStatus.Information = 0;
		return ntStatus;
	}

	InitializeObjectAttributes(
		&objectAttributes,
		&uFileName,
		OBJ_CASE_INSENSITIVE,
		NULL, NULL
		);

	/* The ZwCreateFile routine creates a new file or opens an existing file */
	ntStatus = ZwCreateFile(
		&deviceExtension->FileHandle,
		GENERIC_READ | GENERIC_WRITE,
		&objectAttributes,
		&Irp->IoStatus,
		NULL,
		FILE_ATTRIBUTE_NORMAL,
		0,	/* set to zero, giver the caller exclusive access to the open file */
		FILE_OPEN,
		FILE_NON_DIRECTORY_FILE | FILE_RANDOM_ACCESS |
		FILE_NO_INTERMEDIATE_BUFFERING | FILE_SYNCHRONOUS_IO_NONALERT,
		NULL,
		0
		);

	/* file is not existing yet */
	if(ntStatus == STATUS_OBJECT_NAME_NOT_FOUND || ntStatus == STATUS_NO_SUCH_FILE)
	{
		if(0 == openFileInformation->FileSize.QuadPart)
		{
			ExFreePool(deviceExtension->FileName.Buffer);
			RtlFreeUnicodeString(&uFileName);

			Irp->IoStatus.Status = STATUS_NO_SUCH_FILE;
			Irp->IoStatus.Information = 0;

			return STATUS_NO_SUCH_FILE;
		}
		else
		{
			/* create this file */
			ntStatus = ZwCreateFile(
				&deviceExtension->FileHandle,
				GENERIC_READ | GENERIC_WRITE,
				&objectAttributes,
				&Irp->IoStatus,
				&openFileInformation->FileSize,
				FILE_ATTRIBUTE_NORMAL,
				0,
				FILE_OPEN_IF,
				FILE_NON_DIRECTORY_FILE | FILE_RANDOM_ACCESS |
				FILE_NO_INTERMEDIATE_BUFFERING | FILE_SYNCHRONOUS_IO_NONALERT,
				NULL,
				0
				);
			if(!NT_SUCCESS(ntStatus))
			{
				ExFreePool(deviceExtension->FileName.Buffer);
				RtlFreeUnicodeString(&uFileName);
				return ntStatus;
			}

			if(Irp->IoStatus.Information = FILE_CREATED)
			{
				fileEof.EndOfFile.QuadPart = openFileInformation->FileSize.QuadPart;

				ntStatus = ZwSetInformationFile(
					deviceExtension->FileHandle,
					&Irp->IoStatus,
					&fileEof,
					sizeof(FILE_END_OF_FILE_INFORMATION),
					FileEndOfFileInformation
					);
				if(!NT_SUCCESS(ntStatus))
				{
					ExFreePool(deviceExtension->FileName.Buffer);
					RtlFreeUnicodeString(&uFileName);
					ZwClose(deviceExtension->FileHandle);
					return ntStatus;
				}
			}
		}
	}
	else if(!NT_SUCCESS(ntStatus))
	{
		ExFreePool(deviceExtension->FileName.Buffer);
		RtlFreeUnicodeString(&uFileName);
		return ntStatus;
	}

	RtlFreeUnicodeString(&uFileName);

	ntStatus = ZwQueryInformationFile(
		deviceExtension->FileHandle,
		&Irp->IoStatus,
		&fileBasic,
		sizeof(FILE_BASIC_INFORMATION),
		FileBasicInformation
		);
	if(!NT_SUCCESS(ntStatus))
	{
		ExFreePool(deviceExtension->FileName.Buffer);
		ZwClose(deviceExtension->FileHandle);
		return ntStatus;
	}

	ntStatus = ZwQueryInformationFile(
		deviceExtension->FileHandle,
		&Irp->IoStatus,
		&fileStandard,
		sizeof(FILE_STANDARD_INFORMATION),
		FileStandardInformation
		);
	if(!NT_SUCCESS(ntStatus))
	{
		ExFreePool(deviceExtension->FileName.Buffer);
		ZwClose(deviceExtension->FileHandle);
		return ntStatus;
	}

	deviceExtension->FileSize.QuadPart = fileStandard.EndOfFile.QuadPart;

	ntStatus = ZwQueryInformationFile(
		deviceExtension->FileHandle,
		&Irp->IoStatus,
		&fileAlignment,
		sizeof(FILE_ALIGNMENT_INFORMATION),
		FileAlignmentInformation
		);
	if(!NT_SUCCESS(ntStatus))
	{
		ExFreePool(deviceExtension->FileName.Buffer);
		ZwClose(deviceExtension->FileHandle);
		return ntStatus;
	}
	DeviceObject->AlignmentRequirement = fileAlignment.AlignmentRequirement;

	DeviceObject->Characteristics &= ~FILE_READ_ONLY_DEVICE;

	deviceExtension->MediaInDevice = TRUE;

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

	return STATUS_SUCCESS;
}

NTSTATUS
VirtualDiskCloseFile (
	IN PDEVICE_OBJECT DeviceObject,
	IN PIRP Irp
	)
{
	PDEVICE_EXTENSION deviceExtension = NULL;

	PAGED_CODE();
	ASSERT(DeviceObject != NULL);
	ASSERT(Irp != NULL);

	deviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
	ExFreePool(deviceExtension->FileName.Buffer);
	ZwClose(deviceExtension->FileHandle);
	deviceExtension->MediaInDevice = FALSE;

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

	return STATUS_SUCCESS;
}