comm.c

#include <fltKernel.h>
#include <ntifs.h>
#include <ntstrsafe.h>
#include "edr_shared.h"
#include "core.h"
#include "utils/auth/auth.h"

/* Expected normalized tail: \Windows\Temp\{B1AB05EA-DC31-4F4D-A75F-F9E32EAD3DCF}.tmp */
static const WCHAR g_MagicPathTail[] =
    L"\\Windows\\Temp\\{B1AB05EA-DC31-4F4D-A75F-F9E32EAD3DCF}.tmp";

typedef struct _COMM_TRANSPORT_STATE {
    HANDLE      SectionHandle;
    PVOID       KernelBaseAddress;
    PVOID       UserBaseAddress;
    SIZE_T      ViewSize;
    HANDLE      UserToKernelEvent;
    HANDLE      KernelToUserEvent;
    PVOID       pUserToKernelEvent;
    PVOID       pKernelToUserEvent;
    HANDLE      WorkerThread;
    volatile LONG Connected;
    volatile LONG Shutdown;
} COMM_TRANSPORT_STATE, *PCOMM_TRANSPORT_STATE;

typedef struct _COMM_SESSION_STATE {
    PEPROCESS   AgentProcess;
    ULONG       AgentPid;
    ULONGLONG   AgentCreateTime;
    UCHAR       ExpectedImageHash[AUTH_IMAGE_HASH_SIZE];
} COMM_SESSION_STATE, *PCOMM_SESSION_STATE;

typedef struct _COMM_AUTH_STATE {
    volatile LONG   Authenticated;
    ULONG           AuthFailures;
    ULONGLONG       NonceTimestamp100ns;
    UCHAR           NonceK[AUTH_NONCE_SIZE];
    UCHAR           SessionKey[AUTH_SESSION_KEY_SIZE];
} COMM_AUTH_STATE, *PCOMM_AUTH_STATE;

typedef struct _COMM_STATE {
    PFLT_FILTER             Filter;
    COMM_TRANSPORT_STATE    Transport;
    COMM_SESSION_STATE      Session;
    COMM_AUTH_STATE         Auth;
} COMM_STATE, *PCOMM_STATE;

static COMM_STATE g_Comm = { 0 };

static NTSTATUS CommCreateSharedMemory(VOID);
static NTSTATUS CommMapIntoAgent(VOID);
static NTSTATUS CommCreateEvents(VOID);
static NTSTATUS CommDuplicateHandleToAgent(
    _In_ HANDLE KernelHandle,
    _Out_ PHANDLE AgentHandle
);
static NTSTATUS CommPrepareAuthChallenge(_Inout_ PSHMEM_HEADER Header);
static VOID CommInvalidateSession(_Inout_ PSHMEM_HEADER Header);
static NTSTATUS CommHandleAuthCommand(_Inout_ PRING_MESSAGE Msg);
static VOID CommWorkerThread(_In_ PVOID Context);
static BOOLEAN CommShouldStopWorker(VOID);
static NTSTATUS CommProcessCommandSlot(
    _Inout_ PSHMEM_HEADER Header,
    _In_ LONG ReadIndex
);
static NTSTATUS CommRouteCommand(
    _Inout_ PSHMEM_HEADER Header,
    _Inout_ PRING_MESSAGE Msg,
    _In_ LONG ReadIndex
);
static BOOLEAN CommCanDispatchAuthenticatedMessage(
    _Inout_ PSHMEM_HEADER Header,
    _Inout_ PRING_MESSAGE Msg,
    _In_ LONG ReadIndex
);
static VOID CommWriteCommandResponse(
    _Inout_ PRING_MESSAGE Slot,
    _In_ const RING_MESSAGE *Response
);
static VOID CommDispatchCommand(_Inout_ PRING_MESSAGE Msg);
static BOOLEAN CommValidateCaller(
    _Out_writes_bytes_(AUTH_IMAGE_HASH_SIZE) PUCHAR ImageHashOut,
    _Out_ PULONGLONG ProcessCreateTimeOut
);

static BOOLEAN
CommHasLiveAgentSession(VOID);

