CommModifyKit/driver/DriverEntry.cpp

// DriverEntry.cpp - KMDF 驱动入口
// 职责：初始化 WDF 驱动对象、注册 EvtDeviceAdd（PnP 过滤驱动）
//      创建控制设备 \\.\CommModifyKit 接收 DLL 的 IOCTL
// 注意：本驱动作为串口设备类的 UpperFilters 加载
//      当 PnP Manager 创建串口 PDO 时，调用 EvtDeviceAdd attach 过滤设备
#include <ntddk.h>
#include <wdf.h>
#include <initguid.h>  // 必须在包含定义 GUID 的头文件之前
#include <ntddser.h>   // GUID_DEVINTERFACE_COMPORT
#include "DeviceContext.h"
#include "QueueCallback.h"
#include "ClientConnection.h"
#include "SerialFilter.h"

extern "C" {

DRIVER_INITIALIZE DriverEntry;
EVT_WDF_DRIVER_DEVICE_ADD EvtDeviceAdd;
EVT_WDF_DRIVER_UNLOAD EvtDriverUnload;
EVT_WDF_OBJECT_CONTEXT_CLEANUP EvtDeviceContextCleanup;

} // extern "C"

// PnP 回调使用 C++ 链接（WDF 宏内部处理）
EVT_WDF_DEVICE_PREPARE_HARDWARE EvtDevicePrepareHardware;

namespace commkit_driver {

// 全局驱动对象（仅用于内部日志/调试）
static WDFDRIVER g_wdf_driver = nullptr;

} // namespace commkit_driver

extern "C" NTSTATUS DriverEntry(PDRIVER_OBJECT driver_object,
                                  PUNICODE_STRING registry_path) {
    NTSTATUS status;
    WDF_DRIVER_CONFIG cfg;

    // 初始化全局客户端连接管理器
    commkit_driver::GetClientConnection().Initialize();

    // 初始化 WDF 配置
    WDF_DRIVER_CONFIG_INIT(&cfg, EvtDeviceAdd);
    cfg.EvtDriverUnload = EvtDriverUnload;

    // 创建 WDF 驱动对象
    status = WdfDriverCreate(driver_object, registry_path,
                                WDF_NO_OBJECT_ATTRIBUTES, &cfg,
                                &commkit_driver::g_wdf_driver);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 初始化 ClientConnection（创建控制设备）
    status = commkit_driver::CreateControlDevice(commkit_driver::g_wdf_driver);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    return STATUS_SUCCESS;
}

extern "C" NTSTATUS EvtDeviceAdd(WDFDRIVER driver, PWDFDEVICE_INIT init) {
    UNREFERENCED_PARAMETER(driver);

    NTSTATUS status;

    // 1. 标记为过滤设备
    WdfFdoInitSetFilter(init);

    // 2. 注册 PnP 电源回调（EvtDevicePrepareHardware 解析 COM 编号）
    WDF_PNPPOWER_EVENT_CALLBACKS pnp_callbacks;
    WDF_PNPPOWER_EVENT_CALLBACKS_INIT(&pnp_callbacks);
    pnp_callbacks.EvtDevicePrepareHardware = EvtDevicePrepareHardware;
    WdfDeviceInitSetPnpPowerEventCallbacks(init, &pnp_callbacks);

    // 3. 注册 WDM IRP 预处理回调（拦截 READ/WRITE/DEVICE_CONTROL）
    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchRead,
        IRP_MJ_READ, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchWrite,
        IRP_MJ_WRITE, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchDeviceControl,
        IRP_MJ_DEVICE_CONTROL, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 注册 CREATE/CLOSE/CLEANUP 预处理回调（确保正确转发）
    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchCreate,
        IRP_MJ_CREATE, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchClose,
        IRP_MJ_CLOSE, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = WdfDeviceInitAssignWdmIrpPreprocessCallback(
        init, commkit_driver::SerialFilter::DispatchCleanup,
        IRP_MJ_CLEANUP, nullptr, 0);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 4. 设置宽松 SDDL（匹配 serial.sys 行为）
    // WDF 默认 SDDL 仅 System+Administrators，会导致非 elevated 进程
    // 无法打开 \\.\COMx（即使 Administrator 账户在 filtered token 下也失败）
    // AU=Authenticated Users, WD=Everyone 给读写权限
    static DECLARE_CONST_UNICODE_STRING(
        sddl_filter,
        L"D:P(A;;GA;;;SY)(A;;GA;;;BA)(A;;GRGW;;;AU)(A;;GRGW;;;WD)");
    status = WdfDeviceInitAssignSDDLString(init, &sddl_filter);
    if (!NT_SUCCESS(status)) {
        DbgPrint("[CommModifyKit] WdfDeviceInitAssignSDDLString failed 0x%08X\n", status);
        return status;
    }
    DbgPrint("[CommModifyKit] SDDL assigned OK\n");

    // 5. 创建 WDF 设备
    WDFDEVICE device;
    WDF_OBJECT_ATTRIBUTES attrs;
    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attrs, DEVICE_CONTEXT);
    attrs.EvtCleanupCallback = EvtDeviceContextCleanup;

    status = WdfDeviceCreate(&init, &attrs, &device);
    if (!NT_SUCCESS(status)) {
        DbgPrint("[CommModifyKit] WdfDeviceCreate failed 0x%08X\n", status);
        return status;
    }
    DbgPrint("[CommModifyKit] WdfDeviceCreate OK\n");

    // 5. 初始化设备上下文
    PDEVICE_CONTEXT ctx = DeviceGetContext(device);
    ctx->ComNumber = 0;          // 延迟到 EvtDevicePrepareHardware 解析
    ctx->MonitoringEnabled = FALSE;
    ctx->RingBuffer = nullptr;
    ctx->WdfDevice = device;
    ctx->LowerDevice = WdfDeviceWdmGetAttachedDevice(device);

    // 6. 注册到全局端口表（com_number 暂为 0，PrepareHardware 时更新）
    commkit_driver::GetClientConnection().RegisterFilterDevice(device);

    return STATUS_SUCCESS;
}

