added the use of an active logger count

git-svn-id: svn://svn.code.sf.net/p/processhacker/code@1632 21ef857c-d57f-4fe0-8362-d861dc6d29cd

trunk/KProcessHacker/include/sync.h

@@ -25,6 +25,44 @@
 
 #include "kph.h"
 
+/* General synchronization macros */
+
+/* KphEqualSpin
+ * 
+ * Spins until the first value is equal to the second 
+ * value.
+ */
+FORCEINLINE VOID KphSpinUntilEqual(
+    __inout PLONG Value,
+    __in LONG Value2
+    )
+{
+    while (InterlockedCompareExchange(
+        Value,
+        Value2,
+        Value2
+        ) != Value2)
+        PAUSE();
+}
+
+/* KphNotEqualSpin
+ * 
+ * Spins until the first value is not equal to the second 
+ * value.
+ */
+FORCEINLINE VOID KphSpinUntilNotEqual(
+    __inout PLONG Value,
+    __in LONG Value2
+    )
+{
+    while (InterlockedCompareExchange(
+        Value,
+        Value2,
+        Value2
+        ) == Value2)
+        PAUSE();
+}
+
 /* Spin Locks */
 
 /* KphAcquireBitSpinLock

trunk/KProcessHacker/include/sysservicep.h

@@ -82,6 +82,8 @@ typedef struct _KPHPSS_RESET_BLOCK
 
 #define TAG_EVENT_BLOCK ('BEhP')
 
+#define KPHPSS_EVENT_BLOCK_MAX_SIZE 0x200
+
 #define KPHPSS_EVENT_PROBE_ARGUMENTS_FAILED 0x00000001
 #define KPHPSS_EVENT_COPY_ARGUMENTS_FAILED 0x00000002
 #define KPHPSS_EVENT_KERNEL_MODE 0x00000004

trunk/KProcessHacker/sync.c

@@ -27,6 +27,13 @@ ULONG KphpCountBits(
     __in ULONG_PTR Number
     );
 
+VOID KphpProcessorLockDpc(
+    __in PKDPC Dpc,
+    __in PVOID DeferredContext,
+    __in PVOID SystemArgument1,
+    __in PVOID SystemArgument2
+    );
+
 /* KphfAcquireGuardedLock
  * 
  * Acquires a guarded lock and raises the IRQL to APC_LEVEL.
@@ -76,13 +83,6 @@ VOID FASTCALL KphfReleaseGuardedLock(
     KeLowerIrql(oldIrql);
 }
 
-VOID KphpProcessorLockDpc(
-    __in PKDPC Dpc,
-    __in PVOID DeferredContext,
-    __in PVOID SystemArgument1,
-    __in PVOID SystemArgument2
-    );
-
 /* KphAcquireProcessorLock
  * 
  * Raises the IRQL to DISPATCH_LEVEL and prevents threads from 
@@ -175,12 +175,7 @@ BOOLEAN KphAcquireProcessorLock(
             KeInsertQueueDpc(&ProcessorLock->Dpcs[i], ProcessorLock, NULL);
     
     /* Spinwait for all (other) processors to be acquired. */
-    while (InterlockedCompareExchange(
-        &ProcessorLock->AcquiredProcessors,
-        numberProcessors - 1,
-        numberProcessors - 1
-        ) != numberProcessors - 1)
-        NOTHING;
+    KphSpinUntilEqual(&ProcessorLock->AcquiredProcessors, numberProcessors - 1);
     
     dprintf("KphAcquireProcessorLock: All processors acquired.\n");
     ProcessorLock->Acquired = TRUE;