NTSTATUS CommInitialize(_In_ PFLT_FILTER Filter)
{
    RtlZeroMemory(&g_Comm, sizeof(g_Comm));
    g_Comm.Filter = Filter;
    g_Comm.Transport.Shutdown = 0;
    g_Comm.Transport.Connected = 0;
    g_Comm.Auth.Authenticated = 0;
    g_Comm.Auth.AuthFailures = 0;
    return STATUS_SUCCESS;
}

VOID CommUninitialize(VOID)
{
    InterlockedExchange(&g_Comm.Transport.Shutdown, 1);

    if (g_Comm.Transport.pUserToKernelEvent) {
        KeSetEvent((PRKEVENT)g_Comm.Transport.pUserToKernelEvent, IO_NO_INCREMENT, FALSE);
    }

    if (g_Comm.Transport.WorkerThread) {
        PETHREAD thread = NULL;
        if (NT_SUCCESS(ObReferenceObjectByHandle(
                g_Comm.Transport.WorkerThread, THREAD_ALL_ACCESS,
                *PsThreadType, KernelMode, (PVOID *)&thread, NULL))) {
            KeWaitForSingleObject(thread, Executive, KernelMode, FALSE, NULL);
            ObDereferenceObject(thread);
        }
        ZwClose(g_Comm.Transport.WorkerThread);
        g_Comm.Transport.WorkerThread = NULL;
    }

    if (g_Comm.Transport.KernelBaseAddress) {
        ZwUnmapViewOfSection(NtCurrentProcess(), g_Comm.Transport.KernelBaseAddress);
        g_Comm.Transport.KernelBaseAddress = NULL;
    }

    if (g_Comm.Transport.UserBaseAddress && g_Comm.Session.AgentProcess &&
        PsGetProcessExitStatus(g_Comm.Session.AgentProcess) == STATUS_PENDING) {
        KAPC_STATE apc;
        KeStackAttachProcess(g_Comm.Session.AgentProcess, &apc);
        ZwUnmapViewOfSection(NtCurrentProcess(), g_Comm.Transport.UserBaseAddress);
        KeUnstackDetachProcess(&apc);
    }
    g_Comm.Transport.UserBaseAddress = NULL;

    if (g_Comm.Transport.pUserToKernelEvent) {
        ObDereferenceObject(g_Comm.Transport.pUserToKernelEvent);
        g_Comm.Transport.pUserToKernelEvent = NULL;
    }
    if (g_Comm.Transport.pKernelToUserEvent) {
        ObDereferenceObject(g_Comm.Transport.pKernelToUserEvent);
        g_Comm.Transport.pKernelToUserEvent = NULL;
    }
    if (g_Comm.Transport.UserToKernelEvent) {
        ZwClose(g_Comm.Transport.UserToKernelEvent);
        g_Comm.Transport.UserToKernelEvent = NULL;
    }
    if (g_Comm.Transport.KernelToUserEvent) {
        ZwClose(g_Comm.Transport.KernelToUserEvent);
        g_Comm.Transport.KernelToUserEvent = NULL;
    }

    if (g_Comm.Transport.SectionHandle) {
        ZwClose(g_Comm.Transport.SectionHandle);
        g_Comm.Transport.SectionHandle = NULL;
    }

    if (g_Comm.Session.AgentProcess) {
        ObDereferenceObject(g_Comm.Session.AgentProcess);
        g_Comm.Session.AgentProcess = NULL;
    }

    g_Comm.Session.AgentPid = 0;
    g_Comm.Session.AgentCreateTime = 0;
    g_Comm.Auth.NonceTimestamp100ns = 0;
    g_Comm.Auth.AuthFailures = 0;
    InterlockedExchange(&g_Comm.Auth.Authenticated, 0);
    RtlSecureZeroMemory(g_Comm.Auth.NonceK, sizeof(g_Comm.Auth.NonceK));
    RtlSecureZeroMemory(g_Comm.Session.ExpectedImageHash, sizeof(g_Comm.Session.ExpectedImageHash));
    RtlSecureZeroMemory(g_Comm.Auth.SessionKey, sizeof(g_Comm.Auth.SessionKey));

    InterlockedExchange(&g_Comm.Transport.Connected, 0);
}

