CommModifyKit/driver/QueueCallback.cpp

// QueueCallback.cpp - I/O Queue 回调实现
#include <ntddk.h>
#include <wdf.h>
#include <initguid.h>
#include <wdmsec.h>
#include "QueueCallback.h"
#include "ClientConnection.h"
#include "DeviceContext.h"
#include "../common/CommKitIoctl.h"

namespace commkit_driver {

// SDDL: System+Admins full access, Authenticated Users+Everyone read/write
// 控制设备 \\.\CommModifyKit 需要允许非 elevated 进程也能打开
static const UNICODE_STRING SddlSysAdminFullAccess =
    RTL_CONSTANT_STRING(L"D:P(A;;GA;;;SY)(A;;GA;;;BA)(A;;GRGW;;;AU)(A;;GRGW;;;WD)");

// 专用 Parallel 队列句柄：用于 READ_EVENTS 转发
// WdfDeviceConfigureRequestDispatching 只能按请求类型路由（如 WdfRequestTypeDeviceControl），
// 不能按具体 IOCTL 码路由。改为在 EvtIoDeviceControl 中用 WdfRequestForwardToIoQueue 转发。
static WDFQUEUE g_read_events_queue = nullptr;

NTSTATUS CreateControlDevice(WDFDRIVER driver) {
    NTSTATUS status;
    PWDFDEVICE_INIT init = WdfControlDeviceInitAllocate(
        driver, (PCUNICODE_STRING)&SddlSysAdminFullAccess);
    if (!init) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    // 设置设备名
    DECLARE_UNICODE_STRING_SIZE(device_name, 64);
    RtlInitUnicodeString(&device_name, COMMKIT_DEVICE_NAME);
    status = WdfDeviceInitAssignName(init, &device_name);
    if (!NT_SUCCESS(status)) {
        WdfDeviceInitFree(init);
        return status;
    }

    // 创建设备
    WDFDEVICE control_device;
    WDF_OBJECT_ATTRIBUTES attrs;
    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attrs, DEVICE_CONTEXT);

    status = WdfDeviceCreate(&init, &attrs, &control_device);
    if (!NT_SUCCESS(status)) {
        WdfDeviceInitFree(init);
        return status;
    }

    // 创建符号链接 \\DosDevices\\CommModifyKit → \\.\CommModifyKit
    DECLARE_UNICODE_STRING_SIZE(symbolic_link, 64);
    RtlInitUnicodeString(&symbolic_link, COMMKIT_DEVICE_DOS_NAME);
    status = WdfDeviceCreateSymbolicLink(control_device, &symbolic_link);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 配置 I/O Queue 处理 IOCTL
    // 主队列：Sequential 派发，处理除 READ_EVENTS 外的所有 IOCTL
    // （ATTACH/DETACH/WRITE_PORT/READ_PORT/FREE_ALL 需要串行化避免竞争）
    WDF_IO_QUEUE_CONFIG queue_config;
    WDF_IO_QUEUE_CONFIG_INIT_DEFAULT_QUEUE(&queue_config, WdfIoQueueDispatchSequential);
    queue_config.EvtIoDeviceControl = EvtIoDeviceControl;
    queue_config.EvtIoDefault       = EvtIoDefault;

    WDFQUEUE queue;
    status = WdfIoQueueCreate(control_device, &queue_config,
                                WDF_NO_OBJECT_ATTRIBUTES, &queue);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    // 专用队列：Parallel 派发，仅处理 IOCTL_COMMKIT_READ_EVENTS
    // 独立于主队列，避免 WRITE_PORT/READ_PORT 同步阻塞时 EventLoop 无法读取事件
    // （否则高流量场景下环形缓冲溢出覆盖旧事件）
    WDF_IO_QUEUE_CONFIG read_events_config;
    WDF_IO_QUEUE_CONFIG_INIT(&read_events_config, WdfIoQueueDispatchParallel);
    read_events_config.EvtIoDeviceControl = EvtIoReadEvents;

    WDFQUEUE read_events_queue;
    status = WdfIoQueueCreate(control_device, &read_events_config,
                                WDF_NO_OBJECT_ATTRIBUTES, &read_events_queue);
    if (!NT_SUCCESS(status)) {
        return status;
    }
    // 保存队列句柄供 EvtIoDeviceControl 转发 READ_EVENTS 请求
    g_read_events_queue = read_events_queue;

    // 标记控制设备已就绪
    WdfControlFinishInitializing(control_device);
    return STATUS_SUCCESS;
}

