Windows driver: enhance IRP completion by avoiding inline ones and using CriticalWorkQueue for faster completion dispatch

This change adds more robustness in low memory case and avoid freezes.
It also protects cancel handling with cancel spin lock

src/Driver/EncryptedIoQueue.c

@@ -261,22 +261,6 @@ static void OnItemCompleted (EncryptedIoQueueItem *item, BOOL freeItem)
 		ReleasePoolBuffer (item->Queue, item);
 }
 
-
-static NTSTATUS CompleteOriginalIrp (EncryptedIoQueueItem *item, NTSTATUS status, ULONG_PTR information)
-{
-#ifdef TC_TRACE_IO_QUEUE
-	Dump ("< %I64d [%I64d] %c status=%x info=%I64d\n", item->OriginalIrpOffset, GetElapsedTime (&item->Queue->LastPerformanceCounter), item->Write ? 'W' : 'R', status, (int64) information);
-#endif
-
-	TCCompleteDiskIrp (item->OriginalIrp, status, information);
-
-	item->Status = status;
-	OnItemCompleted (item, TRUE);
-
-	return status;
-}
-
-
 static void AcquireFragmentBuffer (EncryptedIoQueue *queue, uint8 *buffer)
 {
 	NTSTATUS status = STATUS_INVALID_PARAMETER;
@@ -407,18 +391,30 @@ UpdateBuffer(
 static VOID CompleteIrpWorkItemRoutine(PDEVICE_OBJECT DeviceObject, PVOID Context)
 {
 	PCOMPLETE_IRP_WORK_ITEM workItem = (PCOMPLETE_IRP_WORK_ITEM)Context;
-	EncryptedIoQueueItem* item = (EncryptedIoQueueItem * ) workItem->Item;
+	EncryptedIoQueueItem *item = (EncryptedIoQueueItem *)workItem->Item;
-	EncryptedIoQueue* queue = item->Queue;
+	EncryptedIoQueue *queue = (EncryptedIoQueue *)workItem->Queue;
 	KIRQL oldIrql;
 	UNREFERENCED_PARAMETER(DeviceObject);
 
 	__try
 	{
-		// Complete the IRP
+		if (item)
-		TCCompleteDiskIrp(workItem->Irp, workItem->Status, workItem->Information);
+		{
-
+			// Normal path: completion associated with an EncryptedIoQueueItem
-		item->Status = workItem->Status;
+			item->Status = workItem->Status;
-		OnItemCompleted(item, FALSE); // Do not free item here; it will be freed below
+			OnItemCompleted(item, FALSE); // Do not free item here; it will be freed below
+			TCCompleteDiskIrp(workItem->Irp, workItem->Status, workItem->Information);
+		}
+		else
+		{
+			// Raw completion: release lock/counters if requested and finish the IRP
+			if (workItem->ReleaseLockAndCounters && queue)
+			{
+				DecrementOutstandingIoCount(queue);
+				IoReleaseRemoveLock(&queue->RemoveLock, workItem->Irp);
+			}
+			TCCompleteDiskIrp(workItem->Irp, workItem->Status, workItem->Information);
+		}
 	}
 	__finally
 	{
@@ -436,47 +432,94 @@ static VOID CompleteIrpWorkItemRoutine(PDEVICE_OBJECT DeviceObject, PVOID Contex
 		// Release the semaphore to signal that a work item is available
 		KeReleaseSemaphore(&queue->WorkItemSemaphore, IO_DISK_INCREMENT, 1, FALSE);
 
-		// Free the item
+		// Free the queue item if any
-		ReleasePoolBuffer(queue, item);
+		if (item)
+		{
+			ReleasePoolBuffer(queue, item);
+		}
 	}
 }
 
-// Handles the completion of the original IRP.
+// Queues completion of the original IRP via pre-allocated work-item (from an EncryptedIoQueueItem)
-static VOID HandleCompleteOriginalIrp(EncryptedIoQueue* queue, EncryptedIoRequest* request)
+static VOID QueueIrpCompletionFromItem(EncryptedIoQueue *queue,
+									   EncryptedIoQueueItem *item,
+									   NTSTATUS status)
 {
-	NTSTATUS status = KeWaitForSingleObject(&queue->WorkItemSemaphore, Executive, KernelMode, FALSE, NULL);
+	// Must be PASSIVE_LEVEL. We rely on the preallocated pool only.
-	if (queue->ThreadExitRequested)
+	ASSERT(KeGetCurrentIrql() == PASSIVE_LEVEL);
-		return;
 
-	if (!NT_SUCCESS(status))
+	// Block until a work item is available (no inline completion, no alloc).
+	for (;;)
 	{
-		// Handle wait failure: we call the completion routine directly. 
+		NTSTATUS ws = KeWaitForSingleObject(&queue->WorkItemSemaphore,
-		// This is not ideal since it can cause deadlock that we are trying to fix but it is better than losing the IRP.
+											Executive, KernelMode, FALSE, NULL);
-		CompleteOriginalIrp(request->Item, STATUS_INSUFFICIENT_RESOURCES, 0);
+		if (NT_SUCCESS(ws))
+			break;
+		// Non-alertable wait: an error here is unexpected
 	}
-	else
+
+	KIRQL oldIrql;
+	KeAcquireSpinLock(&queue->WorkItemLock, &oldIrql);
+	ASSERT(!IsListEmpty(&queue->FreeWorkItemsList));
+	PLIST_ENTRY freeEntry = RemoveHeadList(&queue->FreeWorkItemsList);
+	KeReleaseSpinLock(&queue->WorkItemLock, oldIrql);
+
+	PCOMPLETE_IRP_WORK_ITEM wi = CONTAINING_RECORD(freeEntry, COMPLETE_IRP_WORK_ITEM, ListEntry);
+
+	KeResetEvent(&queue->NoActiveWorkItemsEvent);
+	InterlockedIncrement(&queue->ActiveWorkItems);
+
+	wi->Irp = item->OriginalIrp;
+	wi->Status = status;
+	wi->Information = NT_SUCCESS(status) ? item->OriginalLength : 0;
+	wi->Item = item;
+	wi->Queue = queue;
+	wi->ReleaseLockAndCounters = FALSE;
+
+	IoQueueWorkItem(wi->WorkItem, CompleteIrpWorkItemRoutine, CriticalWorkQueue, wi);
+}
+
+// Queues completion of a raw IRP not bound to EncryptedIoQueueItem (early error/cancel paths)
+static VOID QueueRawIrpCompletion(EncryptedIoQueue *queue,
+								  PIRP irp, NTSTATUS status, ULONG_PTR info)
+{
+	ASSERT(KeGetCurrentIrql() == PASSIVE_LEVEL);
+
+	for (;;)
 	{
-		// Obtain a work item from the free list.
+		NTSTATUS ws = KeWaitForSingleObject(&queue->WorkItemSemaphore,
-		KIRQL oldIrql;
+											Executive, KernelMode, FALSE, NULL);
-		KeAcquireSpinLock(&queue->WorkItemLock, &oldIrql);
+		if (NT_SUCCESS(ws))
-		PLIST_ENTRY freeEntry = RemoveHeadList(&queue->FreeWorkItemsList);
+			break;
-		KeReleaseSpinLock(&queue->WorkItemLock, oldIrql);
-
-		PCOMPLETE_IRP_WORK_ITEM workItem = CONTAINING_RECORD(freeEntry, COMPLETE_IRP_WORK_ITEM, ListEntry);
-
-		// Increment ActiveWorkItems.
-		InterlockedIncrement(&queue->ActiveWorkItems);
-		KeResetEvent(&queue->NoActiveWorkItemsEvent);
-
-		// Prepare the work item.
-		workItem->Irp = request->Item->OriginalIrp;
-		workItem->Status = request->Item->Status;
-		workItem->Information = NT_SUCCESS(request->Item->Status) ? request->Item->OriginalLength : 0;
-		workItem->Item = request->Item;
-
-		// Queue the work item.
-		IoQueueWorkItem(workItem->WorkItem, CompleteIrpWorkItemRoutine, DelayedWorkQueue, workItem);
 	}
+
+	KIRQL oldIrql;
+	KeAcquireSpinLock(&queue->WorkItemLock, &oldIrql);
+	ASSERT(!IsListEmpty(&queue->FreeWorkItemsList));
+	PLIST_ENTRY freeEntry = RemoveHeadList(&queue->FreeWorkItemsList);
+	KeReleaseSpinLock(&queue->WorkItemLock, oldIrql);
+
+	PCOMPLETE_IRP_WORK_ITEM wi = CONTAINING_RECORD(freeEntry, COMPLETE_IRP_WORK_ITEM, ListEntry);
+
+	KeResetEvent(&queue->NoActiveWorkItemsEvent);
+	InterlockedIncrement(&queue->ActiveWorkItems);
+
+	wi->Irp = irp;
+	wi->Status = status;
+	wi->Information = info;
+	wi->Item = NULL; // raw path: helper will release lock/counters
+	wi->Queue = queue;
+	wi->ReleaseLockAndCounters = TRUE;
+
+	IoQueueWorkItem(wi->WorkItem, CompleteIrpWorkItemRoutine, CriticalWorkQueue, wi);
+}
+
+// Handles the completion of the original IRP.
+static VOID HandleCompleteOriginalIrp(EncryptedIoQueue *queue, EncryptedIoRequest *request)
+{
+	QueueIrpCompletionFromItem(queue,
+							   request->Item,
+							   request->Item->Status);
 }
 
 static VOID CompletionThreadProc(PVOID threadArg)
@@ -770,6 +813,8 @@ static VOID MainThreadProc (PVOID threadArg)
 		{
 			PIRP irp = CONTAINING_RECORD (listEntry, IRP, Tail.Overlay.ListEntry);
 			PIO_STACK_LOCATION irpSp = IoGetCurrentIrpStackLocation (irp);
+			KIRQL irql;
+			BOOLEAN cancelled;
 
 			if (queue->Suspended)
 				KeWaitForSingleObject (&queue->QueueResumedEvent, Executive, KernelMode, FALSE, NULL);
@@ -777,9 +822,8 @@ static VOID MainThreadProc (PVOID threadArg)
 			item = GetPoolBuffer (queue, sizeof (EncryptedIoQueueItem));
 			if (!item)
 			{
-				TCCompleteDiskIrp (irp, STATUS_INSUFFICIENT_RESOURCES, 0);
+				// Defer completion to work item to avoid re-entrant completion on our thread
-				DecrementOutstandingIoCount (queue);
+				QueueRawIrpCompletion(queue, irp, STATUS_INSUFFICIENT_RESOURCES, 0);
-				IoReleaseRemoveLock (&queue->RemoveLock, irp);
 
 				continue;
 			}
@@ -789,10 +833,14 @@ static VOID MainThreadProc (PVOID threadArg)
 			item->TempUserMdl = NULL;
 			item->Status = STATUS_SUCCESS;
 
-			IoSetCancelRoutine (irp, NULL);
+			IoAcquireCancelSpinLock(&irql);
-			if (irp->Cancel)
+			(PDRIVER_CANCEL)IoSetCancelRoutine(irp, NULL);
+			cancelled = irp->Cancel ? TRUE : FALSE;
+			IoReleaseCancelSpinLock(irql);
+			if (cancelled)
 			{
-				CompleteOriginalIrp (item, STATUS_CANCELLED, 0);
+				// Defer cancellation completion
+				QueueIrpCompletionFromItem(queue, item, STATUS_CANCELLED);
 				continue;
 			}
 
@@ -811,7 +859,8 @@ static VOID MainThreadProc (PVOID threadArg)
 				break;
 
 			default:
-				CompleteOriginalIrp (item, STATUS_INVALID_PARAMETER, 0);
+				// Defer completion for invalid parameter
+				QueueIrpCompletionFromItem(queue, item, STATUS_INVALID_PARAMETER);
 				continue;
 			}
 
@@ -832,7 +881,8 @@ static VOID MainThreadProc (PVOID threadArg)
 				hResult = ULongAdd(item->OriginalLength, ENCRYPTION_DATA_UNIT_SIZE, &alignedLength);
 				if (hResult != S_OK)
 				{
-					CompleteOriginalIrp (item, STATUS_INVALID_PARAMETER, 0);
+					// Defer completion for invalid parameter
+					QueueIrpCompletionFromItem(queue, item, STATUS_INVALID_PARAMETER);
 					continue;
 				}
 
@@ -841,7 +891,7 @@ static VOID MainThreadProc (PVOID threadArg)
 				buffer = TCalloc (alignedLength);
 				if (!buffer)
 				{
-					CompleteOriginalIrp (item, STATUS_INSUFFICIENT_RESOURCES, 0);
+					QueueIrpCompletionFromItem (queue, item, STATUS_INSUFFICIENT_RESOURCES);
 					continue;
 				}
 
@@ -861,7 +911,8 @@ static VOID MainThreadProc (PVOID threadArg)
 					if (!NT_SUCCESS(mapStatus))
 					{
 						TCfree (buffer);
-						CompleteOriginalIrp (item, mapStatus, 0);
+						// Defer completion on mapping failure
+						QueueIrpCompletionFromItem(queue, item, mapStatus);
 						continue;
 					}
 
@@ -883,7 +934,8 @@ static VOID MainThreadProc (PVOID threadArg)
 				}
 
 				TCfree (buffer);
-				CompleteOriginalIrp (item, item->Status, NT_SUCCESS (item->Status) ? item->OriginalLength : 0);
+				// Defer completion of misaligned read workaround
+				QueueIrpCompletionFromItem(queue, item, item->Status);
 				continue;
 			}
 
@@ -898,7 +950,8 @@ static VOID MainThreadProc (PVOID threadArg)
 					)
 				)
 			{
-				CompleteOriginalIrp (item, STATUS_INVALID_PARAMETER, 0);
+				// Defer completion for invalid parameter
+				QueueIrpCompletionFromItem(queue, item, STATUS_INVALID_PARAMETER);
 				continue;
 			}
 