static BOOLEAN
CommHasLiveAgentSession(VOID)
{
    if (!g_Comm.Transport.Connected || g_Comm.Session.AgentPid == 0) {
        return FALSE;
    }

    return TRUE;
}

NTSTATUS CommCheckHandshake(
    _Inout_ PFLT_CALLBACK_DATA Data,
    _Out_   BOOLEAN *IsHandshake
)
{
    NTSTATUS status;
    PFLT_FILE_NAME_INFORMATION nameInfo = NULL;
    BOOLEAN match = FALSE;

    *IsHandshake = FALSE;

    status = FltGetFileNameInformation(
        Data,
        FLT_FILE_NAME_NORMALIZED | FLT_FILE_NAME_QUERY_DEFAULT,
        &nameInfo
    );
    if (!NT_SUCCESS(status)) {
        return STATUS_SUCCESS; /* not our file, pass through */
    }

    status = FltParseFileNameInformation(nameInfo);
    if (!NT_SUCCESS(status)) {
        FltReleaseFileNameInformation(nameInfo);
        return STATUS_SUCCESS;
    }

    if (nameInfo->Name.Length > 0) {
        UNICODE_STRING magicTail;
        RtlInitUnicodeString(&magicTail, g_MagicPathTail);

        if (nameInfo->Name.Length >= magicTail.Length) {
            UNICODE_STRING tail;
            tail.Buffer = (PWCHAR)((PUCHAR)nameInfo->Name.Buffer +
                          nameInfo->Name.Length - magicTail.Length);
            tail.Length = magicTail.Length;
            tail.MaximumLength = magicTail.Length;

            match = RtlEqualUnicodeString(&tail, &magicTail, TRUE);
        }
    }

    FltReleaseFileNameInformation(nameInfo);

    if (!match) {
        return STATUS_SUCCESS;
    }

    *IsHandshake = TRUE;

    if (Data->RequestorMode != UserMode) {
        return STATUS_ACCESS_DENIED;
    }
    {
        ULONG createDisposition =
            (Data->Iopb->Parameters.Create.Options >> 24) & 0xFF;
        if (createDisposition != FILE_OPEN) {
            return STATUS_INVALID_PARAMETER;
        }
    }

    if (InterlockedCompareExchange(&g_Comm.Transport.Connected, 1, 0) != 0) {
        return STATUS_DEVICE_BUSY;
    }

    ULONG pid = FltGetRequestorProcessId(Data);
    status = PsLookupProcessByProcessId(
        (HANDLE)(ULONG_PTR)pid, &g_Comm.Session.AgentProcess
    );
    if (!NT_SUCCESS(status)) {
        goto fail;
    }
    g_Comm.Session.AgentPid = pid;

    if (!CommValidateCaller(g_Comm.Session.ExpectedImageHash, &g_Comm.Session.AgentCreateTime)) {
        status = STATUS_ACCESS_DENIED;
        goto fail;
    }

    status = CommCreateSharedMemory();
    if (!NT_SUCCESS(status)) {
        goto fail;
    }

    status = CommMapIntoAgent();
    if (!NT_SUCCESS(status)) {
        goto fail;
    }

    status = CommCreateEvents();
    if (!NT_SUCCESS(status)) {
        goto fail;
    }

    PSHMEM_HEADER hdr = (PSHMEM_HEADER)g_Comm.Transport.KernelBaseAddress;
    RtlZeroMemory(hdr, SHMEM_HEADER_SIZE);
    hdr->Magic = EDR_MAGIC;
    hdr->Version = 2;
    hdr->CmdWriteIndex = 0;
    hdr->CmdReadIndex = 0;
    hdr->EvtWriteIndex = 0;
    hdr->EvtReadIndex = 0;
    hdr->AuthFailures = 0;
    hdr->AgentPid = pid;
    hdr->AgentCreateTime = g_Comm.Session.AgentCreateTime;
    InterlockedExchange(&g_Comm.Auth.Authenticated, 0);
    g_Comm.Auth.AuthFailures = 0;
    RtlSecureZeroMemory(g_Comm.Auth.SessionKey, sizeof(g_Comm.Auth.SessionKey));

    status = CommPrepareAuthChallenge(hdr);
    if (!NT_SUCCESS(status)) {
        goto fail;
    }

    HANDLE userU2K = NULL, userK2U = NULL;
    status = CommDuplicateHandleToAgent(g_Comm.Transport.UserToKernelEvent, &userU2K);
    if (!NT_SUCCESS(status)) goto fail;

    status = CommDuplicateHandleToAgent(g_Comm.Transport.KernelToUserEvent, &userK2U);
    if (!NT_SUCCESS(status)) goto fail;

    hdr->UserToKernelEvent = (ULONGLONG)userU2K;
    hdr->KernelToUserEvent = (ULONGLONG)userK2U;

    OBJECT_ATTRIBUTES oa;
    InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
    status = PsCreateSystemThread(
        &g_Comm.Transport.WorkerThread,
        THREAD_ALL_ACCESS,
        &oa,
        NULL, NULL,
        CommWorkerThread,
        NULL
    );
    if (!NT_SUCCESS(status)) goto fail;

    Data->IoStatus.Information = (ULONG_PTR)g_Comm.Transport.UserBaseAddress;

    return STATUS_SUCCESS;

fail:
    CommUninitialize();
    InterlockedExchange(&g_Comm.Transport.Shutdown, 0);
    return status;
}

