Public/VeraCrypt
1
0
Fork 0
mirror of https://github.com/veracrypt/VeraCrypt.git synced 2025-11-11 02:58:02 -06:00
Code Activity

Windows driver: revert IRP completion overhaul changes until more tests are conducted

Browse Source
This commit is contained in:
Mounir IDRASSI
2025-09-13 15:08:38 +09:00
parent f6f25eec8f
commit 02e7b48836
2 changed files with 125 additions and 315 deletions
Download Patch File Download Diff File Expand all files Collapse all files

432
src/Driver/EncryptedIoQueue.c
View File

@@ -19,7 +19,6 @@
#include "Volumes.h" #include "Volumes.h"
#include <IntSafe.h> #include <IntSafe.h>
#define MAX_WI_RETRIES 8
static void AcquireBufferPoolMutex (EncryptedIoQueue *queue) static void AcquireBufferPoolMutex (EncryptedIoQueue *queue)
{ {
@@ -161,23 +160,19 @@ static void DecrementOutstandingIoCount (EncryptedIoQueue *queue)
static void OnItemCompleted (EncryptedIoQueueItem *item, BOOL freeItem) static void OnItemCompleted (EncryptedIoQueueItem *item, BOOL freeItem)
{ {
EncryptedIoQueue* queue = item->Queue; DecrementOutstandingIoCount (item->Queue);
PIRP originalIrp = item->OriginalIrp; IoReleaseRemoveLock (&item->Queue->RemoveLock, item->OriginalIrp);
if (NT_SUCCESS (item->Status)) if (NT_SUCCESS (item->Status))
{ {
if (item->Write) if (item->Write)
queue->TotalBytesWritten += item->BytesCompleted; item->Queue->TotalBytesWritten += item->OriginalLength;
else else
queue->TotalBytesRead += item->BytesCompleted; item->Queue->TotalBytesRead += item->OriginalLength;
} }
DecrementOutstandingIoCount (queue);
if (freeItem) if (freeItem)
ReleasePoolBuffer (queue, item); ReleasePoolBuffer (item->Queue, item);
// Release the RemoveLock last after we are done touching the queue
IoReleaseRemoveLock (&queue->RemoveLock, originalIrp);
} }
@@ -323,7 +318,6 @@ UpdateBuffer(
return updated; return updated;
} }
// Note: completes the IRP and releases the RemoveLock last (after all queue accesses)
static VOID CompleteIrpWorkItemRoutine(PDEVICE_OBJECT DeviceObject, PVOID Context) static VOID CompleteIrpWorkItemRoutine(PDEVICE_OBJECT DeviceObject, PVOID Context)
{ {
PCOMPLETE_IRP_WORK_ITEM workItem = (PCOMPLETE_IRP_WORK_ITEM)Context; PCOMPLETE_IRP_WORK_ITEM workItem = (PCOMPLETE_IRP_WORK_ITEM)Context;
@@ -332,141 +326,71 @@ static VOID CompleteIrpWorkItemRoutine(PDEVICE_OBJECT DeviceObject, PVOID Contex
KIRQL oldIrql; KIRQL oldIrql;
UNREFERENCED_PARAMETER(DeviceObject); UNREFERENCED_PARAMETER(DeviceObject);
// Complete IRP __try
TCCompleteDiskIrp(workItem->Irp, workItem->Status, workItem->Information);
// Centralized accounting and cleanup
item->Status = workItem->Status;
item->BytesCompleted = workItem->Information;
OnItemCompleted(item, TRUE); // releases RemoveLock last
// Return or free the work item depending on origin
if (workItem->FromPool)
{ {
// Complete the IRP
TCCompleteDiskIrp(workItem->Irp, workItem->Status, workItem->Information);
item->Status = workItem->Status;
OnItemCompleted(item, FALSE); // Do not free item here; it will be freed below
}
__finally
{
// If no active work items remain, signal the event
if (InterlockedDecrement(&queue->ActiveWorkItems) == 0)
{
KeSetEvent(&queue->NoActiveWorkItemsEvent, IO_DISK_INCREMENT, FALSE);
}
// Return the work item to the free list
KeAcquireSpinLock(&queue->WorkItemLock, &oldIrql); KeAcquireSpinLock(&queue->WorkItemLock, &oldIrql);
InsertTailList(&queue->FreeWorkItemsList, &workItem->ListEntry); InsertTailList(&queue->FreeWorkItemsList, &workItem->ListEntry);
KeReleaseSpinLock(&queue->WorkItemLock, oldIrql); KeReleaseSpinLock(&queue->WorkItemLock, oldIrql);
// immediately wake any waiter
KeSetEvent(&queue->WorkItemAvailableEvent, IO_DISK_INCREMENT, FALSE); // Release the semaphore to signal that a work item is available
KeReleaseSemaphore(&queue->WorkItemSemaphore, IO_DISK_INCREMENT, 1, FALSE);
// Free the item
ReleasePoolBuffer(queue, item);
}
}
// Handles the completion of the original IRP.
static VOID HandleCompleteOriginalIrp(EncryptedIoQueue* queue, EncryptedIoRequest* request)
{
NTSTATUS status = KeWaitForSingleObject(&queue->WorkItemSemaphore, Executive, KernelMode, FALSE, NULL);
if (queue->ThreadExitRequested)
return;
if (!NT_SUCCESS(status))
{
// Handle wait failure: we call the completion routine directly.
// This is not ideal since it can cause deadlock that we are trying to fix but it is better than losing the IRP.
CompleteOriginalIrp(request->Item, STATUS_INSUFFICIENT_RESOURCES, 0);
} }
else else
{ {
IoFreeWorkItem(workItem->WorkItem); // Obtain a work item from the free list.
TCfree(workItem); KIRQL oldIrql;
} KeAcquireSpinLock(&queue->WorkItemLock, &oldIrql);
PLIST_ENTRY freeEntry = RemoveHeadList(&queue->FreeWorkItemsList);
KeReleaseSpinLock(&queue->WorkItemLock, oldIrql);
if (InterlockedDecrement(&queue->ActiveWorkItems) == 0) PCOMPLETE_IRP_WORK_ITEM workItem = CONTAINING_RECORD(freeEntry, COMPLETE_IRP_WORK_ITEM, ListEntry);
KeSetEvent(&queue->NoActiveWorkItemsEvent, IO_DISK_INCREMENT, FALSE);
}
// Helper: acquire a completion work item (from pool if available, else elastic allocation) // Increment ActiveWorkItems.
static PCOMPLETE_IRP_WORK_ITEM TryAcquireCompletionWorkItem(EncryptedIoQueue* queue)
{
PCOMPLETE_IRP_WORK_ITEM wi = NULL;
KIRQL irql;
// Try pooled work item
KeAcquireSpinLock(&queue->WorkItemLock, &irql);
if (!IsListEmpty(&queue->FreeWorkItemsList))
{
PLIST_ENTRY e = RemoveHeadList(&queue->FreeWorkItemsList);
wi = CONTAINING_RECORD(e, COMPLETE_IRP_WORK_ITEM, ListEntry);
wi->FromPool = TRUE;
}
KeReleaseSpinLock(&queue->WorkItemLock, irql);
// Elastic fallback: allocate on demand
if (!wi)
{
wi = (PCOMPLETE_IRP_WORK_ITEM)TCalloc(sizeof(COMPLETE_IRP_WORK_ITEM));
if (wi)
{
RtlZeroMemory(wi, sizeof(*wi));
wi->WorkItem = IoAllocateWorkItem(queue->DeviceObject);
if (!wi->WorkItem) { TCfree(wi); wi = NULL; }
else wi->FromPool = FALSE;
}
}
return wi;
}
// Complete the original IRP: try pooled work item, then elastic, else inline.
// Precondition: item->Status and item->BytesCompleted are set by the caller.
static VOID FinalizeOriginalIrp(EncryptedIoQueue* queue, EncryptedIoQueueItem* item)
{
// Try to get a work item (from the preallocated pool first, else elastic alloc)
PCOMPLETE_IRP_WORK_ITEM wi = TryAcquireCompletionWorkItem(queue);
if (wi)
{
InterlockedIncrement(&queue->ActiveWorkItems); InterlockedIncrement(&queue->ActiveWorkItems);
KeResetEvent(&queue->NoActiveWorkItemsEvent); KeResetEvent(&queue->NoActiveWorkItemsEvent);
wi->Irp = item->OriginalIrp; // Prepare the work item.
wi->Status = item->Status; workItem->Irp = request->Item->OriginalIrp;
wi->Information = item->BytesCompleted; workItem->Status = request->Item->Status;
wi->Item = item; workItem->Information = NT_SUCCESS(request->Item->Status) ? request->Item->OriginalLength : 0;
workItem->Item = request->Item;
IoQueueWorkItem(wi->WorkItem, CompleteIrpWorkItemRoutine, DelayedWorkQueue, wi); // Queue the work item.
return; IoQueueWorkItem(workItem->WorkItem, CompleteIrpWorkItemRoutine, DelayedWorkQueue, workItem);
} }
// Final fallback: inline completion (safe OOM path)
TCCompleteDiskIrp(item->OriginalIrp, item->Status, item->BytesCompleted);
OnItemCompleted(item, TRUE);
}
// Handles the completion of the original IRP.
// Returns TRUE if the caller still owns and should free request parameter.
// Returns FALSE if this function requeued request parameter (caller must NOT free it).
static BOOLEAN HandleCompleteOriginalIrp(EncryptedIoQueue* queue, EncryptedIoRequest* request)
{
PCOMPLETE_IRP_WORK_ITEM wi = TryAcquireCompletionWorkItem(queue);
if (!wi)
{
// No work item available. Either try again later, or give up and complete inline.
request->WiRetryCount++;
if (request->WiRetryCount >= MAX_WI_RETRIES)
{
// Final fallback: complete inline on the completion thread (no locks held)
TCCompleteDiskIrp(request->Item->OriginalIrp, request->Item->Status, request->Item->BytesCompleted);
// Complete accounting and release resources (RemoveLock released last inside OnItemCompleted)
OnItemCompleted(request->Item, TRUE);
// We completed: caller should free request
return TRUE;
}
// Requeue and wait for a pooled work item to become available
ExInterlockedInsertTailList(&queue->CompletionThreadQueue,
&request->CompletionListEntry,
&queue->CompletionThreadQueueLock);
KeSetEvent(&queue->CompletionThreadQueueNotEmptyEvent, IO_DISK_INCREMENT, FALSE);
// Prefer waiting on a real signal over blind sleeps (per-request exponential backoff 1...32ms)
ULONG backoff = 1u << min(request->WiRetryCount - 1, 5); // 1,2,4,8,16,32
LARGE_INTEGER duetime; duetime.QuadPart = -(LONGLONG)backoff * 10000;
KeWaitForSingleObject(&queue->WorkItemAvailableEvent, Executive, KernelMode, FALSE, &duetime);
return FALSE; // we requeued so caller must not free request
}
// Queue to system worker thread
InterlockedIncrement(&queue->ActiveWorkItems);
KeResetEvent(&queue->NoActiveWorkItemsEvent);
wi->Irp = request->Item->OriginalIrp;
wi->Status = request->Item->Status;
wi->Information = request->Item->BytesCompleted;
wi->Item = request->Item;
IoQueueWorkItem(wi->WorkItem, CompleteIrpWorkItemRoutine, DelayedWorkQueue, wi);
return TRUE; // caller should free request
} }
static VOID CompletionThreadProc(PVOID threadArg) static VOID CompletionThreadProc(PVOID threadArg)
@@ -474,33 +398,48 @@ static VOID CompletionThreadProc(PVOID threadArg)
EncryptedIoQueue* queue = (EncryptedIoQueue*)threadArg; EncryptedIoQueue* queue = (EncryptedIoQueue*)threadArg;
PLIST_ENTRY listEntry; PLIST_ENTRY listEntry;
EncryptedIoRequest* request; EncryptedIoRequest* request;
UINT64_STRUCT dataUnit;
if (IsEncryptionThreadPoolRunning()) if (IsEncryptionThreadPoolRunning())
KeSetPriorityThread(KeGetCurrentThread(), LOW_REALTIME_PRIORITY); KeSetPriorityThread(KeGetCurrentThread(), LOW_REALTIME_PRIORITY);
for (;;) while (!queue->ThreadExitRequested)
{ {
if (!NT_SUCCESS(KeWaitForSingleObject(&queue->CompletionThreadQueueNotEmptyEvent, Executive, KernelMode, FALSE, NULL))) if (!NT_SUCCESS(KeWaitForSingleObject(&queue->CompletionThreadQueueNotEmptyEvent, Executive, KernelMode, FALSE, NULL)))
continue; continue;
if (queue->ThreadExitRequested)
break;
while ((listEntry = ExInterlockedRemoveHeadList(&queue->CompletionThreadQueue, &queue->CompletionThreadQueueLock))) while ((listEntry = ExInterlockedRemoveHeadList(&queue->CompletionThreadQueue, &queue->CompletionThreadQueueLock)))
{ {
request = CONTAINING_RECORD(listEntry, EncryptedIoRequest, CompletionListEntry); request = CONTAINING_RECORD(listEntry, EncryptedIoRequest, CompletionListEntry);
if (request->EncryptedLength > 0 && NT_SUCCESS(request->Item->Status))
{
ASSERT(request->EncryptedOffset + request->EncryptedLength <= request->Offset.QuadPart + request->Length);
dataUnit.Value = (request->Offset.QuadPart + request->EncryptedOffset) / ENCRYPTION_DATA_UNIT_SIZE;
if (queue->CryptoInfo->bPartitionInInactiveSysEncScope)
dataUnit.Value += queue->CryptoInfo->FirstDataUnitNo.Value;
else if (queue->RemapEncryptedArea)
dataUnit.Value += queue->RemappedAreaDataUnitOffset;
DecryptDataUnits(request->Data + request->EncryptedOffset, &dataUnit, request->EncryptedLength / ENCRYPTION_DATA_UNIT_SIZE, queue->CryptoInfo);
}
// Dump("Read sector %lld count %d\n", request->Offset.QuadPart >> 9, request->Length >> 9);
// Update subst sectors
if((queue->SecRegionData != NULL) && (queue->SecRegionSize > 512)) {
UpdateBuffer(request->Data, queue->SecRegionData, queue->SecRegionSize, request->Offset.QuadPart, request->Length, TRUE);
}
if (request->CompleteOriginalIrp) if (request->CompleteOriginalIrp)
{ {
if (!HandleCompleteOriginalIrp(queue, request)) HandleCompleteOriginalIrp(queue, request);
{
// request was requeued; do not free it now
continue;
}
} }
ReleasePoolBuffer(queue, request); ReleasePoolBuffer(queue, request);
} }
if (queue->ThreadExitRequested)
break;
} }
PsTerminateSystemThread(STATUS_SUCCESS); PsTerminateSystemThread(STATUS_SUCCESS);
@@ -567,9 +506,6 @@ static VOID IoThreadProc (PVOID threadArg)
#endif #endif
// Perform IO request if no preceding request of the item failed // Perform IO request if no preceding request of the item failed
ULONG_PTR bytesXfer = 0; // track per-fragment transfer
request->ActualBytes = 0; // default
request->WiRetryCount = 0;
if (NT_SUCCESS (request->Item->Status)) if (NT_SUCCESS (request->Item->Status))
{ {
if (queue->ThreadBlockReadWrite) if (queue->ThreadBlockReadWrite)
@@ -637,7 +573,6 @@ static VOID IoThreadProc (PVOID threadArg)
else else
request->Item->Status = TCCachedRead (queue, NULL, request->Data, request->Offset, request->Length); request->Item->Status = TCCachedRead (queue, NULL, request->Data, request->Offset, request->Length);
} }
bytesXfer = NT_SUCCESS(request->Item->Status) ? request->Length : 0;
} }
else else
{ {
@@ -648,89 +583,30 @@ static VOID IoThreadProc (PVOID threadArg)
else else
request->Item->Status = TCCachedRead (queue, &ioStatus, request->Data, request->Offset, request->Length); request->Item->Status = TCCachedRead (queue, &ioStatus, request->Data, request->Offset, request->Length);
bytesXfer = (ULONG_PTR) ioStatus.Information;
// If short read on success, mark STATUS_END_OF_FILE but keep bytesXfer (true partial count)
if (NT_SUCCESS (request->Item->Status) && ioStatus.Information != request->Length) if (NT_SUCCESS (request->Item->Status) && ioStatus.Information != request->Length)
request->Item->Status = STATUS_END_OF_FILE; request->Item->Status = STATUS_END_OF_FILE;
} }
// For writes, we can account immediately.
// For reads, defer final accounting until after clamp/copy below.
if (request->Item->Write) {
request->ActualBytes = (ULONG)bytesXfer;
request->Item->BytesCompleted += bytesXfer;
} else {
// Provisional value; will be finalized after decrypt/copy clamp.
request->ActualBytes = (ULONG)bytesXfer;
}
} }
if (request->Item->Write) if (request->Item->Write)
{ {
queue->ReadAheadBufferValid = FALSE; queue->ReadAheadBufferValid = FALSE;
ReleaseFragmentBuffer(queue, request->Data); ReleaseFragmentBuffer (queue, request->Data);
if (request->CompleteOriginalIrp) if (request->CompleteOriginalIrp)
{ {
// call unified path to finish the original IRP HandleCompleteOriginalIrp(queue, request);
FinalizeOriginalIrp(queue, request->Item);
} }
// request itself is no longer needed ReleasePoolBuffer (queue, request);
ReleasePoolBuffer(queue, request);
} }
else else
{ {
BOOL readAhead = FALSE; BOOL readAhead = FALSE;
// Determine how many bytes we actually read for this fragment if (NT_SUCCESS (request->Item->Status))
ULONG copyLen = request->ActualBytes; memcpy (request->OrigDataBufferFragment, request->Data, request->Length);
// Decrypt in IO thread.
if (copyLen > 0 && request->EncryptedLength > 0)
{
UINT64_STRUCT dataUnit;
ULONG encStart = (ULONG)request->EncryptedOffset;
ULONG encInCopy = (copyLen > encStart)
? min((ULONG)copyLen - encStart, (ULONG)request->EncryptedLength)
: 0;
// Trim to full units (avoid returning ciphertext tails)
encInCopy -= (encInCopy % ENCRYPTION_DATA_UNIT_SIZE);
if (copyLen > encStart + encInCopy)
copyLen = encStart + encInCopy;
if (encInCopy > 0)
{
ASSERT((ULONG)request->EncryptedOffset + encInCopy <= copyLen);
dataUnit.Value = (request->Offset.QuadPart + request->EncryptedOffset) / ENCRYPTION_DATA_UNIT_SIZE;
if (queue->CryptoInfo->bPartitionInInactiveSysEncScope)
dataUnit.Value += queue->CryptoInfo->FirstDataUnitNo.Value;
else if (queue->RemapEncryptedArea)
dataUnit.Value += queue->RemappedAreaDataUnitOffset;
DecryptDataUnits(request->Data + request->EncryptedOffset,
&dataUnit,
encInCopy / ENCRYPTION_DATA_UNIT_SIZE,
queue->CryptoInfo);
}
}
// Finalize read accounting after clamp
bytesXfer = copyLen;
request->ActualBytes = (ULONG)bytesXfer;
request->Item->BytesCompleted += bytesXfer;
// Update subst sectors before copying
if (copyLen > 0 && (queue->SecRegionData != NULL) && (queue->SecRegionSize > 512))
{
UpdateBuffer(request->Data, queue->SecRegionData, queue->SecRegionSize,
request->Offset.QuadPart, copyLen, TRUE);
}
if (copyLen > 0)
memcpy(request->OrigDataBufferFragment, request->Data, copyLen);
ReleaseFragmentBuffer (queue, request->Data); ReleaseFragmentBuffer (queue, request->Data);
request->Data = request->OrigDataBufferFragment; request->Data = request->OrigDataBufferFragment;
@@ -744,16 +620,8 @@ static VOID IoThreadProc (PVOID threadArg)
InterlockedIncrement (&queue->OutstandingIoCount); InterlockedIncrement (&queue->OutstandingIoCount);
} }
if (request->CompleteOriginalIrp) ExInterlockedInsertTailList (&queue->CompletionThreadQueue, &request->CompletionListEntry, &queue->CompletionThreadQueueLock);
{ KeSetEvent (&queue->CompletionThreadQueueNotEmptyEvent, IO_DISK_INCREMENT, FALSE);
ExInterlockedInsertTailList (&queue->CompletionThreadQueue, &request->CompletionListEntry, &queue->CompletionThreadQueueLock);
KeSetEvent (&queue->CompletionThreadQueueNotEmptyEvent, IO_DISK_INCREMENT, FALSE);
}
else
{
// Non-final fragment: no need to involve the completion thread
ReleasePoolBuffer(queue, request);
}
if (readAhead) if (readAhead)
{ {
@@ -833,8 +701,8 @@ static VOID MainThreadProc (PVOID threadArg)
item->Queue = queue; item->Queue = queue;
item->OriginalIrp = irp; item->OriginalIrp = irp;
item->Status = STATUS_SUCCESS; item->Status = STATUS_SUCCESS;
item->BytesCompleted = 0;
IoSetCancelRoutine (irp, NULL);
if (irp->Cancel) if (irp->Cancel)
{ {
CompleteOriginalIrp (item, STATUS_CANCELLED, 0); CompleteOriginalIrp (item, STATUS_CANCELLED, 0);
@@ -922,10 +790,7 @@ static VOID MainThreadProc (PVOID threadArg)
} }
TCfree (buffer); TCfree (buffer);
CompleteOriginalIrp (item, item->Status, NT_SUCCESS (item->Status) ? item->OriginalLength : 0);
// Defer completion to avoid inline IoCompleteRequest re-entrancy
item->BytesCompleted = NT_SUCCESS(item->Status) ? item->OriginalLength : 0;
FinalizeOriginalIrp(queue, item);
continue; continue;
} }
@@ -1088,7 +953,7 @@ static VOID MainThreadProc (PVOID threadArg)
if (request->EncryptedLength > 0) if (request->EncryptedLength > 0)
{ {
UINT64_STRUCT dataUnit; UINT64_STRUCT dataUnit;
ASSERT ((ULONG)request->EncryptedOffset + request->EncryptedLength <= request->Length); ASSERT (request->EncryptedOffset + request->EncryptedLength <= request->Offset.QuadPart + request->Length);
dataUnit.Value = (request->Offset.QuadPart + request->EncryptedOffset) / ENCRYPTION_DATA_UNIT_SIZE; dataUnit.Value = (request->Offset.QuadPart + request->EncryptedOffset) / ENCRYPTION_DATA_UNIT_SIZE;
@@ -1227,42 +1092,16 @@ NTSTATUS EncryptedIoQueueResumeFromHold (EncryptedIoQueue *queue)
return STATUS_SUCCESS; return STATUS_SUCCESS;
} }
static VOID FreeCompletionWorkItemPool(EncryptedIoQueue* queue)
{
if (!queue->WorkItemPool) return;
// Drain the list for hygiene
KIRQL irql;
KeAcquireSpinLock(&queue->WorkItemLock, &irql);
while (!IsListEmpty(&queue->FreeWorkItemsList))
(void) RemoveHeadList(&queue->FreeWorkItemsList);
KeReleaseSpinLock(&queue->WorkItemLock, irql);
for (ULONG i = 0; i < queue->MaxWorkItems; ++i)
{
if (queue->WorkItemPool[i].WorkItem)
{
IoFreeWorkItem(queue->WorkItemPool[i].WorkItem);
queue->WorkItemPool[i].WorkItem = NULL;
}
}
TCfree(queue->WorkItemPool);
queue->WorkItemPool = NULL;
}
NTSTATUS EncryptedIoQueueStart (EncryptedIoQueue *queue) NTSTATUS EncryptedIoQueueStart (EncryptedIoQueue *queue)
{ {
NTSTATUS status; NTSTATUS status;
EncryptedIoQueueBuffer *buffer; EncryptedIoQueueBuffer *buffer;
int i, j, preallocatedIoRequestCount, preallocatedItemCount, fragmentSize; int i, j, preallocatedIoRequestCount, preallocatedItemCount, fragmentSize;
KIRQL oldIrql;
preallocatedIoRequestCount = EncryptionIoRequestCount; preallocatedIoRequestCount = EncryptionIoRequestCount;
preallocatedItemCount = EncryptionItemCount; preallocatedItemCount = EncryptionItemCount;
fragmentSize = EncryptionFragmentSize; fragmentSize = EncryptionFragmentSize;
// Clamp preallocation to a sane hard limit
if (preallocatedIoRequestCount > TC_ENC_IO_QUEUE_PREALLOCATED_IO_REQUEST_MAX_COUNT)
preallocatedIoRequestCount = TC_ENC_IO_QUEUE_PREALLOCATED_IO_REQUEST_MAX_COUNT;
queue->StartPending = TRUE; queue->StartPending = TRUE;
queue->ThreadExitRequested = FALSE; queue->ThreadExitRequested = FALSE;
@@ -1364,44 +1203,34 @@ retry_preallocated:
// Initialize the free work item list // Initialize the free work item list
InitializeListHead(&queue->FreeWorkItemsList); InitializeListHead(&queue->FreeWorkItemsList);
KeInitializeSemaphore(&queue->WorkItemSemaphore, EncryptionMaxWorkItems, EncryptionMaxWorkItems);
KeInitializeSpinLock(&queue->WorkItemLock); KeInitializeSpinLock(&queue->WorkItemLock);
KeInitializeEvent(&queue->WorkItemAvailableEvent, SynchronizationEvent, FALSE);
queue->MaxWorkItems = EncryptionMaxWorkItems; queue->MaxWorkItems = EncryptionMaxWorkItems;
// Enforce [1, VC_MAX_WORK_ITEMS] queue->WorkItemPool = (PCOMPLETE_IRP_WORK_ITEM)TCalloc(sizeof(COMPLETE_IRP_WORK_ITEM) * queue->MaxWorkItems);
if (queue->MaxWorkItems == 0) if (!queue->WorkItemPool)
queue->MaxWorkItems = 1;
if (queue->MaxWorkItems > VC_MAX_WORK_ITEMS)
queue->MaxWorkItems = VC_MAX_WORK_ITEMS;
{ {
// Two-phase allocation to avoid leaving stale LIST_ENTRYs on failure goto noMemory;
PCOMPLETE_IRP_WORK_ITEM pool; }
pool = (PCOMPLETE_IRP_WORK_ITEM)TCalloc(sizeof(COMPLETE_IRP_WORK_ITEM) * queue->MaxWorkItems); // Allocate and initialize work items
if (!pool) for (i = 0; i < (int) queue->MaxWorkItems; ++i)
goto noMemory; {
queue->WorkItemPool[i].WorkItem = IoAllocateWorkItem(queue->DeviceObject);
// Allocate all work items first if (!queue->WorkItemPool[i].WorkItem)
for (i = 0; i < (int) queue->MaxWorkItems; ++i)
{ {
pool[i].WorkItem = IoAllocateWorkItem(queue->DeviceObject); // Handle allocation failure
if (!pool[i].WorkItem) // Free previously allocated work items
for (j = 0; j < i; ++j)
{ {
for (j = 0; j < i; ++j) IoFreeWorkItem(queue->WorkItemPool[j].WorkItem);
IoFreeWorkItem(pool[j].WorkItem);
TCfree(pool);
goto noMemory;
} }
TCfree(queue->WorkItemPool);
goto noMemory;
} }
// Publish pool and insert entries into the free list under the spin lock // Insert the work item into the free list
queue->WorkItemPool = pool; ExInterlockedInsertTailList(&queue->FreeWorkItemsList, &queue->WorkItemPool[i].ListEntry, &queue->WorkItemLock);
KeAcquireSpinLock(&queue->WorkItemLock, &oldIrql);
for (i = 0; i < (int) queue->MaxWorkItems; ++i)
{
InsertTailList(&queue->FreeWorkItemsList, &queue->WorkItemPool[i].ListEntry);
}
KeReleaseSpinLock(&queue->WorkItemLock, oldIrql);
} }
queue->ActiveWorkItems = 0; queue->ActiveWorkItems = 0;
@@ -1476,9 +1305,6 @@ err:
FreePoolBuffers (queue); FreePoolBuffers (queue);
if (queue->WorkItemPool)
FreeCompletionWorkItemPool(queue);
queue->StartPending = FALSE; queue->StartPending = FALSE;
return status; return status;
} }
@@ -1497,42 +1323,32 @@ NTSTATUS EncryptedIoQueueStop (EncryptedIoQueue *queue)
Dump ("Queue stopping out=%d\n", queue->OutstandingIoCount); Dump ("Queue stopping out=%d\n", queue->OutstandingIoCount);
queue->ThreadExitRequested = TRUE; queue->ThreadExitRequested = TRUE;
// Wake all threads
KeSetEvent(&queue->MainThreadQueueNotEmptyEvent, IO_DISK_INCREMENT, FALSE);
KeSetEvent(&queue->IoThreadQueueNotEmptyEvent, IO_DISK_INCREMENT, FALSE);
KeSetEvent(&queue->CompletionThreadQueueNotEmptyEvent, IO_DISK_INCREMENT, FALSE);
// Wake any GetPoolBuffer waiters (defensive)
KeSetEvent(&queue->PoolBufferFreeEvent, IO_DISK_INCREMENT, FALSE);
TCStopThread (queue->MainThread, &queue->MainThreadQueueNotEmptyEvent); TCStopThread (queue->MainThread, &queue->MainThreadQueueNotEmptyEvent);
TCStopThread (queue->IoThread, &queue->IoThreadQueueNotEmptyEvent); TCStopThread (queue->IoThread, &queue->IoThreadQueueNotEmptyEvent);
TCStopThread (queue->CompletionThreads[0], &queue->CompletionThreadQueueNotEmptyEvent); TCStopThread (queue->CompletionThreads[0], &queue->CompletionThreadQueueNotEmptyEvent);
TCStopThread (queue->CompletionThreads[1], &queue->CompletionThreadQueueNotEmptyEvent); TCStopThread (queue->CompletionThreads[1], &queue->CompletionThreadQueueNotEmptyEvent);
/* // Wait for active work items to complete
* Wait for all completion work items to finish. KeResetEvent(&queue->NoActiveWorkItemsEvent);
* Dump("Queue stopping active work items=%d\n", queue->ActiveWorkItems);
* Work-item drain protocol: while (InterlockedCompareExchange(&queue->ActiveWorkItems, 0, 0) > 0)
* - queue->ActiveWorkItems is incremented by producers when queuing a completion work item,
* and queue->NoActiveWorkItemsEvent (NotificationEvent) is reset at that time
* (see FinalizeOriginalIrp and HandleCompleteOriginalIrp).
* - CompleteIrpWorkItemRoutine decrements ActiveWorkItems and sets
* NoActiveWorkItemsEvent when the count transitions 1 -> 0.
* - Waiters must not reset the notification event. They simply loop until the count is 0,
* waiting for the final 1 -> 0 transition to signal the event.
*/
Dump("Queue stopping, waiting for active work items (n=%ld)\n",
InterlockedExchangeAdd(&queue->ActiveWorkItems, 0)); // atomic read
for (;;)
{ {
LONG n = InterlockedExchangeAdd(&queue->ActiveWorkItems, 0); // atomic read
if (n == 0)
break;
KeWaitForSingleObject(&queue->NoActiveWorkItemsEvent, Executive, KernelMode, FALSE, NULL); KeWaitForSingleObject(&queue->NoActiveWorkItemsEvent, Executive, KernelMode, FALSE, NULL);
// reset the event again in case multiple work items are completing
KeResetEvent(&queue->NoActiveWorkItemsEvent);
} }
FreeCompletionWorkItemPool(queue); // Free pre-allocated work items
for (ULONG i = 0; i < queue->MaxWorkItems; ++i)
{
if (queue->WorkItemPool[i].WorkItem)
{
IoFreeWorkItem(queue->WorkItemPool[i].WorkItem);
queue->WorkItemPool[i].WorkItem = NULL;
}
}
TCfree(queue->WorkItemPool);
TCfree (queue->FragmentBufferA); TCfree (queue->FragmentBufferA);
TCfree (queue->FragmentBufferB); TCfree (queue->FragmentBufferB);

