windows 内核层简单 map 实现与测试

KernelDataStructs.h

#pragma once
#pragma once  
#include  

// 内存标签定义  
#define KERNEL_MAP_TAG 'KMap'  
#define KERNEL_VECTOR_TAG 'KVec'  

// 安全内存分配宏  
#define SAFE_ALLOC(ptr, size, tag) \
    (ptr) = ExAllocatePoolWithTag(NonPagedPoolNx, (size), (tag)); \
    if (!(ptr)) { \
        return STATUS_INSUFFICIENT_RESOURCES; \
    }  

// 安全释放宏  
#define SAFE_FREE(ptr, tag) \
    if ((ptr)) { \
        ExFreePoolWithTag((ptr), (tag)); \
        (ptr) = NULL; \
    }  

// 函数类型定义  
typedef BOOLEAN(*KeyCompareFunc)(CONST PVOID Key1, CONST PVOID Key2);
typedef ULONG(*HashFunc)(CONST PVOID Key);
typedef VOID(*ElementDestroyFunc)(PVOID Element);

// 通用状态码  
#define KERNEL_MAP_KEY_EXISTS         ((NTSTATUS)0x00000001L)  
#define KERNEL_MAP_KEY_NOT_FOUND      ((NTSTATUS)0x00000002L)  

KernelMap.h

#pragma once  
#include "KernelDataStructs.h"  

// 映射条目结构  
typedef struct _MAP_ENTRY {
    PVOID Key;
    PVOID Value;
    struct _MAP_ENTRY* Next;
} MAP_ENTRY, * PMAP_ENTRY;

// 内核映射结构  
typedef struct _KERNEL_MAP {
    PMAP_ENTRY* Buckets;     // 哈希桶数组  
    ULONG BucketCount;       // 桶的数量  
    ULONG EntryCount;        // 总条目数  
    KSPIN_LOCK MapLock;      // 读写自旋锁  

    // 自定义函数  
    KeyCompareFunc CompareFunc;
    HashFunc HashFunction;
} KERNEL_MAP, * PKERNEL_MAP;

// 函数声明  
NTSTATUS
KernelMapCreate(
    _Out_ PKERNEL_MAP* OutMap,
    _In_ ULONG InitialCapacity,
    _In_opt_ KeyCompareFunc CompareFunc,
    _In_opt_ HashFunc HashFunction
);

NTSTATUS
KernelMapInsert(
    _Inout_ PKERNEL_MAP Map,
    _In_ PVOID Key,
    _In_ PVOID Value
);

NTSTATUS
KernelMapFind(
    _In_ PKERNEL_MAP Map,
    _In_ PVOID Key,
    _Out_ PVOID* OutValue
);

NTSTATUS
KernelMapRemove(
    _Inout_ PKERNEL_MAP Map,
    _In_ PVOID Key
);

VOID
KernelMapDestroy(
    _Inout_ PKERNEL_MAP Map
);

KernelMap.c

#include "KernelMap.h"  

// 默认哈希函数（MurmurHash3）  
ULONG DefaultHashFunction(_In_ CONST PVOID Key) {
    ULONG_PTR KeyValue = (ULONG_PTR)Key;
    KeyValue ^= KeyValue >> 16;
    KeyValue *= 0x85ebca6b;
    KeyValue ^= KeyValue >> 13;
    KeyValue *= 0xc2b2ae35;
    KeyValue ^= KeyValue >> 16;
    return (ULONG)KeyValue;
}

// 默认比较函数  
BOOLEAN DefaultCompareFunction(
    _In_ CONST PVOID Key1,
    _In_ CONST PVOID Key2
) {
    return Key1 == Key2;
}

// 创建映射  
NTSTATUS
KernelMapCreate(
    _Out_ PKERNEL_MAP* OutMap,
    _In_ ULONG InitialCapacity,
    _In_opt_ KeyCompareFunc CompareFunc,
    _In_opt_ HashFunc HashFunction
) {
    PKERNEL_MAP Map;

    // 分配映射结构  
    SAFE_ALLOC(Map, sizeof(KERNEL_MAP), KERNEL_MAP_TAG);

    // 分配桶数组  
    SAFE_ALLOC(Map->Buckets,
        InitialCapacity * sizeof(PMAP_ENTRY),
        KERNEL_MAP_TAG
    );

    // 初始化  
    RtlZeroMemory(Map->Buckets, InitialCapacity * sizeof(PMAP_ENTRY));
    Map->BucketCount = InitialCapacity;
    Map->EntryCount = 0;

    // 设置比较和哈希函数  
    Map->CompareFunc = CompareFunc ? CompareFunc : DefaultCompareFunction;
    Map->HashFunction = HashFunction ? HashFunction : DefaultHashFunction;

    // 初始化自旋锁  
    KeInitializeSpinLock(&Map->MapLock);

    *OutMap = Map;
    return STATUS_SUCCESS;
}

