// Kernel asynchronous I/O request object support.
// This file owns the common wait, completion, cancellation, APC, IOSB, and
// locked-buffer lifetime rules for all async I/O families.

#include <string.h>
#include <kernel/io_async.h>
#include <kernel/ke_event.h>
#include <kernel/ke_apc.h>
#include <kernel/mm_virt.h>
#include <ntlite/syscall.h>

// Releases resources owned by the common async request object.
static void KiIoAsyncRequestDestroy(OBJECT* obj) {
    IO_ASYNC_REQUEST_OBJECT* request =
        CONTAINING_RECORD(obj, IO_ASYNC_REQUEST_OBJECT, Header);

    if (request->CleanupRoutine)
        request->CleanupRoutine(request, request->CleanupContext);

    if (request->WaitEvent) {
        ObDereferenceObject(&request->WaitEvent->Header);
        request->WaitEvent = (struct EVENT*)0;
    }
    if (request->CompletionEvent) {
        ObDereferenceObject(&request->CompletionEvent->Header);
        request->CompletionEvent = (struct EVENT*)0;
    }
    if (request->LockedBuffer) {
        MmUnlockUserBuffer(request->LockedBuffer);
        request->LockedBuffer = (MM_LOCKED_USER_BUFFER*)0;
    }
    if (request->LockedIosb) {
        MmUnlockUserBuffer(request->LockedIosb);
        request->LockedIosb = (MM_LOCKED_USER_BUFFER*)0;
    }
}

// Allocates a typed async request object with the common header first.
__declspec(dllexport) KSTATUS IoCreateAsyncRequest(size_t objectSize,
                                                   IO_ASYNC_REQUEST_CLEANUP_ROUTINE cleanupRoutine,
                                                   void* cleanupContext,
                                                   IO_ASYNC_REQUEST_CANCEL_ROUTINE cancelRoutine,
                                                   void* cancelContext,
                                                   IO_ASYNC_REQUEST_OBJECT** requestOut) {
    if (!requestOut) return NTLSTATUS_INVALID_PARAMETER;
    *requestOut = (IO_ASYNC_REQUEST_OBJECT*)0;
    if (objectSize < sizeof(IO_ASYNC_REQUEST_OBJECT)) return NTLSTATUS_INVALID_PARAMETER;

    IO_ASYNC_REQUEST_OBJECT* request =
        (IO_ASYNC_REQUEST_OBJECT*)ObAllocateObject(ObjectTypeAsyncIo, objectSize);
    if (!request) return NTLSTATUS_NO_MEMORY;

    request->Header.Destroy = KiIoAsyncRequestDestroy;
    LIST_INIT(&request->Link);
    request->State = IoAsyncRequestQueued;
    request->CancelRequested = 0;
    request->IosbSnapshot.Status = NTLSTATUS_PENDING;
    request->IosbSnapshot.Information = 0;
    request->CleanupRoutine = cleanupRoutine;
    request->CleanupContext = cleanupContext;
    request->CancelRoutine = cancelRoutine;
    request->CancelContext = cancelContext;

    KSTATUS status = KeEventCreate(EventManualReset, 0, &request->WaitEvent);
    if (!KSUCCESS(status)) {
        ObDereferenceObject(&request->Header);
        return status;
    }

    *requestOut = request;
    return NTLSTATUS_SUCCESS;
}

// Returns the object header embedded in an async request.
__declspec(dllexport) OBJECT* IoAsyncRequestObject(IO_ASYNC_REQUEST_OBJECT* request) {
    return request ? &request->Header : (OBJECT*)0;
}

// Attaches caller-visible completion state to a new async request.
__declspec(dllexport) void IoInitializeAsyncRequestCompletion(IO_ASYNC_REQUEST_OBJECT* request,
                                                              IO_STATUS_BLOCK* iosb,
                                                              EVENT* completionEvent,
                                                              KAPC* completionApc,
                                                              MM_LOCKED_USER_BUFFER* lockedBuffer,
                                                              MM_LOCKED_USER_BUFFER* lockedIosb) {
    if (!request) return;

    request->CompletionEvent = completionEvent;
    request->CompletionApc = completionApc;
    request->Iosb = iosb;
    request->LockedBuffer = lockedBuffer;
    request->LockedIosb = lockedIosb;
    request->IosbSnapshot.Status = NTLSTATUS_PENDING;
    request->IosbSnapshot.Information = 0;

    if (iosb) {
        iosb->Status = NTLSTATUS_PENDING;
        iosb->Information = 0;
    }
    if (completionEvent)
        ObReferenceObject(&completionEvent->Header);
}

