CommModifyKit/driver/EventRingBuffer.cpp

// EventRingBuffer.cpp - 环形缓冲实现
#include <ntddk.h>
#include <wdf.h>
#include "EventRingBuffer.h"

namespace commkit_driver {

NTSTATUS EventRingBuffer::Initialize() {
    // 分配非分页内存（可在 DISPATCH_LEVEL 访问）
    buffer_ = nullptr;
    head_ = tail_ = count_ = 0;
    KeInitializeSpinLock(&lock_);

    ULONG size = sizeof(COMMKIT_EVENT) * COMMKIT_RING_CAPACITY;
    buffer_ = (COMMKIT_EVENT*)ExAllocatePool2(
        POOL_FLAG_NON_PAGED, size, 'BCMK');
    if (!buffer_) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }
    RtlZeroMemory(buffer_, size);
    head_ = tail_ = count_ = 0;
    return STATUS_SUCCESS;
}

void EventRingBuffer::Cleanup() {
    if (buffer_) {
        ExFreePoolWithTag(buffer_, 'BCMK');
        buffer_ = nullptr;
    }
    count_ = head_ = tail_ = 0;
}

void EventRingBuffer::Push(ULONG sequence, UINT64 timestamp,
                            ULONG com_number, ULONG event_type,
                            ULONG data_size, PVOID data) {
    KIRQL old_irql;
    KeAcquireSpinLock(&lock_, &old_irql);

    if (!buffer_) {
        KeReleaseSpinLock(&lock_, old_irql);
        return;
    }

    // 截断超长数据
    ULONG copy_size = data_size;
    if (copy_size > COMMKIT_MAX_DATA) {
        copy_size = COMMKIT_MAX_DATA;
    }

    // 写入 head 位置
    COMMKIT_EVENT* slot = &buffer_[head_];
    slot->Sequence  = (INT32)sequence;
    slot->TimeStamp = timestamp;
    slot->ComNumber = com_number;
    slot->EventType = event_type;
    slot->DataSize  = copy_size;
    if (copy_size > 0 && data) {
        RtlCopyMemory(slot->Data, data, copy_size);
    }

    // 推进 head
    head_ = (head_ + 1) % COMMKIT_RING_CAPACITY;
    if (count_ == COMMKIT_RING_CAPACITY) {
        // 满了，覆盖最旧：tail 跟随 head
        tail_ = head_;
    } else {
        count_++;
    }

    KeReleaseSpinLock(&lock_, old_irql);
}

LONG EventRingBuffer::PopBatch(PCOMMKIT_EVENT out_buf, LONG max_count) {
    if (!out_buf || max_count <= 0) return 0;

    KIRQL old_irql;
    KeAcquireSpinLock(&lock_, &old_irql);

    if (!buffer_) {
        KeReleaseSpinLock(&lock_, old_irql);
        return 0;
    }

    LONG to_pop = count_ < max_count ? count_ : max_count;
    for (LONG i = 0; i < to_pop; ++i) {
        RtlCopyMemory(&out_buf[i], &buffer_[tail_], sizeof(COMMKIT_EVENT));
        tail_ = (tail_ + 1) % COMMKIT_RING_CAPACITY;
    }
    count_ -= to_pop;

    KeReleaseSpinLock(&lock_, old_irql);
    return to_pop;
}

void EventRingBuffer::Clear() {
    KIRQL old_irql;
    KeAcquireSpinLock(&lock_, &old_irql);
    head_ = tail_ = count_ = 0;
    KeReleaseSpinLock(&lock_, old_irql);
}

} // namespace commkit_driver