devqueue.cpp

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

// DevQueue.cpp -- Custom IRP queuing support
// Copyright (C) 1999 by Walter Oney
// All rights reserved

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

typedef struct _NOTIFY_CONTEXT {
	PQNOTIFYFUNC notify;		// real notification function
	PVOID context;				// context arg for notification function
	LONG count;					// number of busy queues
	} NOTIFY_CONTEXT, *PNOTIFY_CONTEXT;

VOID NotificationCallback(PNOTIFY_CONTEXT ctx);

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

GENERICAPI VOID GENERIC_EXPORT AbortRequests(PDEVQUEUE pdq, NTSTATUS status)
	{							// AbortRequests
	pdq->abortstatus = status;
	CleanupRequests(pdq, NULL, status);
	}							// AbortRequests

GENERICAPI VOID GENERIC_EXPORT AbortAllRequests(PDEVQUEUE* q, ULONG nq, NTSTATUS status)
	{							// AbortAllRequests
	for (ULONG i = 0; i < nq; ++i)
		AbortRequests(q[i], status);
	}							// AbortAllRequests

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

GENERICAPI VOID GENERIC_EXPORT AllowRequests(PDEVQUEUE pdq)
	{							// AllowRequests
	pdq->abortstatus = STATUS_SUCCESS;
	}							// AllowRequests

GENERICAPI VOID GENERIC_EXPORT AllowAllRequests(PDEVQUEUE* q, ULONG nq)
	{							// AllowAllRequests
	for (ULONG i = 0; i < nq; ++i)
		AllowRequests(q[i]);
	}							// AllowAllRequests

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

GENERICAPI NTSTATUS GENERIC_EXPORT AreRequestsBeingAborted(PDEVQUEUE pdq)
	{							// AreRequestsBeingAborted
	return pdq->abortstatus;
	}							// AreRequestsBeingAborted

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

GENERICAPI VOID GENERIC_EXPORT CancelRequest(PDEVQUEUE pdq, PIRP Irp)
	{							// CancelRequest
	KIRQL oldirql = Irp->CancelIrql;

	// Release the global cancel spin lock as soon as possible

	IoReleaseCancelSpinLock(DISPATCH_LEVEL);

	// Acquire our queue-specific queue lock. Note that we stayed at DISPATCH_LEVEL
	// when we released the cancel spin lock

	KeAcquireSpinLockAtDpcLevel(&pdq->lock);

	// (After Peretz & Hanrahan) The IRP is guaranteed to be on *some* queue (maybe a degenerate one),
	// so we unconditionally remove it and complete it.

	RemoveEntryList(&Irp->Tail.Overlay.ListEntry);
	KeReleaseSpinLock(&pdq->lock, oldirql);

	Irp->IoStatus.Status = STATUS_CANCELLED;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	}							// CancelRequest

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

GENERICAPI BOOLEAN GENERIC_EXPORT CheckBusyAndStall(PDEVQUEUE pdq)
	{							// CheckBusyAndStall
	KIRQL oldirql;
	KeAcquireSpinLock(&pdq->lock, &oldirql);
	BOOLEAN busy = pdq->CurrentIrp != NULL;
	if (!busy)
		InterlockedIncrement(&pdq->stallcount);
	KeReleaseSpinLock(&pdq->lock, oldirql);
	return busy;
	}							// CheckBusyAndStall

GENERICAPI BOOLEAN GENERIC_EXPORT CheckAnyBusyAndStall(PDEVQUEUE* q, ULONG nq, PDEVICE_OBJECT fdo)
	{							// CheckAnyBusyAndStall
	ULONG i;

	// Call CheckBusyAndStall for each queue. If one of them is busy,
	// back out by unstalling the queues we stalled.

	for (i = 0; i < nq; ++i)
		if (CheckBusyAndStall(q[i]))
			{					// a queue is busy
			for (--i; (int) i >= 0; --i)
				RestartRequests(q[i], fdo);
			return TRUE;		// indicate at least one queue is busy
			}					// a queue is busy

	// Return FALSE because no queue was busy and all are now stalled

	return FALSE;
	}							// CheckAnyBusyAndStall

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

GENERICAPI VOID GENERIC_EXPORT CleanupRequests(PDEVQUEUE pdq, PFILE_OBJECT fop, NTSTATUS status)
	{							// CleanupRequests
	LIST_ENTRY cancellist;
	InitializeListHead(&cancellist);

	// Create a list of IRPs that belong to the same file object

	KIRQL oldirql;
	KeAcquireSpinLock(&pdq->lock, &oldirql);

	PLIST_ENTRY first = &pdq->head;
	PLIST_ENTRY next;

	for (next = first->Flink; next != first; )
		{						// for each queued IRP
		PIRP Irp = CONTAINING_RECORD(next, IRP, Tail.Overlay.ListEntry);
		PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(Irp);

		// Follow the chain to the next IRP now (so that the next iteration of
		// the loop is properly setup whether we dequeue this IRP or not)

		PLIST_ENTRY current = next;
		next = next->Flink;

		// Skip this IRP if it's not for the same file object as the
		// current IRP_MJ_CLEANUP.

		if (fop && stack->FileObject != fop)
			continue;			// not for same file object

		// (After Hanrahan) Set the CancelRoutine pointer to NULL. If it was
		// already NULL, someone is trying to cancel this IRP right now, so just
		// leave it on the queue and let them do it as soon as we release the spin lock.

		if (!IoSetCancelRoutine(Irp, NULL))
			continue;
		RemoveEntryList(current);
		InsertTailList(&cancellist, current);
		}						// for each queued IRP

	// Release the spin lock. We're about to undertake a potentially time-consuming
	// operation that might conceivably result in a deadlock if we keep the lock.

	KeReleaseSpinLock(&pdq->lock, oldirql);

	// Complete the selected requests.

	while (!IsListEmpty(&cancellist))
		{						// cancel selected requests
		next = RemoveHeadList(&cancellist);
		PIRP Irp = CONTAINING_RECORD(next, IRP, Tail.Overlay.ListEntry);
		Irp->IoStatus.Status = status;
		IoCompleteRequest(Irp, IO_NO_INCREMENT);
		}						// cancel selected requests
	}							// CleanupRequests