@@ -916,7 +969,8 @@ static VOID MainThreadProc (PVOID threadArg)
 
 				if (hResult != S_OK)
 				{
-					CompleteOriginalIrp (item, STATUS_INVALID_PARAMETER, 0);
+					// Defer completion for invalid parameter
+					QueueIrpCompletionFromItem(queue, item, STATUS_INVALID_PARAMETER);
 					continue;
 				}
 				else
@@ -932,7 +986,7 @@ static VOID MainThreadProc (PVOID threadArg)
 						// Do not allow writing to this volume anymore. This is to fake a complete volume
 						// or system failure (otherwise certain kinds of inconsistency within the file
 						// system could indicate that this volume has used hidden volume protection).
-						CompleteOriginalIrp (item, STATUS_INVALID_PARAMETER, 0);
+						QueueIrpCompletionFromItem(queue, item, STATUS_INVALID_PARAMETER);
 						continue;
 					}
 
@@ -946,7 +1000,7 @@ static VOID MainThreadProc (PVOID threadArg)
 						queue->CryptoInfo->bHiddenVolProtectionAction = TRUE;
 
 						// Deny this write operation to prevent the hidden volume from being overwritten
-						CompleteOriginalIrp (item, STATUS_INVALID_PARAMETER, 0);
+						QueueIrpCompletionFromItem(queue, item, STATUS_INVALID_PARAMETER);
 						continue;
 					}
 				}
