diff --git a/UsbDk/RedirectorStrategy.cpp b/UsbDk/RedirectorStrategy.cpp index d677046..edce670 100644 --- a/UsbDk/RedirectorStrategy.cpp +++ b/UsbDk/RedirectorStrategy.cpp @@ -189,6 +189,8 @@ NTSTATUS CUsbDkRedirectorStrategy::PNPPreProcess(PIRP Irp) typedef struct tag_USBDK_REDIRECTOR_REQUEST_CONTEXT { + bool PreprocessingDone; + WDFMEMORY LockedBuffer; ULONG64 EndpointAddress; USB_DK_TRANSFER_TYPE TransferType; @@ -322,37 +324,30 @@ NTSTATUS CUsbDkRedirectorStrategy::IoInCallerContextRWControlTransfer(CRedirecto } size_t bufferLength = static_cast(TransferRequest.bufferLength) - sizeof(WDF_USB_CONTROL_SETUP_PACKET); - if (bufferLength > 0) + + if (context->SetupPacket.Packet.bm.Request.Dir == BMREQUEST_HOST_TO_DEVICE) // write { - if (context->SetupPacket.Packet.bm.Request.Dir == BMREQUEST_HOST_TO_DEVICE) // write - { #pragma warning(push) #pragma warning(disable:4244) //Unsafe conversions on 32 bit - status = WdfRequest.LockUserBufferForRead(static_cast(static_cast(TransferRequest.buffer) + 1), - bufferLength, context->LockedBuffer); + status = WdfRequest.LockUserBufferForRead(static_cast(static_cast(TransferRequest.buffer) + 1), + bufferLength, context->LockedBuffer); #pragma warning(pop) - if (!NT_SUCCESS(status)) - { - TraceEvents(TRACE_LEVEL_ERROR, TRACE_REDIRECTOR, "%!FUNC! LockUserBufferForRead failed %!STATUS!", status); - } - } - else // read + if (!NT_SUCCESS(status)) { -#pragma warning(push) -#pragma warning(disable:4244) //Unsafe conversions on 32 bit - status = WdfRequest.LockUserBufferForWrite(static_cast(static_cast(TransferRequest.buffer) + 1), - bufferLength, context->LockedBuffer); -#pragma warning(pop) - if (!NT_SUCCESS(status)) - { - TraceEvents(TRACE_LEVEL_ERROR, TRACE_REDIRECTOR, "%!FUNC! LockUserBufferForWrite failed %!STATUS!", status); - } + TraceEvents(TRACE_LEVEL_ERROR, TRACE_REDIRECTOR, "%!FUNC! LockUserBufferForRead failed %!STATUS!", status); } } - else + else // read { - context->LockedBuffer = WDF_NO_HANDLE; - status = STATUS_SUCCESS; +#pragma warning(push) +#pragma warning(disable:4244) //Unsafe conversions on 32 bit + status = WdfRequest.LockUserBufferForWrite(static_cast(static_cast(TransferRequest.buffer) + 1), + bufferLength, context->LockedBuffer); +#pragma warning(pop) + if (!NT_SUCCESS(status)) + { + TraceEvents(TRACE_LEVEL_ERROR, TRACE_REDIRECTOR, "%!FUNC! LockUserBufferForWrite failed %!STATUS!", status); + } } return status; @@ -391,6 +386,8 @@ void CUsbDkRedirectorStrategy::IoInCallerContext(WDFDEVICE Device, WDFREQUEST Re } } + WdfRequest.Context()->PreprocessingDone = true; + if (NT_SUCCESS(status)) { CUsbDkFilterStrategy::IoInCallerContext(Device, WdfRequest.Detach()); @@ -526,6 +523,14 @@ void CUsbDkRedirectorStrategy::WritePipe(WDFREQUEST Request) CRedirectorRequest WdfRequest(Request); PUSBDK_REDIRECTOR_REQUEST_CONTEXT Context = WdfRequest.Context(); + if (!Context->PreprocessingDone) + { + //May happen if EvtioInCallerContext wasn't called due to + //lack of memory of other resources + WdfRequest.SetStatus(STATUS_INSUFFICIENT_RESOURCES); + return; + } + switch (Context->TransferType) { case ControlTransferType: @@ -533,7 +538,6 @@ void CUsbDkRedirectorStrategy::WritePipe(WDFREQUEST Request) break; case BulkTransferType: case InterruptTransferType: - if (Context->LockedBuffer != WDF_NO_HANDLE) { m_Target.WritePipeAsync(WdfRequest.Detach(), Context->EndpointAddress, Context->LockedBuffer, [](WDFREQUEST Request, WDFIOTARGET, PWDF_REQUEST_COMPLETION_PARAMS Params, WDFCONTEXT) @@ -555,13 +559,8 @@ void CUsbDkRedirectorStrategy::WritePipe(WDFREQUEST Request) WdfRequest.SetStatus(status); }); } - else - { - WdfRequest.SetStatus(STATUS_INVALID_PARAMETER); - } break; case IsochronousTransferType: - if (Context->LockedBuffer != WDF_NO_HANDLE) { CPreAllocatedWdfMemoryBufferT PacketSizesArray(Context->LockedIsochronousPacketsArray); @@ -572,10 +571,6 @@ void CUsbDkRedirectorStrategy::WritePipe(WDFREQUEST Request) PacketSizesArray.ArraySize(), IsoRWCompletion); } - else - { - WdfRequest.SetStatus(STATUS_INVALID_PARAMETER); - } break; default: TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Error: Wrong transfer type: %d\n", Context->TransferType); @@ -587,6 +582,15 @@ void CUsbDkRedirectorStrategy::ReadPipe(WDFREQUEST Request) { CRedirectorRequest WdfRequest(Request); PUSBDK_REDIRECTOR_REQUEST_CONTEXT Context = WdfRequest.Context(); + + if (!Context->PreprocessingDone) + { + //May happen if EvtioInCallerContext wasn't called due to + //lack of memory of other resources + WdfRequest.SetStatus(STATUS_INSUFFICIENT_RESOURCES); + return; + } + switch (Context->TransferType) { case ControlTransferType: @@ -594,7 +598,6 @@ void CUsbDkRedirectorStrategy::ReadPipe(WDFREQUEST Request) break; case BulkTransferType: case InterruptTransferType: - if (Context->LockedBuffer != WDF_NO_HANDLE) { m_Target.ReadPipeAsync(WdfRequest.Detach(), Context->EndpointAddress, Context->LockedBuffer, [](WDFREQUEST Request, WDFIOTARGET, PWDF_REQUEST_COMPLETION_PARAMS Params, WDFCONTEXT) @@ -616,13 +619,8 @@ void CUsbDkRedirectorStrategy::ReadPipe(WDFREQUEST Request) WdfRequest.SetStatus(status); }); } - else - { - WdfRequest.SetStatus(STATUS_INVALID_PARAMETER); - } break; case IsochronousTransferType: - if (Context->LockedBuffer != WDF_NO_HANDLE) { CPreAllocatedWdfMemoryBufferT PacketSizesArray(Context->LockedIsochronousPacketsArray); @@ -633,10 +631,6 @@ void CUsbDkRedirectorStrategy::ReadPipe(WDFREQUEST Request) PacketSizesArray.ArraySize(), IsoRWCompletion); } - else - { - WdfRequest.SetStatus(STATUS_INVALID_PARAMETER); - } break; default: TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Error: Wrong transfer type: %d\n", Context->TransferType);