RunDown 保护-ExAcquireRundownProtection

#include <ntddk.h>

// 定义全局结构，包含内存指针和 Rundown Protection
typedef struct _PROTECTED_MEMORY {
    EX_RUNDOWN_REF RundownProtection;  // Rundown 保护结构
    PVOID          pMemoryBuffer;      // 受保护的内存区域
    ULONG          MemorySize;         // 内存大小（示例用途）
} PROTECTED_MEMORY, *PPROTECTED_MEMORY;

PROTECTED_MEMORY g_ProtectedMemory;    // 全局受保护内存对象

// 自定义上下文（用于调试跟踪）
typedef struct _MEMORY_ACCESS_CONTEXT {
    ULONG          ThreadId;
    LARGE_INTEGER  AccessTime;
} MEMORY_ACCESS_CONTEXT, *PMEMORY_ACCESS_CONTEXT;

// 初始化受保护的内存区域
NTSTATUS InitializeProtectedMemory()
{
    // 分配非分页内存（示例分配 4096 字节）
    g_ProtectedMemory.pMemoryBuffer = ExAllocatePoolWithTag(
        NonPagedPool,
        4096,
        'PMem'  // 内存标签（自定义四字符标识）
    );

    if (!g_ProtectedMemory.pMemoryBuffer) {
        KdPrint(("Failed to allocate memory!\n"));
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    // 初始化 Rundown Protection
    ExInitializeRundownProtection(&g_ProtectedMemory.RundownProtection);
    g_ProtectedMemory.MemorySize = 4096;
    return STATUS_SUCCESS;
}

// 安全释放内存
VOID SafeFreeProtectedMemory()
{
    // 等待所有线程释放 Rundown 保护
    KdPrint(("Waiting for all memory accesses to complete...\n"));
    ExWaitForRundownProtectionRelease(&g_ProtectedMemory.RundownProtection);

    // 安全释放内存
    if (g_ProtectedMemory.pMemoryBuffer) {
        ExFreePoolWithTag(g_ProtectedMemory.pMemoryBuffer, 'PMem');
        g_ProtectedMemory.pMemoryBuffer = NULL;
        KdPrint(("Memory released successfully.\n"));
    }
}

// 带上下文的获取内存访问权限
BOOLEAN AcquireMemoryAccess(PMEMORY_ACCESS_CONTEXT Context)
{
    // 填充上下文信息（调试用）
    Context->ThreadId = HandleToULong(PsGetCurrentThreadId());
    KeQuerySystemTimePrecise(&Context->AccessTime);

    // 尝试获取 Rundown 保护
    return ExAcquireRundownProtectionEx(
        &g_ProtectedMemory.RundownProtection,
        Context
    );
}

// 释放内存访问权限
VOID ReleaseMemoryAccess(PMEMORY_ACCESS_CONTEXT Context)
{
    ExReleaseRundownProtectionEx(
        &g_ProtectedMemory.RundownProtection,
        Context
    );
}

// 模拟线程访问内存的操作
VOID MemoryAccessThreadRoutine(PVOID StartContext)
{
    MEMORY_ACCESS_CONTEXT ctx = {0};

    if (AcquireMemoryAccess(&ctx)) {
        KdPrint(("[Thread %lu] Accessing memory at %lld\n",
                ctx.ThreadId, ctx.AccessTime.QuadPart));

        // 实际访问内存的代码（示例：写入数据）
        RtlFillMemory(
            g_ProtectedMemory.pMemoryBuffer,
            g_ProtectedMemory.MemorySize,
            0xAA  // 填充示例数据
        );

        // 模拟耗时操作（例如处理数据）
        LARGE_INTEGER delay = { .QuadPart = -10000 * 100 }; // 延迟 100ms
        KeDelayExecutionThread(KernelMode, FALSE, &delay);

        ReleaseMemoryAccess(&ctx);
        KdPrint(("[Thread %lu] Memory access released\n", ctx.ThreadId));
    } else {
        KdPrint(("[Thread %lu] Memory is being freed, access denied\n", ctx.ThreadId));
    }

    PsTerminateSystemThread(STATUS_SUCCESS);
}

// 驱动卸载函数
VOID DriverUnload(PDRIVER_OBJECT DriverObject)
{
    UNREFERENCED_PARAMETER(DriverObject);
    SafeFreeProtectedMemory();
}

// 驱动入口点
NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
{
    UNREFERENCED_PARAMETER(RegistryPath);
    NTSTATUS status;
    HANDLE hThread;
    ULONG i;

    DriverObject->DriverUnload = DriverUnload;

    // 初始化受保护内存
    status = InitializeProtectedMemory();
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 创建 3 个线程模拟并发访问
    for (i = 0; i < 3; i++) {
        status = PsCreateSystemThread(
            &hThread,
            THREAD_ALL_ACCESS,
            NULL,
            NULL,
            NULL,
            MemoryAccessThreadRoutine,
            NULL
        );

        if (!NT_SUCCESS(status)) {
            KdPrint(("Failed to create thread: 0x%X\n", status));
            SafeFreeProtectedMemory();  // 清理资源
            return status;
        }

        ZwClose(hThread);
    }

    return STATUS_SUCCESS;
}