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 and Administrators have full access
static const UNICODE_STRING SddlSysAdminFullAccess = 
    RTL_CONSTANT_STRING(L"D:P(A;;GA;;;SY)(A;;GA;;;BA)");

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
    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;
    }

    // 标记控制设备已就绪
    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) {
            PVOID data = req->Data;
            status = conn.HandleWritePort(req->ComNumber, data, req->DataLen);
        }
        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) {
            PVOID data = req->Data;
            status = conn.HandleReadPort(req->ComNumber, data, req->DataLen);
        }
        break;
    }

    case IOCTL_COMMKIT_READ_EVENTS: {
        PVOID out_buf = nullptr;
        size_t out_len = 0;
        status = WdfRequestRetrieveOutputBuffer(request, 0,
                                                   (PVOID*)&out_buf, &out_len);
        if (NT_SUCCESS(status) && out_buf) {
            ULONG returned = 0;
            status = conn.HandleReadEvents(out_buf, (ULONG)out_len, &returned);
            info = returned;
        }
        break;
    }

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

    default:
        status = STATUS_INVALID_DEVICE_REQUEST;
        break;
    }

    WdfRequestCompleteWithInformation(request, status, info);
}

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

} // namespace commkit_driver