NtSetInformationProcess_Nirvana.cpp

#include <Ntifs.h>

typedef enum _PROCESS_INFORMATION_CLASS {
    ProcessMemoryPriority,
    ProcessMemoryExhaustionInfo,
    ProcessAppMemoryInfo,
    ProcessInPrivateInfo,
    ProcessEDPStateInfo,
    ProcessInformationClassMax
} PROCESS_INFORMATION_CLASS;

BOOLEAN IsCanonical(PVOID address);
PVOID MmGetMaximumUserAddress();
BOOLEAN MmValidateUserCallTarget(PVOID, ULONG_PTR);
VOID FORCEINLINE PspLockProcessExclusive(IN PEPROCESS Process, IN PETHREAD CurrentThread);
VOID FORCEINLINE PspUnlockProcessExclusive(IN PEPROCESS Process, IN PETHREAD CurrentThread);

typedef struct _PROCESS_INSTRUMENTATION_CALLBACK_INFORMATION {
    ULONG Version;
    ULONG Reserved;
    PVOID Callback;
} PROCESS_INSTRUMENTATION_CALLBACK_INFORMATION;

typedef struct _EWOW64PROCESS {
    PVOID Peb;
    USHORT Machine;
} EWOW64PROCESS, *PEWOW64PROCESS;

#define ProcessInstrumentationCallback 40
#define PROCESS_TAG 'yQsP'

LUID SeDebugPrivilege;

NTSTATUS
NTAPI
NtSetInformationProcess(IN HANDLE ProcessHandle,
                        IN PROCESS_INFORMATION_CLASS ProcessInformationClass,
                        IN PVOID ProcessInformation,
                        IN ULONG ProcessInformationLength) {
    ULONG InfoSize = ProcessInformationLength;

    PETHREAD pThread = PsGetCurrentThread();
    KPROCESSOR_MODE AccessMode = ExGetPreviousMode();

    switch ( ProcessInformationClass ) {
        case ProcessInstrumentationCallback:
        {
            if ( (InfoSize - 8) & 0xFFFFFFF7 )
                return STATUS_INFO_LENGTH_MISMATCH;

            NTSTATUS status = STATUS_SUCCESS;
            PROCESS_INSTRUMENTATION_CALLBACK_INFORMATION callback_info = { 0 };
            PVOID cb_address = NULL;

            if ( InfoSize == 8 ) {
                cb_address = callback_info.Callback = (PVOID)*(PULONGLONG)ProcessInformation;
            } else {
                callback_info = *(PROCESS_INSTRUMENTATION_CALLBACK_INFORMATION*)ProcessInformation;
                cb_address = callback_info.Callback;
            }

            if ( callback_info.Reserved )
                return STATUS_INVALID_PARAMETER;

            if ( callback_info.Version != callback_info.Reserved )
                return STATUS_UNKNOWN_REVISION;

            if ( !IsCanonical(cb_address) )
                return STATUS_INVALID_PARAMETER;

            PVOID pObject = NULL;
            status = ObReferenceObjectByHandleWithTag(ProcessHandle,
                                                      0x200,
                                                      *PsProcessType,
                                                      AccessMode,
                                                      PROCESS_TAG,
                                                      &pObject,
                                                      NULL);
            if ( !NT_SUCCESS(status) )
                return status;

            PEPROCESS pCurrentProcess = PsGetCurrentProcess();
            BOOLEAN has_privilege = SeSinglePrivilegeCheck(SeDebugPrivilege, AccessMode);
            PEPROCESS pTargetProcess = (PEPROCESS)pObject;

            if ( !has_privilege && pObject != pCurrentProcess ) {
                ObfDereferenceObjectWithTag(pObject, PROCESS_TAG);
                return STATUS_PRIVILEGE_NOT_HELD;
            }

            if ( !ExAcquireRundownProtection(pTargetProcess->RundownProtect) ) {
                ObfDereferenceObjectWithTag(pTargetProcess, PROCESS_TAG);
                return STATUS_PROCESS_IS_TERMINATING;
            }

            if ( pTargetProcess->WoW64Process ) {
                if ( pCurrentProcess->WoW64Process ) {
                    KAPC_STATE apc_state = { 0 };
                    KeStackAttachProcess(pTargetProcess, &apc_state);

                    if ( callback_info.Callback < MmGetMaximumUserAddress() &&
                         MmValidateUserCallTarget(callback_info.Callback, 1i64) ) {
                        PVOID Peb = NULL;
                        PEWOW64PROCESS wowProcess = (PEWOW64PROCESS)pTargetProcess->WoW64Process;

                        if ( wowProcess )
                            Peb = wowProcess->Peb;

                        Peb[1].Ldr = callback_info.Callback;  // WTF?
                    } else {
                        status = STATUS_INVALID_PARAMETER;
                    }

                    KeUnstackDetachProcess(&apc_state);
LABEL_135:
                    ExReleaseRundownProtection(&pTargetProcess->RundownProtect);
LABEL_136:
                    ObfDereferenceObjectWithTag(pTargetProcess, PROCESS_TAG);
                    return status;
                }
            } else if ( !pCurrentProcess->WoW64Process ) {
                KAPC_STATE apc_state = { 0 };
                KeStackAttachProcess(pTargetProcess, &apc_state);
                PVOID Callback = callback_info.Callback;

                if ( !MmValidateUserCallTarget(callback_info.Callback, 0i64) )
                    status = STATUS_INVALID_PARAMETER;

                KeUnstackDetachProcess(&apc_state);

                if ( NT_SUCCESS(status) ) {
                    PspLockProcessExclusive(pTargetProcess, pThread);

                    pTargetProcess->Pcb.InstrumentationCallback = Callback;

                    PLIST_ENTRY Head = &pTargetProcess->ThreadListHead;
                    PLIST_ENTRY Next = Head->Flink;

                    // walks a list of process threads
                    // yeah, CONTAINING_RECORD
                    while ( Next != Head ) {
                        if ( Callback ) {
                            // set _ETHREAD->_KTHREAD->_DISPATCHER_HEADER->Instrumented
                            _interlockedbittestandset(Next - 420, 0x19u);
                        } else {
                            // unset _ETHREAD->_KTHREAD->_DISPATCHER_HEADER->Instrumented
                            _interlockedbittestandreset(Next - 420, 0x19u);
                        }

                        Next = Next->Flink;
                    }

                    PspUnlockProcessExclusive(pTargetProcess, pThread);
                }

                goto LABEL_135;
            }

            status = STATUS_NOT_SUPPORTED;
            goto LABEL_135;

            break;
        }
    }
}