From f106fa892a55172e5bc6da11ac33d099398ddee8 Mon Sep 17 00:00:00 2001 From: donnaskiez Date: Sun, 12 May 2024 01:27:18 +1000 Subject: [PATCH] small io packet change --- driver/callbacks.c | 4 ++-- driver/common.h | 3 +++ driver/integrity.c | 4 ++-- driver/io.c | 54 +++++++++++++++++++++++++++++++++++++++++++--- driver/io.h | 4 ++-- driver/modules.c | 16 +++++++------- driver/pool.c | 2 +- driver/thread.c | 2 +- 8 files changed, 70 insertions(+), 19 deletions(-) diff --git a/driver/callbacks.c b/driver/callbacks.c index 9023a38..dcb7c97 100644 --- a/driver/callbacks.c +++ b/driver/callbacks.c @@ -747,7 +747,7 @@ ObPreOpCallbackRoutine(_In_ PVOID RegistrationContext, goto end; } - IrpQueueCompletePacket(report, report_size); + IrpQueueSchedulePacket(report, report_size); } end: @@ -942,7 +942,7 @@ EnumHandleCallback(_In_ PHANDLE_TABLE HandleTable, goto end; } - IrpQueueCompletePacket(report, report_size); + IrpQueueSchedulePacket(report, report_size); end: ExUnlockHandleTableEntry(HandleTable, Entry); diff --git a/driver/common.h b/driver/common.h index c1cc372..03a3063 100644 --- a/driver/common.h +++ b/driver/common.h @@ -191,6 +191,8 @@ typedef struct _DEFERRED_REPORTS_LIST { } DEFERRED_REPORTS_LIST, *PDEFERRED_REPORTS_LIST; +#define EVENT_COUNT 5 + typedef struct _IRP_QUEUE_HEAD { LIST_ENTRY queue; volatile UINT32 irp_count; @@ -200,6 +202,7 @@ typedef struct _IRP_QUEUE_HEAD { IO_CSQ csq; KSPIN_LOCK lock; DEFERRED_REPORTS_LIST deferred_reports; + KDPC dpc[EVENT_COUNT]; } IRP_QUEUE_HEAD, *PIRP_QUEUE_HEAD; diff --git a/driver/integrity.c b/driver/integrity.c index 7581d36..ff9ea14 100644 --- a/driver/integrity.c +++ b/driver/integrity.c @@ -892,7 +892,7 @@ ReportInvalidProcessModule(_In_ PPROCESS_MODULE_INFORMATION Module) return; } - IrpQueueCompletePacket(report, report_size); + IrpQueueSchedulePacket(report, report_size); } /* @@ -2238,7 +2238,7 @@ HeartbeatDpcRoutine(_In_ PKDPC Dpc, goto end; } - IrpQueueCompletePacket(packet, packet_size); + IrpQueueSchedulePacket(packet, packet_size); IncrementHeartbeatCounter(config); } diff --git a/driver/io.c b/driver/io.c index 6dca57d..6c809ea 100644 --- a/driver/io.c +++ b/driver/io.c @@ -292,6 +292,7 @@ IrpQueueDeferPacket(_In_ PIRP_QUEUE_HEAD Queue, * * IMPORTANT: All report buffers must be allocated in non paged memory. */ +STATIC NTSTATUS IrpQueueCompletePacket(_In_ PVOID Buffer, _In_ ULONG BufferSize) { @@ -334,6 +335,49 @@ IrpQueueCompletePacket(_In_ PVOID Buffer, _In_ ULONG BufferSize) return status; } +STATIC +VOID +IrpQueueSchedulePacketDpc(_In_ struct _KDPC* Dpc, + _In_opt_ PVOID DeferredContext, + _In_opt_ PVOID SystemArgument1, + _In_opt_ PVOID SystemArgument2) +{ + UNREFERENCED_PARAMETER(Dpc); + UNREFERENCED_PARAMETER(DeferredContext); + + if (!ARGUMENT_PRESENT(SystemArgument1) || + !ARGUMENT_PRESENT(SystemArgument2)) + return; + + PVOID buffer = SystemArgument1; + UINT32 buffer_length = (UINT32)SystemArgument2; + + IrpQueueCompletePacket(buffer, buffer_length); +} + +/* + * Not only does this allow reporting threads to continue execution once the + * report is scheduled (which in some cases such as handle reporting is very + * important for performance reasons), but it allows us to safely report from + * within a region guarded by a mutex as we use a spinlock here (for performance + * reasons, we dont need certain time critical threads being slept whilst we + * wait). + */ +VOID +IrpQueueSchedulePacket(_In_ PVOID Buffer, _In_ UINT32 BufferLength) +{ + PIRP_QUEUE_HEAD queue = GetIrpQueueHead(); + + /* Maybe not the best implementation, but 99.9999% of the time there should + * be a dpc available.*/ + while (TRUE) { + for (UINT32 index = 0; index < EVENT_COUNT; index++) { + if (KeInsertQueueDpc(&queue->dpc[index], Buffer, BufferLength)) + return; + } + } +} + STATIC VOID IrpQueueFreeDeferredPackets() @@ -364,6 +408,10 @@ IrpQueueInitialise() InitializeListHead(&queue->queue); InitializeListHead(&queue->deferred_reports.head); + for (UINT32 index = 0; index < EVENT_COUNT; index++) { + KeInitializeDpc(&queue->dpc[index], IrpQueueSchedulePacketDpc, NULL); + } + status = IoCsqInitialize(&queue->csq, IrpQueueInsert, IrpQueueRemove, @@ -1164,10 +1212,10 @@ DeviceCreate(_In_ PDEVICE_OBJECT DeviceObject, _Inout_ PIRP Irp) UNREFERENCED_PARAMETER(DeviceObject); DEBUG_INFO("Handle to driver opened."); - //NTSTATUS status = ValidatePciDevices(); + // NTSTATUS status = ValidatePciDevices(); - //if (!NT_SUCCESS(status)) - // DEBUG_ERROR("ValidatePciDevices failed with status %x", status); + // if (!NT_SUCCESS(status)) + // DEBUG_ERROR("ValidatePciDevices failed with status %x", status); IoCompleteRequest(Irp, IO_NO_INCREMENT); return Irp->IoStatus.Status; diff --git a/driver/io.h b/driver/io.h index 03cd7c4..5d9cc93 100644 --- a/driver/io.h +++ b/driver/io.h @@ -62,7 +62,7 @@ ValidateIrpInputBuffer(_In_ PIRP Irp, _In_ ULONG RequiredSize); NTSTATUS IrpQueueInitialise(); -NTSTATUS -IrpQueueCompletePacket(_In_ PVOID Buffer, _In_ ULONG BufferSize); +VOID +IrpQueueSchedulePacket(_In_ PVOID Buffer, _In_ UINT32 BufferLength); #endif \ No newline at end of file diff --git a/driver/modules.c b/driver/modules.c index 5adf452..74f76f7 100644 --- a/driver/modules.c +++ b/driver/modules.c @@ -372,7 +372,7 @@ ReportInvalidDriverObject(_In_ PDRIVER_OBJECT Driver, _In_ UINT32 ReportSubType) return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } FORCEINLINE @@ -595,7 +595,7 @@ ReportNmiBlocking() return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } STATIC @@ -632,7 +632,7 @@ ReportMissingCidTableEntry(_In_ PNMI_CONTEXT Context) return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } STATIC @@ -663,7 +663,7 @@ ReportInvalidRipFoundDuringNmi(_In_ PNMI_CONTEXT Context) return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } /* @@ -956,7 +956,7 @@ ReportApcStackwalkViolation(_In_ UINT64 Rip) return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } /* @@ -1292,7 +1292,7 @@ ReportDpcStackwalkViolation(_In_ PDPC_CONTEXT Context, _In_ UINT64 Frame) return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } STATIC @@ -1606,7 +1606,7 @@ ReportDataTableInvalidRoutine(_In_ TABLE_ID TableId, _In_ UINT64 Address) return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } NTSTATUS @@ -1945,7 +1945,7 @@ ReportWin32kBase_DxgInterfaceViolation(_In_ UINT32 TableIndex, return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } STATIC diff --git a/driver/pool.c b/driver/pool.c index 7b3b0bc..4dc4624 100644 --- a/driver/pool.c +++ b/driver/pool.c @@ -745,7 +745,7 @@ FindUnlinkedProcesses() continue; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } end: diff --git a/driver/thread.c b/driver/thread.c index 541eb3c..929f43d 100644 --- a/driver/thread.c +++ b/driver/thread.c @@ -129,7 +129,7 @@ DetectAttachedThreadsProcessCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry, return; } - IrpQueueCompletePacket(report, packet_size); + IrpQueueSchedulePacket(report, packet_size); } VOID