Instantly share code, notes, and snippets.

Embed
What would you like to do?
NTSTATUS __fastcall PspSetCreateProcessNotifyRoutine(PVOID NotifyRoutine, DWORD Flags)
{
BOOL bIsRemove;
BOOL bIsExRoutine;
DWORD UpperFlagBits;
DWORD LdrDataTableEntryFlags;
_EX_CALLBACK_ROUTINE_BLOCK *NewCallBackBlock;
_ETHREAD *CurrentThread;
size_t Index;
_EX_CALLBACK_ROUTINE_BLOCK *CallBackBlock;
// Copy over everything to UpperFlagBits from Flags, but bit 0.
UpperFlagBits = ((DWORD)Flags & 0xFFFFFFFE);
// Check if bit 1 is set. This will only be true if the caller is PsSetCreateProcessNotifyRoutineEx
// or PsSetCreateProcessNotifyRoutineEx2.
bIsExRoutine = (Flags & 2);
// Bit 0 will be set if "Remove" == TRUE from the caller.
bIsRemove = (Flags & 1);
// Bit 0 is set. We want to remove a callback.
if (bIsRemove)
{
// Disable APCs.
CurrentThread = (_ETHREAD *)KeGetCurrentThread();
--CurrentThread->Tcb.KernelApcDisable;
Index = 0;
while ( 1 )
{
CallBackBlock = ExReferenceCallBackBlock(&PspCreateProcessNotifyRoutine[Index]);
if ( CallBackBlock )
{
if ( /* Is it the same routine? */
ExGetCallBackBlockRoutine(CallBackBlock) == NotifyRoutine
/* Is it the same type? e.g. PsSetCreateProcessNotifyRoutineEx vs PsSetCreateProcessNotifyRoutineEx2 vs PsSetCreateProcessNotifyRoutine. */
&& (_DWORD)ExGetCallBackBlockContext(CallBackBlock) == (_DWORD)UpperFlagBits
/* Did we successfully NULL it out? */
&& ExCompareExchangeCallBack(&PspCreateProcessNotifyRoutine[Index], NULL, CallBackBlock) )
{
// Decrement global count.
if ( bIsExRoutine )
_InterlockedDecrement(&PspCreateProcessNotifyRoutineExCount);
else
_InterlockedDecrement(&PspCreateProcessNotifyRoutineCount);
ExDereferenceCallBackBlock(&PspCreateProcessNotifyRoutine[Index], CallBackBlock);
KiLeaveCriticalRegionUnsafe(CurrentThread);
ExWaitForCallBacks(CallBackBlock);
ExFreePoolWithTag(CallBackBlock, 0);
return STATUS_SUCCESS;
}
ExDereferenceCallBackBlock(&PspCreateProcessNotifyRoutine[Index], CallBackBlock);
}
Index++;
// Maximum callbacks == 64
if ( Index >= 0x40 )
{
KiLeaveCriticalRegionUnsafe(CurrentThread);
return STATUS_PROCEDURE_NOT_FOUND; // Could not find entry to remove.
}
}
}
else // We want to add a callback.
{
if ( bIsExRoutine )
LdrDataTableEntryFlags = 0x20; // "Ex" routine must have _KLDR_DATA_TABLE_ENTRY.IntegrityCheck bit set.
else
LdrDataTableEntryFlags = 0;
if ( !MmVerifyCallbackFunctionCheckFlags(NotifyRoutine, LdrDataTableEntryFlags) )
return STATUS_ACCESS_DENIED;
// Allocate new data structure.
NewCallBackBlock = ExAllocateCallBack(NotifyRoutine, UpperFlagBits);
if ( !NewCallBackBlock )
return STATUS_INSUFFICIENT_RESOURCES;
Index = 0;
while ( !ExCompareExchangeCallBack(&PspCreateProcessNotifyRoutine[Index], NewCallBackBlock, NULL) )
{
Index++;
if ( Index >= 0x40 )
{
// No space for callbacks.
ExFreePoolWithTag(NewCallBackBlock, 0);
return STATUS_INVALID_PARAMETER;
}
}
// Increment global counters.
if ( bIsExRoutine )
{
_InterlockedIncrement(&PspCreateProcessNotifyRoutineExCount);
if ( !(PspNotifyEnableMask & 4) )
_interlockedbittestandset(&PspNotifyEnableMask, 2u); // Have "Ex" callbacks.
}
else
{
_InterlockedIncrement(&PspCreateProcessNotifyRoutineCount);
if ( !(PspNotifyEnableMask & 2) )
_interlockedbittestandset(&PspNotifyEnableMask, 1u); // Have base-type of callbacks.
}
return STATUS_SUCCESS;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment