/*
 * usbshare - IOCTL handling and URB forwarding
 */

#include "usbshare.h"

static VOID UsbShareCompleteTransfer(
    _In_ WDFREQUEST Request,
    _In_ WDFIOTARGET Target,
    _In_ PWDF_REQUEST_COMPLETION_PARAMS Params,
    _In_ WDFCONTEXT Context);

static NTSTATUS UsbShareHandleClaim(_In_ PDEVICE_CONTEXT Context, _In_ WDFREQUEST Request);
static NTSTATUS UsbShareHandleGetDescriptors(_In_ PDEVICE_CONTEXT Context, _In_ WDFREQUEST Request);
static NTSTATUS UsbShareHandleSubmit(_In_ PDEVICE_CONTEXT Context, _In_ WDFREQUEST Request);
static NTSTATUS UsbShareHandleCancel(_In_ PDEVICE_CONTEXT Context, _In_ WDFREQUEST Request);
static NTSTATUS UsbShareHandleSetInterface(_In_ PDEVICE_CONTEXT Context, _In_ WDFREQUEST Request);
static NTSTATUS UsbShareHandleClearHalt(_In_ PDEVICE_CONTEXT Context, _In_ WDFREQUEST Request);

VOID
UsbShareEvtIoDeviceControl(
    _In_ WDFQUEUE Queue,
    _In_ WDFREQUEST Request,
    _In_ size_t OutputBufferLength,
    _In_ size_t InputBufferLength,
    _In_ ULONG IoControlCode
)
{
    WDFDEVICE device = WdfIoQueueGetDevice(Queue);
    PDEVICE_CONTEXT context = GetDeviceContext(device);
    NTSTATUS status;

    UNREFERENCED_PARAMETER(OutputBufferLength);
    UNREFERENCED_PARAMETER(InputBufferLength);

    switch (IoControlCode) {
    case IOCTL_USBSHARE_CLAIM:
        status = UsbShareHandleClaim(context, Request);
        break;

    case IOCTL_USBSHARE_RELEASE:
        UsbShareReleaseClaim(context, WdfRequestGetFileObject(Request));
        status = STATUS_SUCCESS;
        break;

    case IOCTL_USBSHARE_GET_DESCRIPTORS:
        status = UsbShareHandleGetDescriptors(context, Request);
        break;

    case IOCTL_USBSHARE_SUBMIT:
        status = UsbShareHandleSubmit(context, Request);
        /*
         * A submitted transfer completes asynchronously; the completion
         * routine owns the request from here.
         */
        if (status == STATUS_PENDING) {
            return;
        }
        break;

    case IOCTL_USBSHARE_CANCEL:
        status = UsbShareHandleCancel(context, Request);
        break;

    case IOCTL_USBSHARE_SET_INTERFACE:
        status = UsbShareHandleSetInterface(context, Request);
        break;

    case IOCTL_USBSHARE_CLEAR_HALT:
        status = UsbShareHandleClearHalt(context, Request);
        break;

    case IOCTL_USBSHARE_RESET:
        status = WdfUsbTargetDeviceResetPortSynchronously(context->UsbDevice);
        break;

    default:
        /*
         * Not ours. As a filter we must pass it on rather than fail it —
         * some other component in the stack may be waiting for the answer.
         */
        {
            WDF_REQUEST_SEND_OPTIONS options;
            WDF_REQUEST_SEND_OPTIONS_INIT(&options, WDF_REQUEST_SEND_OPTION_SEND_AND_FORGET);
            WdfRequestFormatRequestUsingCurrentType(Request);
            if (!WdfRequestSend(Request, WdfDeviceGetIoTarget(device), &options)) {
                WdfRequestComplete(Request, WdfRequestGetStatus(Request));
            }
            return;
        }
    }

    WdfRequestComplete(Request, status);
}

