nt-lite/code/kernel/ke_queue.c

// Kernel message queue primitive. This file owns bounded queue storage and queue timeout behavior.

#include <stdint.h>
#include <string.h>
#include <kernel/ke_queue.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_time.h>
#include <kernel/ke_irql.h>
#include <kernel/ob_object.h>

#define QUEUE_TIMEOUT_INFINITE UINT64_MAX

// Destroys a queue and releases its private state events.
static void KiQueueDestroy(OBJECT* obj) {
    QUEUE* queue = (QUEUE*)obj;
    if (!queue)
        return;
    if (queue->NotEmptyEvent) {
        ObDereferenceObject(&queue->NotEmptyEvent->Header);
        queue->NotEmptyEvent = (EVENT*)0;
    }
    if (queue->NotFullEvent) {
        ObDereferenceObject(&queue->NotFullEvent->Header);
        queue->NotFullEvent = (EVENT*)0;
    }
}

// Updates waitable state after one message has been inserted.
static void KiQueueNoteInsertedLocked(QUEUE* queue, uint32_t previousCount) {
    if (!queue || !queue->NotEmptyEvent || !queue->NotFullEvent)
        return;

    if (previousCount == 0)
        KeSignalEvent(queue->NotEmptyEvent);
    if (queue->Count == QUEUE_DEPTH)
        KeClearEvent(queue->NotFullEvent);
}

// Updates waitable state after one message has been removed.
static void KiQueueNoteRemovedLocked(QUEUE* queue, uint32_t previousCount) {
    if (!queue || !queue->NotEmptyEvent || !queue->NotFullEvent)
        return;

    if (previousCount == QUEUE_DEPTH)
        KeSignalEvent(queue->NotFullEvent);
    if (queue->Count == 0)
        KeClearEvent(queue->NotEmptyEvent);
}

// Computes an absolute monotonic deadline for a finite timeout.
static uint64_t KiQueueTimeoutDeadline(uint64_t timeoutMilliseconds) {
    return KeMonotonicDeadlineAfterMilliseconds(timeoutMilliseconds);
}

// Computes the remaining wait time for a finite timeout deadline.
static uint64_t KiQueueRemainingMilliseconds(uint64_t deadlineMicroseconds) {
    return KeMonotonicMillisecondsUntil(deadlineMicroseconds);
}

// Waits for a queue state event using either an infinite or finite timeout.
static KSTATUS KiQueueWaitForEvent(QUEUE* queue, EVENT* event, uint64_t timeoutMilliseconds, uint64_t deadlineMicroseconds) {
    if (!queue || !event)
        return NTLSTATUS_INVALID_PARAMETER;

    uint64_t waitMilliseconds = timeoutMilliseconds;
    if (timeoutMilliseconds != QUEUE_TIMEOUT_INFINITE) {
        waitMilliseconds = KiQueueRemainingMilliseconds(deadlineMicroseconds);
        if (waitMilliseconds == 0)
            return NTLSTATUS_TIMEOUT;
    }

    EVENT* events[1] = { event };
    return KeWaitForMultipleEvents(events, 1, 0, waitMilliseconds, (uint32_t*)0);
}

// Creates a fixed-depth queue with waitable occupancy and capacity events.
KSTATUS KeQueueCreate(QUEUE** out) {
    if (!out) return NTLSTATUS_INVALID_PARAMETER;
    *out = (QUEUE*)0;

    QUEUE* q = (QUEUE*)ObAllocateObject(ObjectTypeQueue, sizeof(QUEUE));
    if (!q)  return NTLSTATUS_NO_MEMORY;
    q->Header.Destroy = KiQueueDestroy;
    q->Lock  = (SPINLOCK)SPINLOCK_INIT;
    q->Head  = q->Tail = q->Count = 0;

    KSTATUS s = KeEventCreate(EventManualReset, 0, &q->NotEmptyEvent);
    if (!KSUCCESS(s)) {
        ObDereferenceObject(&q->Header);
        return s;
    }

    s = KeEventCreate(EventManualReset, 1, &q->NotFullEvent);
    if (!KSUCCESS(s)) {
        ObDereferenceObject(&q->Header);
        return s;
    }

    *out = q;
    return NTLSTATUS_SUCCESS;
}

// Sends one fixed-size queue message, waiting according to the timeout contract.
KSTATUS KeSendMessage(QUEUE* queue, const QUEUE_MESSAGE* msg, uint64_t timeoutMilliseconds) {
    if (!queue || !msg) return NTLSTATUS_INVALID_PARAMETER;
    if (timeoutMilliseconds != 0 && !KeCanWait()) return NTLSTATUS_INVALID_IRQL;

    uint64_t deadlineMicroseconds = 0;
    if (timeoutMilliseconds != 0 && timeoutMilliseconds != QUEUE_TIMEOUT_INFINITE)
        deadlineMicroseconds = KiQueueTimeoutDeadline(timeoutMilliseconds);

    while (true) {
        uintptr_t flags = KeAcquireSpinlock(&queue->Lock);
        if (queue->Count < QUEUE_DEPTH) {
            uint32_t previousCount = queue->Count;
            queue->Ring[queue->Tail] = *msg;
            queue->Tail = (queue->Tail + 1) % QUEUE_DEPTH;
            queue->Count++;
            KiQueueNoteInsertedLocked(queue, previousCount);
            KeReleaseSpinlock(&queue->Lock, flags);
            return NTLSTATUS_SUCCESS;
        }
        KeClearEvent(queue->NotFullEvent);
        KeReleaseSpinlock(&queue->Lock, flags);

        if (timeoutMilliseconds == 0)
            return NTLSTATUS_WOULD_BLOCK;

        KSTATUS s = KiQueueWaitForEvent(queue, queue->NotFullEvent, timeoutMilliseconds, deadlineMicroseconds);
        if (s == NTLSTATUS_TIMEOUT || KFAILED(s))
            return s;
    }
}

// Receives one fixed-size queue message, waiting according to the timeout contract.
KSTATUS KeReceiveMessage(QUEUE* queue, QUEUE_MESSAGE* msg, uint64_t timeoutMilliseconds) {
    if (!queue || !msg) return NTLSTATUS_INVALID_PARAMETER;
    if (timeoutMilliseconds != 0 && !KeCanWait()) return NTLSTATUS_INVALID_IRQL;

    uint64_t deadlineMicroseconds = 0;
    if (timeoutMilliseconds != 0 && timeoutMilliseconds != QUEUE_TIMEOUT_INFINITE)
        deadlineMicroseconds = KiQueueTimeoutDeadline(timeoutMilliseconds);

    while (true) {
        uintptr_t flags = KeAcquireSpinlock(&queue->Lock);
        if (queue->Count > 0) {
            uint32_t previousCount = queue->Count;
            *msg = queue->Ring[queue->Head];
            queue->Head = (queue->Head + 1) % QUEUE_DEPTH;
            queue->Count--;
            KiQueueNoteRemovedLocked(queue, previousCount);
            KeReleaseSpinlock(&queue->Lock, flags);
            return NTLSTATUS_SUCCESS;
        }
        KeClearEvent(queue->NotEmptyEvent);
        KeReleaseSpinlock(&queue->Lock, flags);

        if (timeoutMilliseconds == 0)
            return NTLSTATUS_WOULD_BLOCK;

        KSTATUS s = KiQueueWaitForEvent(queue, queue->NotEmptyEvent, timeoutMilliseconds, deadlineMicroseconds);
        if (s == NTLSTATUS_TIMEOUT || KFAILED(s))
            return s;
    }
}