yttrium-virtio-gpu/Balloon/sys/queue.c

#include "precomp.h"

#if defined(EVENT_TRACING)
#include "queue.tmh"
#endif

#ifdef USE_BALLOON_SERVICE

EVT_WDF_IO_QUEUE_IO_WRITE BalloonIoWrite;
EVT_WDF_IO_QUEUE_IO_STOP BalloonIoStop;
EVT_WDF_REQUEST_CANCEL BalloonEvtRequestCancel;

#ifdef ALLOC_PRAGMA
#pragma alloc_text(PAGE, BalloonQueueInitialize)
#endif

NTSTATUS
BalloonQueueInitialize(IN WDFDEVICE Device)
{
    NTSTATUS status = STATUS_SUCCESS;
    WDF_IO_QUEUE_CONFIG queueConfig = {0};

    PAGED_CODE();

    WDF_IO_QUEUE_CONFIG_INIT_DEFAULT_QUEUE(&queueConfig, WdfIoQueueDispatchSequential);
    queueConfig.EvtIoWrite = BalloonIoWrite;
    queueConfig.EvtIoStop = BalloonIoStop;

    status = WdfIoQueueCreate(Device, &queueConfig, WDF_NO_OBJECT_ATTRIBUTES, WDF_NO_HANDLE);

    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_PNP, "WdfIoQueueCreate failed 0x%08x\n", status);
    }
    return status;
}

VOID BalloonIoWrite(IN WDFQUEUE Queue, IN WDFREQUEST Request, IN size_t Length)
{
    PVOID buffer = NULL;
    size_t buffSize;
    PDEVICE_CONTEXT devCtx = NULL;
    NTSTATUS status = STATUS_SUCCESS;

    TraceEvents(TRACE_LEVEL_INFORMATION,
                DBG_HW_ACCESS,
                "BalloonIoWrite Called! Queue 0x%p, Request 0x%p Length %d\n",
                Queue,
                Request,
                Length);

    devCtx = GetDeviceContext(WdfIoQueueGetDevice(Queue));

    if (Length < sizeof(BALLOON_STAT))
    {
        TraceEvents(TRACE_LEVEL_INFORMATION,
                    DBG_HW_ACCESS,
                    "BalloonIoWrite Buffer Length too small %d, expected is %d\n",
                    Length,
                    sizeof(BALLOON_STAT));
        WdfRequestComplete(Request, STATUS_BUFFER_TOO_SMALL);
        return;
    }

    status = WdfRequestRetrieveInputBuffer(Request, sizeof(BALLOON_STAT), &buffer, &buffSize);
    if (!NT_SUCCESS(status) || (buffer == NULL))
    {
        TraceEvents(TRACE_LEVEL_INFORMATION,
                    DBG_HW_ACCESS,
                    "BalloonIoWrite Could not get request memory buffer 0x%08x\n",
                    status);
        WdfRequestCompleteWithInformation(Request, status, 0L);
        return;
    }

    if (devCtx->HandleWriteRequest)
    {
        Length = min(buffSize, sizeof(BALLOON_STAT) * VIRTIO_BALLOON_S_NR);
#if 0
        {
            size_t i;
            PBALLOON_STAT stat = (PBALLOON_STAT)buffer;
            for (i = 0; i < Length / sizeof(BALLOON_STAT); ++i)
            {
                TraceEvents(TRACE_LEVEL_INFORMATION, DBG_HW_ACCESS,
                    "tag = %d, val = %08I64X\n\n", stat[i].tag, stat[i].val);
            }
        }
#endif
        RtlCopyMemory(devCtx->MemStats, buffer, Length);
        WdfRequestCompleteWithInformation(Request, status, Length);
        TraceEvents(TRACE_LEVEL_INFORMATION,
                    DBG_HW_ACCESS,
                    "WdfRequestCompleteWithInformation Called! Queue 0x%p, Request 0x%p Length %d Status 0x%08x\n",
                    Queue,
                    Request,
                    Length,
                    status);

        devCtx->HandleWriteRequest = FALSE;
        BalloonMemStats(WdfIoQueueGetDevice(Queue));
    }
    else
    {
        status = WdfRequestMarkCancelableEx(Request, BalloonEvtRequestCancel);
        if (!NT_SUCCESS(status))
        {
            TraceEvents(TRACE_LEVEL_ERROR, DBG_READ, "WdfRequestMarkCancelableEx failed: 0x%08x\n", status);
            WdfRequestCompleteWithInformation(Request, status, 0L);
            return;
        }

        devCtx->PendingWriteRequest = Request;
    }
}

VOID BalloonIoStop(IN WDFQUEUE Queue, IN WDFREQUEST Request, IN ULONG ActionFlags)
{
    UNREFERENCED_PARAMETER(Queue);

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_READ, "--> %!FUNC! Request: %p", Request);

    if (ActionFlags & WdfRequestStopRequestCancelable)
    {
        NTSTATUS status = WdfRequestUnmarkCancelable(Request);
        __analysis_assume(status != STATUS_NOT_SUPPORTED);
        if (status == STATUS_CANCELLED)
        {
            WdfRequestStopAcknowledge(Request, FALSE);
            goto end_io_stop;
        }
    }

    if (ActionFlags & WdfRequestStopActionSuspend)
    {
        WdfRequestStopAcknowledge(Request, TRUE);
    }
    else if (ActionFlags & WdfRequestStopActionPurge)
    {
        WdfRequestComplete(Request, STATUS_CANCELLED);
    }
    else
    {
        WdfRequestComplete(Request, STATUS_UNSUCCESSFUL);
    }

end_io_stop:
    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_READ, "<-- %!FUNC!");
}

VOID BalloonEvtRequestCancel(IN WDFREQUEST Request)
{
    PDEVICE_CONTEXT devCtx = GetDeviceContext(WdfIoQueueGetDevice(WdfRequestGetIoQueue(Request)));

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_READ, "--> BalloonEvtRequestCancel Cancelled Request: %p\n", Request);

    NT_ASSERT(devCtx->PendingWriteRequest == Request);
    devCtx->PendingWriteRequest = NULL;
    WdfRequestComplete(Request, STATUS_CANCELLED);

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_READ, "<-- BalloonEvtRequestCancel\n");
}

#endif // USE_BALLOON_SERVICE