reg.c

/*
 * reg.c — Registry monitoring and confirm-mode control.
 *
 * Uses CmRegisterCallbackEx to intercept registry operations.
 * Agent submits paths to a confirm watchlist; operations on those
 * paths are suspended until Agent approves or 5s timeout (deny).
 *
 * Mirrors the fs.c confirm watchlist design.
 */

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

#define REG_MAX_CONFIRM_PATHS   64
#define REG_CONFIRM_PATH_LEN    520

typedef struct _REG_CONFIRM_PATH {
    BOOLEAN     Active;
    WCHAR       Path[REG_CONFIRM_PATH_LEN];
    USHORT      PathLength;
} REG_CONFIRM_PATH, *PREG_CONFIRM_PATH;

static REG_CONFIRM_PATH  g_RegConfirmPaths[REG_MAX_CONFIRM_PATHS] = { 0 };
static EX_PUSH_LOCK      g_RegConfirmPathsLock;

#define REG_MAX_PENDING         32
#define REG_CONFIRM_TIMEOUT     (-5LL * 10 * 1000 * 1000)   /* 5 seconds */

typedef struct _REG_PENDING_OP {
    BOOLEAN     InUse;
    ULONG       RequestId;
    KEVENT      CompletionEvent;
    BOOLEAN     Approved;
} REG_PENDING_OP, *PREG_PENDING_OP;

static REG_PENDING_OP   g_RegPendingOps[REG_MAX_PENDING] = { 0 };
static KSPIN_LOCK        g_RegPendingLock;
static volatile LONG     g_RegNextRequestId = 1;

static LARGE_INTEGER     g_RegCookie = { 0 };
static volatile LONG     g_RegInitialized = 0;

static NTSTATUS RegCallback(
    _In_ PVOID CallbackContext,
    _In_opt_ PVOID Argument1,
    _In_opt_ PVOID Argument2
);
static BOOLEAN RegIsOurProcess(VOID);
static BOOLEAN RegMatchesConfirmPath(_In_ PUNICODE_STRING KeyPath);
static BOOLEAN RegPathMatchesPattern(
    _In_ PUNICODE_STRING Pattern,
    _In_ PUNICODE_STRING Path
);
static NTSTATUS RegAddConfirmPathLocked(
    _In_reads_(PathChars) PWCHAR Path,
    _In_ USHORT PathChars
);
static NTSTATUS RegRemoveConfirmPathLocked(
    _In_reads_(PathChars) PWCHAR Path,
    _In_ USHORT PathChars
);
static NTSTATUS RegApplyConfirmPathBatchLocked(
    _In_reads_bytes_(DataLength) PUCHAR Data,
    _In_ ULONG DataLength,
    _In_ BOOLEAN Add
);
static PREG_PENDING_OP RegAllocPendingOp(_Out_ PULONG RequestId);
static VOID RegFreePendingOp(_In_ PREG_PENDING_OP Op);
static NTSTATUS RegWaitForConfirm(
    _In_ ULONG OpType,
    _In_ PUNICODE_STRING KeyPath,
    _Out_ BOOLEAN *Approved
);
static NTSTATUS RegCompleteConfirm(_In_ ULONG RequestId, _In_ BOOLEAN Approve);

static BOOLEAN RegIsExemptProcessId(_In_ ULONG ProcessId);

NTSTATUS RegModuleInitialize(VOID)
{
    NTSTATUS status;
    UNICODE_STRING altitude = RTL_CONSTANT_STRING(L"370020.1");

    ExInitializePushLock(&g_RegConfirmPathsLock);
    RtlZeroMemory(g_RegConfirmPaths, sizeof(g_RegConfirmPaths));
    KeInitializeSpinLock(&g_RegPendingLock);
    RtlZeroMemory(g_RegPendingOps, sizeof(g_RegPendingOps));
    g_RegNextRequestId = 1;

    status = CmRegisterCallbackEx(
        RegCallback,
        &altitude,
        g_DriverObject,
        NULL,
        &g_RegCookie,
        NULL
    );
    if (!NT_SUCCESS(status)) return status;

    InterlockedExchange(&g_RegInitialized, 1);
    return STATUS_SUCCESS;
}

