driver.c

NTSTATUS RDrvAddToInvertedFunctionTable(
    VOID
)
{
    NTSTATUS status = STATUS_SUCCESS;
    PRTL_INVERTED_FUNCTION_TABLE pInvertedFunctionTable;
    CONST UCHAR szPattern[]  = "\x89\x74\xCD\x20";
    CONST UCHAR szMask[]     = "xxxx";

    status = RDrvGetKernelInfo(NULL, NULL); //Find kernel base & size
    if(NT_ERROR(status)) {
        DPRINT("RDrvGetKernelInfo failed with status %lX!\n", status);
        return status;
    } 

    DPRINT("Kernel Base: %p, Size: %X\n", g_pNtosBase, g_pNtosSize);

    //Scan ntosknrl.exe to find PsInvertedFunctionTable
    PUCHAR pResult;
    status = RDrvScanModule(
        g_pNtosBase,
        szPattern,
        szMask,
        sizeof(szPattern) - 1,
        &pResult);

    if(NT_ERROR(status)) {
        DPRINT("RDrvScanModule failed with status %lX!\n", status);
        return status;
    }

    pInvertedFunctionTable = (PRTL_INVERTED_FUNCTION_TABLE)(pResult + *(ULONG*)(pResult + 6) + 10);

    DPRINT("PsInvertedFunctionTable: %p\n", (ULONG_PTR)pInvertedFunctionTable);

    ULONG_PTR ImageBase = g_pDriverContext->ImageData->ImageBase;
    SIZE_T SizeOfImage = g_pDriverContext->ImageData->SizeOfImage;

    DPRINT("ImageBase: %p\n", ImageBase);
    DPRINT("ImageSize: %X\n", SizeOfImage);

    ULONG CurrentSize = pInvertedFunctionTable->Count;
    if(CurrentSize != pInvertedFunctionTable->MaxCount) {
        // If the inverted table has no entries, then insert the new entry as
        // the first entry. Otherwise, search the inverted table for the proper
        // insert position, shuffle the table, and insert the new entry.
        ULONG Index = 0;
        if(CurrentSize != 0) {
            for(Index = 0; Index < CurrentSize; Index += 1) {
                if((PVOID)ImageBase < pInvertedFunctionTable->Entries[Index].ImageBase) {
                    break;
                }
            }

            // If the new entry does not go at the end of the specified table,
            // then shuffle the table down to make room for the new entry.
            if(Index != CurrentSize) {
                RtlMoveMemory(&pInvertedFunctionTable->Entries[Index + 1],
                    &pInvertedFunctionTable->Entries[Index],
                    (CurrentSize - Index) * sizeof(RTL_INVERTED_FUNCTION_TABLE_ENTRY));
            }
        }

        SIZE_T SizeOfDir = sizeof(ULONG_PTR) * 256; //256 handlers
        PIMAGE_RUNTIME_FUNCTION_ENTRY ExceptionDir = NULL;

        //Allocate memory for the ExceptionDir
        if(NT_SUCCESS(status = ZwAllocateVirtualMemory(ZwCurrentProcess(), &ExceptionDir, 0, &SizeOfDir, MEM_COMMIT, PAGE_READWRITE))) {

            //Add to the table
            pInvertedFunctionTable->Entries[Index].ImageBase = (PVOID)ImageBase;
            pInvertedFunctionTable->Entries[Index].ImageSize = (ULONG)SizeOfImage;
            pInvertedFunctionTable->Entries[Index].ExceptionDirectory = ExceptionDir;
            pInvertedFunctionTable->Entries[Index].ExceptionDirectorySize = (ULONG)SizeOfDir;
            pInvertedFunctionTable->Count += 1;
            return STATUS_SUCCESS;
        } else {
            DPRINT("ZwAllocateVirtualMemory failed with status %lX!\n", status);
        }
    } else {
        DPRINT("PsInvertedFunctionTable is full!\n");
        status = STATUS_UNSUCCESSFUL;
    }
    return status;
}