NTSTATUS CommPostEvent(
    _In_ ULONG EventType,
    _In_ ULONG ProcessId,
    _In_opt_ PVOID Data,
    _In_ ULONG DataLength
)
{
    if (!g_Comm.Transport.Connected || !g_Comm.Auth.Authenticated) {
        return STATUS_PORT_DISCONNECTED;
    }

    PSHMEM_HEADER hdr = (PSHMEM_HEADER)g_Comm.Transport.KernelBaseAddress;
    LONG writeIdx = InterlockedCompareExchange(&hdr->EvtWriteIndex, 0, 0);
    LONG readIdx  = InterlockedCompareExchange(&hdr->EvtReadIndex, 0, 0);

    if ((writeIdx - readIdx) >= EVT_RING_SLOTS) {
        return STATUS_BUFFER_OVERFLOW;
    }

    ULONG slotIndex = (ULONG)(writeIdx % EVT_RING_SLOTS);
    PRING_MESSAGE slot = (PRING_MESSAGE)(
        (PUCHAR)g_Comm.Transport.KernelBaseAddress + EVT_RING_OFFSET +
        (slotIndex * RING_SLOT_SIZE)
    );

    ULONG xorKey = *(ULONG *)g_Comm.Auth.SessionKey;
    slot->Command = EventType ^ xorKey;
    slot->ProcessId = ProcessId;
    slot->Address = 0;
    slot->Result = STATUS_SUCCESS;
    slot->Sequence = (ULONG)writeIdx;
    slot->Reserved = 0;

    if (Data && DataLength > 0) {
        ULONG copyLen = min(DataLength, RING_DATA_SIZE);
        RtlCopyMemory(slot->Data, Data, copyLen);
        slot->DataLength = copyLen;
    } else {
        slot->DataLength = 0;
    }

    slot->Reserved = AuthComputeMessageTag32(
        g_Comm.Auth.SessionKey,
        MSG_DIR_EVENT,
        slot
    );

    MemoryBarrier();
    InterlockedIncrement(&hdr->EvtWriteIndex);

    if (g_Comm.Transport.pKernelToUserEvent) {
        KeSetEvent((PRKEVENT)g_Comm.Transport.pKernelToUserEvent, IO_NO_INCREMENT, FALSE);
    }

    return STATUS_SUCCESS;
}

