CommModifyKit/driver/ClientConnection.cpp

// ClientConnection.cpp - 用户态连接管理实现
#include <ntddk.h>
#include <wdf.h>
#include "ClientConnection.h"
#include "EventRingBuffer.h"
#include "SerialFilter.h"  // g_sequence_counter
#include "../common/CommKitIoctl.h"
#include "../common/CommKitEvents.h"

namespace commkit_driver {

// 全局实例
ClientConnection g_ClientConnection;

void ClientConnection::Initialize() {
    ports_count_ = 0;
    control_device_ = nullptr;
    callback_registered_ = FALSE;
    KeInitializeSpinLock(&ports_lock_);
    RtlZeroMemory(ports_table_, sizeof(ports_table_));
}

void ClientConnection::Cleanup() {
    // 释放所有端口的环形缓冲
    KIRQL old_irql;
    KeAcquireSpinLock(&ports_lock_, &old_irql);
    for (ULONG i = 0; i < ports_count_; ++i) {
        PDEVICE_CONTEXT ctx = ports_table_[i].Context;
        if (ctx && ctx->RingBuffer) {
            EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
            ring->Cleanup();
            ExFreePoolWithTag(ring, 'RBCK');
            ctx->RingBuffer = nullptr;
        }
    }
    ports_count_ = 0;
    KeReleaseSpinLock(&ports_lock_, old_irql);
}

void ClientConnection::RegisterFilterDevice(WDFDEVICE device) {
    KIRQL old_irql;
    KeAcquireSpinLock(&ports_lock_, &old_irql);

    if (ports_count_ < 256) {
        PDEVICE_CONTEXT ctx = DeviceGetContext(device);
        ports_table_[ports_count_].ComNumber = 0;  // 延迟到 UpdatePortComNumber
        ports_table_[ports_count_].Context = ctx;
        ports_count_++;
        ctx->ComNumber = 0;
        ctx->MonitoringEnabled = FALSE;
        ctx->WdfDevice = device;
    }

    KeReleaseSpinLock(&ports_lock_, old_irql);
}

void ClientConnection::UpdatePortComNumber(WDFDEVICE device, ULONG com_number) {
    KIRQL old_irql;
    KeAcquireSpinLock(&ports_lock_, &old_irql);

    for (ULONG i = 0; i < ports_count_; ++i) {
        if (ports_table_[i].Context &&
            ports_table_[i].Context->WdfDevice == device) {
            ports_table_[i].ComNumber = com_number;
            ports_table_[i].Context->ComNumber = com_number;
            break;
        }
    }

    KeReleaseSpinLock(&ports_lock_, old_irql);
}

void ClientConnection::UnregisterFilterDevice(ULONG com_number) {
    KIRQL old_irql;
    KeAcquireSpinLock(&ports_lock_, &old_irql);

    for (ULONG i = 0; i < ports_count_; ++i) {
        if (ports_table_[i].ComNumber == com_number) {
            // 移动最后一个元素到当前位置
            PDEVICE_CONTEXT ctx = ports_table_[i].Context;
            if (ctx && ctx->RingBuffer) {
                EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
                ring->Cleanup();
                ExFreePoolWithTag(ring, 'RBCK');
                ctx->RingBuffer = nullptr;
            }
            ports_table_[i] = ports_table_[ports_count_ - 1];
            ports_count_--;
            break;
        }
    }

    KeReleaseSpinLock(&ports_lock_, old_irql);
}

PDEVICE_CONTEXT ClientConnection::FindPortContext(ULONG com_number) {
    PDEVICE_CONTEXT result = nullptr;
    KIRQL old_irql;
    KeAcquireSpinLock(&ports_lock_, &old_irql);
    for (ULONG i = 0; i < ports_count_; ++i) {
        if (ports_table_[i].ComNumber == com_number) {
            result = ports_table_[i].Context;
            break;
        }
    }
    KeReleaseSpinLock(&ports_lock_, old_irql);
    return result;
}

NTSTATUS ClientConnection::HandleRegisterCallback() {
    callback_registered_ = TRUE;
    return STATUS_SUCCESS;
}

NTSTATUS ClientConnection::HandleAttachPort(ULONG com_number) {
    // 拒绝 com_number==0：EvtDevicePrepareHardware 解析失败时 ComNumber 保持 0，
    // 若允许 attach 会误匹配第一个未解析出编号的设备
    if (com_number == 0) {
        DbgPrint("[CommModifyKit] HandleAttachPort: com_number=0 rejected\n");
        return STATUS_INVALID_PARAMETER;
    }

    PDEVICE_CONTEXT ctx = FindPortContext(com_number);
    if (!ctx) {
        DbgPrint("[CommModifyKit] HandleAttachPort: COM%u NOT in ports_table (count=%u)\n",
                 com_number, ports_count_);
        // 打印 ports_table_ 里所有已注册的 COM 编号，便于诊断
        KIRQL old_irql;
        KeAcquireSpinLock(&ports_lock_, &old_irql);
        for (ULONG i = 0; i < ports_count_; ++i) {
            DbgPrint("[CommModifyKit]   ports_table[%u].ComNumber=%u\n",
                     i, ports_table_[i].ComNumber);
        }
        KeReleaseSpinLock(&ports_lock_, old_irql);
        return STATUS_DEVICE_DOES_NOT_EXIST;
    }
    DbgPrint("[CommModifyKit] HandleAttachPort: COM%u found in ports_table\n", com_number);

    // 如果未分配环形缓冲，先分配
    if (!ctx->RingBuffer) {
        EventRingBuffer* ring = (EventRingBuffer*)ExAllocatePool2(
            POOL_FLAG_NON_PAGED, sizeof(EventRingBuffer), 'RBCK');
        if (!ring) {
            return STATUS_INSUFFICIENT_RESOURCES;
        }
        // placement new 等价：手动调用构造
        RtlZeroMemory(ring, sizeof(EventRingBuffer));
        NTSTATUS status = ring->Initialize();
        if (!NT_SUCCESS(status)) {
            ring->Cleanup();
            ExFreePoolWithTag(ring, 'RBCK');
            return status;
        }
        ctx->RingBuffer = ring;
    }

    // 推入 OP_OPEN 事件（使用全局共享序列号，保证所有事件类型 Sequence 唯一递增）
    EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
    LONG cur = InterlockedIncrement(&g_sequence_counter);
    UINT64 timestamp = KeQueryInterruptTime();  // 100ns 单位
    ring->Push((ULONG)cur, timestamp, com_number, COMMKIT_OP_OPEN, 0, nullptr);

    ctx->MonitoringEnabled = TRUE;
    return STATUS_SUCCESS;
}

NTSTATUS ClientConnection::HandleDetachPort(ULONG com_number) {
    PDEVICE_CONTEXT ctx = FindPortContext(com_number);
    if (!ctx) {
        return STATUS_DEVICE_DOES_NOT_EXIST;
    }

    ctx->MonitoringEnabled = FALSE;

    // 推入 OP_CLOSE 事件（使用全局共享序列号）
    if (ctx->RingBuffer) {
        EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
        LONG cur = InterlockedIncrement(&g_sequence_counter);
        UINT64 timestamp = KeQueryInterruptTime();
        ring->Push((ULONG)cur, timestamp, com_number, COMMKIT_OP_CLOSE, 0, nullptr);
    }
    return STATUS_SUCCESS;
}

NTSTATUS ClientConnection::HandleWritePort(ULONG com_number, PVOID data, ULONG len) {
    PDEVICE_CONTEXT ctx = FindPortContext(com_number);
    if (!ctx || !ctx->LowerDevice) {
        return STATUS_DEVICE_DOES_NOT_EXIST;
    }

    // 构造同步写 IRP 发送给下层串口设备
    KEVENT event;
    KeInitializeEvent(&event, NotificationEvent, FALSE);

    IO_STATUS_BLOCK io_status = {};
    PIRP irp = IoBuildSynchronousFsdRequest(
        IRP_MJ_WRITE, ctx->LowerDevice, data, len, nullptr, &event, &io_status);
    if (!irp) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    NTSTATUS status = IoCallDriver(ctx->LowerDevice, irp);
    if (status == STATUS_PENDING) {
        KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, nullptr);
        status = io_status.Status;
    }