VOID RegModuleUninitialize(VOID)
{
    if (!InterlockedCompareExchange(&g_RegInitialized, 0, 1)) return;

    if (g_RegCookie.QuadPart != 0) {
        CmUnRegisterCallback(g_RegCookie);
        g_RegCookie.QuadPart = 0;
    }

    {
        KIRQL oldIrql;
        KeAcquireSpinLock(&g_RegPendingLock, &oldIrql);
        for (ULONG i = 0; i < REG_MAX_PENDING; i++) {
            if (g_RegPendingOps[i].InUse) {
                g_RegPendingOps[i].Approved = FALSE;
                KeSetEvent(&g_RegPendingOps[i].CompletionEvent,
                           IO_NO_INCREMENT, FALSE);
            }
        }
        KeReleaseSpinLock(&g_RegPendingLock, oldIrql);
    }
}

NTSTATUS RegHandleConfirmCommand(_Inout_ PRING_MESSAGE Msg, ULONG Cmd)
{
    switch (Cmd) {
    case CMD_REG_APPROVE: {
        if (Msg->DataLength < sizeof(ULONG)) return STATUS_INVALID_PARAMETER;
        ULONG reqId = *(PULONG)Msg->Data;
        return RegCompleteConfirm(reqId, TRUE);
    }
    case CMD_REG_DENY: {
        if (Msg->DataLength < sizeof(ULONG)) return STATUS_INVALID_PARAMETER;
        ULONG reqId = *(PULONG)Msg->Data;
        return RegCompleteConfirm(reqId, FALSE);
    }
    case CMD_REG_ADD_CONFIRM_PATH: {
        if (Msg->DataLength < sizeof(WCHAR) ||
            (Msg->DataLength % sizeof(WCHAR)) != 0) {
            return STATUS_INVALID_PARAMETER;
        }

        KeEnterCriticalRegion();
        ExAcquirePushLockExclusive(&g_RegConfirmPathsLock);

        NTSTATUS st = RegApplyConfirmPathBatchLocked(
            Msg->Data,
            Msg->DataLength,
            TRUE
        );

        ExReleasePushLockExclusive(&g_RegConfirmPathsLock);
        KeLeaveCriticalRegion();
        return st;
    }
    case CMD_REG_REMOVE_CONFIRM_PATH: {
        if (Msg->DataLength < sizeof(WCHAR) ||
            (Msg->DataLength % sizeof(WCHAR)) != 0) {
            return STATUS_INVALID_PARAMETER;
        }

        KeEnterCriticalRegion();
        ExAcquirePushLockExclusive(&g_RegConfirmPathsLock);

        NTSTATUS st = RegApplyConfirmPathBatchLocked(
            Msg->Data,
            Msg->DataLength,
            FALSE
        );

        ExReleasePushLockExclusive(&g_RegConfirmPathsLock);
        KeLeaveCriticalRegion();
        return st;
    }
    case CMD_REG_QUERY_CONFIRM_PATHS: {
        typedef struct { WCHAR Path[REG_CONFIRM_PATH_LEN]; } REG_PATH_ENTRY;
        ULONG maxEntries = RING_DATA_SIZE / sizeof(REG_PATH_ENTRY);
        REG_PATH_ENTRY *out = (REG_PATH_ENTRY *)Msg->Data;
        ULONG count = 0;

        KeEnterCriticalRegion();
        ExAcquirePushLockShared(&g_RegConfirmPathsLock);

        for (ULONG i = 0; i < REG_MAX_CONFIRM_PATHS && count < maxEntries; i++) {
            if (!g_RegConfirmPaths[i].Active) continue;
            RtlZeroMemory(out[count].Path, sizeof(out[count].Path));
            RtlCopyMemory(out[count].Path, g_RegConfirmPaths[i].Path,
                          g_RegConfirmPaths[i].PathLength);
            count++;
        }

        ExReleasePushLockShared(&g_RegConfirmPathsLock);
        KeLeaveCriticalRegion();

        Msg->DataLength = count * sizeof(REG_PATH_ENTRY);
        return STATUS_SUCCESS;
    }
    default:
        return STATUS_INVALID_PARAMETER;
    }
}

static BOOLEAN RegIsOurProcess(VOID)
{
    ULONG callerPid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId();
    return RegIsExemptProcessId(callerPid);
}

static BOOLEAN RegIsExemptProcessId(_In_ ULONG ProcessId)
{
    if (ProcessId == 4) {
        return TRUE;
    }

    return CommIsAgentProcessId(ProcessId);
}

static NTSTATUS RegGetObjectName(
    _In_ PVOID Object,
    _Out_ PUNICODE_STRING *Name
)
{
    NTSTATUS status;

    status = CmCallbackGetKeyObjectIDEx(
        &g_RegCookie, Object, NULL, Name, 0
    );
    return status;
}