@@ -956,7 +1010,7 @@ static VOID MainThreadProc (PVOID threadArg)
 			{
 				// Prevent inappropriately designed software from damaging important data that may be out of sync with the backup on the Rescue Disk (such as the end of the encrypted area).
 				Dump ("Preventing write to the system encryption key data area\n");
-				CompleteOriginalIrp (item, STATUS_MEDIA_WRITE_PROTECTED, 0);
+				QueueIrpCompletionFromItem(queue, item, STATUS_MEDIA_WRITE_PROTECTED);
 				continue;
 			}
 			else if (item->Write && IsHiddenSystemRunning()
@@ -964,7 +1018,7 @@ static VOID MainThreadProc (PVOID threadArg)
 				 || RegionsOverlap (item->OriginalOffset.QuadPart, item->OriginalOffset.QuadPart + item->OriginalLength - 1, GetBootDriveLength(), _I64_MAX)))
 			{
 				Dump ("Preventing write to boot loader or host protected area\n");
-				CompleteOriginalIrp (item, STATUS_MEDIA_WRITE_PROTECTED, 0);
+				QueueIrpCompletionFromItem(queue, item, STATUS_MEDIA_WRITE_PROTECTED);
 				continue;
 			} 
 			else if (item->Write
@@ -973,7 +1027,7 @@ static VOID MainThreadProc (PVOID threadArg)
 			{
 				// Prevent inappropriately designed software from damaging important data
 				Dump ("Preventing write to the system GPT area\n");
-				CompleteOriginalIrp (item, STATUS_MEDIA_WRITE_PROTECTED, 0);
+				QueueIrpCompletionFromItem(queue, item, STATUS_MEDIA_WRITE_PROTECTED);
 				continue;
 			}
 			
@@ -986,7 +1040,7 @@ static VOID MainThreadProc (PVOID threadArg)
 				&item->TempUserMdl);
 			if (!NT_SUCCESS(mapStatus))
 			{
-				CompleteOriginalIrp (item, mapStatus, 0);
+				QueueIrpCompletionFromItem (queue, item, mapStatus);
 				continue;
 			}
 
