/*
 * usbshare - driver entry, device setup and claim lifecycle
 */

#include "usbshare.h"

NTSTATUS
DriverEntry(
    _In_ PDRIVER_OBJECT DriverObject,
    _In_ PUNICODE_STRING RegistryPath
)
{
    WDF_DRIVER_CONFIG config;
    NTSTATUS status;

    WDF_DRIVER_CONFIG_INIT(&config, UsbShareEvtDeviceAdd);

    status = WdfDriverCreate(DriverObject, RegistryPath,
                             WDF_NO_OBJECT_ATTRIBUTES, &config, WDF_NO_HANDLE);
    if (!NT_SUCCESS(status)) {
        KdPrint(("usbshare: WdfDriverCreate failed 0x%x\n", status));
    }

    return status;
}

NTSTATUS
UsbShareEvtDeviceAdd(
    _In_ WDFDRIVER Driver,
    _Inout_ PWDFDEVICE_INIT DeviceInit
)
{
    NTSTATUS status;
    WDFDEVICE device;
    WDF_OBJECT_ATTRIBUTES attributes;
    WDF_PNPPOWER_EVENT_CALLBACKS pnpCallbacks;
    WDF_FILEOBJECT_CONFIG fileConfig;
    WDF_IO_QUEUE_CONFIG queueConfig;
    PDEVICE_CONTEXT context;
    WDFQUEUE queue;

    UNREFERENCED_PARAMETER(Driver);

    /*
     * Declaring ourselves a filter is what makes this driver safe to attach
     * to arbitrary devices: the framework then forwards every request we do
     * not explicitly handle to the driver below, so a device we know nothing
     * about keeps working exactly as before.
     */
    WdfFdoInitSetFilter(DeviceInit);

    WDF_PNPPOWER_EVENT_CALLBACKS_INIT(&pnpCallbacks);
    pnpCallbacks.EvtDevicePrepareHardware = UsbShareEvtDevicePrepareHardware;
    WdfDeviceInitSetPnpPowerEventCallbacks(DeviceInit, &pnpCallbacks);

    /*
     * File create and close callbacks give us the claim lifecycle: a claim is
     * tied to a handle, so when the client exits — cleanly or not — the
     * kernel closes the handle and the device goes back to its class driver.
     * Without this a crashed client would leave hardware unusable until
     * reboot.
     */
    WDF_FILEOBJECT_CONFIG_INIT(&fileConfig,
                               UsbShareEvtDeviceFileCreate,
                               UsbShareEvtFileClose,
                               WDF_NO_EVENT_CALLBACK); /* no cleanup callback */
    WdfDeviceInitSetFileObjectConfig(DeviceInit, &fileConfig,
                                     WDF_NO_OBJECT_ATTRIBUTES);

    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attributes, DEVICE_CONTEXT);

    status = WdfDeviceCreate(&DeviceInit, &attributes, &device);
    if (!NT_SUCCESS(status)) {
        KdPrint(("usbshare: WdfDeviceCreate failed 0x%x\n", status));
        return status;
    }

    context = GetDeviceContext(device);
    RtlZeroMemory(context, sizeof(DEVICE_CONTEXT));

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

    status = WdfSpinLockCreate(&attributes, &context->ClaimLock);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = WdfSpinLockCreate(&attributes, &context->PendingLock);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    status = WdfCollectionCreate(&attributes, &context->PendingTransfers);
    if (!NT_SUCCESS(status)) {
        return status;
    }

    /*
     * Default queue. Requests we do not recognise are forwarded down by the
     * framework because this is a filter device.
     */
    WDF_IO_QUEUE_CONFIG_INIT_DEFAULT_QUEUE(&queueConfig, WdfIoQueueDispatchParallel);
    queueConfig.EvtIoDeviceControl = UsbShareEvtIoDeviceControl;
    queueConfig.EvtIoInternalDeviceControl = UsbShareEvtIoInternalDeviceControl;
    queueConfig.EvtIoDefault = UsbShareEvtIoDefault;

    status = WdfIoQueueCreate(device, &queueConfig, WDF_NO_OBJECT_ATTRIBUTES, &queue);
    if (!NT_SUCCESS(status)) {
        KdPrint(("usbshare: WdfIoQueueCreate failed 0x%x\n", status));
        return status;
    }

    /* Publish the interface so user mode can find this device. */
    status = WdfDeviceCreateDeviceInterface(device, &GUID_DEVINTERFACE_USBSHARE, NULL);
    if (!NT_SUCCESS(status)) {
        KdPrint(("usbshare: WdfDeviceCreateDeviceInterface failed 0x%x\n", status));
        return status;
    }

    return STATUS_SUCCESS;
}