static BOOLEAN RegMatchesConfirmPath(_In_ PUNICODE_STRING KeyPath)
{
    BOOLEAN matched = FALSE;

    KeEnterCriticalRegion();
    ExAcquirePushLockShared(&g_RegConfirmPathsLock);

    for (ULONG i = 0; i < REG_MAX_CONFIRM_PATHS; i++) {
        if (!g_RegConfirmPaths[i].Active) continue;

        UNICODE_STRING watchPath;
        watchPath.Buffer = g_RegConfirmPaths[i].Path;
        watchPath.Length = g_RegConfirmPaths[i].PathLength;
        watchPath.MaximumLength = sizeof(g_RegConfirmPaths[i].Path);

        if (RegPathMatchesPattern(&watchPath, KeyPath)) {
            matched = TRUE;
            break;
        }
    }

    ExReleasePushLockShared(&g_RegConfirmPathsLock);
    KeLeaveCriticalRegion();
    return matched;
}

static BOOLEAN RegPathMatchesPattern(
    _In_ PUNICODE_STRING Pattern,
    _In_ PUNICODE_STRING Path
)
{
    if (!Pattern || !Path || Pattern->Length == 0 || Path->Length == 0) {
        return FALSE;
    }

    if (RtlEqualUnicodeString(Pattern, Path, TRUE)) {
        return TRUE;
    }

    __try {
        return FsRtlIsNameInExpression(Pattern, Path, TRUE, NULL);
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        return FALSE;
    }
}

static NTSTATUS RegAddConfirmPathLocked(
    _In_reads_(PathChars) PWCHAR Path,
    _In_ USHORT PathChars
)
{
    USHORT pathBytes = (USHORT)(PathChars * sizeof(WCHAR));
    if (PathChars == 0 || pathBytes > (REG_CONFIRM_PATH_LEN - 1) * sizeof(WCHAR)) {
        return STATUS_BUFFER_OVERFLOW;
    }

    UNICODE_STRING target;
    target.Buffer = Path;
    target.Length = pathBytes;
    target.MaximumLength = pathBytes;

    for (ULONG i = 0; i < REG_MAX_CONFIRM_PATHS; i++) {
        if (!g_RegConfirmPaths[i].Active) continue;

        UNICODE_STRING existing;
        existing.Buffer = g_RegConfirmPaths[i].Path;
        existing.Length = g_RegConfirmPaths[i].PathLength;
        existing.MaximumLength = sizeof(g_RegConfirmPaths[i].Path);
        if (RtlEqualUnicodeString(&target, &existing, TRUE)) {
            return STATUS_SUCCESS;
        }
    }

    for (ULONG i = 0; i < REG_MAX_CONFIRM_PATHS; i++) {
        if (!g_RegConfirmPaths[i].Active) {
            RtlZeroMemory(g_RegConfirmPaths[i].Path, sizeof(g_RegConfirmPaths[i].Path));
            RtlCopyMemory(g_RegConfirmPaths[i].Path, Path, pathBytes);
            g_RegConfirmPaths[i].PathLength = pathBytes;
            g_RegConfirmPaths[i].Active = TRUE;
            return STATUS_SUCCESS;
        }
    }

    return STATUS_INSUFFICIENT_RESOURCES;
}

static NTSTATUS RegRemoveConfirmPathLocked(
    _In_reads_(PathChars) PWCHAR Path,
    _In_ USHORT PathChars
)
{
    USHORT pathBytes = (USHORT)(PathChars * sizeof(WCHAR));
    if (PathChars == 0) {
        return STATUS_INVALID_PARAMETER;
    }

    UNICODE_STRING target;
    target.Buffer = Path;
    target.Length = pathBytes;
    target.MaximumLength = pathBytes;

    BOOLEAN removed = FALSE;
    for (ULONG i = 0; i < REG_MAX_CONFIRM_PATHS; i++) {
        if (!g_RegConfirmPaths[i].Active) continue;

        UNICODE_STRING existing;
        existing.Buffer = g_RegConfirmPaths[i].Path;
        existing.Length = g_RegConfirmPaths[i].PathLength;
        existing.MaximumLength = sizeof(g_RegConfirmPaths[i].Path);
        if (RtlEqualUnicodeString(&target, &existing, TRUE)) {
            g_RegConfirmPaths[i].Active = FALSE;
            removed = TRUE;
        }
    }

    return removed ? STATUS_SUCCESS : STATUS_NOT_FOUND;
}

