readwrite.cpp

usb驱动及其测试程序（IRP包的循环传递演示）

// Read/Write request processors for loopback driver
// Copyright (C) 1999 by Walter Oney
// All rights reserved

#include "stddcls.h"
#include "driver.h"

#ifdef DBG
	#define MSGUSBSTRING(d,s,i) { \
		UNICODE_STRING sd; \
		if (i && NT_SUCCESS(GetStringDescriptor(d,i,&sd))) { \
			DbgPrint(s, sd.Buffer); \
			RtlFreeUnicodeString(&sd); \
		}}
#else
	#define MSGUSBSTRING(d,i,s)
#endif

struct _RWCONTEXT : public _URB
	{							// struct _RWCONTEXT
	ULONG_PTR va;				// virtual address for next segment of transfer
	ULONG length;				// length remaining to transfer
	PMDL mdl;					// partial MDL
	ULONG numxfer;				// cumulate transfer count
	BOOLEAN multistage;			// TRUE if multistage transfer
	};							// struct _RWCONTEXT

typedef struct _RWCONTEXT RWCONTEXT, *PRWCONTEXT;

NTSTATUS OnReadWriteComplete(PDEVICE_OBJECT fdo, PIRP Irp, PRWCONTEXT ctx);

#define MAXTRANSFER PAGE_SIZE

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

VOID AbortPipe(PDEVICE_OBJECT fdo, USBD_PIPE_HANDLE hpipe)
	{							// AbortPipe
	PAGED_CODE();
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	URB urb;

	urb.UrbHeader.Length = (USHORT) sizeof(_URB_PIPE_REQUEST);
	urb.UrbHeader.Function = URB_FUNCTION_ABORT_PIPE;
	urb.UrbPipeRequest.PipeHandle = hpipe;

	NTSTATUS status = SendAwaitUrb(fdo, &urb);
	if (!NT_SUCCESS(status))
		KdPrint((DRIVERNAME " - Error %X in AbortPipe\n", status));
	}							// AbortPipe

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS DispatchCreate(PDEVICE_OBJECT fdo, PIRP Irp)
	{							// DispatchCreate
	PAGED_CODE();
	KdPrint((DRIVERNAME " - IRP_MJ_CREATE\n"));
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);

	// Claim the remove lock in Win2K so that removal waits until the
	// handle closes. Don't do this in Win98, however, because this
	// device might be removed by surprise with handles open, whereupon
	// we'll deadlock in HandleRemoveDevice waiting for a close that
	// can never happen because we can't run the user-mode code that
	// would do the close.

	NTSTATUS status;
	if (win98)
		status = STATUS_SUCCESS;
	else 
		status = IoAcquireRemoveLock(&pdx->RemoveLock, stack->FileObject);

	if (NT_SUCCESS(status))
		InterlockedIncrement(&pdx->handles);
	return CompleteRequest(Irp, status, 0);
	}							// DispatchCreate

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS DispatchClose(PDEVICE_OBJECT fdo, PIRP Irp)
	{							// DispatchClose
	PAGED_CODE();
	KdPrint((DRIVERNAME " - IRP_MJ_CLOSE\n"));
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
	InterlockedDecrement(&pdx->handles);
	
	// Release the remove lock to match the acquisition done in DispatchCreate

	if (!win98)
		IoReleaseRemoveLock(&pdx->RemoveLock, stack->FileObject);

	return CompleteRequest(Irp, STATUS_SUCCESS, 0);
	}							// DispatchClose

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS GetStringDescriptor(PDEVICE_OBJECT fdo, UCHAR istring, PUNICODE_STRING s)
	{							// GetStringDescriptor
	NTSTATUS status;
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	URB urb;

	UCHAR data[256];			// maximum-length buffer

	// If this is the first time here, read string descriptor zero and arbitrarily select
	// the first language identifer as the one to use in subsequent get-descriptor calls.

	if (!pdx->langid)
		{						// determine default language id
		UsbBuildGetDescriptorRequest(&urb, sizeof(_URB_CONTROL_DESCRIPTOR_REQUEST), USB_STRING_DESCRIPTOR_TYPE,
			0, 0, data, NULL, sizeof(data), NULL);
		status = SendAwaitUrb(fdo, &urb);
		if (!NT_SUCCESS(status))
			return status;
		pdx->langid = *(LANGID*)(data + 2);
		}						// determine default language id

	// Fetch the designated string descriptor.

	UsbBuildGetDescriptorRequest(&urb, sizeof(_URB_CONTROL_DESCRIPTOR_REQUEST), USB_STRING_DESCRIPTOR_TYPE,
		istring, pdx->langid, data, NULL, sizeof(data), NULL);
	status = SendAwaitUrb(fdo, &urb);
	if (!NT_SUCCESS(status))
		return status;

	ULONG nchars = (data[0] - 2) / 2;
	PWSTR p = (PWSTR) ExAllocatePool(PagedPool, data[0]);
	if (!p)
		return STATUS_INSUFFICIENT_RESOURCES;

	memcpy(p, data + 2, nchars*2);
	p[nchars] = 0;

	s->Length = (USHORT) (2 * nchars);
	s->MaximumLength = (USHORT) ((2 * nchars) + 2);
	s->Buffer = p;

	return STATUS_SUCCESS;
	}							// GetStringDescriptor