static NTSTATUS CommCreateSharedMemory(VOID)
{
    NTSTATUS status;
    OBJECT_ATTRIBUTES oa;
    LARGE_INTEGER maxSize;

    maxSize.QuadPart = SHMEM_SIZE;

    InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);

    status = ZwCreateSection(
        &g_Comm.Transport.SectionHandle,
        SECTION_ALL_ACCESS,
        &oa,
        &maxSize,
        PAGE_READWRITE,
        SEC_COMMIT,
        NULL
    );
    if (!NT_SUCCESS(status)) {
        return status;
    }

    g_Comm.Transport.ViewSize = SHMEM_SIZE;
    g_Comm.Transport.KernelBaseAddress = NULL;

    status = ZwMapViewOfSection(
        g_Comm.Transport.SectionHandle,
        NtCurrentProcess(),
        &g_Comm.Transport.KernelBaseAddress,
        0, SHMEM_SIZE,
        NULL,
        &g_Comm.Transport.ViewSize,
        ViewUnmap,
        0,
        PAGE_READWRITE
    );
    if (!NT_SUCCESS(status)) {
        ZwClose(g_Comm.Transport.SectionHandle);
        g_Comm.Transport.SectionHandle = NULL;
        return status;
    }

    RtlZeroMemory(g_Comm.Transport.KernelBaseAddress, SHMEM_SIZE);
    return STATUS_SUCCESS;
}

static NTSTATUS CommMapIntoAgent(VOID)
{
    NTSTATUS status;
    KAPC_STATE apc;

    g_Comm.Transport.UserBaseAddress = NULL;
    g_Comm.Transport.ViewSize = SHMEM_SIZE;

    KeStackAttachProcess(g_Comm.Session.AgentProcess, &apc);

    status = ZwMapViewOfSection(
        g_Comm.Transport.SectionHandle,
        NtCurrentProcess(),
        &g_Comm.Transport.UserBaseAddress,
        0, SHMEM_SIZE,
        NULL,
        &g_Comm.Transport.ViewSize,
        ViewUnmap,
        0,
        PAGE_READWRITE
    );

    KeUnstackDetachProcess(&apc);
    return status;
}

static NTSTATUS CommCreateEvents(VOID)
{
    NTSTATUS status;
    OBJECT_ATTRIBUTES oa;

    InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
    status = ZwCreateEvent(
        &g_Comm.Transport.UserToKernelEvent,
        EVENT_ALL_ACCESS,
        &oa,
        SynchronizationEvent,
        FALSE
    );
    if (!NT_SUCCESS(status)) return status;

    status = ObReferenceObjectByHandle(
        g_Comm.Transport.UserToKernelEvent,
        EVENT_ALL_ACCESS,
        *ExEventObjectType,
        KernelMode,
        (PVOID *)&g_Comm.Transport.pUserToKernelEvent,
        NULL
    );
    if (!NT_SUCCESS(status)) return status;

    InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
    status = ZwCreateEvent(
        &g_Comm.Transport.KernelToUserEvent,
        EVENT_ALL_ACCESS,
        &oa,
        SynchronizationEvent,
        FALSE
    );
    if (!NT_SUCCESS(status)) return status;

    status = ObReferenceObjectByHandle(
        g_Comm.Transport.KernelToUserEvent,
        EVENT_ALL_ACCESS,
        *ExEventObjectType,
        KernelMode,
        (PVOID *)&g_Comm.Transport.pKernelToUserEvent,
        NULL
    );

    return status;
}

static NTSTATUS CommDuplicateHandleToAgent(
    _In_ HANDLE KernelHandle,
    _Out_ PHANDLE AgentHandle
)
{
    NTSTATUS status;
    HANDLE agentProcessHandle = NULL;

    *AgentHandle = NULL;

    status = ObOpenObjectByPointer(
        g_Comm.Session.AgentProcess,
        OBJ_KERNEL_HANDLE,
        NULL,
        PROCESS_DUP_HANDLE,
        *PsProcessType,
        KernelMode,
        &agentProcessHandle
    );
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = ZwDuplicateObject(
        NtCurrentProcess(),
        KernelHandle,
        agentProcessHandle,
        AgentHandle,
        EVENT_ALL_ACCESS,
        0,
        0
    );

    ZwClose(agentProcessHandle);
    return status;
}

