Balloon: Check ActionFlags before request unmark

Signed-off-by: Kostiantyn Kostiuk <konstantin@daynix.com>
virtio-win · Aug 31, 2022 · 663b6dd · 663b6dd
1 parent 6266327
commit 663b6dd
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/Balloon/sys/queue.c b/Balloon/sys/queue.c
@@ -121,11 +121,18 @@ VOID BalloonIoStop(IN WDFQUEUE Queue,
     TraceEvents(TRACE_LEVEL_VERBOSE, DBG_READ,
         "--> %!FUNC! Request: %p", Request);
 
-    if (WdfRequestUnmarkCancelable(Request) == STATUS_CANCELLED)
+    if (ActionFlags & WdfRequestStopRequestCancelable)
     {
-        WdfRequestStopAcknowledge(Request, FALSE);
+        NTSTATUS status = WdfRequestUnmarkCancelable(Request);
+        __analysis_assume(status != STATUS_NOT_SUPPORTED);
+        if (status == STATUS_CANCELLED)
+        {
+            WdfRequestStopAcknowledge(Request, FALSE);
+            goto end_io_stop;
+        }
     }
-    else if (ActionFlags & WdfRequestStopActionSuspend)
+
+    if (ActionFlags & WdfRequestStopActionSuspend)
     {
         WdfRequestStopAcknowledge(Request, TRUE);
     }
@@ -138,6 +145,7 @@ VOID BalloonIoStop(IN WDFQUEUE Queue,
         WdfRequestComplete(Request, STATUS_UNSUCCESSFUL);
     }
 
+end_io_stop:
     TraceEvents(TRACE_LEVEL_VERBOSE, DBG_READ, "<-- %!FUNC!");
 }