static NTSTATUS RegApplyConfirmPathBatchLocked(
    _In_reads_bytes_(DataLength) PUCHAR Data,
    _In_ ULONG DataLength,
    _In_ BOOLEAN Add
)
{
    if (!Data || DataLength < sizeof(WCHAR) || (DataLength % sizeof(WCHAR)) != 0) {
        return STATUS_INVALID_PARAMETER;
    }

    PWCHAR list = (PWCHAR)Data;
    ULONG totalChars = DataLength / sizeof(WCHAR);
    ULONG pos = 0;
    BOOLEAN sawEntry = FALSE;
    ULONG successCount = 0;
    NTSTATUS firstError = STATUS_SUCCESS;

    while (pos < totalChars) {
        ULONG start = pos;
        while (pos < totalChars && list[pos] != L'\0') {
            pos++;
        }

        ULONG lenChars = pos - start;
        if (lenChars > 0) {
            NTSTATUS st;
            sawEntry = TRUE;
            st = Add
                ? RegAddConfirmPathLocked(&list[start], (USHORT)lenChars)
                : RegRemoveConfirmPathLocked(&list[start], (USHORT)lenChars);

            if (NT_SUCCESS(st)) {
                successCount++;
            } else if (firstError == STATUS_SUCCESS) {
                firstError = st;
            }
        }

        if (pos < totalChars && list[pos] == L'\0') {
            pos++;
        }
    }

    if (!sawEntry) {
        return STATUS_INVALID_PARAMETER;
    }
    if (successCount > 0) {
        return STATUS_SUCCESS;
    }
    if (firstError != STATUS_SUCCESS) {
        return firstError;
    }
    return Add ? STATUS_INSUFFICIENT_RESOURCES : STATUS_NOT_FOUND;
}

static PREG_PENDING_OP RegAllocPendingOp(_Out_ PULONG RequestId)
{
    KIRQL oldIrql;
    KeAcquireSpinLock(&g_RegPendingLock, &oldIrql);

    for (ULONG i = 0; i < REG_MAX_PENDING; i++) {
        if (!g_RegPendingOps[i].InUse) {
            g_RegPendingOps[i].InUse = TRUE;
            g_RegPendingOps[i].Approved = FALSE;
            g_RegPendingOps[i].RequestId =
                (ULONG)InterlockedIncrement(&g_RegNextRequestId);
            KeInitializeEvent(&g_RegPendingOps[i].CompletionEvent,
                              SynchronizationEvent, FALSE);
            *RequestId = g_RegPendingOps[i].RequestId;
            KeReleaseSpinLock(&g_RegPendingLock, oldIrql);
            return &g_RegPendingOps[i];
        }
    }

    KeReleaseSpinLock(&g_RegPendingLock, oldIrql);
    *RequestId = 0;
    return NULL;
}

static VOID RegFreePendingOp(_In_ PREG_PENDING_OP Op)
{
    KIRQL oldIrql;
    KeAcquireSpinLock(&g_RegPendingLock, &oldIrql);
    Op->InUse = FALSE;
    KeReleaseSpinLock(&g_RegPendingLock, oldIrql);
}

static NTSTATUS RegCompleteConfirm(_In_ ULONG RequestId, _In_ BOOLEAN Approve)
{
    KIRQL oldIrql;
    KeAcquireSpinLock(&g_RegPendingLock, &oldIrql);

    for (ULONG i = 0; i < REG_MAX_PENDING; i++) {
        if (g_RegPendingOps[i].InUse &&
            g_RegPendingOps[i].RequestId == RequestId) {
            g_RegPendingOps[i].Approved = Approve;
            KeSetEvent(&g_RegPendingOps[i].CompletionEvent,
                       IO_NO_INCREMENT, FALSE);
            KeReleaseSpinLock(&g_RegPendingLock, oldIrql);
            return STATUS_SUCCESS;
        }
    }

    KeReleaseSpinLock(&g_RegPendingLock, oldIrql);
    return STATUS_NOT_FOUND;
}