VOID EvtIoDeviceControl(WDFQUEUE queue, WDFREQUEST request,
                          size_t output_buffer_length, size_t input_buffer_length,
                          ULONG io_control_code) {
    UNREFERENCED_PARAMETER(queue);
    UNREFERENCED_PARAMETER(input_buffer_length);
    UNREFERENCED_PARAMETER(output_buffer_length);

    NTSTATUS status = STATUS_SUCCESS;
    ULONG info = 0;
    auto& conn = GetClientConnection();

    switch (io_control_code) {
    case IOCTL_COMMKIT_REGISTER_CALLBACK:
        status = conn.HandleRegisterCallback();
        break;

    case IOCTL_COMMKIT_ATTACH_PORT: {
        COMMKIT_PORT_REQUEST* req = nullptr;
        size_t buf_len = 0;
        status = WdfRequestRetrieveInputBuffer(request, sizeof(*req),
                                                  (PVOID*)&req, &buf_len);
        if (NT_SUCCESS(status) && req) {
            status = conn.HandleAttachPort(req->ComNumber);
        }
        break;
    }

    case IOCTL_COMMKIT_DETACH_PORT: {
        COMMKIT_PORT_REQUEST* req = nullptr;
        size_t buf_len = 0;
        status = WdfRequestRetrieveInputBuffer(request, sizeof(*req),
                                                  (PVOID*)&req, &buf_len);
        if (NT_SUCCESS(status) && req) {
            status = conn.HandleDetachPort(req->ComNumber);
        }
        break;
    }

    case IOCTL_COMMKIT_WRITE_PORT: {
        COMMKIT_DATA_REQUEST* req = nullptr;
        size_t buf_len = 0;
        status = WdfRequestRetrieveInputBuffer(request, sizeof(*req),
                                                  (PVOID*)&req, &buf_len);
        if (NT_SUCCESS(status) && req) {
            // 校验 DataLen 不超过实际缓冲区容量，防止越界读取
            ULONG header_size = FIELD_OFFSET(COMMKIT_DATA_REQUEST, Data);
            ULONG data_capacity = (buf_len > header_size)
                                     ? (ULONG)(buf_len - header_size) : 0;
            if (req->DataLen > data_capacity) {
                status = STATUS_INVALID_PARAMETER;
            } else {
                PVOID data = req->Data;
                ULONG bytes_written = 0;
                status = conn.HandleWritePort(req->ComNumber, data, req->DataLen,
                                                &bytes_written);
                // 更新 DataLen 为实际写入字节数，用户态 DLL 可直接读取
                req->DataLen = bytes_written;
                // METHOD_BUFFERED: info = 从 SystemBuffer 拷贝回用户缓冲的字节数
                // = 头部(ComNumber+DataLen) + 实际数据
                info = header_size + bytes_written;
            }
        }
        break;
    }

    case IOCTL_COMMKIT_READ_PORT: {
        COMMKIT_DATA_REQUEST* req = nullptr;
        size_t buf_len = 0;
        status = WdfRequestRetrieveInputBuffer(request, sizeof(*req),
                                                  (PVOID*)&req, &buf_len);
        if (NT_SUCCESS(status) && req) {
            ULONG header_size = FIELD_OFFSET(COMMKIT_DATA_REQUEST, Data);
            ULONG data_capacity = (buf_len > header_size)
                                     ? (ULONG)(buf_len - header_size) : 0;
            if (req->DataLen > data_capacity) {
                status = STATUS_INVALID_PARAMETER;
            } else {
                PVOID data = req->Data;
                ULONG bytes_read = 0;
                status = conn.HandleReadPort(req->ComNumber, data, req->DataLen,
                                               &bytes_read);
                // 更新 DataLen 为实际读取字节数，用户态 DLL 可直接读取
                req->DataLen = bytes_read;
                // METHOD_BUFFERED: info = 从 SystemBuffer 拷贝回用户缓冲的字节数
                // = 头部(ComNumber+DataLen) + 实际数据
                info = header_size + bytes_read;
            }
        }
        break;
    }

    case IOCTL_COMMKIT_READ_EVENTS:
        // 转发到专用 Parallel 队列，避免被 WRITE_PORT/READ_PORT 的同步阻塞拖慢
        if (g_read_events_queue) {
            NTSTATUS fwd = WdfRequestForwardToIoQueue(request, g_read_events_queue);
            if (NT_SUCCESS(fwd)) {
                return;  // 转发成功，请求已移至专用队列，不再完成
            }
            // 转发失败，按错误完成
        }
        status = STATUS_INVALID_DEVICE_REQUEST;
        break;

    case IOCTL_COMMKIT_FREE_ALL:
        status = conn.HandleFreeAll();
        break;

    default:
        status = STATUS_INVALID_DEVICE_REQUEST;
        break;
    }

    WdfRequestCompleteWithInformation(request, status, info);
}

// 专用队列回调：处理 IOCTL_COMMKIT_READ_EVENTS
// 独立于主队列，即使 WRITE_PORT/READ_PORT 阻塞，EventLoop 仍可读取事件
VOID EvtIoReadEvents(WDFQUEUE queue, WDFREQUEST request,
                       size_t output_buffer_length, size_t input_buffer_length,
                       ULONG io_control_code) {
    UNREFERENCED_PARAMETER(queue);
    UNREFERENCED_PARAMETER(input_buffer_length);
    UNREFERENCED_PARAMETER(output_buffer_length);

    if (io_control_code != IOCTL_COMMKIT_READ_EVENTS) {
        WdfRequestComplete(request, STATUS_INVALID_DEVICE_REQUEST);
        return;
    }

    PVOID out_buf = nullptr;
    size_t out_len = 0;
    NTSTATUS status = WdfRequestRetrieveOutputBuffer(request, 0,
                                                        (PVOID*)&out_buf, &out_len);
    if (NT_SUCCESS(status) && out_buf) {
        ULONG returned = 0;
        status = GetClientConnection().HandleReadEvents(out_buf, (ULONG)out_len, &returned);
        WdfRequestCompleteWithInformation(request, status, returned);
    } else {
        WdfRequestComplete(request, STATUS_INVALID_PARAMETER);
    }
}

VOID EvtIoDefault(WDFQUEUE queue, WDFREQUEST request) {
    UNREFERENCED_PARAMETER(queue);
    WdfRequestComplete(request, STATUS_INVALID_DEVICE_REQUEST);
}

} // namespace commkit_driver