general/toaster/toastmon: stop iotarget/timer on d0exit

Bucket IDs:
 - 0x10d_5_vrf_wdftoastmon!toastmon_readrequestcompletionroutine
 - AV_VRF_Wdf01000!FxRequestBase::Vf_VerifyDispose
 - AV_VRF_wdftoastmon!Toastmon_EvtTimerPostRequests

WDF verifier breaks:
 - Driver is trying to delete WDFREQUEST while it is still active on WDFIOTARGET
 - Calling WdfTimerStart when the timer object is running down will lead to a crash
 - WDFREQUEST already sent to a target

3 files changed, 149 insertions, 63 deletions

diff --git a/general/toaster/toastDrv/kmdf/toastmon/toastmon.H b/general/toaster/toastDrv/kmdf/toastmon/toastmon.H
index 62a986a8..68b6fdce 100644
--- a/general/toaster/toastDrv/kmdf/toastmon/toastmon.H
+++ b/general/toaster/toastDrv/kmdf/toastmon/toastmon.H
@@ -44,14 +44,6 @@ typedef struct  _TARGET_DEVICE_INFO {
     WDFREQUEST        WriteRequest;
     WDFTIMER          TimerForPostingRequests;
 
-    //
-    // Set to TRUE while the target is opened. Will be set to FALSE at query remove (for a graceful remove)
-    // or removal complete (surprise removal of the target). Can be set back to TRUE if the graceful remove
-    // fails and the query remove is canceled.
-    //
-    // Guarded by DeviceExtension->TargetDeviceCollectionLock
-    //
-    BOOLEAN           Opened;
 } TARGET_DEVICE_INFO, *PTARGET_DEVICE_INFO;
 
 WDF_DECLARE_CONTEXT_TYPE_WITH_NAME(TARGET_DEVICE_INFO, GetTargetDeviceInfo)
@@ -71,6 +63,8 @@ WDF_DECLARE_CONTEXT_TYPE_WITH_NAME(TIMER_CONTEXT, GetTimerContext)
 DRIVER_INITIALIZE                       DriverEntry;
 EVT_WDF_DRIVER_DEVICE_ADD               ToastMon_EvtDeviceAdd;
 EVT_WDF_DEVICE_CONTEXT_CLEANUP          ToastMon_EvtDeviceContextCleanup;
+EVT_WDF_DEVICE_D0_ENTRY                 ToastMon_EvtDeviceD0Entry;
+EVT_WDF_DEVICE_D0_EXIT                  ToastMon_EvtDeviceD0Exit;
 EVT_WDF_IO_TARGET_QUERY_REMOVE          ToastMon_EvtIoTargetQueryRemove;
 EVT_WDF_IO_TARGET_REMOVE_CANCELED       ToastMon_EvtIoTargetRemoveCanceled;
 EVT_WDF_IO_TARGET_REMOVE_COMPLETE       ToastMon_EvtIoTargetRemoveComplete;
diff --git a/general/toaster/toastDrv/kmdf/toastmon/toastmon.c b/general/toaster/toastDrv/kmdf/toastmon/toastmon.c
index f140e9d7..c011c73e 100644
--- a/general/toaster/toastDrv/kmdf/toastmon/toastmon.c
+++ b/general/toaster/toastDrv/kmdf/toastmon/toastmon.c
@@ -38,6 +38,7 @@ Environment:
 #pragma alloc_text (INIT, DriverEntry)
 #pragma alloc_text (PAGE, ToastMon_EvtDeviceAdd)
 #pragma alloc_text (PAGE, ToastMon_EvtDeviceContextCleanup)
+#pragma alloc_text (PAGE, ToastMon_EvtDeviceD0Exit)
 #pragma alloc_text (PAGE, ToastMon_PnpNotifyInterfaceChange)
 #pragma alloc_text (PAGE, ToastMon_EvtIoTargetQueryRemove)
 #pragma alloc_text (PAGE, ToastMon_EvtIoTargetRemoveCanceled)
