VOID
NTAPI
ExInitializeCallBack(
- IN PEX_CALLBACK Callback
+ IN OUT PEX_CALLBACK Callback
+);
+
+PEX_CALLBACK_ROUTINE_BLOCK
+NTAPI
+ExAllocateCallBack(
+ IN PEX_CALLBACK_FUNCTION Function,
+ IN PVOID Context
+);
+
+VOID
+NTAPI
+ExFreeCallBack(
+ IN PEX_CALLBACK_ROUTINE_BLOCK CallbackRoutineBlock
+);
+
+BOOLEAN
+NTAPI
+ExCompareExchangeCallBack (
+ IN OUT PEX_CALLBACK CallBack,
+ IN PEX_CALLBACK_ROUTINE_BLOCK NewBlock,
+ IN PEX_CALLBACK_ROUTINE_BLOCK OldBlock
+);
+
+PEX_CALLBACK_ROUTINE_BLOCK
+NTAPI
+ExReferenceCallBackBlock(
+ IN OUT PEX_CALLBACK CallBack
+);
+
+VOID
+NTAPI
+ExDereferenceCallBackBlock(
+ IN OUT PEX_CALLBACK CallBack,
+ IN PEX_CALLBACK_ROUTINE_BLOCK CallbackRoutineBlock
+);
+
+PEX_CALLBACK_FUNCTION
+NTAPI
+ExGetCallBackBlockRoutine(
+ IN PEX_CALLBACK_ROUTINE_BLOCK CallbackRoutineBlock
+);
+
+PVOID
+NTAPI
+ExGetCallBackBlockContext(
+ IN PEX_CALLBACK_ROUTINE_BLOCK CallbackRoutineBlock
+);
+
+VOID
+NTAPI
+ExWaitForCallBacks(
+ IN PEX_CALLBACK_ROUTINE_BLOCK CallbackRoutineBlock
);
/* Rundown Functions ********************************************************/
return ExSystemExceptionFilter();
}
+/* CALLBACKS *****************************************************************/
+
+VOID
+FORCEINLINE
+ExDoCallBack(IN OUT PEX_CALLBACK Callback,
+ IN PVOID Context,
+ IN PVOID Argument1,
+ IN PVOID Argument2)
+{
+ PEX_CALLBACK_ROUTINE_BLOCK CallbackRoutineBlock;
+ PEX_CALLBACK_FUNCTION Function;
+
+ /* Reference the block */
+ CallbackRoutineBlock = ExReferenceCallBackBlock(Callback);
+ if (CallbackRoutineBlock)
+ {
+ /* Get the function */
+ Function = ExGetCallBackBlockRoutine(CallbackRoutineBlock);
+
+ /* Do the callback */
+ Function(Context, Argument1, Argument2);
+
+ /* Now dereference it */
+ ExDereferenceCallBackBlock(Callback, CallbackRoutineBlock);
+ }
+}
+
/* RUNDOWN *******************************************************************/
#ifdef _WIN64
IN BOOLEAN Create)
{
ULONG i;
- CLIENT_ID Cid = CurrentThread->Cid;
- /* Loop the notify routines */
- for (i = 0; i < PspThreadNotifyRoutineCount; i++)
+ /* Check if we have registered routines */
+ if (PspThreadNotifyRoutineCount)
{
- /* Call it */
- PspThreadNotifyRoutine[i](Cid.UniqueProcess, Cid.UniqueThread, Create);
+ /* Loop callbacks */
+ for (i = 0; i < PSP_MAX_CREATE_THREAD_NOTIFY; i++)
+ {
+ /* Do the callback */
+ ExDoCallBack(&PspThreadNotifyRoutine[i],
+ CurrentThread->Cid.UniqueProcess,
+ CurrentThread->Cid.UniqueThread,
+ (PVOID)(ULONG_PTR)Create);
+ }
}
}
IN BOOLEAN Create)
{
ULONG i;
- HANDLE ProcessId = (HANDLE)CurrentProcess->UniqueProcessId;
- HANDLE ParentId = CurrentProcess->InheritedFromUniqueProcessId;
- /* Loop the notify routines */
- for(i = 0; i < PSP_MAX_CREATE_PROCESS_NOTIFY; ++i)
+ /* Check if we have registered routines */
+ if (PspProcessNotifyRoutineCount)
{
- /* Make sure it exists */
- if(PspProcessNotifyRoutine[i])
+ /* Loop callbacks */
+ for (i = 0; i < PSP_MAX_CREATE_PROCESS_NOTIFY; i++)
{
- /* Call it */
- PspProcessNotifyRoutine[i](ParentId, ProcessId, Create);
+ /* Do the callback */
+ ExDoCallBack(&PspProcessNotifyRoutine[i],
+ CurrentProcess->InheritedFromUniqueProcessId,
+ (PVOID)(ULONG_PTR)Create,
+ NULL);
}
}
}
/* Loop the notify routines */
for (i = 0; i < PSP_MAX_LOAD_IMAGE_NOTIFY; ++ i)
{
- /* Make sure it exists */
- if (PspLoadImageNotifyRoutine[i])
- {
- /* Call it */
- PspLoadImageNotifyRoutine[i](FullImageName, ProcessId, ImageInfo);
- }
+ /* Do the callback */
+ ExDoCallBack(&PspLoadImageNotifyRoutine[i],
+ FullImageName,
+ ProcessId,
+ ImageInfo);
}
}
/* GLOBALS *******************************************************************/
-BOOLEAN PsImageNotifyEnabled = TRUE;
-ULONG PspThreadNotifyRoutineCount;
-PCREATE_THREAD_NOTIFY_ROUTINE
-PspThreadNotifyRoutine[PSP_MAX_CREATE_THREAD_NOTIFY];
-PCREATE_PROCESS_NOTIFY_ROUTINE
-PspProcessNotifyRoutine[PSP_MAX_CREATE_PROCESS_NOTIFY];
-PLOAD_IMAGE_NOTIFY_ROUTINE
-PspLoadImageNotifyRoutine[PSP_MAX_LOAD_IMAGE_NOTIFY];
+BOOLEAN PsImageNotifyEnabled = FALSE;
+ULONG PspThreadNotifyRoutineCount, PspProcessNotifyRoutineCount;
+ULONG PspLoadImageNotifyRoutineCount;
+EX_CALLBACK PspThreadNotifyRoutine[PSP_MAX_CREATE_THREAD_NOTIFY];
+EX_CALLBACK PspProcessNotifyRoutine[PSP_MAX_CREATE_PROCESS_NOTIFY];
+EX_CALLBACK PspLoadImageNotifyRoutine[PSP_MAX_LOAD_IMAGE_NOTIFY];
PLEGO_NOTIFY_ROUTINE PspLegoNotifyRoutine;
/* PUBLIC FUNCTIONS **********************************************************/
IN BOOLEAN Remove)
{
ULONG i;
+ PEX_CALLBACK_ROUTINE_BLOCK CallBack;
+ PAGED_CODE();
- /* Check if it's a removal or addition */
+ /* Check if we're removing */
if (Remove)
{
- /* Loop the routines */
- for(i = 0; i < PSP_MAX_CREATE_PROCESS_NOTIFY; i++)
+ /* Loop all the routines */
+ for (i = 0; i < PSP_MAX_CREATE_PROCESS_NOTIFY; i++)
{
- /* Check for a match */
- if (PspProcessNotifyRoutine[i] == NotifyRoutine)
+ /* Reference the callback block */
+ CallBack = ExReferenceCallBackBlock(&PspProcessNotifyRoutine[i]);
+ if (!CallBack) continue;
+
+ /* Check it this is a matching block */
+ if (ExGetCallBackBlockRoutine(CallBack) != (PVOID)NotifyRoutine)
+ {
+ /* It's not, try the next one */
+ continue;
+ }
+
+ /* It is, clear the current routine */
+ if (ExCompareExchangeCallBack(&PspProcessNotifyRoutine[i],
+ NULL,
+ CallBack))
{
- /* Remove and return */
- PspProcessNotifyRoutine[i] = NULL;
+ /* Decrement the number of routines */
+ InterlockedDecrement(&PspProcessNotifyRoutineCount);
+
+ /* Dereference the block */
+ ExDereferenceCallBackBlock(&PspProcessNotifyRoutine[i],
+ CallBack);
+
+ /* Wait for actice callbacks */
+ ExWaitForCallBacks(CallBack);
+
+ /* Free the callback and exit */
+ ExFreeCallBack (CallBack);
return STATUS_SUCCESS;
}
+
+ /* Dereference the block */
+ ExDereferenceCallBackBlock(&PspProcessNotifyRoutine[i],
+ CallBack);
}
+
+ /* We didn't find any matching block */
+ return STATUS_PROCEDURE_NOT_FOUND;
}
else
{
- /* Loop the routines */
- for(i = 0; i < PSP_MAX_CREATE_PROCESS_NOTIFY; i++)
+ /* Allocate a callback */
+ CallBack = ExAllocateCallBack((PVOID)NotifyRoutine, NULL);
+ if (!CallBack) return STATUS_INSUFFICIENT_RESOURCES;
+
+ /* Loop all callbacks */
+ for (i = 0; i < PSP_MAX_CREATE_PROCESS_NOTIFY; i++)
{
- /* Find an empty one */
- if (!PspProcessNotifyRoutine[i])
+ /* Add this routine if it's an empty slot */
+ if (ExCompareExchangeCallBack(&PspProcessNotifyRoutine[i],
+ CallBack,
+ NULL))
{
- /* Add it */
- PspProcessNotifyRoutine[i] = NotifyRoutine;
+ /* Found and inserted into an empty slot, return */
+ InterlockedIncrement(&PspProcessNotifyRoutineCount);
return STATUS_SUCCESS;
}
}
- }
- /* Nothing found */
- return Remove ? STATUS_PROCEDURE_NOT_FOUND : STATUS_INVALID_PARAMETER;
+ /* We didn't find a free slot, free the callback and fail */
+ ExFreeCallBack(CallBack);
+ return STATUS_INVALID_PARAMETER;
+ }
}
/*
PsRemoveLoadImageNotifyRoutine(IN PLOAD_IMAGE_NOTIFY_ROUTINE NotifyRoutine)
{
ULONG i;
+ PEX_CALLBACK_ROUTINE_BLOCK CallBack;
+ PAGED_CODE();
- /* Loop the routines */
- for(i = 0; i < PSP_MAX_LOAD_IMAGE_NOTIFY; i++)
+ /* Loop all callbacks */
+ for (i = 0; i < PSP_MAX_LOAD_IMAGE_NOTIFY; i++)
{
- /* Check for a match */
- if (PspLoadImageNotifyRoutine[i] == NotifyRoutine)
+ /* Reference this slot */
+ CallBack = ExReferenceCallBackBlock(&PspLoadImageNotifyRoutine[i]);
+ if (CallBack)
{
- /* Remove and return */
- PspLoadImageNotifyRoutine[i] = NULL;
- return STATUS_SUCCESS;
+ /* Check for a match */
+ if (ExGetCallBackBlockRoutine(CallBack) == (PVOID)NotifyRoutine)
+ {
+ /* Try removing it if it matches */
+ if (ExCompareExchangeCallBack(&PspLoadImageNotifyRoutine[i],
+ NULL,
+ CallBack))
+ {
+ /* We removed it, now dereference the block */
+ InterlockedDecrement(&PspLoadImageNotifyRoutineCount);
+ ExDereferenceCallBackBlock(&PspLoadImageNotifyRoutine[i],
+ CallBack);
+
+ /* Wait for active callbacks */
+ ExWaitForCallBacks(CallBack);
+
+ /* Free the callback and return */
+ ExFreeCallBack(CallBack);
+ return STATUS_SUCCESS;
+ }
+ }
+
+ /* Dereference the callback */
+ ExDereferenceCallBackBlock(&PspLoadImageNotifyRoutine[i], CallBack);
}
}
- /* Nothing found */
- return STATUS_INVALID_PARAMETER;
+ /* Nothing found to remove */
+ return STATUS_PROCEDURE_NOT_FOUND;
}
/*
PsSetLoadImageNotifyRoutine(IN PLOAD_IMAGE_NOTIFY_ROUTINE NotifyRoutine)
{
ULONG i;
+ PEX_CALLBACK_ROUTINE_BLOCK CallBack;
+ PAGED_CODE();
- /* Loop the routines */
+ /* Allocate a callback */
+ CallBack = ExAllocateCallBack((PVOID)NotifyRoutine, NULL);
+ if (!CallBack) return STATUS_INSUFFICIENT_RESOURCES;
+
+ /* Loop callbacks */
for (i = 0; i < PSP_MAX_LOAD_IMAGE_NOTIFY; i++)
{
- /* Find an empty one */
- if (!PspLoadImageNotifyRoutine[i])
+ /* Add this entry if the slot is empty */
+ if (ExCompareExchangeCallBack(&PspLoadImageNotifyRoutine[i],
+ CallBack,
+ NULL))
{
- /* Add it */
- PspLoadImageNotifyRoutine[i] = NotifyRoutine;
+ /* Return success */
+ InterlockedIncrement(&PspLoadImageNotifyRoutineCount);
+ PsImageNotifyEnabled = TRUE;
return STATUS_SUCCESS;
}
}
- /* Nothing found */
- return STATUS_INVALID_PARAMETER;
+ /* No free space found, fail */
+ ExFreeCallBack(CallBack);
+ return STATUS_INSUFFICIENT_RESOURCES;
}
/*
PsRemoveCreateThreadNotifyRoutine(IN PCREATE_THREAD_NOTIFY_ROUTINE NotifyRoutine)
{
ULONG i;
+ PEX_CALLBACK_ROUTINE_BLOCK CallBack;
+ PAGED_CODE();
- /* Loop the routines */
- for(i = 0; i < PSP_MAX_CREATE_THREAD_NOTIFY; i++)
+ /* Loop all callbacks */
+ for (i = 0; i < PSP_MAX_CREATE_THREAD_NOTIFY; i++)
{
- /* Check for a match */
- if (PspThreadNotifyRoutine[i] == NotifyRoutine)
+ /* Reference this slot */
+ CallBack = ExReferenceCallBackBlock(&PspThreadNotifyRoutine[i]);
+ if (CallBack)
{
- /* Remove and return */
- PspThreadNotifyRoutine[i] = NULL;
- return STATUS_SUCCESS;
+ /* Check for a match */
+ if (ExGetCallBackBlockRoutine(CallBack) == (PVOID)NotifyRoutine)
+ {
+ /* Try removing it if it matches */
+ if (ExCompareExchangeCallBack(&PspThreadNotifyRoutine[i],
+ NULL,
+ CallBack))
+ {
+ /* We removed it, now dereference the block */
+ InterlockedDecrement(&PspThreadNotifyRoutineCount);
+ ExDereferenceCallBackBlock(&PspThreadNotifyRoutine[i],
+ CallBack);
+
+ /* Wait for active callbacks */
+ ExWaitForCallBacks(CallBack);
+
+ /* Free the callback and return */
+ ExFreeCallBack(CallBack);
+ return STATUS_SUCCESS;
+ }
+ }
+
+ /* Dereference the callback */
+ ExDereferenceCallBackBlock(&PspThreadNotifyRoutine[i], CallBack);
}
}
- /* Nothing found */
- return STATUS_INVALID_PARAMETER;
+ /* Nothing found to remove */
+ return STATUS_PROCEDURE_NOT_FOUND;
}
/*
NTAPI
PsSetCreateThreadNotifyRoutine(IN PCREATE_THREAD_NOTIFY_ROUTINE NotifyRoutine)
{
- /* Make sure we didn't register too many */
- if (PspThreadNotifyRoutineCount >= PSP_MAX_CREATE_THREAD_NOTIFY)
+ ULONG i;
+ PEX_CALLBACK_ROUTINE_BLOCK CallBack;
+ PAGED_CODE();
+
+ /* Allocate a callback */
+ CallBack = ExAllocateCallBack((PVOID)NotifyRoutine, NULL);
+ if (!CallBack) return STATUS_INSUFFICIENT_RESOURCES;
+
+ /* Loop callbacks */
+ for (i = 0; i < PSP_MAX_CREATE_THREAD_NOTIFY; i++)
{
- /* Fail */
- return STATUS_INSUFFICIENT_RESOURCES;
+ /* Add this entry if the slot is empty */
+ if (ExCompareExchangeCallBack(&PspThreadNotifyRoutine[i],
+ CallBack,
+ NULL))
+ {
+ /* Return success */
+ InterlockedIncrement(&PspThreadNotifyRoutineCount);
+ return STATUS_SUCCESS;
+ }
}
- /* Register this one */
- PspThreadNotifyRoutine[PspThreadNotifyRoutineCount++] = NotifyRoutine;
- return STATUS_SUCCESS;
+ /* No free space found, fail */
+ ExFreeCallBack(CallBack);
+ return STATUS_INSUFFICIENT_RESOURCES;
}
/* EOF */