NTSTATUS
UsbShareEvtDevicePrepareHardware(
    _In_ WDFDEVICE Device,
    _In_ WDFCMRESLIST ResourcesRaw,
    _In_ WDFCMRESLIST ResourcesTranslated
)
{
    NTSTATUS status;
    PDEVICE_CONTEXT context = GetDeviceContext(Device);
    WDF_USB_DEVICE_CREATE_CONFIG createConfig;
    USB_DEVICE_DESCRIPTOR deviceDescriptor;
    WDF_USB_DEVICE_SELECT_CONFIG_PARAMS configParams;
    UCHAR i;

    UNREFERENCED_PARAMETER(ResourcesRaw);
    UNREFERENCED_PARAMETER(ResourcesTranslated);

    /* PrepareHardware can run more than once across power transitions. */
    if (context->UsbDevice != NULL) {
        return STATUS_SUCCESS;
    }

    WDF_USB_DEVICE_CREATE_CONFIG_INIT(&createConfig, USBD_CLIENT_CONTRACT_VERSION_602);

    status = WdfUsbTargetDeviceCreateWithParameters(Device, &createConfig,
                                                    WDF_NO_OBJECT_ATTRIBUTES,
                                                    &context->UsbDevice);
    if (!NT_SUCCESS(status)) {
        KdPrint(("usbshare: WdfUsbTargetDeviceCreateWithParameters failed 0x%x\n", status));
        return status;
    }

    WdfUsbTargetDeviceGetDeviceDescriptor(context->UsbDevice, &deviceDescriptor);

    context->Info.VendorId = deviceDescriptor.idVendor;
    context->Info.ProductId = deviceDescriptor.idProduct;
    context->Info.BcdDevice = deviceDescriptor.bcdDevice;
    context->Info.DeviceClass = deviceDescriptor.bDeviceClass;
    context->Info.DeviceSubClass = deviceDescriptor.bDeviceSubClass;
    context->Info.DeviceProtocol = deviceDescriptor.bDeviceProtocol;
    context->Info.NumConfigurations = deviceDescriptor.bNumConfigurations;
    context->Info.ConfigurationValue = 1;

    /*
     * Select a configuration so pipe handles become available.
     *
     * This is the part most likely to need adjusting: on a device the class
     * driver has already configured, selecting again may be redundant or
     * disruptive. A more careful implementation would query the current
     * configuration first and only select if none is active.
     */
    WDF_USB_DEVICE_SELECT_CONFIG_PARAMS_INIT_SINGLE_INTERFACE(&configParams);

    status = WdfUsbTargetDeviceSelectConfig(context->UsbDevice,
                                            WDF_NO_OBJECT_ATTRIBUTES,
                                            &configParams);
    if (!NT_SUCCESS(status)) {
        KdPrint(("usbshare: WdfUsbTargetDeviceSelectConfig failed 0x%x\n", status));
        /*
         * Not fatal: without pipes only control transfers work, but the
         * filter must not break the device for the class driver either way.
         */
        return STATUS_SUCCESS;
    }

    context->UsbInterface = configParams.Types.SingleInterface.ConfiguredUsbInterface;

    /* Map pipes by full endpoint address. */
    {
        BYTE pipeCount = configParams.Types.SingleInterface.NumberConfiguredPipes;
        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;

                /*
                 * Let short reads through. Without this a transfer that
                 * returns fewer bytes than requested fails, which is normal
                 * and expected for interrupt endpoints.
                 */
                WdfUsbTargetPipeSetNoMaximumPacketSizeCheck(pipe);
            }
        }
    }

    return STATUS_SUCCESS;
}

VOID
UsbShareEvtDeviceFileCreate(
    _In_ WDFDEVICE Device,
    _In_ WDFREQUEST Request,
    _In_ WDFFILEOBJECT FileObject
)
{
    UNREFERENCED_PARAMETER(Device);
    UNREFERENCED_PARAMETER(FileObject);

    /* Opening the handle is always allowed; claiming is a separate step. */
    WdfRequestComplete(Request, STATUS_SUCCESS);
}

VOID
UsbShareEvtFileClose(
    _In_ WDFFILEOBJECT FileObject
)
{
    WDFDEVICE device = WdfFileObjectGetDevice(FileObject);
    PDEVICE_CONTEXT context = GetDeviceContext(device);

    /*
     * The safety net: if this handle held the claim, give the device back.
     * This runs whether the client exited cleanly or was killed.
     */
    UsbShareReleaseClaim(context, FileObject);
}

BOOLEAN
UsbShareIsClaimed(
    _In_ PDEVICE_CONTEXT Context
)
{
    BOOLEAN claimed;

    WdfSpinLockAcquire(Context->ClaimLock);
    claimed = Context->Claimed;
    WdfSpinLockRelease(Context->ClaimLock);

    return claimed;
}

VOID
UsbShareReleaseClaim(
    _In_ PDEVICE_CONTEXT Context,
    _In_opt_ WDFFILEOBJECT Owner
)
{
    BOOLEAN released = FALSE;

    WdfSpinLockAcquire(Context->ClaimLock);

    /*
     * With an owner given, only that owner may release — otherwise closing an
     * unrelated handle would hand the device back while a client is using it.
     */
    if (Context->Claimed && (Owner == NULL || Context->ClaimOwner == Owner)) {
        Context->Claimed = FALSE;
        Context->ClaimOwner = NULL;
        released = TRUE;
    }

    WdfSpinLockRelease(Context->ClaimLock);

    if (released) {
        KdPrint(("usbshare: device released\n"));
    }
}