    // 注：此 IRP 直接发给 LowerDevice，绕过 SerialFilter::DispatchWrite，
    // 不会触发 OnWriteComplete 捕获。需手动将写入数据 Push 为 OP_WRITE 事件，
    // 否则监控端看不到注入的写流量（监控数据不完整）
    if (NT_SUCCESS(status) && ctx->MonitoringEnabled && ctx->RingBuffer) {
        ULONG written = (ULONG)io_status.Information;
        if (written > 0 && data) {
            LONG seq = InterlockedIncrement(&g_sequence_counter);
            UINT64 timestamp = KeQueryInterruptTime();
            EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
            ring->Push((ULONG)seq, timestamp, com_number, COMMKIT_OP_WRITE,
                        written, data);
        }
    }

    return status;
}

NTSTATUS ClientConnection::HandleReadPort(ULONG com_number, PVOID data, ULONG len) {
    PDEVICE_CONTEXT ctx = FindPortContext(com_number);
    if (!ctx || !ctx->LowerDevice) {
        return STATUS_DEVICE_DOES_NOT_EXIST;
    }

    // 构造同步读 IRP 发送给下层串口设备
    KEVENT event;
    KeInitializeEvent(&event, NotificationEvent, FALSE);

    IO_STATUS_BLOCK io_status = {};
    PIRP irp = IoBuildSynchronousFsdRequest(
        IRP_MJ_READ, ctx->LowerDevice, data, len, nullptr, &event, &io_status);
    if (!irp) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    NTSTATUS status = IoCallDriver(ctx->LowerDevice, irp);
    if (status == STATUS_PENDING) {
        KeWaitForSingleObject(&event, Executive, KernelMode, FALSE, nullptr);
        status = io_status.Status;
    }

    // 注：此 IRP 直接发给 LowerDevice，绕过 SerialFilter::DispatchRead，
    // 不会触发 OnReadComplete 捕获。需手动将读取数据 Push 为 OP_READ 事件，
    // 否则读取的数据既不返回调用者（METHOD_BUFFERED + info=0）也不进 RingBuffer，完全丢失
    if (NT_SUCCESS(status) && ctx->MonitoringEnabled && ctx->RingBuffer) {
        ULONG read_bytes = (ULONG)io_status.Information;
        if (read_bytes > 0 && data) {
            LONG seq = InterlockedIncrement(&g_sequence_counter);
            UINT64 timestamp = KeQueryInterruptTime();
            EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
            ring->Push((ULONG)seq, timestamp, com_number, COMMKIT_OP_READ,
                        read_bytes, data);
        }
    }

    return status;
}