@@ -241,12 +236,7 @@ VOID KphReleaseProcessorLock(
     InterlockedExchange(&ProcessorLock->ReleaseSignal, 1);
     
     /* Spinwait for all acquired processors to be released. */
-    while (InterlockedCompareExchange(
-        &ProcessorLock->AcquiredProcessors,
-        0,
-        0
-        ))
-        NOTHING;
+    KphSpinUntilEqual(&ProcessorLock->AcquiredProcessors, 0);
     
     dprintf("KphReleaseProcessorLock: All processors released.\n");
     
@@ -313,12 +303,7 @@ VOID KphpProcessorLockDpc(
     InterlockedIncrement(&processorLock->AcquiredProcessors);
     
     /* Spin until we get the signal to release the processor. */
-    while (!InterlockedCompareExchange(
-        &processorLock->ReleaseSignal,
-        1,
-        1
-        ))
-        NOTHING;
+    KphSpinUntilNotEqual(&processorLock->ReleaseSignal, 0);
     
     /* Decrease the number of acquired processors. */
     InterlockedDecrement(&processorLock->AcquiredProcessors);

trunk/KProcessHacker/sysservice.c

@@ -31,21 +31,24 @@
 
 #include "include/sysservicep.h"
 #include "include/hook.h"
+#include "include/sync.h"
 #include "include/trace.h"
 
 extern PDRIVER_OBJECT KphDriverObject;
 
 FAST_MUTEX KphSsMutex;
 /* Whether system service logging has been initialized. */
-BOOLEAN KphSsInitialized;
+BOOLEAN KphSsInitialized = FALSE;
 /* The KiFastCallEntry hook. */
 KPH_HOOK KphSsKiFastCallEntryHook;
+/* The number of active loggers. */
+ULONG KphSsNumberOfActiveLoggers = 0;
 
 PKPH_OBJECT_TYPE KphSsClientEntryType;
 PKPH_OBJECT_TYPE KphSsProcessEntryType;
 
+FAST_MUTEX KphSsProcessListMutex;
 LIST_ENTRY KphSsProcessListHead;
-EX_RUNDOWN_REF KphSsRundownProtect;
 
 /* KphSsLogInit
  * 
@@ -58,6 +61,7 @@ NTSTATUS KphSsLogInit()
     /* Initialize the process list. */
     InitializeListHead(&KphSsProcessListHead);
     ExInitializeFastMutex(&KphSsMutex);
+    ExInitializeFastMutex(&KphSsProcessListMutex);
     
     /* Initialize the object types. */
     status = KphCreateObjectType(
@@ -104,9 +108,6 @@ NTSTATUS KphSsLogStart()
         return STATUS_UNSUCCESSFUL;
     }
     
-    /* (Re-)initialize rundown protection. */
-    ExInitializeRundownProtection(&KphSsRundownProtect);
-    
     /* Hook KiFastCallEntry. Logging will start from now. */
     KphInitializeHook(
         &KphSsKiFastCallEntryHook,
@@ -152,8 +153,8 @@ NTSTATUS KphSsLogStop()
         return status;
     }
     
-    /* Wait for all loggers to finish. */
-    ExWaitForRundownProtectionRelease(&KphSsRundownProtect);
+    /* Spin until the logger count reaches 0. */
+    KphSpinUntilEqual(&KphSsNumberOfActiveLoggers, 0);
     
     KphSsInitialized = FALSE;
     
@@ -165,7 +166,9 @@ NTSTATUS KphSsLogStop()
 /* KphSsCreateClientEntry
  * 
  * Creates a client entry which describes a client of the 
- * system service logger. Clients receieve system service log events.
+ * system service logger. Clients receieve system service log events. 
+ * Note that a client may have several process entries associated 
+ * with it.
  * 
  * ClientEntry: A variable which receives a pointer to the client entry.
  * ProcessHandle: A handle to the client process, with PROCESS_VM_WRITE 
@@ -363,9 +366,9 @@ NTSTATUS KphSsCreateProcessEntry(
     processEntry->TargetProcess = processObject;
     processEntry->Flags = Flags;
     
-    ExAcquireFastMutex(&KphSsMutex);
+    ExAcquireFastMutex(&KphSsProcessListMutex);
     InsertHeadList(&KphSsProcessListHead, &processEntry->ProcessListEntry);
-    ExReleaseFastMutex(&KphSsMutex);
+    ExReleaseFastMutex(&KphSsProcessListMutex);
     
     *ProcessEntry = processEntry;
     
@@ -386,9 +389,9 @@ VOID NTAPI KphpSsProcessEntryDeleteProcedure(
     
     KphDereferenceObject(processEntry->Client);
     
-    ExAcquireFastMutex(&KphSsMutex);
+    ExAcquireFastMutex(&KphSsProcessListMutex);
     RemoveEntryList(&processEntry->ProcessListEntry);
-    ExReleaseFastMutex(&KphSsMutex);
+    ExReleaseFastMutex(&KphSsProcessListMutex);
 }
 
 /* KphpSsCreateEventBlock
@@ -450,6 +453,10 @@ NTSTATUS KphpSsCreateEventBlock(
     traceSize = capturedFrames * sizeof(PVOID);
     eventBlockSize = sizeof(KPHPSS_EVENT_BLOCK) + argumentsSize + traceSize;
     
+    /* Check if the event block is too large. */
+    if (eventBlockSize > KPHPSS_EVENT_BLOCK_MAX_SIZE)
+        return STATUS_UNSUCCESSFUL;
+    
     /* Allocate the event block. */
     eventBlock = ExAllocatePoolWithTag(PagedPool, eventBlockSize, TAG_EVENT_BLOCK);
     
@@ -763,7 +770,7 @@ VOID NTAPI KphpSsLogSystemServiceCall(
         return;
     }
     
-    ExAcquireFastMutex(&KphSsMutex);
+    ExAcquireFastMutex(&KphSsProcessListMutex);
     
     currentListEntry = KphSsProcessListHead.Flink;
     processEntryCount = 0;
@@ -790,7 +797,7 @@ VOID NTAPI KphpSsLogSystemServiceCall(
         currentListEntry = currentListEntry->Flink;
     }
     
-    ExReleaseFastMutex(&KphSsMutex);
+    ExReleaseFastMutex(&KphSsProcessListMutex);
     
     /* If we didn't find any process entries, don't bother creating the 
      * event block.
@@ -899,13 +906,17 @@ __declspec(naked) VOID NTAPI KphpSsNewKiFastCallEntry()
         mov     cl, [ebx+eax] /* ecx = size of the arguments, in bytes. */
         shr     ecx, 2 /* divide by 2 to get the number of arguments (all ULONGs) */
         
-        /* Call the KiFastCallEntry proc. */
+        /* Call the KiFastCallEntry proc while maintaining the logger count 
+         * so that the driver doesn't get unloaded while we're executing.
+         */
         push    esi /* Thread */
         push    edi /* ServiceTable */
         push    ecx /* NumberOfArguments */
         push    edx /* Arguments */
         push    eax /* Number */
+        lock inc dword ptr KphSsNumberOfActiveLoggers
         call    KphpSsLogSystemServiceCall
+        lock dec dword ptr KphSsNumberOfActiveLoggers
         
         /* Restore the registers and resume execution in KiFastCallEntry. */
         pop     eax