static NTSTATUS
UsbShareHandleClaim(
    _In_ PDEVICE_CONTEXT Context,
    _In_ WDFREQUEST Request
)
{
    NTSTATUS status;
    PUSBSHARE_DEVICE_INFO info;
    WDFFILEOBJECT fileObject = WdfRequestGetFileObject(Request);

    status = WdfRequestRetrieveOutputBuffer(Request, sizeof(USBSHARE_DEVICE_INFO),
                                            (PVOID *)&info, NULL);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    WdfSpinLockAcquire(Context->ClaimLock);

    if (Context->Claimed && Context->ClaimOwner != fileObject) {
        WdfSpinLockRelease(Context->ClaimLock);
        return STATUS_DEVICE_BUSY;
    }

    Context->Claimed = TRUE;
    Context->ClaimOwner = fileObject;

    WdfSpinLockRelease(Context->ClaimLock);

    *info = Context->Info;
    WdfRequestSetInformation(Request, sizeof(USBSHARE_DEVICE_INFO));

    KdPrint(("usbshare: device claimed (%04x:%04x)\n", info->VendorId, info->ProductId));
    return STATUS_SUCCESS;
}

static NTSTATUS
UsbShareHandleGetDescriptors(
    _In_ PDEVICE_CONTEXT Context,
    _In_ WDFREQUEST Request
)
{
    NTSTATUS status;
    PVOID buffer;
    size_t bufferLength;

    if (Context->Descriptors == NULL) {
        status = UsbShareBuildDescriptorBlob(Context);
        if (!NT_SUCCESS(status)) {
            return status;
        }
    }

    status = WdfRequestRetrieveOutputBuffer(Request, 1, &buffer, &bufferLength);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    if (bufferLength < Context->DescriptorsLength) {
        /* Report the needed size so the caller can retry. */
        WdfRequestSetInformation(Request, Context->DescriptorsLength);
        return STATUS_BUFFER_TOO_SMALL;
    }

    RtlCopyMemory(buffer, Context->Descriptors, Context->DescriptorsLength);
    WdfRequestSetInformation(Request, Context->DescriptorsLength);

    return STATUS_SUCCESS;
}

/*
 * Builds the descriptor blob: device descriptor followed by every
 * configuration descriptor, matching what Linux returns when reading a
 * usbdevfs file. The client parses both with the same code.
 */