8
src/Driver/EncryptedIoQueue.h
View File

@@ -46,7 +46,6 @@ typedef struct _COMPLETE_IRP_WORK_ITEM
ULONG_PTR Information; ULONG_PTR Information;
void* Item; void* Item;
LIST_ENTRY ListEntry; // For managing free work items LIST_ENTRY ListEntry; // For managing free work items
BOOLEAN FromPool; // TRUE if taken from prealloc pool
} COMPLETE_IRP_WORK_ITEM, * PCOMPLETE_IRP_WORK_ITEM; } COMPLETE_IRP_WORK_ITEM, * PCOMPLETE_IRP_WORK_ITEM;
typedef struct typedef struct
@@ -141,13 +140,11 @@ typedef struct
PCOMPLETE_IRP_WORK_ITEM WorkItemPool; PCOMPLETE_IRP_WORK_ITEM WorkItemPool;
ULONG MaxWorkItems; ULONG MaxWorkItems;
LIST_ENTRY FreeWorkItemsList; LIST_ENTRY FreeWorkItemsList;
KSEMAPHORE WorkItemSemaphore;
KSPIN_LOCK WorkItemLock; KSPIN_LOCK WorkItemLock;
volatile LONG ActiveWorkItems; volatile LONG ActiveWorkItems;
KEVENT NoActiveWorkItemsEvent; KEVENT NoActiveWorkItemsEvent;
// signaled whenever a pooled work item returns to the free list
KEVENT WorkItemAvailableEvent;
} EncryptedIoQueue; } EncryptedIoQueue;
@@ -159,7 +156,6 @@ typedef struct
ULONG OriginalLength; ULONG OriginalLength;
LARGE_INTEGER OriginalOffset; LARGE_INTEGER OriginalOffset;
NTSTATUS Status; NTSTATUS Status;
ULONG_PTR BytesCompleted; // actual bytes transferred across all fragments
#ifdef TC_TRACE_IO_QUEUE #ifdef TC_TRACE_IO_QUEUE
LARGE_INTEGER OriginalIrpOffset; LARGE_INTEGER OriginalIrpOffset;
@@ -179,8 +175,6 @@ typedef struct
ULONG EncryptedLength; ULONG EncryptedLength;
uint8 *Data; uint8 *Data;
uint8 *OrigDataBufferFragment; uint8 *OrigDataBufferFragment;
ULONG ActualBytes; // actual bytes transferred for this fragment (0 on failure)
UCHAR WiRetryCount; // count how many times we failed to get a work item
LIST_ENTRY ListEntry; LIST_ENTRY ListEntry;
LIST_ENTRY CompletionListEntry; LIST_ENTRY CompletionListEntry;