@@ -1010,7 +1064,7 @@ static VOID MainThreadProc (PVOID threadArg)
 				if (!request)
 				{
 					InterlockedDecrement(&queue->IoThreadPendingRequestCount);
-					CompleteOriginalIrp (item, STATUS_INSUFFICIENT_RESOURCES, 0);
+					QueueIrpCompletionFromItem (queue, item, STATUS_INSUFFICIENT_RESOURCES);
 					break;
 				}
 				request->Item = item;
@@ -1311,19 +1365,18 @@ retry_preallocated:
 		goto noMemory;
 	}
 
-	// Allocate and initialize work items
+	// Initialize the work item pool
+	for (j = 0; j < (int) queue->MaxWorkItems; ++j)
+	{
+		queue->WorkItemPool[j].WorkItem = NULL;
+	}
+
+	// Allocate work items
 	for (i = 0; i < (int) queue->MaxWorkItems; ++i)
 	{
 		queue->WorkItemPool[i].WorkItem = IoAllocateWorkItem(queue->DeviceObject);
 		if (!queue->WorkItemPool[i].WorkItem)
 		{
-			// Handle allocation failure
-			// Free previously allocated work items
-			for (j = 0; j < i; ++j)
-			{
-				IoFreeWorkItem(queue->WorkItemPool[j].WorkItem);
-			}
-			TCfree(queue->WorkItemPool);
 			goto noMemory;
 		}
 
@@ -1384,6 +1437,19 @@ noMemory:
 	status = STATUS_INSUFFICIENT_RESOURCES;
 
 err:
+    if (queue->WorkItemPool)
+	{
+        for (i = 0; i < (int) queue->MaxWorkItems; ++i)
+		{
+            if (queue->WorkItemPool[i].WorkItem)
+			{
+                IoFreeWorkItem(queue->WorkItemPool[i].WorkItem);
+                queue->WorkItemPool[i].WorkItem = NULL;
+            }
+        }
+        TCfree(queue->WorkItemPool);
+        queue->WorkItemPool = NULL;
+    }
 	if (queue->FragmentBufferA)
 		TCfree (queue->FragmentBufferA);
 	if (queue->FragmentBufferB)

src/Driver/EncryptedIoQueue.h

@@ -44,11 +44,13 @@ typedef struct _COMPLETE_IRP_WORK_ITEM
 	PIRP Irp;
 	NTSTATUS Status;
 	ULONG_PTR Information;
-	void* Item;
+	void* Item; // NULL for raw completion
+	struct _EncryptedIoQueueStruct* Queue; // Owning queue
+	BOOLEAN ReleaseLockAndCounters; // TRUE for raw completion: release RemoveLock and decrement OutstandingIoCount
 	LIST_ENTRY ListEntry; // For managing free work items
 } COMPLETE_IRP_WORK_ITEM, * PCOMPLETE_IRP_WORK_ITEM;
 
-typedef struct
+typedef struct _EncryptedIoQueueStruct
 {
 	PDEVICE_OBJECT DeviceObject;