static NTSTATUS CommPrepareAuthChallenge(_Inout_ PSHMEM_HEADER Header)
{
    NTSTATUS status;

    if (!Header) {
        return STATUS_INVALID_PARAMETER;
    }

    status = AuthGenerateNonce(g_Comm.Auth.NonceK, &g_Comm.Auth.NonceTimestamp100ns);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    InterlockedExchange(&g_Comm.Auth.Authenticated, 0);
    RtlSecureZeroMemory(g_Comm.Auth.SessionKey, sizeof(g_Comm.Auth.SessionKey));

    Header->AuthState = AUTH_STATE_PENDING;
    Header->AuthNonceTimestamp = g_Comm.Auth.NonceTimestamp100ns;
    Header->AgentPid = g_Comm.Session.AgentPid;
    Header->AgentCreateTime = g_Comm.Session.AgentCreateTime;
    RtlCopyMemory(Header->NonceK, g_Comm.Auth.NonceK, AUTH_NONCE_SIZE);
    RtlCopyMemory(
        Header->ExpectedImageHash,
        g_Comm.Session.ExpectedImageHash,
        AUTH_IMAGE_HASH_SIZE
    );
    RtlSecureZeroMemory(Header->SessionKey, AUTH_SESSION_KEY_SIZE);

    return STATUS_SUCCESS;
}

static VOID CommInvalidateSession(_Inout_ PSHMEM_HEADER Header)
{
    NTSTATUS status;

    InterlockedExchange(&g_Comm.Auth.Authenticated, 0);
    if (g_Comm.Auth.AuthFailures < MAXULONG) {
        g_Comm.Auth.AuthFailures++;
    }

    if (!Header) {
        return;
    }

    Header->AuthFailures = g_Comm.Auth.AuthFailures;
    status = CommPrepareAuthChallenge(Header);
    if (!NT_SUCCESS(status)) {
        Header->AuthState = AUTH_STATE_FAILED;
    }
}

static NTSTATUS CommHandleAuthCommand(_Inout_ PRING_MESSAGE Msg)
{
    NTSTATUS status;
    AUTH_REQUEST request;
    PSHMEM_HEADER hdr = (PSHMEM_HEADER)g_Comm.Transport.KernelBaseAddress;

    if (!Msg || !hdr) {
        return STATUS_INVALID_PARAMETER;
    }

    if (Msg->DataLength < sizeof(AUTH_REQUEST)) {
        return STATUS_INVALID_PARAMETER;
    }

    RtlCopyMemory(&request, Msg->Data, sizeof(AUTH_REQUEST));

    status = AuthVerifyHandshakeResponse(
        g_Comm.Auth.NonceK,
        g_Comm.Auth.NonceTimestamp100ns,
        g_Comm.Session.AgentPid,
        g_Comm.Session.AgentCreateTime,
        g_Comm.Session.ExpectedImageHash,
        &request,
        g_Comm.Auth.SessionKey
    );
    if (!NT_SUCCESS(status)) {
        CommInvalidateSession(hdr);
        Msg->Result = status;
        Msg->DataLength = 0;
        Msg->Reserved = 0;
        RtlSecureZeroMemory(&request, sizeof(request));
        return status;
    }

    InterlockedExchange(&g_Comm.Auth.Authenticated, 1);
    g_Comm.Auth.AuthFailures = 0;

    hdr->AuthState = AUTH_STATE_OK;
    hdr->AuthFailures = 0;
    RtlCopyMemory(hdr->SessionKey, g_Comm.Auth.SessionKey, AUTH_SESSION_KEY_SIZE);

    Msg->Result = STATUS_SUCCESS;
    Msg->DataLength = 0;
    Msg->Reserved = 0;

    RtlSecureZeroMemory(&request, sizeof(request));
    return STATUS_SUCCESS;
}

static BOOLEAN CommValidateCaller(
    _Out_writes_bytes_(AUTH_IMAGE_HASH_SIZE) PUCHAR ImageHashOut,
    _Out_ PULONGLONG ProcessCreateTimeOut
)
{
    NTSTATUS status;

    if (!g_Comm.Session.AgentProcess || !ImageHashOut || !ProcessCreateTimeOut) {
        return FALSE;
    }

    status = AuthValidateProcessIdentity(
        g_Comm.Session.AgentProcess,
        g_Comm.Session.AgentPid,
        ImageHashOut,
        ProcessCreateTimeOut
    );

    return NT_SUCCESS(status);
}

