nt!PspSetCreateProcessNotifyRoutine - Windows 10 (15063.0.amd64fre.rs2_release.170317-1834).cpp

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;
  }
}