Fix potential deadlock due to unregistering callouts during a PowerState callback (#408)79745

Moves the unregister of callouts into a workitem routine. Additionally fixed an issue which could potentially leak an IOWorkItem

5 files changed, 210 insertions, 46 deletions

diff --git a/network/trans/WFPSampler/sys/Framework_Events.cpp b/network/trans/WFPSampler/sys/Framework_Events.cpp
index cd2290f6..baf59cc8 100644
--- a/network/trans/WFPSampler/sys/Framework_Events.cpp
+++ b/network/trans/WFPSampler/sys/Framework_Events.cpp
@@ -104,6 +104,7 @@ VOID EventCleanupDriverObject(_In_ WDFOBJECT driverObject)
    Notes:                                                                                       <br>
                                                                                                 <br>
    MSDN_Ref: HTTP://MSDN.Microsoft.com/En-US/Library/Windows/Hardware/FF540840.aspx             <br>
+             HTTP://MSDN.Microsoft.com/En-US/Library/Windows/Hardware/FF549133.aspx             <br>
 */
 _IRQL_requires_min_(PASSIVE_LEVEL)
 _IRQL_requires_max_(PASSIVE_LEVEL)
@@ -159,11 +160,25 @@ VOID EventCleanupDeviceObject(_In_ WDFOBJECT deviceObject)
 
    FwpmBfeStateUnsubscribeChanges(g_bfeSubscriptionHandle);
 
-   UnregisterPowerStateChangeCallback(&g_deviceExtension);
-
    if(g_pNDISPoolData)
       KrnlHlprNDISPoolDataDestroy(&g_pNDISPoolData);
 
+   UnregisterPowerStateChangeCallback(&g_deviceExtension);
+
+   if(g_pPowerStateExitIOWorkItem)
+   {
+      IoFreeWorkItem(g_pPowerStateExitIOWorkItem);
+
+      g_pPowerStateExitIOWorkItem = 0;
+   }
+
+   if(g_pPowerStateEnterIOWorkItem)
+   {
+      IoFreeWorkItem(g_pPowerStateEnterIOWorkItem);
+
+      g_pPowerStateEnterIOWorkItem = 0;
+   }
+
 #if DBG
    
    DbgPrintEx(DPFLTR_IHVNETWORK_ID,
diff --git a/network/trans/WFPSampler/sys/Framework_PowerStates.cpp b/network/trans/WFPSampler/sys/Framework_PowerStates.cpp
index ed67b11f..c615f8a1 100644
--- a/network/trans/WFPSampler/sys/Framework_PowerStates.cpp
+++ b/network/trans/WFPSampler/sys/Framework_PowerStates.cpp
@@ -20,6 +20,118 @@
 #include "Framework_WFPSamplerCalloutDriver.h" /// .
 #include "Framework_PowerStates.tmh"           /// $(OBJ_PATH)\$(O)\
 
+KGUARDED_MUTEX guardedMutex = {0};
+
+/**
+ @framework_function="ActOnPowerStateEnter"
+
+   Purpose:  Passive function to handle entering a power managed state.  Prior to dropping into a power managed   <br>
+             state, unregister the callouts.  This prevents the driver from potentially taking any more           <br>
+             references on NBLs, and instead causes the filters to return BLOCK.  This should also allow adequate <br>
+             time to drain currently queued NBLs.                                                                 <br>
+                                                                                                                  <br>
+   Notes:                                                                                                         <br>    
+                                                                                                                  <br>
+   MSDN_Ref: https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/content/wdm/nc-wdm-io_workitem_routine <br>
+*/
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_
+_Function_class_(IO_WORKITEM_ROUTINE)
+VOID ActOnPowerStateEnter(_In_ PDEVICE_OBJECT pDeviceObject,
+                          _Inout_opt_ PVOID pContext)
+{
+#if DBG
+
+   DbgPrintEx(DPFLTR_IHVNETWORK_ID,
+              DPFLTR_INFO_LEVEL,
+              " ---> ActOnPowerStateEnter()\n");
+
+#endif /// DBG
+
+   UNREFERENCED_PARAMETER(pDeviceObject);
+   UNREFERENCED_PARAMETER(pContext);
+
+   KeAcquireGuardedMutex(&guardedMutex);
+
+   if(g_calloutsRegistered == TRUE)
+   {
+      NTSTATUS status = STATUS_SUCCESS;
+
+      status = KrnlHlprExposedCalloutsUnregister();
+      HLPR_BAIL_ON_FAILURE(status);
+
+      g_calloutsRegistered = FALSE;
+   }
+
+   HLPR_BAIL_LABEL:
+
+   KeReleaseGuardedMutex(&guardedMutex);
+
+#if DBG
+
+   DbgPrintEx(DPFLTR_IHVNETWORK_ID,
+              DPFLTR_INFO_LEVEL,
+              " <--- ActOnPowerStateEnter()\n");
+
+#endif /// DBG
+
+   return;
+}
+
+/**
+ @framework_function="ActOnPowerStateExit"
+
+   Purpose:  Passive function to handle exiting a power managed state.  When coming out of a power managed state, <br>
+             register the callouts so normal processing of NBLs will commence.                                    <br>
+                                                                                                                  <br>
+   Notes:                                                                                                         <br>    
+                                                                                                                  <br>
+   MSDN_Ref: https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/content/wdm/nc-wdm-io_workitem_routine <br>
+*/
+_IRQL_requires_(PASSIVE_LEVEL)
+_IRQL_requires_same_
+_Function_class_(IO_WORKITEM_ROUTINE)
+VOID ActOnPowerStateExit(_In_ PDEVICE_OBJECT pDeviceObject,
+                         _Inout_opt_ PVOID pContext)
+{
+#if DBG
+
+   DbgPrintEx(DPFLTR_IHVNETWORK_ID,
+              DPFLTR_INFO_LEVEL,
+              " ---> ActOnPowerStateExit()\n");
+
+#endif /// DBG
+
+   UNREFERENCED_PARAMETER(pDeviceObject);
+   UNREFERENCED_PARAMETER(pContext);
+
+   KeAcquireGuardedMutex(&guardedMutex);
+
+   if(g_calloutsRegistered == FALSE)
+   {
+      NTSTATUS status = STATUS_SUCCESS;
+
+      status = KrnlHlprExposedCalloutsRegister();
+      HLPR_BAIL_ON_FAILURE(status);
+
+      g_calloutsRegistered = TRUE;
+   }
+
+   HLPR_BAIL_LABEL:
+
+   KeReleaseGuardedMutex(&guardedMutex);
+
+#if DBG
+
+   DbgPrintEx(DPFLTR_IHVNETWORK_ID,
+              DPFLTR_INFO_LEVEL,
+              " <--- ActOnPowerStateExit()\n");
+
+#endif /// DBG
+
+   return;
+}
+
 /**
  @framework_function="PowerStateCallback"
  
@@ -47,36 +159,23 @@ VOID PowerStateCallback(_In_ VOID* pCallbackContext,
 
    if(pPowerStateEvent == (VOID*)PO_CB_SYSTEM_STATE_LOCK)
    {
-      NTSTATUS status = STATUS_SUCCESS;
-
       if(pEventSpecifics)
       {
          /// entering the ON state (S0), so return operation to normal
-         if(g_calloutsRegistered == FALSE)
-         {
-            status = KrnlHlprExposedCalloutsRegister();
-            HLPR_BAIL_ON_FAILURE(status);
-
-            g_calloutsRegistered = TRUE;
-         }
+         IoQueueWorkItem(g_pPowerStateExitIOWorkItem,
+                         ActOnPowerStateExit,
+                         DelayedWorkQueue,
+                         0);
       }
       else
       {
-         /// leaving the ON state (S0) to sleep (S1/S2/S3) or hibernate (S4)
-         /// Unregister the callouts so no more injection will take place.  By default, the filters 
-         /// invoking the callouts will return block.
-         if(g_calloutsRegistered == TRUE)
-         {
-            status = KrnlHlprExposedCalloutsUnregister();
-            HLPR_BAIL_ON_FAILURE(status);
-
-            g_calloutsRegistered = FALSE;
-         }
+         IoQueueWorkItem(g_pPowerStateEnterIOWorkItem,
+                         ActOnPowerStateEnter,
+                         DelayedWorkQueue,
+                         0);
       }
    }
 
-   HLPR_BAIL_LABEL:
-
 #if DBG
    
    DbgPrintEx(DPFLTR_IHVNETWORK_ID,
@@ -89,7 +188,7 @@ VOID PowerStateCallback(_In_ VOID* pCallbackContext,
 }
 
 /**
- @framework_function="RegisterPowerStateChangeCallback"
+ @framework_function="UnregisterPowerStateChangeCallback"
  
    Purpose:  Unregister a callback that handled notifications of power state changes.           <br>
                                                                                                 <br>
@@ -168,6 +267,8 @@ NTSTATUS RegisterPowerStateChangeCallback(_Inout_ DEVICE_EXTENSION* pDeviceExten
    OBJECT_ATTRIBUTES objectAttributes = {0};
    UNICODE_STRING    unicodeString    = {0};
 
+   KeInitializeGuardedMutex(&guardedMutex);
+
    RtlInitUnicodeString(&unicodeString,
                         L"\\Callback\\PowerState");
 
diff --git a/network/trans/WFPSampler/sys/Framework_WFPSamplerCalloutDriver.cpp b/network/trans/WFPSampler/sys/Framework_WFPSamplerCalloutDriver.cpp
index 88824779..88369c81 100644
--- a/network/trans/WFPSampler/sys/Framework_WFPSamplerCalloutDriver.cpp
+++ b/network/trans/WFPSampler/sys/Framework_WFPSamplerCalloutDriver.cpp
@@ -32,13 +32,15 @@ extern "C"
    DRIVER_INITIALIZE DriverEntry;
 };
 
-PDEVICE_OBJECT         g_pWDMDevice            = 0;
-NDIS_POOL_DATA*        g_pNDISPoolData         = 0;
-BOOLEAN                g_calloutsRegistered    = FALSE;
-HANDLE                 g_bfeSubscriptionHandle = 0;
-SERIALIZATION_LIST     g_bsiSerializationList  = {0};
-WFPSAMPLER_DEVICE_DATA g_WFPSamplerDeviceData  = {0};
-DEVICE_EXTENSION       g_deviceExtension       = {0};
+PDEVICE_OBJECT         g_pWDMDevice                 = 0;
+PIO_WORKITEM           g_pPowerStateEnterIOWorkItem = 0;
+PIO_WORKITEM           g_pPowerStateExitIOWorkItem  = 0;
+NDIS_POOL_DATA*        g_pNDISPoolData              = 0;
+BOOLEAN                g_calloutsRegistered         = FALSE;
+HANDLE                 g_bfeSubscriptionHandle      = 0;
+SERIALIZATION_LIST     g_bsiSerializationList       = {0};
+WFPSAMPLER_DEVICE_DATA g_WFPSamplerDeviceData       = {0};
+DEVICE_EXTENSION       g_deviceExtension            = {0};
 KSPIN_LOCK             g_bpeSpinLock;
 WDFDRIVER              g_WDFDriver;
 WDFDEVICE              g_WDFDevice;
@@ -405,6 +407,14 @@ NTSTATUS PrvDriverDeviceAdd(_In_ WDFDRIVER* pWDFDriver)
    HLPR_BAIL_ON_NULL_POINTER_WITH_STATUS(g_pWDMDevice,
                                          status);
 
+   g_pPowerStateEnterIOWorkItem = IoAllocateWorkItem(g_pWDMDevice);
+   HLPR_BAIL_ON_NULL_POINTER_WITH_STATUS(g_pPowerStateEnterIOWorkItem,
+                                         status);
+
+   g_pPowerStateExitIOWorkItem = IoAllocateWorkItem(g_pWDMDevice);
+   HLPR_BAIL_ON_NULL_POINTER_WITH_STATUS(g_pPowerStateExitIOWorkItem,
+                                         status);
+
    status = RegisterPowerStateChangeCallback(&g_deviceExtension);
    HLPR_BAIL_ON_FAILURE(status);
 
diff --git a/network/trans/WFPSampler/sys/Framework_WFPSamplerCalloutDriver.h b/network/trans/WFPSampler/sys/Framework_WFPSamplerCalloutDriver.h
index 52b03284..b8902383 100644
--- a/network/trans/WFPSampler/sys/Framework_WFPSamplerCalloutDriver.h
+++ b/network/trans/WFPSampler/sys/Framework_WFPSamplerCalloutDriver.h
@@ -77,6 +77,9 @@ typedef struct DEVICE_EXTENSION_
 
 extern PDEVICE_OBJECT g_pWDMDevice;
 
+extern PIO_WORKITEM g_pPowerStateEnterIOWorkItem;
+extern PIO_WORKITEM g_pPowerStateExitIOWorkItem;
+
 extern WDFDRIVER g_WDFDriver;
 extern WDFDEVICE g_WDFDevice;
 
diff --git a/network/trans/WFPSampler/syslib/HelperFunctions_WorkItems.cpp b/network/trans/WFPSampler/syslib/HelperFunctions_WorkItems.cpp
index 4c47f3b3..23f4a6df 100644
--- a/network/trans/WFPSampler/syslib/HelperFunctions_WorkItems.cpp
+++ b/network/trans/WFPSampler/syslib/HelperFunctions_WorkItems.cpp
@@ -843,9 +843,16 @@ NTSTATUS KrnlHlprWorkItemQueue(_In_ PDEVICE_OBJECT pWDMDevice,
 
    HLPR_BAIL_LABEL:
 
-   if(status != STATUS_SUCCESS &&
-      pWorkItemData)
-      KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+   if(status != STATUS_SUCCESS)
+   {
+      if(pWorkItemData)
+         KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+      else
+      {
+         if(pIOWorkItem)
+            IoFreeWorkItem(pIOWorkItem);
+      }
+   }
 
 #if DBG
    
@@ -929,9 +936,16 @@ NTSTATUS KrnlHlprWorkItemQueue(_In_ PDEVICE_OBJECT pWDMDevice,
 
    HLPR_BAIL_LABEL:
 
-   if(status != STATUS_SUCCESS &&
-      pWorkItemData)
-      KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+   if(status != STATUS_SUCCESS)
+   {
+      if(pWorkItemData)
+         KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+      else
+      {
+         if(pIOWorkItem)
+            IoFreeWorkItem(pIOWorkItem);
+      }
+   }
 
 #if DBG
    
@@ -1015,9 +1029,16 @@ NTSTATUS KrnlHlprWorkItemQueue(_In_ PDEVICE_OBJECT pWDMDevice,
 
    HLPR_BAIL_LABEL:
 
-   if(status != STATUS_SUCCESS &&
-      pWorkItemData)
-      KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+   if(status != STATUS_SUCCESS)
+   {
+      if(pWorkItemData)
+         KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+      else
+      {
+         if(pIOWorkItem)
+            IoFreeWorkItem(pIOWorkItem);
+      }
+   }
 
 #if DBG
    
@@ -1098,9 +1119,16 @@ NTSTATUS KrnlHlprWorkItemQueue(_In_ PDEVICE_OBJECT pWDMDevice,
 
    HLPR_BAIL_LABEL:
 
-   if(status != STATUS_SUCCESS &&
-      pWorkItemData)
-      KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+   if(status != STATUS_SUCCESS)
+   {
+      if(pWorkItemData)
+         KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+      else
+      {
+         if(pIOWorkItem)
+            IoFreeWorkItem(pIOWorkItem);
+      }
+   }
 
 #if DBG
    
@@ -1181,9 +1209,16 @@ NTSTATUS KrnlHlprWorkItemQueue(_In_ PDEVICE_OBJECT pWDMDevice,
 
    HLPR_BAIL_LABEL:
 
-   if(status != STATUS_SUCCESS &&
-      pWorkItemData)
-      KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+   if(status != STATUS_SUCCESS)
+   {
+      if(pWorkItemData)
+         KrnlHlprWorkItemDataDestroy(&pWorkItemData);
+      else
+      {
+         if(pIOWorkItem)
+            IoFreeWorkItem(pIOWorkItem);
+      }
+   }
 
 #if DBG