readwrite.cpp

Programming the Microsoft Windows driver model.2nd 随书光盘。内有很多作者送的实用工具和随书源码。WDM编程

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

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

#if 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

	// The following fields are used for error recovery only. Making this
	// structure bigger to contain them simply saves extra complication in the
	// code when we decide we need to recover the error

	PIO_WORKITEM rcitem;		// work item created for recovery
	PIRP Irp;					// the main IRP that we're going to fail
	};							// struct _RWCONTEXT

typedef struct _RWCONTEXT RWCONTEXT, *PRWCONTEXT;

ULONG GetStatus(PDEVICE_OBJECT fdo);
NTSTATUS OnReadWriteComplete(PDEVICE_OBJECT fdo, PIRP Irp, PRWCONTEXT ctx);
VOID OnCancelReadWrite(IN PDEVICE_OBJECT fdo, IN PIRP Irp);
VOID RecoverFromError(PDEVICE_OBJECT fdo, PRWCONTEXT ctx);

///////////////////////////////////////////////////////////////////////////////
// AbortPipe is called as part of an attempt to recover after an I/O error to
// abort pending requests for a given pipe.

#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);

	NTSTATUS status = STATUS_SUCCESS;
	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);
	return CompleteRequest(Irp, STATUS_SUCCESS, 0);
	}							// DispatchClose

///////////////////////////////////////////////////////////////////////////////
// I (WCO) revised DispatchReadWrite in SP-7 to use a queue. The
// code that used to be in this routine was moved to StartIo. A queue is
// needed to avoid mixing segments from different IRPs in the calls
// down to the bus driver.

#pragma PAGEDCODE

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

	// This device can only handle up to 4096 bytes in one transfer

	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);
	if (stack->Parameters.Read.Length > 4096)
		{
		KdPrint((DRIVERNAME " - Failing read/write bigger than 4096 bytes\n"));
		return CompleteRequest(Irp, STATUS_INVALID_PARAMETER, 0);
		}

	IoMarkIrpPending(Irp);
	StartPacket(&pdx->dqReadWrite, fdo, Irp, OnCancelReadWrite);
	return STATUS_PENDING;
	}							// DispatchReadWrite

#pragma LOCKEDCODE

VOID OnCancelReadWrite(IN PDEVICE_OBJECT fdo, IN PIRP Irp)
	{							// OnCancelReadWrite
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	CancelRequest(&pdx->dqReadWrite, Irp);
	}							// OnCancelReadWrite

///////////////////////////////////////////////////////////////////////////////
// GetStatus is called during an attempt to recover after an error to determine
// the USBD status for the device.

#pragma PAGEDCODE

ULONG GetStatus(PDEVICE_OBJECT fdo)
	{							// GetStatus
	PAGED_CODE();
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;

	KEVENT event;
	KeInitializeEvent(&event, NotificationEvent, FALSE);
	IO_STATUS_BLOCK iostatus;
	ULONG portstatus;

	PIRP Irp = IoBuildDeviceIoControlRequest(IOCTL_INTERNAL_USB_GET_PORT_STATUS,
		pdx->LowerDeviceObject, NULL, 0, NULL, 0, TRUE, &event, &iostatus);
	if (!Irp)
		return 0;
	IoGetNextIrpStackLocation(Irp)->Parameters.Others.Argument1 = (PVOID) &portstatus;

	NTSTATUS status = IoCallDriver(pdx->LowerDeviceObject, Irp);
	if (status == STATUS_PENDING)
		{
		KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, NULL);
		status = iostatus.Status;
		}
	
	if (!NT_SUCCESS(status))
		{
		KdPrint((DRIVERNAME " - Error %X trying to reset device\n", status));
		return 0;
		}

	return portstatus;
	}							// GetStatus

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

#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] - sizeof(WCHAR)) / sizeof(WCHAR);
	if (nchars > 127)
		nchars = 127;
	PWSTR p = (PWSTR) ExAllocatePool(PagedPool, (nchars + 1) * sizeof(WCHAR));
	if (!p)
		return STATUS_INSUFFICIENT_RESOURCES;

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

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

	return STATUS_SUCCESS;
	}							// GetStringDescriptor

///////////////////////////////////////////////////////////////////////////////
// This is the completion routine for IRP_MJ_READ and IRP_MJ_WRITE requests.