// Updates the lifecycle state owned by the request's operation family.
__declspec(dllexport) void IoAsyncRequestSetState(IO_ASYNC_REQUEST_OBJECT* request,
                                                  IO_ASYNC_REQUEST_STATE state) {
    if (request) request->State = state;
}

// Returns the last lifecycle state published by the operation family.
__declspec(dllexport) IO_ASYNC_REQUEST_STATE IoAsyncRequestGetState(IO_ASYNC_REQUEST_OBJECT* request) {
    return request ? request->State : IoAsyncRequestCancelled;
}

// Marks a request as cancelled so providers can observe it before completion.
__declspec(dllexport) void IoAsyncRequestRequestCancel(IO_ASYNC_REQUEST_OBJECT* request) {
    if (request) request->CancelRequested = 1;
}

// Returns whether cancellation has been requested for the async request.
__declspec(dllexport) uint8_t IoAsyncRequestIsCancelRequested(IO_ASYNC_REQUEST_OBJECT* request) {
    return request ? request->CancelRequested : 0;
}

// Returns the wait event owned by an async request.
__declspec(dllexport) EVENT* IoGetAsyncRequestWaitEvent(IO_ASYNC_REQUEST_OBJECT* request) {
    return request ? request->WaitEvent : (EVENT*)0;
}

// Waits for an async request to complete according to wait flags.
__declspec(dllexport) KSTATUS IoWaitAsyncRequest(IO_ASYNC_REQUEST_OBJECT* request,
                                                 uint32_t flags) {
    if (!request || !request->WaitEvent) return NTLSTATUS_INVALID_PARAMETER;
    if (flags & ~SYS_WAIT_FLAG_ALERTABLE) return NTLSTATUS_INVALID_PARAMETER;
    return KeWaitForEvent(request->WaitEvent,
                          (flags & SYS_WAIT_FLAG_ALERTABLE) ? 1 : 0);
}

// Cancels an async request through its operation-family cancel routine.
__declspec(dllexport) KSTATUS IoCancelAsyncRequest(IO_ASYNC_REQUEST_OBJECT* request) {
    if (!request) return NTLSTATUS_INVALID_PARAMETER;
    if (request->CancelRoutine)
        return request->CancelRoutine(request, request->CancelContext);

    request->CancelRequested = 1;
    if (request->State == IoAsyncRequestCompleted ||
        request->State == IoAsyncRequestCancelled)
        return NTLSTATUS_OBJECT_DEAD;
    return NTLSTATUS_PENDING;
}

// Completes an async request and publishes all caller-visible completion state.
__declspec(dllexport) KSTATUS IoCompleteAsyncRequest(IO_ASYNC_REQUEST_OBJECT* request,
                                                     KSTATUS status,
                                                     size_t information) {
    if (!request) return NTLSTATUS_INVALID_PARAMETER;

    request->IosbSnapshot.Status = status;
    request->IosbSnapshot.Information = information;
    request->State = (status == NTLSTATUS_CANCELLED) ? IoAsyncRequestCancelled
                                                   : IoAsyncRequestCompleted;

    if (request->Iosb) {
        request->Iosb->Status = status;
        request->Iosb->Information = information;
    }

    // User APCs run against the same IOSB state that waiters will observe.
    if (request->CompletionApc) {
        if (!KeInsertQueueApc(request->CompletionApc,
                              request->Iosb,
                              (void*)(uintptr_t)status)) {
            if (request->CompletionApc->RundownRoutine)
                request->CompletionApc->RundownRoutine(request->CompletionApc);
        }
    }

    if (request->CompletionEvent)
        KeSignalEvent(request->CompletionEvent);
    if (request->WaitEvent)
        KeSignalEvent(request->WaitEvent);

    if (request->LockedBuffer) {
        MmUnlockUserBuffer(request->LockedBuffer);
        request->LockedBuffer = (MM_LOCKED_USER_BUFFER*)0;
    }
    if (request->LockedIosb) {
        MmUnlockUserBuffer(request->LockedIosb);
        request->LockedIosb = (MM_LOCKED_USER_BUFFER*)0;
    }

    return NTLSTATUS_SUCCESS;
}

// Copies the last completion snapshot out of an async request.
__declspec(dllexport) KSTATUS IoQueryAsyncRequest(IO_ASYNC_REQUEST_OBJECT* request,
                                                  IO_STATUS_BLOCK* iosbOut) {
    if (!request || !iosbOut) return NTLSTATUS_INVALID_PARAMETER;
    iosbOut->Status = request->IosbSnapshot.Status;
    iosbOut->Information = request->IosbSnapshot.Information;
    return NTLSTATUS_SUCCESS;
}