nt-lite/code/kernel/ke_irp.c

// I/O request packet infrastructure. This file owns IRP dispatch, queuing, cancellation, and remove-lock helpers.

#include <ntlite/atomic.h>
#include <stdbool.h>
#include <string.h>
#include <intrin.h>
#include <kernel/ke_irp.h>
#include <kernel/ke_event.h>
#include <kernel/ke_semaphore.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_apc.h>
#include <kernel/ob_object.h>
#include <kernel/mm_virt.h>
#include <ntlite/ddi/io_driver.h>

// IRP_QUEUE operations.

__declspec(dllexport) void IrpQueueInit(IRP_QUEUE* q, struct SEMAPHORE* sem) {
    q->Lock      = (SPINLOCK)SPINLOCK_INIT;
    q->Semaphore = sem;
    LIST_INIT(&q->Head);
}

// Implements the irp Queue Enqueue operation.
__declspec(dllexport) void IrpQueueEnqueue(IRP_QUEUE* q, IRP* irp) {
    uintptr_t flags = KeAcquireSpinlock(&q->Lock);
    if (irp->Flags & IRP_FLAG_PRIORITY)
        ListInsertHead(&q->Head, &irp->Link);
    else
        ListInsertTail(&q->Head, &irp->Link);
    KeReleaseSpinlock(&q->Lock, flags);
    KeSignalSemaphore((SEMAPHORE*)q->Semaphore);
}

// Implements the irp Queue Dequeue operation.
__declspec(dllexport) IRP* IrpQueueDequeue(IRP_QUEUE* q) {
    uintptr_t flags = KeAcquireSpinlock(&q->Lock);
    LIST_ENTRY* node = ListPopHead(&q->Head);
    KeReleaseSpinlock(&q->Lock, flags);
    if (!node) return (IRP*)0;
    return CONTAINING_RECORD(node, IRP, Link);
}

// Cancels every queued IRP and invokes an optional ownership callback.
__declspec(dllexport) void IrpQueueCancelAllWithCallback(IRP_QUEUE* q,
                                                         KSTATUS status,
                                                         IRP_QUEUE_CANCEL_CALLBACK callback,
                                                         void* context) {
    if (!q) return;
    while (true) {
        uintptr_t flags = KeAcquireSpinlock(&q->Lock);
        LIST_ENTRY* node = ListPopHead(&q->Head);
        KeReleaseSpinlock(&q->Lock, flags);
        if (!node) break;

        IRP* irp = CONTAINING_RECORD(node, IRP, Link);
        irp->Cancel = true;
        irp->Flags |= IRP_FLAG_CANCELLED;
        if (callback)
            callback(irp, context);
        IoCompleteRequest(irp, status, 0);
    }
}

// Implements the irp Queue Cancel All operation.
__declspec(dllexport) void IrpQueueCancelAll(IRP_QUEUE* q, KSTATUS status) {
    IrpQueueCancelAllWithCallback(q, status, (IRP_QUEUE_CANCEL_CALLBACK)0, (void*)0);
}

// Remove locks.

__declspec(dllexport) KSTATUS IoInitializeRemoveLock(IO_REMOVE_LOCK* lock) {
    if (!lock) return NTLSTATUS_INVALID_PARAMETER;
    memset(lock, 0, sizeof(*lock));
    lock->Lock = (SPINLOCK)SPINLOCK_INIT;
    lock->IoCount = 1; // sentinel reference released by IoReleaseRemoveLockAndWait
    KSTATUS st = KeEventCreate(EventManualReset, 0, &lock->RemovedEvent);
    if (!KSUCCESS(st)) return st;
    return NTLSTATUS_SUCCESS;
}

// Releases remove Lock.
__declspec(dllexport) void IoDestroyRemoveLock(IO_REMOVE_LOCK* lock) {
    if (!lock) return;
    if (lock->RemovedEvent) {
        ObDereferenceObject(&lock->RemovedEvent->Header);
        lock->RemovedEvent = (EVENT*)0;
    }
}