NTSTATUS ClientConnection::HandleReadEvents(PVOID out_buf, ULONG out_size, PULONG returned) {
    *returned = 0;
    if (!callback_registered_) {
        return STATUS_DEVICE_NOT_READY;
    }

    PCOMMKIT_EVENT events = (PCOMMKIT_EVENT)out_buf;
    ULONG max_count = out_size / sizeof(COMMKIT_EVENT);
    if (max_count == 0) {
        return STATUS_BUFFER_TOO_SMALL;
    }

    ULONG total_popped = 0;
    KIRQL old_irql;
    KeAcquireSpinLock(&ports_lock_, &old_irql);

    // 按端口顺序弹出事件到输出缓冲
    for (ULONG i = 0; i < ports_count_ && total_popped < max_count; ++i) {
        PDEVICE_CONTEXT ctx = ports_table_[i].Context;
        if (!ctx || !ctx->RingBuffer || !ctx->MonitoringEnabled) {
            continue;
        }
        EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
        LONG popped = ring->PopBatch(events + total_popped, max_count - total_popped);
        total_popped += (ULONG)popped;
    }

    KeReleaseSpinLock(&ports_lock_, old_irql);

    // 弹出后按 Sequence 升序排序（锁外执行，不持自旋锁做 O(n²) 操作）
    // 各端口 RingBuffer 内部 Sequence 是递增的，但跨端口弹出后顺序被打乱
    // （如 port0 的 seq=100-109 可能排在 port1 的 seq=95-99 之前）
    // 使用插入排序：事件批量小（通常 ≤16），O(n²) 开销可忽略
    if (total_popped > 1) {
        for (ULONG i = 1; i < total_popped; ++i) {
            COMMKIT_EVENT key = events[i];
            LONG j = (LONG)i - 1;
            while (j >= 0 && events[j].Sequence > key.Sequence) {
                events[j + 1] = events[j];
                j--;
            }
            events[j + 1] = key;
        }
    }

    *returned = total_popped * sizeof(COMMKIT_EVENT);
    return STATUS_SUCCESS;
}

NTSTATUS ClientConnection::HandleFreeAll() {
    KIRQL old_irql;
    KeAcquireSpinLock(&ports_lock_, &old_irql);
    for (ULONG i = 0; i < ports_count_; ++i) {
        PDEVICE_CONTEXT ctx = ports_table_[i].Context;
        if (ctx) {
            ctx->MonitoringEnabled = FALSE;
            if (ctx->RingBuffer) {
                ((EventRingBuffer*)ctx->RingBuffer)->Clear();
            }
        }
    }
    KeReleaseSpinLock(&ports_lock_, old_irql);
    callback_registered_ = FALSE;
    return STATUS_SUCCESS;
}

} // namespace commkit_driver