From a79fc8493d96dc9531ea153726e31c1f5ab99c63 Mon Sep 17 00:00:00 2001
From: Nodir Temirkhodjaev <nodir.temir@gmail.com>
Date: Mon, 13 Feb 2023 15:39:51 +0300
Subject: [PATCH] Driver: Reset Cancel Routine of cancelled IRP

---
 src/driver/fortbuf.c | 6 ++++--
 src/driver/fortdev.c | 8 +++++---
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/driver/fortbuf.c b/src/driver/fortbuf.c
index 4d7a1a0d..57326786 100644
--- a/src/driver/fortbuf.c
+++ b/src/driver/fortbuf.c
@@ -264,8 +264,10 @@ FORT_API NTSTATUS fort_buffer_cancel_pending(PFORT_BUFFER buf, PIRP irp, ULONG_P
 
     *info = 0;
 
+    /* Cancel routines are called at IRQL = DISPATCH_LEVEL */
+
     KLOCK_QUEUE_HANDLE lock_queue;
-    KeAcquireInStackQueuedSpinLock(&buf->lock, &lock_queue);
+    KeAcquireInStackQueuedSpinLockAtDpcLevel(&buf->lock, &lock_queue);
     if (irp == buf->irp) {
         buf->irp = NULL;
         buf->out_len = 0;
@@ -277,7 +279,7 @@ FORT_API NTSTATUS fort_buffer_cancel_pending(PFORT_BUFFER buf, PIRP irp, ULONG_P
             status = STATUS_SUCCESS;
         }
     }
-    KeReleaseInStackQueuedSpinLock(&lock_queue);
+    KeReleaseInStackQueuedSpinLockFromDpcLevel(&lock_queue);
 
     return status;
 }
diff --git a/src/driver/fortdev.c b/src/driver/fortdev.c
index fea200a5..e0a7daa2 100644
--- a/src/driver/fortdev.c
+++ b/src/driver/fortdev.c
@@ -117,6 +117,7 @@ static void fort_device_cancel_pending(PDEVICE_OBJECT device, PIRP irp)
 
     const NTSTATUS status = fort_buffer_cancel_pending(&g_device->buffer, irp, &info);
 
+    IoSetCancelRoutine(irp, NULL);
     IoReleaseCancelSpinLock(irp->CancelIrql); /* before IoCompleteRequest()! */
 
     fort_request_complete_info(irp, status, info);
@@ -177,12 +178,13 @@ static NTSTATUS fort_device_control_getlog(PVOID out, ULONG out_len, PIRP irp, U
         const NTSTATUS status = fort_buffer_xmove(&g_device->buffer, irp, out, out_len, info);
 
         if (status == STATUS_PENDING) {
-            KIRQL cirq;
-
             IoMarkIrpPending(irp);
 
+            KIRQL cirq;
             IoAcquireCancelSpinLock(&cirq);
-            IoSetCancelRoutine(irp, fort_device_cancel_pending);
+            {
+                IoSetCancelRoutine(irp, fort_device_cancel_pending);
+            }
             IoReleaseCancelSpinLock(cirq);
         }
         return status;