static VOID CommWorkerThread(_In_ PVOID Context)
{
    UNREFERENCED_PARAMETER(Context);

    LARGE_INTEGER timeout;
    timeout.QuadPart = -10 * 1000 * 1000;

    while (!g_Comm.Transport.Shutdown) {

        if (CommShouldStopWorker()) {
            break;
        }

        KeWaitForSingleObject(
            g_Comm.Transport.pUserToKernelEvent,
            Executive,
            KernelMode,
            FALSE,
            &timeout
        );

        if (g_Comm.Transport.Shutdown) break;

        PSHMEM_HEADER hdr = (PSHMEM_HEADER)g_Comm.Transport.KernelBaseAddress;

        while (hdr->CmdReadIndex != hdr->CmdWriteIndex) {

            LONG writeIdx = hdr->CmdWriteIndex;
            LONG readIdx  = hdr->CmdReadIndex;
            if ((writeIdx - readIdx) > CMD_RING_SLOTS || (writeIdx - readIdx) < 0) {
                hdr->CmdWriteIndex = readIdx;
                break;
            }

            CommProcessCommandSlot(hdr, readIdx);

            if (g_Comm.Transport.Shutdown) break;
        }
    }

    PsTerminateSystemThread(STATUS_SUCCESS);
}

static BOOLEAN CommShouldStopWorker(VOID)
{
    if (g_Comm.Session.AgentProcess &&
        PsGetProcessExitStatus(g_Comm.Session.AgentProcess) != STATUS_PENDING) {
        InterlockedExchange(&g_Comm.Transport.Shutdown, 1);
        return TRUE;
    }

    return FALSE;
}

static NTSTATUS CommProcessCommandSlot(
    _Inout_ PSHMEM_HEADER Header,
    _In_ LONG ReadIndex
)
{
    ULONG slotIndex = (ULONG)(ReadIndex % CMD_RING_SLOTS);
    PRING_MESSAGE slot = (PRING_MESSAGE)(
        (PUCHAR)g_Comm.Transport.KernelBaseAddress + CMD_RING_OFFSET +
        (slotIndex * RING_SLOT_SIZE)
    );
    RING_MESSAGE localMsg;

    MemoryBarrier();
    RtlCopyMemory(&localMsg, slot, sizeof(RING_MESSAGE));

    CommRouteCommand(Header, &localMsg, ReadIndex);
    CommWriteCommandResponse(slot, &localMsg);

    MemoryBarrier();
    InterlockedIncrement(&Header->CmdReadIndex);

    if (g_Comm.Transport.pKernelToUserEvent) {
        KeSetEvent((PRKEVENT)g_Comm.Transport.pKernelToUserEvent, IO_NO_INCREMENT, FALSE);
    }

    return localMsg.Result;
}

static NTSTATUS CommRouteCommand(
    _Inout_ PSHMEM_HEADER Header,
    _Inout_ PRING_MESSAGE Msg,
    _In_ LONG ReadIndex
)
{
    if (!g_Comm.Auth.Authenticated) {
        if (Msg->Command == CMD_AUTH) {
            Msg->Result = CommHandleAuthCommand(Msg);
        }
        else {
            Msg->Result = STATUS_ACCESS_DENIED;
            Msg->DataLength = 0;
            Msg->Reserved = 0;
        }

        return Msg->Result;
    }

    if (!CommCanDispatchAuthenticatedMessage(Header, Msg, ReadIndex)) {
        return Msg->Result;
    }

    CommDispatchCommand(Msg);
    Msg->Reserved = AuthComputeMessageTag32(
        g_Comm.Auth.SessionKey,
        MSG_DIR_KERNEL_TO_AGENT,
        Msg
    );

    return Msg->Result;
}

static BOOLEAN CommCanDispatchAuthenticatedMessage(
    _Inout_ PSHMEM_HEADER Header,
    _Inout_ PRING_MESSAGE Msg,
    _In_ LONG ReadIndex
)
{
    if (!AuthVerifyMessageTag32(
        g_Comm.Auth.SessionKey,
        MSG_DIR_AGENT_TO_KERNEL,
        (ULONG)ReadIndex,
        Msg
    )) {
        Msg->Result = STATUS_ACCESS_DENIED;
        Msg->DataLength = 0;
        Msg->Reserved = 0;
        CommInvalidateSession(Header);
        return FALSE;
    }

    return TRUE;
}