// Implements the acquire Remove Lock operation.
__declspec(dllexport) KSTATUS IoAcquireRemoveLock(IO_REMOVE_LOCK* lock, IRP* irp) {
    (void)irp;
    if (!lock) return NTLSTATUS_INVALID_PARAMETER;
    uintptr_t flags = KeAcquireSpinlock(&lock->Lock);
    if (lock->Removing) {
        KeReleaseSpinlock(&lock->Lock, flags);
        return NTLSTATUS_OBJECT_DEAD;
    }
    NtlAtomicIncrement32(&lock->IoCount);
    KeReleaseSpinlock(&lock->Lock, flags);
    return NTLSTATUS_SUCCESS;
}

// Releases remove Lock.
__declspec(dllexport) void IoReleaseRemoveLock(IO_REMOVE_LOCK* lock, IRP* irp) {
    (void)irp;
    if (!lock) return;
    int32_t left = NtlAtomicDecrement32(&lock->IoCount);
    if (left == 0 && lock->Removing && lock->RemovedEvent)
        KeSignalEvent(lock->RemovedEvent);
}

// Releases remove Lock And Wait.
__declspec(dllexport) void IoReleaseRemoveLockAndWait(IO_REMOVE_LOCK* lock, IRP* removeIrp) {
    (void)removeIrp;
    if (!lock) return;

    uintptr_t flags = KeAcquireSpinlock(&lock->Lock);
    lock->Removing = true;
    int32_t left = NtlAtomicDecrement32(&lock->IoCount); // release sentinel
    bool done = (left == 0);
    KeReleaseSpinlock(&lock->Lock, flags);

    if (done) {
        if (lock->RemovedEvent) KeSignalEvent(lock->RemovedEvent);
        return;
    }

    if (lock->RemovedEvent)
        KeWaitForEvent(lock->RemovedEvent, 0);
}

// IRP helpers.

static size_t IoIrpSize(uint8_t stackCount) {
    if (stackCount == 0) stackCount = 1;
    return sizeof(IRP) + ((size_t)stackCount - 1u) * sizeof(IO_STACK_LOCATION);
}

// Creates irp.
__declspec(dllexport) KSTATUS IoAllocateIrp(uint8_t stackCount, IRP** outIrp) {
    if (!outIrp) return NTLSTATUS_INVALID_PARAMETER;
    if (stackCount == 0) stackCount = 1;
    IRP* irp = (IRP*)MmAllocatePoolTag(IoIrpSize(stackCount), MM_POOL_TAG('i', 'r', 'p', ' '));
    if (!irp) return NTLSTATUS_NO_MEMORY;
    IoInitializeIrp(irp, stackCount);
    *outIrp = irp;
    return NTLSTATUS_SUCCESS;
}

// Releases irp.
__declspec(dllexport) void IoFreeIrp(IRP* irp) {
    if (irp) MmFreePool(irp);
}

// Initializes irp.
__declspec(dllexport) void IoInitializeIrp(IRP* irp, uint8_t stackCount) {
    if (!irp) return;
    if (stackCount == 0) stackCount = 1;
    memset(irp, 0, IoIrpSize(stackCount));
    LIST_INIT(&irp->Link);
    irp->Status = NTLSTATUS_FAILURE;
    irp->Information = 0;
    irp->CancelLock = (SPINLOCK)SPINLOCK_INIT;
    irp->StackCount = stackCount;
    irp->CurrentLocation = 0;
}

// Returns current Irp Stack Location.
__declspec(dllexport) IO_STACK_LOCATION* IoGetCurrentIrpStackLocation(IRP* irp) {
    if (!irp || irp->CurrentLocation >= irp->StackCount) return (IO_STACK_LOCATION*)0;
    return &irp->Stack[irp->CurrentLocation];
}

// Returns next Irp Stack Location.
__declspec(dllexport) IO_STACK_LOCATION* IoGetNextIrpStackLocation(IRP* irp) {
    if (!irp || irp->CurrentLocation + 1u >= irp->StackCount) return (IO_STACK_LOCATION*)0;
    return &irp->Stack[irp->CurrentLocation + 1u];
}