///////////////////////////////////////////////////////////////////////////////

#pragma LOCKEDCODE

NTSTATUS OnReadWriteComplete(PDEVICE_OBJECT fdo, PIRP Irp, PRWCONTEXT ctx)
	{							// OnReadWriteComplete
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	BOOLEAN read = (ctx->UrbBulkOrInterruptTransfer.TransferFlags & USBD_TRANSFER_DIRECTION_IN) != 0;
	ctx->numxfer += ctx->UrbBulkOrInterruptTransfer.TransferBufferLength;

#if DBG
	USBD_STATUS urbstatus = URB_STATUS(ctx);
	if (!USBD_SUCCESS(urbstatus))
		KdPrint(("LOOPBACK -  %s failed with USBD status %X\n", read ? "Read" : "Write", urbstatus));
#endif

	// If this stage completed without error, resubmit the URB to perform the
	// next stage

	NTSTATUS status = Irp->IoStatus.Status;
	if (NT_SUCCESS(status) && ctx->length)
		{						// start next stage

		// Calculate length of next stage of the transfer

		ULONG seglen = ctx->length;
		if (seglen > MAXTRANSFER)

	#if (MAXTRANSFER >= PAGE_SIZE)
			seglen = (ULONG_PTR) PAGE_ALIGN(ctx->va) + PAGE_SIZE - ctx->va;
	#else
			seglen = MAXTRANSFER;
	#endif

		// We're now in arbitrary thread context, so the virtual address of the
		// user buffer is currently meaningless. IoBuildPartialMdl copies physical
		// page numbers from the original IRP's probed-and-locked MDL, however,
		// so our process context doesn't matter.

		IoBuildPartialMdl(Irp->MdlAddress, ctx->mdl, (PVOID) ctx->va, seglen);

		// Reinitialize the URB

		ctx->UrbBulkOrInterruptTransfer.TransferBufferLength = seglen;

		// The "next" stack location is the one belonging to the driver
		// underneath us. It's been mostly set to zero and must be reinitialized.

		PIO_STACK_LOCATION stack = IoGetNextIrpStackLocation(Irp);
		stack->MajorFunction = IRP_MJ_INTERNAL_DEVICE_CONTROL;
		stack->Parameters.Others.Argument1 = (PVOID) (PURB) ctx;
		stack->Parameters.DeviceIoControl.IoControlCode = IOCTL_INTERNAL_USB_SUBMIT_URB;
		IoSetCompletionRoutine(Irp, (PIO_COMPLETION_ROUTINE) OnReadWriteComplete,
			(PVOID) ctx, TRUE, TRUE, TRUE);

		ctx->va += seglen;
		ctx->length -= seglen;

		// Pass the recycled IRP down and ignore the return code from IoCallDriver. If
		// this stage gets completed synchronously, this very completion routine will
		// already have been called recursively to deal with it. If not, the IRP is
		// now active again. In either case, we want the current invocation IoCompleteRequest
		// to stop working on this IRP.

		IoCallDriver(pdx->LowerDeviceObject, Irp);
		return STATUS_MORE_PROCESSING_REQUIRED; // halt completion process in its tracks!
		}						// start next stage

	// The request is complete now

	if (NT_SUCCESS(status))
		Irp->IoStatus.Information = ctx->numxfer;
	else
		{					// had an error
		if (read)
			InterlockedIncrement(&pdx->inerror);
		else
			InterlockedIncrement(&pdx->outerror);

		KdPrint(("LOOPBACK - %s finished with error %X\n", read ? "Read" : "Write", status));
		}					// had an error

	ExFreePool(ctx->mdl);
	ExFreePool(ctx);
	IoReleaseRemoveLock(&pdx->RemoveLock, Irp);

	return status;
	}							// OnReadWriteComplete