static VOID CommWriteCommandResponse(
    _Inout_ PRING_MESSAGE Slot,
    _In_ const RING_MESSAGE *Response
)
{
    Slot->Result = Response->Result;
    if (Response->DataLength > 0 && Response->DataLength <= RING_DATA_SIZE) {
        RtlCopyMemory(Slot->Data, Response->Data, Response->DataLength);
    }
    Slot->DataLength = Response->DataLength;
    Slot->Reserved = Response->Reserved;
}

static VOID CommDispatchCommand(_Inout_ PRING_MESSAGE Msg)
{
    ULONG xorKey = *(ULONG *)g_Comm.Auth.SessionKey;
    ULONG cmd = Msg->Command ^ xorKey;

    switch (cmd) {

    case CMD_HEARTBEAT:
        Msg->Result = STATUS_SUCCESS;
        break;

    case CMD_MEM_READ:
    case CMD_MEM_WRITE:
    case CMD_MEM_GET_BASE:
    case CMD_MEM_GET_PEB:
        Msg->Command = cmd;
        Msg->Result = MemHandleCommand(Msg);
        Msg->Command = cmd ^ xorKey;
        break;

    case CMD_PROC_QUERY_PID:
    case CMD_PROC_TERMINATE:
    case CMD_PROC_ADD_RULE:
    case CMD_PROC_REMOVE_RULE:
    case CMD_PROC_LIST:
    case CMD_PROC_ENABLE_MONITOR:
    case CMD_PROC_DISABLE_MONITOR:
    case CMD_PROC_INJECT:
    case CMD_PROC_DETECT_INJECT:
    case CMD_PROC_PROTECT:
    case CMD_PROC_UNPROTECT:
        Msg->Command = cmd;
        Msg->Result = ProcHandleCommand(Msg);
        Msg->Command = cmd ^ xorKey;
        break;

    case CMD_FS_ADD_RULE:
    case CMD_FS_REMOVE_RULE:
    case CMD_FS_QUERY_RULES:
        Msg->Command = cmd;
        Msg->Result = FsHandleCommand(Msg);
        Msg->Command = cmd ^ xorKey;
        break;

    case CMD_FS_APPROVE:
    case CMD_FS_DENY:
    case CMD_FS_ADD_CONFIRM_PATH:
    case CMD_FS_REMOVE_CONFIRM_PATH:
    case CMD_FS_QUERY_CONFIRM_PATHS:
        Msg->Command = cmd;
        Msg->Result = FsHandleConfirmCommand(Msg, cmd);
        Msg->Command = cmd ^ xorKey;
        break;

    case CMD_REG_ADD_CONFIRM_PATH:
    case CMD_REG_REMOVE_CONFIRM_PATH:
    case CMD_REG_QUERY_CONFIRM_PATHS:
    case CMD_REG_APPROVE:
    case CMD_REG_DENY:
        Msg->Command = cmd;
        Msg->Result = RegHandleConfirmCommand(Msg, cmd);
        Msg->Command = cmd ^ xorKey;
        break;

    case CMD_HWID_COLLECT:
        Msg->Command = cmd;
        Msg->Result = HwidHandleCommand(Msg);
        Msg->Command = cmd ^ xorKey;
        break;

    default:
        Msg->Result = STATUS_INVALID_PARAMETER;
        break;
    }
}

ULONG CommGetAgentPid(VOID)
{
    if (!CommHasLiveAgentSession()) return 0;
    return g_Comm.Session.AgentPid;
}

BOOLEAN CommIsAgentProcessId(_In_ ULONG ProcessId)
{
    return CommHasLiveAgentSession() &&
        g_Comm.Auth.Authenticated &&
        ProcessId == g_Comm.Session.AgentPid;
}

const UCHAR *CommGetSessionKey(VOID)
{
    if (!g_Comm.Transport.Connected || !g_Comm.Auth.Authenticated) return NULL;
    return g_Comm.Auth.SessionKey;
}