CommModifyKit/driver/SerialFilter.cpp

// SerialFilter.cpp - 串口 IRP 拦截实现
// KMDF 过滤驱动模式：通过 WdfDeviceInitAssignWdmIrpPreprocessCallback 注册
// 在预处理回调中注册完成例程，捕获 READ/WRITE 数据
#include <ntddk.h>
#include <wdf.h>
#include "SerialFilter.h"
#include "EventRingBuffer.h"
#include "ClientConnection.h"
#include "DeviceContext.h"
#include "../common/CommKitEvents.h"

namespace commkit_driver {

// 全局原子序列号（所有事件类型共享，保证 Sequence 字段全局唯一递增）
LONG g_sequence_counter = 0;

// 获取当前时间戳（100ns 单位）
static UINT64 GetCurrentTimeStamp() {
    return KeQueryInterruptTime();
}

void SerialFilter::CaptureIrpData(PDEVICE_CONTEXT ctx, PIRP irp, ULONG event_type) {
    if (!ctx || !ctx->MonitoringEnabled || !ctx->RingBuffer) {
        return;
    }

    PIO_STACK_LOCATION sl = IoGetCurrentIrpStackLocation(irp);
    if (!sl) return;

    // 仅在 IRP 成功完成时捕获
    if (!NT_SUCCESS(irp->IoStatus.Status)) {
        return;
    }

    ULONG data_size = 0;
    PVOID data_ptr = nullptr;

    // 根据 IRP 类型提取数据
    if (event_type == COMMKIT_OP_READ) {
        // 实际读取的字节数（IRP 完成后 Information = 读取字节数）
        data_size = (ULONG)irp->IoStatus.Information;
        if (data_size == 0) return;
        // 数据可能在 SystemBuffer / MdlAddress / UserBuffer
        if (irp->MdlAddress) {
            data_ptr = MmGetSystemAddressForMdlSafe(irp->MdlAddress, NormalPagePriority);
        } else if (irp->AssociatedIrp.SystemBuffer) {
            data_ptr = irp->AssociatedIrp.SystemBuffer;
        }
    } else if (event_type == COMMKIT_OP_WRITE) {
        data_size = (ULONG)irp->IoStatus.Information;
        if (data_size == 0) {
            data_size = sl->Parameters.Write.Length;
        }
        if (data_size == 0) return;
        if (irp->MdlAddress) {
            data_ptr = MmGetSystemAddressForMdlSafe(irp->MdlAddress, NormalPagePriority);
        } else if (irp->AssociatedIrp.SystemBuffer) {
            data_ptr = irp->AssociatedIrp.SystemBuffer;
        }
    }

    if (!data_ptr || data_size == 0) {
        return;
    }

    // 推入环形缓冲
    LONG seq = InterlockedIncrement(&g_sequence_counter);
    auto* ring = (EventRingBuffer*)ctx->RingBuffer;
    ring->Push((ULONG)seq, GetCurrentTimeStamp(),
                ctx->ComNumber, event_type,
                data_size, data_ptr);
}

NTSTATUS SerialFilter::OnReadComplete(PDEVICE_OBJECT dev, PIRP irp, PVOID ctx) {
    UNREFERENCED_PARAMETER(dev);
    PDEVICE_CONTEXT device_ctx = (PDEVICE_CONTEXT)ctx;
    if (device_ctx) {
        CaptureIrpData(device_ctx, irp, COMMKIT_OP_READ);
    }
    // 下层驱动返回 STATUS_PENDING 时需标记 IRP 为 pending
    if (irp->PendingReturned) {
        IoMarkIrpPending(irp);
    }
    return STATUS_SUCCESS;
}

NTSTATUS SerialFilter::OnWriteComplete(PDEVICE_OBJECT dev, PIRP irp, PVOID ctx) {
    UNREFERENCED_PARAMETER(dev);
    PDEVICE_CONTEXT device_ctx = (PDEVICE_CONTEXT)ctx;
    if (device_ctx) {
        CaptureIrpData(device_ctx, irp, COMMKIT_OP_WRITE);
    }
    if (irp->PendingReturned) {
        IoMarkIrpPending(irp);
    }
    return STATUS_SUCCESS;
}

NTSTATUS SerialFilter::DispatchRead(WDFDEVICE dev, PIRP irp) {
    PDEVICE_CONTEXT ctx = DeviceGetContext(dev);

    if (ctx && ctx->MonitoringEnabled && ctx->RingBuffer) {
        IoSetCompletionRoutine(irp, OnReadComplete, ctx, TRUE, TRUE, TRUE);
    }

    return WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
}

NTSTATUS SerialFilter::DispatchWrite(WDFDEVICE dev, PIRP irp) {
    PDEVICE_CONTEXT ctx = DeviceGetContext(dev);

    if (ctx && ctx->MonitoringEnabled && ctx->RingBuffer) {
        IoSetCompletionRoutine(irp, OnWriteComplete, ctx, TRUE, TRUE, TRUE);
    }

    return WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
}

NTSTATUS SerialFilter::DispatchDeviceControl(WDFDEVICE dev, PIRP irp) {
    // DeviceControl 透传，不拦截（串口 IOCTL 由应用层处理）
    return WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
}

NTSTATUS SerialFilter::DispatchCreate(WDFDEVICE dev, PIRP irp) {
    // CREATE/CLOSE/CLEANUP 直接转发，不做任何处理
    DbgPrint("[CommModifyKit] DispatchCreate entered, IRP=%p\n", irp);
    NTSTATUS status = WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
    DbgPrint("[CommModifyKit] DispatchCreate returned 0x%08X\n", status);
    return status;
}

NTSTATUS SerialFilter::DispatchClose(WDFDEVICE dev, PIRP irp) {
    return WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
}

NTSTATUS SerialFilter::DispatchCleanup(WDFDEVICE dev, PIRP irp) {
    return WdfDeviceWdmDispatchPreprocessedIrp(dev, irp);
}

NTSTATUS SerialFilter::CreateAndAttach(WDFDRIVER driver,
                                         PDEVICE_OBJECT serial_pdo,
                                         ULONG com_number,
                                         WDFDEVICE* out_filter_device) {
    // 不再使用：过滤设备由 EvtDeviceAdd 自动创建
    UNREFERENCED_PARAMETER(driver);
    UNREFERENCED_PARAMETER(serial_pdo);
    UNREFERENCED_PARAMETER(com_number);
    *out_filter_device = nullptr;
    return STATUS_NOT_IMPLEMENTED;
}

} // namespace commkit_driver