// 插入映射项  
NTSTATUS
KernelMapInsert(
    _Inout_ PKERNEL_MAP Map,
    _In_ PVOID Key,
    _In_ PVOID Value
) {
    KIRQL OldIrql;
    ULONG HashValue;
    PMAP_ENTRY NewEntry, Current;

    // 计算哈希值  
    HashValue = Map->HashFunction(Key) % Map->BucketCount;

    // 分配新条目  
    SAFE_ALLOC(NewEntry, sizeof(MAP_ENTRY), KERNEL_MAP_TAG);

    // 初始化新条目  
    NewEntry->Key = Key;
    NewEntry->Value = Value;
    NewEntry->Next = NULL;

    // 加锁插入  
    KeAcquireSpinLock(&Map->MapLock, &OldIrql);

    // 检查是否已存在相同键  
    Current = Map->Buckets[HashValue];
    while (Current) {
        if (Map->CompareFunc(Current->Key, Key)) {
            // 键已存在  
            KeReleaseSpinLock(&Map->MapLock, OldIrql);
            SAFE_FREE(NewEntry, KERNEL_MAP_TAG);
            return KERNEL_MAP_KEY_EXISTS;
        }
        Current = Current->Next;
    }

    // 头插法处理冲突  
    NewEntry->Next = Map->Buckets[HashValue];
    Map->Buckets[HashValue] = NewEntry;
    Map->EntryCount++;

    KeReleaseSpinLock(&Map->MapLock, OldIrql);
    return STATUS_SUCCESS;
}

// 查找映射项  
NTSTATUS
KernelMapFind(
    _In_ PKERNEL_MAP Map,
    _In_ PVOID Key,
    _Out_ PVOID* OutValue
) {
    KIRQL OldIrql;
    ULONG HashValue;
    PMAP_ENTRY Current;

    // 计算哈希值  
    HashValue = Map->HashFunction(Key) % Map->BucketCount;

    // 加锁查找  
    KeAcquireSpinLock(&Map->MapLock, &OldIrql);

    // 遍历对应桶的链表  
    Current = Map->Buckets[HashValue];
    while (Current) {
        if (Map->CompareFunc(Current->Key, Key)) {
            *OutValue = Current->Value;
            KeReleaseSpinLock(&Map->MapLock, OldIrql);
            return STATUS_SUCCESS;
        }
        Current = Current->Next;
    }

    KeReleaseSpinLock(&Map->MapLock, OldIrql);
    return KERNEL_MAP_KEY_NOT_FOUND;
}

// 移除映射项  
NTSTATUS
KernelMapRemove(
    _Inout_ PKERNEL_MAP Map,
    _In_ PVOID Key
) {
    KIRQL OldIrql;
    ULONG HashValue;
    PMAP_ENTRY Current, Previous = NULL;

    // 计算哈希值  
    HashValue = Map->HashFunction(Key) % Map->BucketCount;

    // 加锁移除  
    KeAcquireSpinLock(&Map->MapLock, &OldIrql);

    // 遍历对应桶的链表  
    Current = Map->Buckets[HashValue];
    while (Current) {
        if (Map->CompareFunc(Current->Key, Key)) {
            if (Previous) {
                Previous->Next = Current->Next;
            }
            else {
                Map->Buckets[HashValue] = Current->Next;
            }

            SAFE_FREE(Current, KERNEL_MAP_TAG);
            Map->EntryCount--;

            KeReleaseSpinLock(&Map->MapLock, OldIrql);
            return STATUS_SUCCESS;
        }
        Previous = Current;
        Current = Current->Next;
    }

    KeReleaseSpinLock(&Map->MapLock, OldIrql);
    return KERNEL_MAP_KEY_NOT_FOUND;
}

// 销毁映射  
VOID
KernelMapDestroy(
    _Inout_ PKERNEL_MAP Map
) {
    ULONG i;
    PMAP_ENTRY Current, Next;

    // 释放每个桶的链表  
    for (i = 0; i < Map->BucketCount; i++) {
        Current = Map->Buckets[i];
        while (Current) {
            Next = Current->Next;
            SAFE_FREE(Current, KERNEL_MAP_TAG);
            Current = Next;
        }
    }

    // 释放桶数组和映射结构  
    SAFE_FREE(Map->Buckets, KERNEL_MAP_TAG);
    SAFE_FREE(Map, KERNEL_MAP_TAG);
}

