CommModifyKit/driver/ClientConnection.cpp

// ClientConnection.cpp - 用户态连接管理实现
#include <ntddk.h>
#include <wdf.h>
#include "ClientConnection.h"
#include "EventRingBuffer.h"
#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) {
    PDEVICE_CONTEXT ctx = FindPortContext(com_number);
    if (!ctx) {
        return STATUS_DEVICE_DOES_NOT_EXIST;
    }

    // 如果未分配环形缓冲，先分配
    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 事件（序列号由全局计数器递增）
    EventRingBuffer* ring = (EventRingBuffer*)ctx->RingBuffer;
    static LONG seq = 0;
    LONG cur = InterlockedIncrement(&seq);
    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;
        static LONG seq2 = 0;
        LONG cur = InterlockedIncrement(&seq2);
        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;
    }

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

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

    *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