queue.c

该源码是用DDK编写的WDM驱动程序

// queue.c
//
// Generated by C DriverWizard 3.1.0 (Build 1722)
// Requires DDK Only
// File created on 7/11/2007
//

#include "pch.h"
#ifdef CHARSAMPLE_DDK_WMI_TRACE
#include "queue.tmh"
#endif

///////////////////////////////////////////////////////////////////////////////////////////////////
// Cancel Safe Driver Managed IRP Queue
///////////////////////////////////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CharSample_DDKInitializeQueue 
//      Sets up a driver managed IRP queue.
//
//  Arguments:
//      IN  Queue
//              An instance of our queue structure
//
//      IN  StartIoRoutine
//              Routine where queue IRPs are sent to be processed
//
//      IN  DeviceObject
//              Device object for our driver
//
//      IN  bUseCharSample_DDKStartIoDpc
//              Flag to indicate that the queue should use queue a DPC for 
//              calling StartIo if StartIo recursion is possible
//
//  Return Value:
//      none
//
VOID CharSample_DDKInitializeQueue(
    IN  PCHARSAMPLE_DDK_QUEUE          Queue,
    IN  PCHARSAMPLE_DDK_QUEUE_STARTIO     StartIoRoutine,
    IN  PDEVICE_OBJECT          DeviceObject,
    IN  BOOLEAN                 bUseCharSample_DDKStartIoDpc
    )
{
    // must provide StartIo routine
    ASSERT(StartIoRoutine != NULL);

    // save off the user info
    Queue->StartIoRoutine = StartIoRoutine;
    Queue->DeviceObject = DeviceObject;
    Queue->bUseCharSample_DDKStartIoDpc = bUseCharSample_DDKStartIoDpc;

    // queues are created in a stalled state
    // Start device will unstall them
    Queue->StallCount = 1;

    // initialize our queue lock
    KeInitializeSpinLock(&Queue->QueueLock);

    // initialize our IRP list
    InitializeListHead(&Queue->IrpQueue);

    // initialize our CharSample_DDKStartIoDpc
    if (bUseCharSample_DDKStartIoDpc)
    {
        KeInitializeDpc(&Queue->CharSample_DDKStartIoDpc, CharSample_DDKStartIoDpc, Queue);
    }    

    // initialize stop event
    KeInitializeEvent(&Queue->StopEvent, NotificationEvent, FALSE);

    Queue->ErrorStatus = STATUS_SUCCESS;

    return;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CharSample_DDKQueueIrp
//      Inserts an IRP into the queue if the queue is busy, or sends IRP
//      to StartIo routine if the queue is not busy.
//
//  Arguments:
//      IN  Queue
//              An instance of our queue structure
//
//      IN  Irp
//              The IRP to add to the queue
//
//  Return Value:
//      Status
//
NTSTATUS CharSample_DDKQueueIrp(
    IN  PCHARSAMPLE_DDK_QUEUE    Queue,
    IN  PIRP            Irp
    )
{
    NTSTATUS            status;
    KIRQL               oldIrql;
    PCHARSAMPLE_DDK_DEVICE_EXTENSION   deviceExtension;

    deviceExtension = (PCHARSAMPLE_DDK_DEVICE_EXTENSION)Queue->DeviceObject->DeviceExtension;

    // grab the queue protection
    KeAcquireSpinLock(&Queue->QueueLock, &oldIrql);

    // If the queue has been invalidated, complete the IRP
    if (Queue->ErrorStatus != STATUS_SUCCESS)
    {
        status = Queue->ErrorStatus;

        // drop the queue protection
        KeReleaseSpinLock(&Queue->QueueLock, oldIrql);

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

        IoCompleteRequest(Irp, IO_NO_INCREMENT);

        return status;
    }

    // see if the queue is busy
    if ((Queue->CurrentIrp == NULL) && (Queue->StallCount == 0))
    {
        // mark irp pending since STATUS_PENDING is returned
        IoMarkIrpPending(Irp);
        status = STATUS_PENDING;

        // set the current IRP
        Queue->CurrentIrp = Irp;

        // drop the queue protection
        // raise our IRQL before calling StartIo
        KeReleaseSpinLockFromDpcLevel(&Queue->QueueLock);

        // call the user's StartIo routine
        Queue->StartIoRoutine(Queue->DeviceObject, Queue->CurrentIrp);

        // drop our IRQL back
        KeLowerIrql(oldIrql);

        return status;
    }

    // put our queue pointer into the IRP
    Irp->Tail.Overlay.DriverContext[0] = Queue;

    // queue the IRP
    InsertTailList(&Queue->IrpQueue, &Irp->Tail.Overlay.ListEntry);

    // insert our queue cancel routine into the IRP
    IoSetCancelRoutine(Irp, CharSample_DDKQueueCancelRoutine);

    // Make sure the IRP was not cancelled before 
    // we inserted our cancel routine
    if (Irp->Cancel)
    {
        // If the IRP was cancelled after we put in our cancel routine we
        // will get back NULL here and we will let the cancel routine handle
        // the IRP.  If NULL is not returned here then we know the IRP was 
        // cancelled before we inserted the queue cancel routine, and we
        // need to call our cancel routine to handle the IRP.
        if (IoSetCancelRoutine(Irp, NULL) != NULL)
        {
            RemoveEntryList(&Irp->Tail.Overlay.ListEntry);

            // drop the queue protection
            KeReleaseSpinLock(&Queue->QueueLock, oldIrql);

            // cancel the IRP
            status = STATUS_CANCELLED;
            Irp->IoStatus.Status = status;
            Irp->IoStatus.Information = 0;
            IoCompleteRequest(Irp, IO_NO_INCREMENT);

            return status;
        }
    }

    // mark irp pending since STATUS_PENDING is returned
    IoMarkIrpPending(Irp);
    status = STATUS_PENDING;

    // drop the queue protection
    KeReleaseSpinLock(&Queue->QueueLock, oldIrql);

    return status;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CharSample_DDKStartNext
//      Pulls the next available IRP from the queue and sends it to 
//      StartIo for processing.
//
//  Arguments:
//      IN  Queue
//              An instance of our queue structure
//
//  Return Value:
//      none
//
VOID CharSample_DDKStartNext(
    IN  PCHARSAMPLE_DDK_QUEUE    Queue
    )
{
    KIRQL       oldIrql;
    PLIST_ENTRY entry;

    // grab the queue protection
    KeAcquireSpinLock(&Queue->QueueLock, &oldIrql);

    // set the current IRP pointer to NULL
    Queue->CurrentIrp = NULL;

    // check if there are entries in the queue
    if (IsListEmpty(&Queue->IrpQueue) || (Queue->StallCount > 0) || (Queue->ErrorStatus != STATUS_SUCCESS))
    {
        // set event that queue is stalled
        KeSetEvent(&Queue->StopEvent, IO_NO_INCREMENT, FALSE);
        KeReleaseSpinLock(&Queue->QueueLock, oldIrql);

        return;
    }

    // get a new IRP from the queue
    for (entry = Queue->IrpQueue.Flink; entry != &Queue->IrpQueue; entry = entry->Flink)
    {
        // get the IRP from the list entry
        Queue->CurrentIrp = CONTAINING_RECORD(entry, IRP, Tail.Overlay.ListEntry);

        ASSERT(Queue->CurrentIrp->Type == IO_TYPE_IRP);

        // if we found an IRP, pull the queue cancel routine out of it
        // See if the IRP is canceled
        if (IoSetCancelRoutine(Queue->CurrentIrp, NULL) == NULL)
        {
            // cancel routine already has a hold on this IRP, 
            // just go on to the next one
            Queue->CurrentIrp = NULL;
        }
        else
        {
            // Found a usable IRP, pull the entry out of the list
            RemoveEntryList(entry);
            break;
        }
    }

    if (Queue->CurrentIrp == NULL)
    {
        // set event that queue is stalled
        KeSetEvent(&Queue->StopEvent, IO_NO_INCREMENT, FALSE);
        KeReleaseSpinLock(&Queue->QueueLock, oldIrql);
        return;
    }

    // found an IRP call StartIo

    // Determine if we need to queue a DPC or not.
    // We only use the DPC if the user specified to protect 
    // against StartIo recursion (bUseCharSample_DDKStartIoDpc), and the 
    // queue has multiple entries. If there are not multiple 
    // entries in the queue we won't recurse anyway.
    if (Queue->bUseCharSample_DDKStartIoDpc && !IsListEmpty(&Queue->IrpQueue))
    {
        // drop the queue protection
        KeReleaseSpinLock(&Queue->QueueLock, oldIrql);

        // queue our StartIo DPC to prevent StartIo recursion
        KeInsertQueueDpc(&Queue->CharSample_DDKStartIoDpc, NULL, NULL);
    }
    else
    {
        // drop the queue protection
        // raise our IRQL before calling StartIo
        KeReleaseSpinLockFromDpcLevel(&Queue->QueueLock);

        // call the user's StartIo routine
        Queue->StartIoRoutine(Queue->DeviceObject, Queue->CurrentIrp);

        // drop our IRQL back
        KeLowerIrql(oldIrql);
    }

    return;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CharSample_DDKFlushQueue
//      Cancels all IRPs in the queue, or all IRPs in the queue related to a 
//      particular open handle.
//
//  Arguments:
//      IN  Queue
//              An instance of our queue structure
//
//      IN  FileObject
//              If NULL all IRPs in queue are cancelled, if non-NULL then all
//              IRPs related to this file object are cancelled.
//
//  Return Value:
//      none
//
VOID CharSample_DDKFlushQueue(
    IN  PCHARSAMPLE_DDK_QUEUE    Queue,
    IN  PFILE_OBJECT        FileObject
    )
{
    PLIST_ENTRY         entry;
    PLIST_ENTRY         nextEntry;
    PIRP                irp;
    PIO_STACK_LOCATION  irpStack;
    KIRQL               oldIrql;
    LIST_ENTRY          cancelList;

    // initialize our cancel list
    InitializeListHead(&cancelList);

    // grab the queue protection
    KeAcquireSpinLock(&Queue->QueueLock, &oldIrql);

    // Look at the first entry in the queue
    entry = Queue->IrpQueue.Flink;
    while (entry != &Queue->IrpQueue)
    {
        // get the next list entry
        nextEntry = entry->Flink;

        // Get the IRP from the entry
        irp = CONTAINING_RECORD(entry, IRP, Tail.Overlay.ListEntry);

        ASSERT(irp->Type == IO_TYPE_IRP);

        // Determine if we need to pull out of queue
        if (FileObject != NULL)
        {
            // get the current IRP stack location from the IRP
            irpStack = IoGetCurrentIrpStackLocation(irp);

            if (irpStack->FileObject != FileObject)
            {
                // go to the next entry
                entry = nextEntry;

                // We are not flushing this IRP
                continue;
            }
        }

        // Attempt to cancel the IRP
        if (IoSetCancelRoutine(irp, NULL) == NULL)
        {
            // go to the next entry
            entry = nextEntry;

            // cancel routine already has this IRP,
            // just go on
            continue;
        }

        // pull the IRP from the queue
        RemoveEntryList(entry);

        InsertTailList(&cancelList, entry);

        // go to the next entry
        entry = nextEntry;
    }

    // drop the queue protection
    KeReleaseSpinLock(&Queue->QueueLock, oldIrql);

    // Now clear out our cancel list
    while (!IsListEmpty(&cancelList))
    {
        // Get the first entry on the list
        entry = RemoveHeadList(&cancelList);

        // Get the IRP for that entry
        irp = CONTAINING_RECORD(entry, IRP, Tail.Overlay.ListEntry);

        // Cancel the IRP
        irp->IoStatus.Status = STATUS_CANCELLED;
        irp->IoStatus.Information = 0;
        IoCompleteRequest(irp, IO_NO_INCREMENT);
    }

    return;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CharSample_DDKInvalidateQueue
//      Stops queue from receiving anymore IRPs, all IRPs are completed upon
//      receipt
//
//  Arguments:
//      IN  Queue
//              An instance of our queue structure
//
//  Return Value:
//      none
//
VOID CharSample_DDKInvalidateQueue(
    IN  PCHARSAMPLE_DDK_QUEUE    Queue,
    IN  NTSTATUS                  ErrorStatus   
    )
{
    // indicate the queue is shutdown
    Queue->ErrorStatus = TRUE;

    // flush all requests from the queue
    CharSample_DDKFlushQueue(Queue, NULL);

    return;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CharSample_DDKPauseQueue
//      Stops queue from sending anymore IRPs to StartIo to be processed
//
//  Arguments:
//      IN  Queue
//              An instance of our queue structure
//
//  Return Value:
//      none
//
VOID CharSample_DDKPauseQueue(
    IN  PCHARSAMPLE_DDK_QUEUE    Queue
    )
{
    KIRQL   oldIrql;
    BOOLEAN bBusy;

    KeAcquireSpinLock(&Queue->QueueLock, &oldIrql);

    // indicate the queue is paused
    InterlockedIncrement(&Queue->StallCount);

    bBusy = Queue->CurrentIrp != NULL;
    if (bBusy)
    {
        // reset stop event
        KeClearEvent(&Queue->StopEvent);
    }

    KeReleaseSpinLock(&Queue->QueueLock, oldIrql);

    if (bBusy)
    {
        KeWaitForSingleObject(&Queue->StopEvent, Executive, KernelMode, FALSE, NULL);
    }

    return;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CharSample_DDKRestartQueue
//      Restarts IRP processing for a queue that was paused 
//      using CharSample_DDKPauseQueue()
//      or invalidated using CharSample_DDKInvalidateQueue()
//
//  Arguments:
//      IN  Queue
//              An instance of our queue structure
//
//  Return Value:
//      none
//
VOID CharSample_DDKRestartQueue(
    IN  PCHARSAMPLE_DDK_QUEUE    Queue
    )
{
    // if the queue is stalled or invalid, restart it
    if (InterlockedDecrement(&Queue->StallCount) == 0)
    {
        CharSample_DDKStartNext(Queue);
    }

    return;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CharSample_DDKStartIoDpc
//      DPC routine used to call StartIo, if bUseCharSample_DDKStartIoDpc is specified and
//      StartIo recursion is possible.
//
//  Arguments:
//      IN  Dpc
//              DPC object
//
//      IN  Context
//              An instance of our queue structure
//
//      IN  Unused1
//              Not used
//
//      IN  Unused2
//              Not used
//
//  Return Value:
//      none
//
VOID CharSample_DDKStartIoDpc(
    IN  PKDPC   Dpc,
    IN  PVOID   Context,
    IN  PVOID   Unused1,
    IN  PVOID   Unused2
    )
{
    PCHARSAMPLE_DDK_QUEUE  queue;

    queue = (PCHARSAMPLE_DDK_QUEUE)Context;

    // call the user's StartIo routine
    queue->StartIoRoutine(queue->DeviceObject, queue->CurrentIrp);

    return;
}

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