多核发dpc安全inline hook

VOID 
OpSafeInlineHook(PVOID TargetAddress, PVOID ReadyOpCode, ULONG OpCodeLength)
{
    PMDL MdlFuncAddress;

    ASSERT(TargetAddress && ReadyOpCode && OpCodeLength);

    if (ScmMapVirtualAddress(TargetAddress, 0x400, &MdlFuncAddress)) 
    {
        WPOFF();
        RtlCopyMemory(TargetAddress, ReadyOpCode, OpCodeLength);
        WPON();
        ScmUnmapVirtualAddress(MdlFuncAddress);
    }
}

VOID SafeHookDpcRoutine (
    __in struct _KDPC *Dpc,
    __in_opt PDPC_CONTEXT DeferredContext,
    __in_opt PVOID SystemArgument1,
    __in_opt PVOID SystemArgument2
    )
{
    InterlockedIncrement(&DeferredContext->LockedProcessors);
    do {
        __asm   pause;
    } while (DeferredContext->ReleaseFlag == FALSE);
    InterlockedDecrement(&DeferredContext->LockedProcessors);
}

BOOL ScHeSafeInlineHook(PVOID TargetAddress, PVOID ReadyOpCode, ULONG OpCodeLength)
{
    BOOL result = FALSE;
    DPC_CONTEXT DpcContext;
    KAFFINITY OrigAffinity;
    UNICODE_STRING NameString;
    CCHAR CurrentProcessor;
    CCHAR Processor;
    PKDPC Dpc;
    ULONG i;
    KIRQL OrigIrql;
    pFnKeSetAffinityThread KeSetAffinityThread = NULL;
    
    RtlInitUnicodeString(&NameString, L"KeSetAffinityThread");
    KeSetAffinityThread = (pFnKeSetAffinityThread)MmGetSystemRoutineAddress(&NameString);

    OrigAffinity = KeSetAffinityThread(KeGetCurrentThread(), 1); 
    OrigIrql = KeRaiseIrqlToDpcLevel();

    if (KeNumberProcessors > 1) {

        CurrentProcessor = (CCHAR)KeGetCurrentProcessorNumber();
        DpcContext.Dpcs = ExAllocatePoolWithTag(NonPagedPool, KeNumberProcessors * sizeof(KDPC), MEM_TAG);
        DpcContext.LockedProcessors = 1;
        DpcContext.ReleaseFlag = FALSE;

        for (Processor = 0; Processor < KeNumberProcessors; Processor++)
        {
            if (Processor == CurrentProcessor)  continue;
            Dpc = &DpcContext.Dpcs[Processor];
            KeInitializeDpc(Dpc, SafeHookDpcRoutine, &DpcContext);
            KeSetTargetProcessorDpc(Dpc, Processor);
            KeInsertQueueDpc(Dpc, NULL, NULL);
        }

        for (i = 0; i < 0x800000; i++) {
            __asm   pause;
            if (DpcContext.LockedProcessors == (ULONG)KeNumberProcessors) break;
        }
        
        if (DpcContext.LockedProcessors != (ULONG)KeNumberProcessors) {
            KdPrint(("[ScSafeInlineHook] Failed to insert dpc to other processors"));
            DpcContext.ReleaseFlag = TRUE;
            for (Processor = 0; Processor < KeNumberProcessors; Processor++) 
            {
                if (Processor != CurrentProcessor) {
                    KeRemoveQueueDpc(&DpcContext.Dpcs[Processor]);
                }
            }
        } else {
            KdPrint(("[ScSafeInlineHook] Insert dpc succeed, now start inline hook"));
            OpSafeInlineHook(TargetAddress, ReadyOpCode, OpCodeLength);
            result = TRUE;
            DpcContext.ReleaseFlag = TRUE;  
        }
        do {
            __asm   pause;
        } while (DpcContext.LockedProcessors != 1);

        ExFreePoolWithTag(DpcContext.Dpcs, MEM_TAG);

    } else {

        OpSafeInlineHook(TargetAddress, ReadyOpCode, OpCodeLength);
        result = TRUE;
    }
    KeLowerIrql(OrigIrql);
    KeSetAffinityThread(KeGetCurrentThread(), OrigAffinity); 
    return result;
}