测试用例

NTSTATUS TestMapOperations() {
    PKERNEL_MAP TestMap;
    PVOID Value;
    NTSTATUS Status;
    ULONG i;

    // 创建映射  
    Status = KernelMapCreate(&TestMap, 16, NULL, NULL);
    if (!NT_SUCCESS(Status)) {
        DbgPrint("Map Creation Failed: %x\n", Status);
        return Status;
    }

    // 测试1：单个元素插入和查找  
    DbgPrint("Test 1: Single Element Insertion and Lookup\n");
    Status = KernelMapInsert(TestMap, (PVOID)1, (PVOID)100);
    if (!NT_SUCCESS(Status)) {
        DbgPrint("Map Insert Failed: %x\n", Status);
        KernelMapDestroy(TestMap);
        return Status;
    }

    Status = KernelMapFind(TestMap, (PVOID)1, &Value);
    if (!NT_SUCCESS(Status) || Value != (PVOID)100) {
        DbgPrint("Map Find Failed or Incorrect Value: %x\n", Status);
        KernelMapDestroy(TestMap);
        return Status;
    }
    DbgPrint("Single Element Test Passed. Value: %p\n", Value);

    // 测试2：重复插入同一个键  
    DbgPrint("\nTest 2: Duplicate Key Insertion\n");
    Status = KernelMapInsert(TestMap, (PVOID)1, (PVOID)200);
    if (Status != KERNEL_MAP_KEY_EXISTS) {
        DbgPrint("Duplicate Key Insertion Should Fail: %x\n", Status);
        KernelMapDestroy(TestMap);
        return STATUS_UNSUCCESSFUL;
    }
    DbgPrint("Duplicate Key Insertion Test Passed\n");

    // 测试3：删除存在的元素  
    DbgPrint("\nTest 3: Remove Existing Element\n");
    Status = KernelMapRemove(TestMap, (PVOID)1);
    if (!NT_SUCCESS(Status)) {
        DbgPrint("Map Remove Failed: %x\n", Status);
        KernelMapDestroy(TestMap);
        return Status;
    }

    // 验证元素已删除  
    Status = KernelMapFind(TestMap, (PVOID)1, &Value);
    if (Status != KERNEL_MAP_KEY_NOT_FOUND) {
        DbgPrint("Element Should Not Be Found After Removal: %x\n", Status);
        KernelMapDestroy(TestMap);
        return STATUS_UNSUCCESSFUL;
    }
    DbgPrint("Remove Existing Element Test Passed\n");

    // 测试4：删除不存在的元素  
    DbgPrint("\nTest 4: Remove Non-Existent Element\n");
    Status = KernelMapRemove(TestMap, (PVOID)999);
    if (Status != KERNEL_MAP_KEY_NOT_FOUND) {
        DbgPrint("Remove Non-Existent Element Should Fail: %x\n", Status);
        KernelMapDestroy(TestMap);
        return STATUS_UNSUCCESSFUL;
    }
    DbgPrint("Remove Non-Existent Element Test Passed\n");

    // 测试5：多元素随机插入和删除  
    DbgPrint("\nTest 5: Multiple Element Random Insertion and Removal\n");
    for (i = 0; i < 100; i++) {
        // 使用伪随机数作为键和值  
        PVOID Key = (PVOID)((i * 17) % 1000);
        PVOID Val = (PVOID)((i * 23) % 5000);

        Status = KernelMapInsert(TestMap, Key, Val);
        if (!NT_SUCCESS(Status) && Status != KERNEL_MAP_KEY_EXISTS) {
            DbgPrint("Random Insertion Failed for Key %p: %x\n", Key, Status);
            KernelMapDestroy(TestMap);
            return Status;
        }
    }

    // 随机查找和删除  
    for (i = 0; i < 50; i++) {
        PVOID Key = (PVOID)((i * 17) % 1000);
        PVOID Val;

        Status = KernelMapFind(TestMap, Key, &Val);
        if (NT_SUCCESS(Status)) {
            Status = KernelMapRemove(TestMap, Key);
            if (!NT_SUCCESS(Status)) {
                DbgPrint("Random Removal Failed for Key %p: %x\n", Key, Status);
                KernelMapDestroy(TestMap);
                return Status;
            }
        }
    }

    DbgPrint("Multiple Element Random Test Passed\n");

    // 销毁映射  
    KernelMapDestroy(TestMap);

    DbgPrint("All Map Tests Completed Successfully\n");
    return STATUS_SUCCESS;
}