// Implements the skip Current Irp Stack Location operation.
__declspec(dllexport) void IoSkipCurrentIrpStackLocation(IRP* irp) {
    if (!irp) return;
    if (irp->CurrentLocation + 1u < irp->StackCount)
        irp->CurrentLocation++;
}

// Copies current Irp Stack Location To Next.
__declspec(dllexport) void IoCopyCurrentIrpStackLocationToNext(IRP* irp) {
    IO_STACK_LOCATION* cur = IoGetCurrentIrpStackLocation(irp);
    IO_STACK_LOCATION* next = IoGetNextIrpStackLocation(irp);
    if (!cur || !next) return;
    memcpy(next, cur, sizeof(*next));
    irp->CurrentLocation++;
}

// Updates completion Routine.
__declspec(dllexport) void IoSetCompletionRoutine(IRP* irp, IRP_COMPLETION_FN fn, void* ctx) {
    if (!irp) return;
    irp->CompletionFn = fn;
    irp->CompletionCtx = ctx;
}

// Updates cancel Routine.
__declspec(dllexport) IRP_CANCEL_ROUTINE IoSetCancelRoutine(IRP* irp, IRP_CANCEL_ROUTINE routine) {
    if (!irp) return (IRP_CANCEL_ROUTINE)0;
    uintptr_t flags = KeAcquireSpinlock(&irp->CancelLock);
    IRP_CANCEL_ROUTINE old = irp->CancelRoutine;
    irp->CancelRoutine = routine;
    KeReleaseSpinlock(&irp->CancelLock, flags);
    return old;
}

// Cancels irp.
__declspec(dllexport) void IoCancelIrp(IRP* irp) {
    if (!irp) return;

    uintptr_t flags = KeAcquireSpinlock(&irp->CancelLock);
    irp->Cancel = true;
    irp->Flags |= IRP_FLAG_CANCELLED;
    IRP_CANCEL_ROUTINE routine = irp->CancelRoutine;
    irp->CancelRoutine = (IRP_CANCEL_ROUTINE)0;
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    DEVICE_OBJECT* dev = ios ? ios->DeviceObject : (DEVICE_OBJECT*)0;
    KeReleaseSpinlock(&irp->CancelLock, flags);

    if (routine)
        routine(dev, irp);
}

__declspec(dllexport) void IoBuildReadWriteIrp(IRP* irp,
                                               IO_MAJOR_FUNCTION major,
                                               uint64_t byteOffset,
                                               uint32_t length,
                                               void* buffer) {
    if (!irp) return;
    IoInitializeIrp(irp, 1);
    irp->UserBuffer = buffer;
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios) return;
    ios->MajorFunction = major;
    ios->MinorFunction = 0;
    ios->Parameters.ReadWrite.ByteOffset = byteOffset;
    ios->Parameters.ReadWrite.Length = length;
}

__declspec(dllexport) void IoBuildPnpIrp(IRP* irp,
                                         IO_PNP_MINOR_FUNCTION minor) {
    if (!irp) return;
    uint8_t stackCount = irp->StackCount ? irp->StackCount : 1;
    IoInitializeIrp(irp, stackCount);
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios) return;
    ios->MajorFunction = IRP_MJ_PNP;
    ios->MinorFunction = (uint8_t)minor;
}

__declspec(dllexport) void IoBuildDeviceIoControlIrp(IRP* irp,
                                                     IO_MAJOR_FUNCTION major,
                                                     uint32_t ioControlCode,
                                                     void* inputBuffer,
                                                     uint32_t inputLength,
                                                     void* outputBuffer,
                                                     uint32_t outputLength) {
    if (!irp) return;
    uint8_t stackCount = irp->StackCount ? irp->StackCount : 1;
    IoInitializeIrp(irp, stackCount);
    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    if (!ios) return;
    ios->MajorFunction = major;
    ios->MinorFunction = 0;
    ios->Parameters.DeviceIoControl.IoControlCode = ioControlCode;
    ios->Parameters.DeviceIoControl.InputBuffer = inputBuffer;
    ios->Parameters.DeviceIoControl.InputLength = inputLength;
    ios->Parameters.DeviceIoControl.OutputBuffer = outputBuffer;
    ios->Parameters.DeviceIoControl.OutputLength = outputLength;
}