GENERICAPI VOID GENERIC_EXPORT CleanupAllRequests(PDEVQUEUE* q, ULONG nq, PFILE_OBJECT fop, NTSTATUS status)
	{							// CleanupAllRequests
	for (ULONG i = 0; i < nq; ++i)
		CleanupRequests(q[i], fop, status);
	}							// CleanupAllRequests

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

GENERICAPI PIRP GENERIC_EXPORT GetCurrentIrp(PDEVQUEUE pdq)
	{							// GetCurrentIrp
	return pdq->CurrentIrp;
	}							// GetCurrentIrp

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

GENERICAPI VOID GENERIC_EXPORT InitializeQueue(PDEVQUEUE pdq, PDRIVER_STARTIO StartIo)
	{							// InitializeQueue
	InitializeListHead(&pdq->head);
	KeInitializeSpinLock(&pdq->lock);
	pdq->StartIo = StartIo;
	pdq->stallcount = 1;
	pdq->CurrentIrp = NULL;
	KeInitializeEvent(&pdq->evStop, NotificationEvent, FALSE);
	pdq->abortstatus = (NTSTATUS) 0;
	pdq->notify = NULL;
	pdq->notifycontext = 0;
	}							// InitializeQueue

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

VOID NotificationCallback(PNOTIFY_CONTEXT ctx)
	{							// NotificationCallback
	if (InterlockedDecrement(&ctx->count) > 0)
		return;

	(*ctx->notify)(ctx->context);

	ExFreePool(ctx);
	}							// NotificationCallback

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

GENERICAPI VOID GENERIC_EXPORT RestartRequests(PDEVQUEUE pdq, PDEVICE_OBJECT fdo)
	{							// RestartRequests

	// The original version of this routine called StartNextPacket to restart the
	// queue. Reader Sink Ho pointed out a race condition, such that an intervening
	// call to StartPacket in another thread or on another CPU would cause StartNextPacket
	// to end up dequeuing a second IRP.

	KIRQL oldirql;
	KeAcquireSpinLock(&pdq->lock, &oldirql);

	if (InterlockedDecrement(&pdq->stallcount) > 0)
		{						// still stalled
		KeReleaseSpinLock(&pdq->lock, oldirql);
		return;
		}						// still stalled

	// Dequeue and start the IRP at the head of the list. See the comments in
	// StartNextPacket for an explanation of the cancel logic.

	while (!pdq->stallcount && !pdq->CurrentIrp && !pdq->abortstatus && !IsListEmpty(&pdq->head))
		{						// start first queued IRP
		PLIST_ENTRY next = RemoveHeadList(&pdq->head);
		PIRP Irp = CONTAINING_RECORD(next, IRP, Tail.Overlay.ListEntry);

		if (!IoSetCancelRoutine(Irp, NULL))
			{					// IRP being cancelled right now
			ASSERT(Irp->Cancel);	// else CancelRoutine shouldn't be NULL!
			InitializeListHead(&Irp->Tail.Overlay.ListEntry);
			continue;			// with "start first queued IRP"
			}					// IRP being cancelled right now

		pdq->CurrentIrp = Irp;
		KeReleaseSpinLockFromDpcLevel(&pdq->lock);
		(*pdq->StartIo)(fdo, Irp);
		KeLowerIrql(oldirql);
		return;
		}						// start first queued IRP

	// No IRPs need to be started (or else all queued IRPs were being cancelled)

	KeReleaseSpinLock(&pdq->lock, oldirql);
	}							// RestartRequests

GENERICAPI VOID GENERIC_EXPORT RestartAllRequests(PDEVQUEUE* q, ULONG nq, PDEVICE_OBJECT fdo)
	{							// RestartAllRequests
	for (ULONG i = 0; i < nq; ++i)
		RestartRequests(q[i], fdo);
	}							// RestartAllRequests

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

GENERICAPI VOID GENERIC_EXPORT StallRequests(PDEVQUEUE pdq)
	{							// StallRequests
	InterlockedIncrement(&pdq->stallcount);
	}							// StallRequests

GENERICAPI VOID GENERIC_EXPORT StallAllRequests(PDEVQUEUE* q, ULONG nq)
	{							// StallAllRequests
	for (ULONG i = 0; i < nq; ++i)
		StallRequests(q[i]);
	}							// StallAllRequests

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

GENERICAPI NTSTATUS GENERIC_EXPORT StallRequestsAndNotify(PDEVQUEUE pdq, PQNOTIFYFUNC notify, PVOID context)
	{							// StallRequestsAndNotify
	NTSTATUS status;