NTSTATUS
UsbShareBuildDescriptorBlob(
    _In_ PDEVICE_CONTEXT Context
)
{
    NTSTATUS status;
    USB_DEVICE_DESCRIPTOR deviceDescriptor;
    PUCHAR blob = NULL;
    ULONG blobSize = 0;
    ULONG offset;
    UCHAR configIndex;

    WdfUsbTargetDeviceGetDeviceDescriptor(Context->UsbDevice, &deviceDescriptor);

    /* First pass: total up the sizes. */
    blobSize = sizeof(USB_DEVICE_DESCRIPTOR);

    for (configIndex = 0; configIndex < deviceDescriptor.bNumConfigurations; configIndex++) {
        USHORT configSize = 0;

        status = WdfUsbTargetDeviceRetrieveConfigDescriptor(Context->UsbDevice, NULL, &configSize);
        if (status != STATUS_BUFFER_TOO_SMALL && !NT_SUCCESS(status)) {
            return status;
        }
        blobSize += configSize;

        /*
         * Only configuration 0 can be retrieved through this API; devices
         * with several configurations would need a raw control transfer per
         * configuration. They are rare enough to leave for later.
         */
        break;
    }

    blob = (PUCHAR)ExAllocatePool2(POOL_FLAG_NON_PAGED, blobSize, USBSHARE_POOL_TAG);
    if (blob == NULL) {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    RtlCopyMemory(blob, &deviceDescriptor, sizeof(USB_DEVICE_DESCRIPTOR));
    offset = sizeof(USB_DEVICE_DESCRIPTOR);

    {
        USHORT configSize = (USHORT)(blobSize - offset);
        status = WdfUsbTargetDeviceRetrieveConfigDescriptor(Context->UsbDevice,
                                                            blob + offset,
                                                            &configSize);
        if (!NT_SUCCESS(status)) {
            ExFreePoolWithTag(blob, USBSHARE_POOL_TAG);
            return status;
        }
    }

    Context->Descriptors = blob;
    Context->DescriptorsLength = blobSize;

    return STATUS_SUCCESS;
}

static NTSTATUS
UsbShareHandleSubmit(
    _In_ PDEVICE_CONTEXT Context,
    _In_ WDFREQUEST Request
)
{
    NTSTATUS status;
    PUSBSHARE_TRANSFER transfer;
    size_t inputLength;
    PUCHAR payload;
    WDFMEMORY urbMemory;
    PURB urb;
    WDFUSBPIPE pipe;
    PREQUEST_CONTEXT reqContext;
    WDF_OBJECT_ATTRIBUTES attributes;
    WDFIOTARGET target;

    if (!UsbShareIsClaimed(Context)) {
        return STATUS_INVALID_DEVICE_STATE;
    }

    status = WdfRequestRetrieveInputBuffer(Request, sizeof(USBSHARE_TRANSFER),
                                           (PVOID *)&transfer, &inputLength);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    if (inputLength < sizeof(USBSHARE_TRANSFER) + transfer->BufferLength) {
        return STATUS_BUFFER_TOO_SMALL;
    }
    payload = (PUCHAR)transfer + sizeof(USBSHARE_TRANSFER);

    /* Attach a context so a later cancel can find this request. */
    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attributes, REQUEST_CONTEXT);
    status = WdfObjectAllocateContext(Request, &attributes, (PVOID *)&reqContext);
    if (!NT_SUCCESS(status)) {
        return status;
    }
    reqContext->TransferId = transfer->Id;
    reqContext->ExpectedLength = transfer->BufferLength;

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    attributes.ParentObject = Request;

    if (transfer->Type == USBSHARE_TRANSFER_CONTROL) {
        status = WdfUsbTargetDeviceCreateUrb(Context->UsbDevice, &attributes,
                                             &urbMemory, &urb);
        if (!NT_SUCCESS(status)) {
            return status;
        }

        /*
         * The setup packet arrives in USB wire order and is copied verbatim.
         * Reinterpreting the fields here would only introduce a chance to get
         * the endianness wrong.
         */
        UsbBuildVendorRequest(urb,
                              URB_FUNCTION_VENDOR_DEVICE,
                              sizeof(struct _URB_CONTROL_VENDOR_OR_CLASS_REQUEST),
                              (transfer->Direction == USBSHARE_DIR_IN)
                                  ? USBD_TRANSFER_DIRECTION_IN : 0,
                              0,
                              transfer->Setup[0],  /* bmRequestType */
                              transfer->Setup[1],  /* bRequest */
                              *(USHORT *)&transfer->Setup[2],  /* wValue */
                              *(USHORT *)&transfer->Setup[4],  /* wIndex */
                              payload,
                              NULL,
                              transfer->BufferLength,
                              NULL);
    } else {
        pipe = Context->Pipes[transfer->EndpointAddress];
        if (pipe == NULL) {
            return STATUS_INVALID_PARAMETER;
        }

        status = WdfUsbTargetDeviceCreateUrb(Context->UsbDevice, &attributes,
                                             &urbMemory, &urb);
        if (!NT_SUCCESS(status)) {
            return status;
        }

        /*
         * Bulk and interrupt share one URB function; the pipe handle decides
         * which it actually is.
         */
        urb->UrbBulkOrInterruptTransfer.Hdr.Length =
            sizeof(struct _URB_BULK_OR_INTERRUPT_TRANSFER);
        urb->UrbBulkOrInterruptTransfer.Hdr.Function =
            URB_FUNCTION_BULK_OR_INTERRUPT_TRANSFER;
        urb->UrbBulkOrInterruptTransfer.PipeHandle = WdfUsbTargetPipeWdmGetPipeHandle(pipe);
        urb->UrbBulkOrInterruptTransfer.TransferBuffer = payload;
        urb->UrbBulkOrInterruptTransfer.TransferBufferLength = transfer->BufferLength;
        urb->UrbBulkOrInterruptTransfer.TransferBufferMDL = NULL;
        urb->UrbBulkOrInterruptTransfer.UrbLink = NULL;
        urb->UrbBulkOrInterruptTransfer.TransferFlags =
            (transfer->Direction == USBSHARE_DIR_IN)
                ? (USBD_TRANSFER_DIRECTION_IN | USBD_SHORT_TRANSFER_OK)
                : 0;
    }

    reqContext->UrbMemory = urbMemory;
    reqContext->Urb = urb;

    target = WdfUsbTargetDeviceGetIoTarget(Context->UsbDevice);

    status = WdfUsbTargetDeviceFormatRequestForUrb(Context->UsbDevice, Request,
                                                   urbMemory, NULL);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    WdfRequestSetCompletionRoutine(Request, UsbShareCompleteTransfer, Context);

    /* Track it so a cancel can find it. */
    WdfSpinLockAcquire(Context->PendingLock);
    WdfCollectionAdd(Context->PendingTransfers, Request);
    WdfSpinLockRelease(Context->PendingLock);

    if (!WdfRequestSend(Request, target, WDF_NO_SEND_OPTIONS)) {
        status = WdfRequestGetStatus(Request);

        WdfSpinLockAcquire(Context->PendingLock);
        WdfCollectionRemove(Context->PendingTransfers, Request);
        WdfSpinLockRelease(Context->PendingLock);

        return status;
    }

    return STATUS_PENDING;
}