static NTSTATUS RegWaitForConfirm(
    _In_ ULONG OpType,
    _In_ PUNICODE_STRING KeyPath,
    _Out_ BOOLEAN *Approved
)
{
    NTSTATUS status;
    ULONG requestId = 0;

    *Approved = FALSE;

    PREG_PENDING_OP op = RegAllocPendingOp(&requestId);
    if (!op) return STATUS_INSUFFICIENT_RESOURCES;

    ULONG pathBytes = (ULONG)min(KeyPath->Length,
        RING_DATA_SIZE - sizeof(REG_CONFIRM_REQUEST));
    ULONG totalLen = sizeof(REG_CONFIRM_REQUEST) + pathBytes;
    PUCHAR evtBuf = (PUCHAR)ExAllocatePool2(POOL_FLAG_NON_PAGED, totalLen, 'fCgR');
    if (!evtBuf) {
        RegFreePendingOp(op);
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    PREG_CONFIRM_REQUEST req = (PREG_CONFIRM_REQUEST)evtBuf;
    req->RequestId = requestId;
    req->OperationType = OpType;
    req->CallerPid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId();
    req->PathLength = pathBytes;
    RtlCopyMemory((PUCHAR)req + sizeof(REG_CONFIRM_REQUEST),
                   KeyPath->Buffer, pathBytes);

    status = CommPostEvent(
        EVT_REG_CONFIRM_REQ,
        req->CallerPid,
        evtBuf,
        totalLen
    );
    ExFreePoolWithTag(evtBuf, 'fCgR');
    if (!NT_SUCCESS(status)) {
        RegFreePendingOp(op);
        return status;
    }

    LARGE_INTEGER timeout;
    timeout.QuadPart = REG_CONFIRM_TIMEOUT;
    status = KeWaitForSingleObject(
        &op->CompletionEvent, Executive, KernelMode,
        FALSE, &timeout
    );

    if (status == STATUS_TIMEOUT) {
        *Approved = FALSE;
    }
    else if (NT_SUCCESS(status)) {
        *Approved = op->Approved;
    }
    else {
        *Approved = FALSE;
    }

    RegFreePendingOp(op);
    return STATUS_SUCCESS;
}

static NTSTATUS RegCallback(
    _In_ PVOID CallbackContext,
    _In_opt_ PVOID Argument1,
    _In_opt_ PVOID Argument2
)
{
    UNREFERENCED_PARAMETER(CallbackContext);

    if (!g_RegInitialized) return STATUS_SUCCESS;
    if (RegIsOurProcess()) return STATUS_SUCCESS;
    if (Argument1 == NULL || Argument2 == NULL) return STATUS_SUCCESS;

    REG_NOTIFY_CLASS notifyClass = (REG_NOTIFY_CLASS)(ULONG_PTR)Argument1;
    ULONG opType = 0;
    PVOID keyObject = NULL;

    switch (notifyClass) {
    case RegNtPreSetValueKey: {
        PREG_SET_VALUE_KEY_INFORMATION info =
            (PREG_SET_VALUE_KEY_INFORMATION)Argument2;
        keyObject = info->Object;
        opType = REG_OP_SET_VALUE;
        break;
    }
    case RegNtPreDeleteValueKey: {
        PREG_DELETE_VALUE_KEY_INFORMATION info =
            (PREG_DELETE_VALUE_KEY_INFORMATION)Argument2;
        keyObject = info->Object;
        opType = REG_OP_DELETE_VALUE;
        break;
    }
    case RegNtPreDeleteKey: {
        PREG_DELETE_KEY_INFORMATION info =
            (PREG_DELETE_KEY_INFORMATION)Argument2;
        keyObject = info->Object;
        opType = REG_OP_DELETE_KEY;
        break;
    }
    case RegNtPreRenameKey: {
        PREG_RENAME_KEY_INFORMATION info =
            (PREG_RENAME_KEY_INFORMATION)Argument2;
        keyObject = info->Object;
        opType = REG_OP_RENAME_KEY;
        break;
    }
    case RegNtPreCreateKeyEx: {
        PREG_CREATE_KEY_INFORMATION info =
            (PREG_CREATE_KEY_INFORMATION)Argument2;
        keyObject = info->RootObject;
        opType = REG_OP_CREATE_KEY;
        break;
    }
    default:
        return STATUS_SUCCESS;
    }

    if (!keyObject) return STATUS_SUCCESS;

    PUNICODE_STRING keyName = NULL;
    NTSTATUS status = RegGetObjectName(keyObject, &keyName);
    if (!NT_SUCCESS(status) || !keyName) return STATUS_SUCCESS;

    if (RegMatchesConfirmPath(keyName)) {
        BOOLEAN approved = FALSE;
        RegWaitForConfirm(opType, keyName, &approved);
        CmCallbackReleaseKeyObjectIDEx(keyName);
        if (!approved) {
            return STATUS_ACCESS_DENIED;
        }
        return STATUS_SUCCESS;
    }

    CmCallbackReleaseKeyObjectIDEx(keyName);
    return STATUS_SUCCESS;
}