@@ -135,6 +136,7 @@ Return Value:
 {
     WDF_OBJECT_ATTRIBUTES           attributes;
     NTSTATUS                        status = STATUS_SUCCESS;
+    WDF_PNPPOWER_EVENT_CALLBACKS    pnpCallbacks;
     WDFDEVICE                       device;
     PDEVICE_EXTENSION               deviceExtension;
 
@@ -152,6 +154,12 @@ Return Value:
 
     attributes.EvtCleanupCallback = ToastMon_EvtDeviceContextCleanup;
 
+    WDF_PNPPOWER_EVENT_CALLBACKS_INIT(&pnpCallbacks);
+    pnpCallbacks.EvtDeviceD0Entry = ToastMon_EvtDeviceD0Entry;
+    pnpCallbacks.EvtDeviceD0Exit  = ToastMon_EvtDeviceD0Exit;
+
+    WdfDeviceInitSetPnpPowerEventCallbacks(DeviceInit, &pnpCallbacks);
+
     //
     // Create a framework device object.This call will inturn create
     // a WDM deviceobject, attach to the lower stack and set the
@@ -290,6 +298,105 @@ Return Value:
 
 _Use_decl_annotations_
 NTSTATUS
+ToastMon_EvtDeviceD0Entry(
+    WDFDEVICE               Device,
+    WDF_POWER_DEVICE_STATE  PreviousState
+    )
+/*++
+
+Routine Description:
+
+    EvtDeviceD0Entry performs operations that are needed when the driver's
+    device enters the D0 power state.
+
+    EvtDeviceD0Entry should not be marked as Pageable.
+
+--*/
+{
+    PDEVICE_EXTENSION   deviceExtension;
+    PTARGET_DEVICE_INFO targetDeviceInfo;
+    WDFIOTARGET         ioTarget;
+    ULONG               count, i;
+
+    UNREFERENCED_PARAMETER(PreviousState);
+
+    KdPrint( ("ToastMon_EvtDeviceD0Entry\n"));
+
+    deviceExtension = GetDeviceExtension((WDFDEVICE)Device);
+
+    //
+    // (Re-)Start IoTarget and timers
+    //
+
+    WdfWaitLockAcquire(deviceExtension->TargetDeviceCollectionLock, NULL);
+
+    count = WdfCollectionGetCount(deviceExtension->TargetDeviceCollection);
+    for (i = 0; i < count; i ++) {
+
+        ioTarget = WdfCollectionGetItem(deviceExtension->TargetDeviceCollection, i);
+        WdfIoTargetStart(ioTarget);
+
+        targetDeviceInfo = GetTargetDeviceInfo(ioTarget);
+        WdfTimerStart(targetDeviceInfo->TimerForPostingRequests,
+                                        WDF_REL_TIMEOUT_IN_MS(1));
+    }
+
+    WdfWaitLockRelease(deviceExtension->TargetDeviceCollectionLock);
+
+    return STATUS_SUCCESS;
+}
+
+_Use_decl_annotations_
+NTSTATUS
+ToastMon_EvtDeviceD0Exit(
+    WDFDEVICE               Device,
+    WDF_POWER_DEVICE_STATE  TargetState
+    )
+/*++
+
+Routine Description:
+
+    EvtDeviceD0Exit performs operations that are needed when the driver's
+    device leaves the D0 power state.
+
+--*/
+{
+    PDEVICE_EXTENSION   deviceExtension;
+    PTARGET_DEVICE_INFO targetDeviceInfo;
+    WDFIOTARGET         ioTarget;
+    ULONG               count, i;
+
+    UNREFERENCED_PARAMETER(TargetState);
+
+    PAGED_CODE();
+
+    KdPrint( ("ToastMon_EvtDeviceD0Exit\n"));
+
+    deviceExtension = GetDeviceExtension((WDFDEVICE)Device);
+
+    //
+    // Stop IoTarget and timers
+    //
+
+    WdfWaitLockAcquire(deviceExtension->TargetDeviceCollectionLock, NULL);
+
+    count = WdfCollectionGetCount(deviceExtension->TargetDeviceCollection);
+    for (i = 0; i < count; i ++) {
+
+        ioTarget = WdfCollectionGetItem(deviceExtension->TargetDeviceCollection, i);
+        WdfIoTargetStop(ioTarget, WdfIoTargetWaitForSentIoToComplete);
+
+        targetDeviceInfo = GetTargetDeviceInfo(ioTarget);
+        WdfTimerStop(targetDeviceInfo->TimerForPostingRequests, TRUE);
+    }
+
+    WdfWaitLockRelease(deviceExtension->TargetDeviceCollectionLock);
+
+    return STATUS_SUCCESS;
+}
+
+_Use_decl_annotations_
+NTSTATUS
 ToastMon_PnpNotifyInterfaceChange(
     PVOID NotificationStruct,
     PVOID Context
@@ -567,13 +674,6 @@ Return Value:
     WdfTimerStart(targetDeviceInfo->TimerForPostingRequests,
                                         WDF_REL_TIMEOUT_IN_MS(1));
 
-    //
-    // Can be set outside of holding TargetDeviceCollectionLock because we are
-    // creating the target in a notification callback and all subsequent state changes notifications
-    // can only happen after the arrival notification returns back to the kernel
-    //
-    targetDeviceInfo->Opened = TRUE;
-
     *Target = ioTarget;
 
     return status;
@@ -606,7 +706,6 @@ Return Value:
 --*/
 {
     PTARGET_DEVICE_INFO         targetDeviceInfo;
-    WDFWAITLOCK                 targetDeviceCollectionLock;
 
     PAGED_CODE();
 
@@ -620,15 +719,6 @@ Return Value:
 
     WdfTimerStop(targetDeviceInfo->TimerForPostingRequests, TRUE);
 
-    targetDeviceCollectionLock = targetDeviceInfo->DeviceExtension->TargetDeviceCollectionLock;
-
-    //
-    // The target is being query removed, set Opened to FALSE to match this state change.
-    //
-    WdfWaitLockAcquire(targetDeviceCollectionLock, NULL);
-    targetDeviceInfo->Opened = FALSE;
-    WdfWaitLockRelease(targetDeviceCollectionLock);
-
     WdfIoTargetCloseForQueryRemove(IoTarget);
 
     return STATUS_SUCCESS;
@@ -657,7 +747,6 @@ Return Value:
 --*/
 {
     PTARGET_DEVICE_INFO         targetDeviceInfo;
-    WDFWAITLOCK                 targetDeviceCollectionLock;
     WDF_IO_TARGET_OPEN_PARAMS   openParams;
     NTSTATUS                    status;
 
@@ -680,16 +769,6 @@ Return Value:
         return;
     }
 
-    targetDeviceCollectionLock = targetDeviceInfo->DeviceExtension->TargetDeviceCollectionLock;
-
-    //
-    // The query remove has failed and the target has been successfully reopened. Set Opened
-    // back to TRUE to reflect the state change.
-    //
-    WdfWaitLockAcquire(targetDeviceCollectionLock, NULL);
-    targetDeviceInfo->Opened = TRUE;
-    WdfWaitLockRelease(targetDeviceCollectionLock);
-
     //
     // Restart the timer.
     //
@@ -733,14 +812,11 @@ Return Value:
     WdfTimerStop(targetDeviceInfo->TimerForPostingRequests, TRUE);
 
     //
-    // Remove the target device from the collection and set Opened to FALSE to match
-    // the state change (in the case of a surprise removal of the target,
-    // ToastMon_EvtIoTargetQueryRemove is not called so Opened is still TRUE).
+    // Remove the target device from the collection
     //
     WdfWaitLockAcquire(deviceExtension->TargetDeviceCollectionLock, NULL);
 
     WdfCollectionRemove(deviceExtension->TargetDeviceCollection, IoTarget);
-    targetDeviceInfo->Opened = FALSE;
 
     WdfWaitLockRelease(deviceExtension->TargetDeviceCollectionLock);
 
@@ -827,7 +903,17 @@ Return Value:
 
     targetInfo = GetTargetDeviceInfo(IoTarget);
 
-    request = targetInfo->ReadRequest;
+    //
+    // Clear the ReadRequest field in the context to avoid
+    // being reposted even before the reqeuest completes.
+    // This will be reset in the complete routine when the request completes.
+    //
+    request = InterlockedExchangePointer(&targetInfo->ReadRequest, NULL);
+    if (request == NULL) {
+        status = STATUS_INVALID_DEVICE_REQUEST;
+        KdPrint(("Request is already posted\n", status));
+        goto exit;
+    }
 
     //
     // Allocate memory for read. Ideally I should have allocated the memory upfront along
@@ -848,7 +934,7 @@ Return Value:
 
     if (!NT_SUCCESS(status)) {
         KdPrint(("WdfMemoryCreate failed 0x%x\n", status));
-        return status;
+        goto exit;
     }
 
     status = WdfIoTargetFormatRequestForRead(
@@ -859,26 +945,23 @@ Return Value:
                                 NULL); // OutputBufferOffset
    if (!NT_SUCCESS(status)) {
         KdPrint(("WdfIoTargetFormatRequestForRead failed 0x%x\n", status));
-        return status;
+        goto exit;
     }
 
     WdfRequestSetCompletionRoutine(request,
                    Toastmon_ReadRequestCompletionRoutine,
                    targetInfo);
 
-    //
-    // Clear the ReadRequest field in the context to avoid
-    // being reposted even before the reqeuest completes.
-    // This will be reset in the complete routine when the request completes.
-    //
-    targetInfo->ReadRequest = NULL;
-
     if(WdfRequestSend(request, IoTarget, WDF_NO_SEND_OPTIONS) == FALSE) {
         status = WdfRequestGetStatus(request);
         KdPrint(("WdfRequestSend failed 0x%x\n", status));
-        targetInfo->ReadRequest = request;
+        goto exit;
     }
 
+exit:
+    if (!NT_SUCCESS(status)) {
+        targetInfo->ReadRequest = request;
+    }
     return status;
 }
 
@@ -907,7 +990,17 @@ Return Value:
 
     targetInfo = GetTargetDeviceInfo(IoTarget);
 
-    request = targetInfo->WriteRequest;
+    //
+    // Clear the Write field in the context to avoid
+    // being reposted even before the reqeuest completes.
+    // This will be reset in the complete routine when the request completes.
+    //
+    request = InterlockedExchangePointer(&targetInfo->WriteRequest, NULL);
+    if (request == NULL) {
+        status = STATUS_INVALID_DEVICE_REQUEST;
+        KdPrint(("Request is already posted\n", status));
+        goto exit;
+    }
 
     //
     // Allocate memory for write. Ideally I should have allocated the memory upfront along
@@ -928,7 +1021,7 @@ Return Value:
 
     if (!NT_SUCCESS(status)) {
         KdPrint(("WdfMemoryCreate failed 0x%x\n", status));
-        return status;
+        goto exit;
     }
 
     status = WdfIoTargetFormatRequestForWrite(
@@ -939,26 +1032,23 @@ Return Value:
                                 NULL); // OutputBufferOffset
    if (!NT_SUCCESS(status)) {
         KdPrint(("WdfIoTargetFormatRequestForWrite failed 0x%x\n", status));
-        return status;
+        goto exit;
     }
 
     WdfRequestSetCompletionRoutine(request,
                    Toastmon_WriteRequestCompletionRoutine,
                    targetInfo);
 
-    //
-    // Clear the WriteRequest field in the context to avoid
-    // being reposted even before the reqeuest completes.
-    // This will be reset in the complete routine when the request completes.
-    //
-    targetInfo->WriteRequest = NULL;
-
     if(WdfRequestSend(request, IoTarget, WDF_NO_SEND_OPTIONS) == FALSE) {
         status = WdfRequestGetStatus(request);
         KdPrint(("WdfRequestSend failed 0x%x\n", status));
-        targetInfo->WriteRequest = request;
+        goto exit;
     }
 
+exit:
+    if (!NT_SUCCESS(status)) {
+        targetInfo->WriteRequest = request;
+    }
     return status;
 }
 
diff --git a/general/toaster/toastDrv/kmdf/toastmon/wmi.c b/general/toaster/toastDrv/kmdf/toastmon/wmi.c
index 8d8b4406..eb0e8e70 100644
--- a/general/toaster/toastDrv/kmdf/toastmon/wmi.c
+++ b/general/toaster/toastDrv/kmdf/toastmon/wmi.c
@@ -178,6 +178,7 @@ Return Value:
     PDEVICE_EXTENSION       deviceExt = Context;
     WDFCOLLECTION           hCollection = deviceExt->TargetDeviceCollection;
     WDFIOTARGET             ioTarget;
+    WDF_IO_TARGET_STATE     ioTargetState;
     ULONG                   i;
 
     PAGED_CODE();
@@ -193,7 +194,8 @@ Return Value:
         // The WdfIoTargetWdmGetXxxDeviceObject APIs can only be called while the target is opened, otherwise
         // they can return undefined values.
         //
-        if (GetTargetDeviceInfo(ioTarget)->Opened == FALSE) {
+        ioTargetState = WdfIoTargetGetState(ioTarget);
+        if (ioTargetState != WdfIoTargetStarted) {
             KdPrint(("WDFIOTARGET %p not in an opened state.\n", ioTarget));
             continue;
         }