static VOID
UsbShareCompleteTransfer(
    _In_ WDFREQUEST Request,
    _In_ WDFIOTARGET Target,
    _In_ PWDF_REQUEST_COMPLETION_PARAMS Params,
    _In_ WDFCONTEXT CompletionContext
)
{
    PDEVICE_CONTEXT context = (PDEVICE_CONTEXT)CompletionContext;
    PREQUEST_CONTEXT reqContext = GetRequestContext(Request);
    PUSBSHARE_TRANSFER_RESULT result;
    NTSTATUS status;
    size_t outputLength;
    ULONG transferred = 0;
    PUCHAR outPayload;

    UNREFERENCED_PARAMETER(Target);

    WdfSpinLockAcquire(context->PendingLock);
    WdfCollectionRemove(context->PendingTransfers, Request);
    WdfSpinLockRelease(context->PendingLock);

    if (reqContext->Urb != NULL) {
        transferred = reqContext->Urb->UrbBulkOrInterruptTransfer.TransferBufferLength;
    }

    status = WdfRequestRetrieveOutputBuffer(Request, sizeof(USBSHARE_TRANSFER_RESULT),
                                            (PVOID *)&result, &outputLength);
    if (!NT_SUCCESS(status)) {
        WdfRequestComplete(Request, status);
        return;
    }

    result->Id = reqContext->TransferId;
    result->Status = Params->IoStatus.Status;
    result->UsbdStatus = (reqContext->Urb != NULL)
        ? reqContext->Urb->UrbHeader.Status : 0;
    result->ActualLength = transferred;

    /*
     * Copy the received payload after the result header, but only as much as
     * the output buffer holds — a device may return more than expected.
     */
    if (transferred > 0 && outputLength > sizeof(USBSHARE_TRANSFER_RESULT)) {
        ULONG room = (ULONG)(outputLength - sizeof(USBSHARE_TRANSFER_RESULT));
        ULONG copy = (transferred < room) ? transferred : room;

        outPayload = (PUCHAR)result + sizeof(USBSHARE_TRANSFER_RESULT);
        RtlCopyMemory(outPayload,
                      reqContext->Urb->UrbBulkOrInterruptTransfer.TransferBuffer,
                      copy);

        WdfRequestSetInformation(Request, sizeof(USBSHARE_TRANSFER_RESULT) + copy);
    } else {
        WdfRequestSetInformation(Request, sizeof(USBSHARE_TRANSFER_RESULT));
    }

    /*
     * Always complete successfully: the transfer's own outcome travels in the
     * result structure. Failing the IOCTL would lose the distinction between
     * "the ioctl did not work" and "the device stalled".
     */
    WdfRequestComplete(Request, STATUS_SUCCESS);
}