#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((DRIVERNAME " - %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 && !Irp->Cancel)
		{						// start next stage

		// Calculate length of next stage of the transfer

		ULONG seglen = ctx->length;
		if (seglen > pdx->maxtransfer)
			seglen = pdx->maxtransfer;

		// 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.

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

		// See the comment about this in DispatchReadWrite...

		if (!GenericGetSystemAddressForMdl(mdl))
			{						// can't map transfer segment
			KdPrint((DRIVERNAME " - Can't map memory for read or write\n"));
			status = STATUS_INSUFFICIENT_RESOURCES;
			Irp->IoStatus.Status = STATUS_INSUFFICIENT_RESOURCES;	// okay to change because DispatchReadWrite returned STATUS_PENDING
			goto FinishCompletion;
			}						// can't map transfer segment

		// 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

FinishCompletion:				// label reached if MDL locking fails
	if (NT_SUCCESS(status))
		Irp->IoStatus.Information = ctx->numxfer;
	else
		{					// had an error
		if (read)
			InterlockedIncrement(&pdx->inerror);
		else
			InterlockedIncrement(&pdx->outerror);

		KdPrint((DRIVERNAME " - %s finished with error %X\n", read ? "Read" : "Write", status));

		if (status != STATUS_CANCELLED)
			{					// recover after I/O error

			// Attempt to recover after an error. Whether or not the attempt succeeds, this IRP
			// is a goner and will eventually get completed with its current error status. Since
			// StartIo is serializing IRPs, we needn't worry about anyone else trying to read or
			// write the device while this is going on. We have to call RecoverFromError at PASSIVE_LEVEL,
			// though. If we were tricky and did that directly from here, we would need to
			// use IoSetCompletionRoutineEx because whatever unload protection our caller has
			// done expires as soon as the main IRP completes (and RecoverFromError completes the
			// main IRP.) Therefore, always queue a work item, which independently protects us
			// from unloading.

			ctx->rcitem = IoAllocateWorkItem(fdo);
			if (!ctx->rcitem)
				KdPrint((DRIVERNAME " - Can't allocate work item for error recovery\n"));
			else
				{			// try to recover
				ctx->Irp = Irp;
				IoQueueWorkItem(ctx->rcitem, (PIO_WORKITEM_ROUTINE) RecoverFromError, CriticalWorkQueue, (PVOID) ctx);
				return STATUS_MORE_PROCESSING_REQUIRED;	// defer completion a while longer
				}			// try to recover
			}					// recover after I/O error
		}					// had an error

	IoFreeMdl(ctx->mdl);
	ExFreePool(ctx);
	StartNextPacket(&pdx->dqReadWrite, fdo);
	IoReleaseRemoveLock(&pdx->RemoveLock, Irp);

	return STATUS_SUCCESS;		// allow IRP to continue completing
	}							// OnReadWriteComplete

///////////////////////////////////////////////////////////////////////////////
// RecoverFromError is called in a system worker thread to attempt recovery
// after a device error.

#pragma PAGEDCODE

VOID RecoverFromError(PDEVICE_OBJECT fdo, PRWCONTEXT ctx)
	{							// RecoverFromError
	PAGED_CODE();
	PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) fdo->DeviceExtension;
	BOOLEAN read = (ctx->UrbBulkOrInterruptTransfer.TransferFlags & USBD_TRANSFER_DIRECTION_IN) != 0;

	KdPrint((DRIVERNAME " - Beginning error recovery\n"));

	// Determine the status of our device

	ULONG portstatus = GetStatus(fdo);

	// If the port is connected but not enabled, try resetting it

	USBD_PIPE_HANDLE hpipe = read ? pdx->hinpipe : pdx->houtpipe;
	if (!(portstatus & USBD_PORT_ENABLED) && (portstatus & USBD_PORT_CONNECTED))
		{						// reset port
		KdPrint((DRIVERNAME " - Error recovery aborting pipe\n"));
		AbortPipe(fdo, hpipe);

		KdPrint((DRIVERNAME " - Error recovery resetting device\n"));
		ResetDevice(fdo);
		}						// reset port

	// Now reset the failed endpoint

	ResetPipe(fdo, hpipe);

	KdPrint((DRIVERNAME " - Finished with error recovery\n"));

	IoFreeWorkItem(ctx->rcitem);

	// Finish completing the failed IRP. Note that our completion routine will
	// not be called this time, so we have to do all the cleanup here

	PIRP Irp = ctx->Irp;
	IoFreeMdl(ctx->mdl);
	ExFreePool(ctx);
	StartNextPacket(&pdx->dqReadWrite, fdo);
	IoReleaseRemoveLock(&pdx->RemoveLock, Irp);
	IoCompleteRequest(Irp, IO_NO_INCREMENT);	// resume completion of failed IRP
	}							// RecoverFromError

///////////////////////////////////////////////////////////////////////////////
// ResetDevice is called during an attempt to recover after an error in order to
// reset the device.

#pragma PAGEDCODE