Windows driver: overhaul IRP completion path in EncryptedIoQueue, add dual completion threads, precise byte accounting & safer UpdateBuffer

Major changes:
- Added pooled + elastic work item model with retry/backoff (MAX_WI_RETRIES). removed semaphore usage.
- Introduced two completion threads to reduce contention and latency under heavy IO.
- Added BytesCompleted (per IRP) and ActualBytes (per fragment) for correct short read/write accounting. total read/write stats now reflect real transferred bytes instead of requested length.
- Moved decryption of read fragments into IO thread. completion threads now only finalize IRPs (reduces race window and simplifies flow).
- Deferred final IRP completion via FinalizeOriginalIrp to avoid inline IoCompleteRequest re-entrancy. added safe OOM inline fallback.
- Implemented work item pool drain & orderly shutdown (ActiveWorkItems + NoActiveWorkItemsEvent) with robust stop protocol.
- Replaced semaphore-based work item acquisition with spin lock + free list + event (WorkItemAvailableEvent). added exponential backoff for transient exhaustion.
- Added elastic (on-demand) work item allocation with pool vs dynamic origin tracking (FromPool).
- Added FreeCompletionWorkItemPool() for symmetric cleanup; ensured all threads are explicitly awakened during stop.
- Added second completion thread replacing single CompletionThread.
- Hardened UpdateBuffer: fixed parameter name typo, added bounds/overflow checks using IntSafe (ULongLongAdd), validated Count, guarded sector end computation.
- Fixed GPT/system region write protection logic to pass correct length instead of end offset.
- Ensured ASSERTs use fragment‑relative bounds (cast + length) and avoided mixed 64/32 comparisons.
- Added MAX_WI_RETRIES constant. added WiRetryCount field in EncryptedIoRequest.
- Ensured RemoveLock is released only after all queue/accounting updates (OnItemCompleted).
- Reset/read-ahead logic preserved. read-ahead trigger now based on actual completion & zero pending fragment count.
- General refactoring, clearer separation of concerns (TryAcquireCompletionWorkItem / FinalizeOriginalIrp / HandleCompleteOriginalIrp).

Safety / correctness improvements:
- Accurate short read handling (STATUS_END_OF_FILE with true byte count).
- Eliminated risk of double free or premature RemoveLock release on completion paths.
- Prevented potential overflow in sector end arithmetic.
- Reduced contention and potential deadlock scenarios present with previous semaphore wait path.
This commit is contained in:
Mounir IDRASSI
2025-09-07 23:58:35 +09:00
parent 55adda7504
commit 062b385a69
2 changed files with 350 additions and 128 deletions

View File

@@ -46,6 +46,7 @@ typedef struct _COMPLETE_IRP_WORK_ITEM
ULONG_PTR Information;
void* Item;
LIST_ENTRY ListEntry; // For managing free work items
BOOLEAN FromPool; // TRUE if taken from prealloc pool
} COMPLETE_IRP_WORK_ITEM, * PCOMPLETE_IRP_WORK_ITEM;
typedef struct
@@ -86,8 +87,8 @@ typedef struct
KSPIN_LOCK IoThreadQueueLock;
KEVENT IoThreadQueueNotEmptyEvent;
// Completion thread
PKTHREAD CompletionThread;
// Completion threads
PKTHREAD CompletionThreads[2]; // two threads to handle completions
LIST_ENTRY CompletionThreadQueue;
KSPIN_LOCK CompletionThreadQueueLock;
KEVENT CompletionThreadQueueNotEmptyEvent;
@@ -140,11 +141,16 @@ typedef struct
PCOMPLETE_IRP_WORK_ITEM WorkItemPool;
ULONG MaxWorkItems;
LIST_ENTRY FreeWorkItemsList;
KSEMAPHORE WorkItemSemaphore;
KSPIN_LOCK WorkItemLock;
// Backoff (ms) when overflow work-item alloc fails. grows up to 32ms, reset on success
ULONG OverflowBackoffMs;
volatile LONG ActiveWorkItems;
KEVENT NoActiveWorkItemsEvent;
// signaled whenever a pooled work item returns to the free list
KEVENT WorkItemAvailableEvent;
} EncryptedIoQueue;
@@ -156,6 +162,7 @@ typedef struct
ULONG OriginalLength;
LARGE_INTEGER OriginalOffset;
NTSTATUS Status;
ULONG_PTR BytesCompleted; // actual bytes transferred across all fragments
#ifdef TC_TRACE_IO_QUEUE
LARGE_INTEGER OriginalIrpOffset;
@@ -175,6 +182,8 @@ typedef struct
ULONG EncryptedLength;
uint8 *Data;
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 CompletionListEntry;