///////////////////////////////////////////////////////////////////////////////

#pragma PAGEDCODE

NTSTATUS DispatchReadWrite(PDEVICE_OBJECT fdo, PIRP Irp)
	{							// ReadWrite
	PAGED_CODE();
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

	NTSTATUS status = IoAcquireRemoveLock(&pdx->RemoveLock, Irp);
	if (!NT_SUCCESS(status))
		return CompleteRequest(Irp, status, 0);
		
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
	BOOLEAN read = stack->MajorFunction == IRP_MJ_READ;

	USBD_PIPE_HANDLE hpipe = read ? pdx->hinpipe : pdx->houtpipe;
	LONG haderr;
	if (read)
		haderr = InterlockedExchange(&pdx->inerror, 0);
	else
		haderr = InterlockedExchange(&pdx->outerror, 0);
	if (haderr && !NT_SUCCESS(ResetPipe(fdo, hpipe)))
		ResetDevice(fdo);

	// Our strategy will be to do the transfer in stages up to MAXTRANSFER bytes
	// long using a single URB that we resubmit during the completion routine.
	// Note that chained URBs via UrbLink are not supported in either Win98 or Win2K.

	PRWCONTEXT ctx = (PRWCONTEXT) ExAllocatePool(NonPagedPool, sizeof(RWCONTEXT));
	if (!ctx)
		{
		KdPrint((DRIVERNAME " - Can't allocate memory for context structure\n"));
		IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
		return CompleteRequest(Irp, STATUS_INSUFFICIENT_RESOURCES, 0);
		}
	RtlZeroMemory(ctx, sizeof(RWCONTEXT));

	ULONG length = Irp->MdlAddress ?  MmGetMdlByteCount(Irp->MdlAddress) : 0;
	if (!length)
		{						// zero-length read
		IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
		return CompleteRequest(Irp, STATUS_SUCCESS, 0);
		}						// zero-length read
	ULONG_PTR va = (ULONG_PTR) MmGetMdlVirtualAddress(Irp->MdlAddress);

	ULONG urbflags =  USBD_SHORT_TRANSFER_OK | (read ? USBD_TRANSFER_DIRECTION_IN : USBD_TRANSFER_DIRECTION_OUT);

	// Calculate the length of the segment we'll transfer in the first stage.
	// If we're using PAGE_SIZE as our maximum transfer length (which would be the
	// default used by USBD), do a short initial transfer to get us up to a page
	// boundary.

	ULONG seglen = length;
	if (seglen > MAXTRANSFER)

#if (MAXTRANSFER >= PAGE_SIZE)
		seglen = (ULONG_PTR) PAGE_ALIGN(va) + PAGE_SIZE - va;
#else
		seglen = MAXTRANSFER;
#endif

	// Allocate an MDL for each segment of the transfer. The length parameter here
	// is chosen so that the MDL has the capacity to describe two physical pages,
	// which is the most that any stage will ever need.

	PMDL mdl = IoAllocateMdl((PVOID) va, PAGE_SIZE, FALSE, FALSE, NULL);
	if (!mdl)
		{					// can't allocate MDL
		KdPrint((DRIVERNAME " - Can't allocate memory for MDL\n"));
		ExFreePool(ctx);
		IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
		return CompleteRequest(Irp, STATUS_INSUFFICIENT_RESOURCES, 0);