NTSTATUS EvtDevicePrepareHardware(WDFDEVICE device,
                                     WDFCMRESLIST resources,
                                     WDFCMRESLIST resources_translated) {
    UNREFERENCED_PARAMETER(resources);
    UNREFERENCED_PARAMETER(resources_translated);

    NTSTATUS status;
    PDEVICE_CONTEXT ctx = DeviceGetContext(device);

    // 从 PnP 属性解析 COM 编号
    // 方式：获取 DevicePropertyFriendlyName 或 DevicePropertyDeviceDescription，
    // 匹配 "COM%d" 模式
    WDFMEMORY memory;
    status = WdfDeviceAllocAndQueryProperty(
        device, DevicePropertyFriendlyName, NonPagedPoolNx, 0, &memory);
    if (!NT_SUCCESS(status)) {
        // 尝试 DeviceDescription 作为后备
        status = WdfDeviceAllocAndQueryProperty(
            device, DevicePropertyDeviceDescription, NonPagedPoolNx, 0, &memory);
    }

    if (NT_SUCCESS(status)) {
        PWCHAR name = (PWCHAR)WdfMemoryGetBuffer(memory, nullptr);
        if (name) {
            // 打印 FriendlyName/DeviceDescription 前 40 字符用于诊断
            DbgPrint("[CommModifyKit] FriendlyName='%ws'\n", name);
            // 解析 "COMx" 模式（可能出现在 "COMx (port)" 或 "通信端口 (COMx)" 等格式中）
            ULONG name_len = (ULONG)wcslen(name);
            for (ULONG i = 0; i + 3 <= name_len; i++) {
                if ((name[i] == L'C' || name[i] == L'c') &&
                    (name[i+1] == L'O' || name[i+1] == L'o') &&
                    (name[i+2] == L'M' || name[i+2] == L'm')) {
                    if (name[i+3] >= L'0' && name[i+3] <= L'9') {
                        ULONG num = 0;
                        ULONG j = i + 3;
                        while (j < name_len && name[j] >= L'0' && name[j] <= L'9') {
                            num = num * 10 + (name[j] - L'0');
                            j++;
                        }
                        if (num > 0 && num <= 256) {
                            ctx->ComNumber = num;
                            break;
                        }
                    }
                }
            }
        }
        WdfObjectDelete(memory);
    }

    // 后备方案：FriendlyName/DeviceDescription 均未解析出 COM 编号时，
    // 从 PDO 名称查询 HKLM\HARDWARE\DEVICEMAP\SERIALCOMM 获取 "COMx" 映射
    if (ctx->ComNumber == 0) {
        WDFMEMORY pdo_memory;
        status = WdfDeviceAllocAndQueryProperty(
            device, DevicePropertyPhysicalDeviceObjectName,
            NonPagedPoolNx, 0, &pdo_memory);
        if (NT_SUCCESS(status)) {
            PWCHAR pdo_name = (PWCHAR)WdfMemoryGetBuffer(pdo_memory, nullptr);
            if (pdo_name && pdo_name[0] != L'\0') {
                UNICODE_STRING key_path;
                RtlInitUnicodeString(&key_path,
                    L"\\Registry\\Machine\\HARDWARE\\DEVICEMAP\\SERIALCOMM");
                OBJECT_ATTRIBUTES obj_attr;
                InitializeObjectAttributes(&obj_attr, &key_path,
                    OBJ_CASE_INSENSITIVE, nullptr, nullptr);

                HANDLE hkey;
                status = ZwOpenKey(&hkey, KEY_READ, &obj_attr);
                if (NT_SUCCESS(status)) {
                    UNICODE_STRING value_name;
                    RtlInitUnicodeString(&value_name, pdo_name);

                    UCHAR result_buf[256] = {0};
                    ULONG result_len = 0;
                    status = ZwQueryValueKey(hkey, &value_name,
                        KeyValuePartialInformation,
                        result_buf, sizeof(result_buf), &result_len);
                    if (NT_SUCCESS(status)) {
                        KEY_VALUE_PARTIAL_INFORMATION* kv_info =
                            (KEY_VALUE_PARTIAL_INFORMATION*)result_buf;
                        if (kv_info->Type == REG_SZ && kv_info->DataLength >= sizeof(WCHAR)) {
                            PWCHAR com_name = (PWCHAR)kv_info->Data;
                            ULONG com_len = kv_info->DataLength / sizeof(WCHAR);
                            // 解析 "COMx"
                            if (com_len >= 4 &&
                                (com_name[0] == L'C' || com_name[0] == L'c') &&
                                (com_name[1] == L'O' || com_name[1] == L'o') &&
                                (com_name[2] == L'M' || com_name[2] == L'm') &&
                                com_name[3] >= L'0' && com_name[3] <= L'9') {
                                ULONG num = 0;
                                ULONG j = 3;
                                while (j < com_len &&
                                       com_name[j] >= L'0' && com_name[j] <= L'9') {
                                    num = num * 10 + (com_name[j] - L'0');
                                    j++;
                                }
                                if (num > 0 && num <= 256) {
                                    ctx->ComNumber = num;
                                }
                            }
                        }
                    }
                    ZwClose(hkey);
                }
            }
            WdfObjectDelete(pdo_memory);
        }
    }

    // 后备方案 2：查询本设备注册的 GUID_DEVINTERFACE_COMPORT 接口
    // 虚拟串口可能不在 SERIALCOMM 中，但会注册 COMPORT 接口
    // 注意：必须传入本设备 PDO，不能传 nullptr 枚举所有设备后用 IoGetDeviceObjectPointer 打开，
    //       因为那会发送 IRP_MJ_CREATE 到正在初始化的设备栈，导致死锁
    if (ctx->ComNumber == 0) {
        PDEVICE_OBJECT currentPdo = WdfDeviceWdmGetPhysicalDevice(device);
        if (currentPdo) {
            PWSTR interfaceList = nullptr;
            status = IoGetDeviceInterfaces(
                &GUID_DEVINTERFACE_COMPORT,
                currentPdo,  // 仅查询本设备的接口
                DEVICE_INTERFACE_INCLUDE_NONACTIVE,
                &interfaceList);
            if (NT_SUCCESS(status) && interfaceList) {
                PWCHAR p = interfaceList;
                while (*p && ctx->ComNumber == 0) {
                    ULONG len = (ULONG)wcslen(p);
                    // 从符号链接路径解析 COM 编号（如 "\\?\COM12"）
                    for (ULONG i = 0; i + 3 <= len; i++) {
                        if ((p[i] == L'C' || p[i] == L'c') &&
                            (p[i+1] == L'O' || p[i+1] == L'o') &&
                            (p[i+2] == L'M' || p[i+2] == L'm') &&
                            p[i+3] >= L'0' && p[i+3] <= L'9') {
                            ULONG num = 0;
                            ULONG j = i + 3;
                            while (j < len && p[j] >= L'0' && p[j] <= L'9') {
                                num = num * 10 + (p[j] - L'0');
                                j++;
                            }
                            if (num > 0 && num <= 256) {
                                ctx->ComNumber = num;
                                break;
                            }
                        }
                    }
                    p += len + 1;
                }
                ExFreePool(interfaceList);
            }
        }
    }

    // 更新端口表中的 COM 编号
    commkit_driver::GetClientConnection().UpdatePortComNumber(device, ctx->ComNumber);

    DbgPrint("[CommModifyKit] EvtDevicePrepareHardware OK, ComNumber=%u\n", ctx->ComNumber);
    return STATUS_SUCCESS;
}

// 设备上下文清理回调：设备移除时注销端口表条目
VOID EvtDeviceContextCleanup(WDFOBJECT object) {
    WDFDEVICE device = (WDFDEVICE)object;
    PDEVICE_CONTEXT ctx = DeviceGetContext(device);
    if (ctx && ctx->ComNumber != 0) {
        commkit_driver::GetClientConnection().UnregisterFilterDevice(ctx->ComNumber);
    }
}

extern "C" VOID EvtDriverUnload(WDFDRIVER driver) {
    UNREFERENCED_PARAMETER(driver);
    commkit_driver::GetClientConnection().Cleanup();
    commkit_driver::g_wdf_driver = nullptr;
}