static NTSTATUS
UsbShareHandleCancel(
    _In_ PDEVICE_CONTEXT Context,
    _In_ WDFREQUEST Request
)
{
    NTSTATUS status;
    PUSBSHARE_CANCEL cancel;
    ULONG i, count;
    WDFREQUEST target = NULL;

    status = WdfRequestRetrieveInputBuffer(Request, sizeof(USBSHARE_CANCEL),
                                           (PVOID *)&cancel, NULL);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    WdfSpinLockAcquire(Context->PendingLock);

    count = WdfCollectionGetCount(Context->PendingTransfers);
    for (i = 0; i < count; i++) {
        WDFREQUEST candidate = (WDFREQUEST)WdfCollectionGetItem(Context->PendingTransfers, i);
        PREQUEST_CONTEXT candidateContext = GetRequestContext(candidate);

        if (candidateContext != NULL && candidateContext->TransferId == cancel->Id) {
            target = candidate;
            break;
        }
    }

    WdfSpinLockRelease(Context->PendingLock);

    if (target == NULL) {
        /* Already finished. Not an error: the caller gets its result anyway. */
        return STATUS_SUCCESS;
    }

    WdfRequestCancelSentRequest(target);
    return STATUS_SUCCESS;
}

static NTSTATUS
UsbShareHandleSetInterface(
    _In_ PDEVICE_CONTEXT Context,
    _In_ WDFREQUEST Request
)
{
    NTSTATUS status;
    PUSBSHARE_SET_INTERFACE params;
    WDF_USB_INTERFACE_SELECT_SETTING_PARAMS settingParams;

    if (!UsbShareIsClaimed(Context)) {
        return STATUS_INVALID_DEVICE_STATE;
    }

    status = WdfRequestRetrieveInputBuffer(Request, sizeof(USBSHARE_SET_INTERFACE),
                                           (PVOID *)&params, NULL);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    if (Context->UsbInterface == NULL) {
        return STATUS_INVALID_DEVICE_STATE;
    }

    /*
     * Going through the framework rather than sending a raw SET_INTERFACE is
     * essential: the USB stack has to re-open the pipes and, for isochronous
     * endpoints, reserve bandwidth. A raw control transfer changes the device
     * without telling the stack, after which every later transfer fails.
     */
    WDF_USB_INTERFACE_SELECT_SETTING_PARAMS_INIT_SETTING(&settingParams,
                                                         params->AlternateSetting);

    status = WdfUsbInterfaceSelectSetting(Context->UsbInterface,
                                          WDF_NO_OBJECT_ATTRIBUTES,
                                          &settingParams);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    /* Pipe handles change with the setting, so rebuild the map. */
    RtlZeroMemory(Context->Pipes, sizeof(Context->Pipes));
    {
        BYTE pipeCount = WdfUsbInterfaceGetNumConfiguredPipes(Context->UsbInterface);
        BYTE i;

        for (i = 0; i < pipeCount; i++) {
            WDF_USB_PIPE_INFORMATION pipeInfo;
            WDFUSBPIPE pipe;

            WDF_USB_PIPE_INFORMATION_INIT(&pipeInfo);
            pipe = WdfUsbInterfaceGetConfiguredPipe(Context->UsbInterface, i, &pipeInfo);
            if (pipe != NULL) {
                Context->Pipes[pipeInfo.EndpointAddress] = pipe;
                WdfUsbTargetPipeSetNoMaximumPacketSizeCheck(pipe);
            }
        }
    }

    return STATUS_SUCCESS;
}

static NTSTATUS
UsbShareHandleClearHalt(
    _In_ PDEVICE_CONTEXT Context,
    _In_ WDFREQUEST Request
)
{
    NTSTATUS status;
    PUSBSHARE_CLEAR_HALT params;
    WDFUSBPIPE pipe;

    if (!UsbShareIsClaimed(Context)) {
        return STATUS_INVALID_DEVICE_STATE;
    }

    status = WdfRequestRetrieveInputBuffer(Request, sizeof(USBSHARE_CLEAR_HALT),
                                           (PVOID *)&params, NULL);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    pipe = Context->Pipes[params->EndpointAddress];
    if (pipe == NULL) {
        return STATUS_INVALID_PARAMETER;
    }

    return WdfUsbTargetPipeResetSynchronously(pipe, WDF_NO_HANDLE, NULL);
}