// Completion/submission.

__declspec(dllexport) KSTATUS IoCompleteRequest(IRP* irp, KSTATUS status, size_t information) {
    if (!irp) return NTLSTATUS_INVALID_PARAMETER;
    if (irp->Flags & IRP_FLAG_COMPLETED)
        return NTLSTATUS_ALREADY_EXISTS;

    irp->Flags |= IRP_FLAG_COMPLETED;
    irp->Status = status;
    irp->Information = information;

    // Completion callbacks may recycle an embedded IRP, while a synchronous
    // waiter may release stack storage as soon as its event is signaled.
    // Snapshot every completion target before either ownership handoff.
    KAPC* completionApc = (KAPC*)irp->CompletionApc;
    EVENT* completionEvent = (EVENT*)irp->CompletionEvent;
    IRP_COMPLETION_FN completionFn = irp->CompletionFn;
    void* completionCtx = irp->CompletionCtx;

    IO_STACK_LOCATION* ios = IoGetCurrentIrpStackLocation(irp);
    DEVICE_OBJECT* dev = ios ? ios->DeviceObject : (DEVICE_OBJECT*)0;
    if (dev) {
        NtlAtomicDecrement32(&dev->OutstandingIrps);
        IoDereferenceDeviceObject(dev);
        ios->DeviceObject = (DEVICE_OBJECT*)0;
    }

    if (completionApc)
        KeInsertQueueApc(completionApc, irp, (void*)(uintptr_t)status);

    // The callback owns the final callback-based handoff. Do not dereference
    // the IRP after it returns because embedded requests may be freed there.
    if (completionFn)
        completionFn(irp, completionCtx);

    // The event is the final synchronous ownership handoff. A waiter may
    // invalidate a stack IRP immediately after this signal becomes visible.
    if (completionEvent)
        KeSignalEvent(completionEvent);

    return NTLSTATUS_SUCCESS;
}

// Submits irp.
__declspec(dllexport) KSTATUS IoSubmitIrp(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return NTLSTATUS_INVALID_PARAMETER;
    irp->Status = NTLSTATUS_FAILURE;
    irp->Information = 0;
    irp->Flags &= ~IRP_FLAG_COMPLETED;
    return IoCallDriver(device, irp);
}

// Submits irp Sync.
__declspec(dllexport) KSTATUS IoSubmitIrpSync(DEVICE_OBJECT* device, IRP* irp) {
    if (!device || !irp) return NTLSTATUS_INVALID_PARAMETER;
    if (!KeCanWait()) return NTLSTATUS_INVALID_IRQL;

    EVENT* ev = (EVENT*)0;
    KSTATUS st = KeEventCreate(EventAutoReset, 0, &ev);
    if (!KSUCCESS(st)) return st;

    irp->CompletionEvent = ev;

    // Drivers may retain this IRP, its transfer buffer, and the caller's
    // kernel stack until completion. Defer process termination across that
    // ownership interval rather than letting teardown free live storage.
    KeEnterTerminationDeferral();
    st = IoSubmitIrp(device, irp);
    if (st == NTLSTATUS_PENDING) {
        do {
            st = KeWaitForEvent(ev, 0);
        } while (st != NTLSTATUS_SUCCESS);
        st = irp->Status;
    }

    // Completion events are signaled only after IoCompleteRequest has finished
    // every operation that can dereference a synchronous caller's stack IRP.
    irp->CompletionEvent = (EVENT*)0;
    ObDereferenceObject(&ev->Header);
    KeLeaveTerminationDeferral();
    return st;
}