多核发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;

}