RedirectorStrategy: Ensure there are no short buffers on read

To avoid "BABBLE" conditions. Signed-off-by: Dmitry Fleytman <dfleytma@redhat.com>
daynix · Dec 24, 2015 · 32aadcb · 32aadcb
1 parent f322815
commit 32aadcb
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 4 deletions.
diff --git a/UsbDk/RedirectorStrategy.cpp b/UsbDk/RedirectorStrategy.cpp
@@ -189,6 +189,9 @@ typedef struct tag_USBDK_REDIRECTOR_REQUEST_CONTEXT
 
     WDFMEMORY LockedIsochronousPacketsArray;
     WDFMEMORY LockedIsochronousResultsArray;
+
+    WDFMEMORY ShadowBuffer;
+    PVOID     ShadowBufferPtr;
 } USBDK_REDIRECTOR_REQUEST_CONTEXT, *PUSBDK_REDIRECTOR_REQUEST_CONTEXT;
 
 class CRedirectorRequest : public CWdfRequest
@@ -608,6 +611,52 @@ void CUsbDkRedirectorStrategy::TraceTransferError(const CRedirectorRequest &WdfR
                 DataBuffer.Size());
 }
 
+bool CUsbDkRedirectorStrategy::SupplyMaxPacketSizeShadowBuffer(CRedirectorRequest &WdfRequest,
+                                                               PUSBDK_REDIRECTOR_REQUEST_CONTEXT RequestContext)
+{
+    CPreAllocatedWdfMemoryBuffer OrigBuffer(RequestContext->LockedBuffer);
+    auto MaxPacketSize = m_Target.GetPipeMaxPacketSize(RequestContext->EndpointAddress);
+
+    if (OrigBuffer.Size() < MaxPacketSize)
+    {
+        TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_REDIRECTOR,
+                    "%!FUNC! Going to allocate shadow buffer of %llu bytes intead of %llu bytes for pipe %llu",
+                    MaxPacketSize, OrigBuffer.Size(), RequestContext->EndpointAddress);
+
+        auto status = WdfMemoryCreate(WDF_NO_OBJECT_ATTRIBUTES,
+                                      NonPagedPool,
+                                      'MBSU',
+                                      MaxPacketSize,
+                                      &RequestContext->ShadowBuffer,
+                                      &RequestContext->ShadowBufferPtr);
+        if (!NT_SUCCESS(status))
+        {
+            TraceEvents(TRACE_LEVEL_ERROR, TRACE_REDIRECTOR,
+                        "%!FUNC! Failed to allocate shadow buffer of %llu bytes: %!STATUS!", MaxPacketSize, status);
+
+            WdfRequest.SetStatus(status);
+            return false;
+        }
+    }
+
+    return true;
+}
+
+void CUsbDkRedirectorStrategy::RetireMaxPacketSizeShadowBuffer(PUSBDK_REDIRECTOR_REQUEST_CONTEXT RequestContext)
+{
+    if (RequestContext->ShadowBufferPtr != nullptr)
+    {
+        ASSERT(RequestContext->Direction == UsbDkTransferDirection::Read);
+
+        ASSERT((RequestContext->TransferType == BulkTransferType) ||
+               (RequestContext->TransferType == InterruptTransferType));
+
+        CPreAllocatedWdfMemoryBuffer OrigBuffer(RequestContext->LockedBuffer);
+        RtlCopyMemory(OrigBuffer.Ptr(), RequestContext->ShadowBufferPtr, OrigBuffer.Size());
+        WdfObjectDelete(RequestContext->ShadowBuffer);
+    }
+}
+
 void CUsbDkRedirectorStrategy::ReadPipe(WDFREQUEST Request)
 {
     CRedirectorRequest WdfRequest(Request);
@@ -631,7 +680,14 @@ void CUsbDkRedirectorStrategy::ReadPipe(WDFREQUEST Request)
     case BulkTransferType:
     case InterruptTransferType:
         {
-            m_Target.ReadPipeAsync(WdfRequest.Detach(), Context->EndpointAddress, Context->LockedBuffer,
+            if (!SupplyMaxPacketSizeShadowBuffer(WdfRequest, Context))
+            {
+                break;
+            }
+
+            auto DataBuffer = (Context->ShadowBufferPtr != nullptr) ? Context->ShadowBuffer : Context->LockedBuffer;
+
+            m_Target.ReadPipeAsync(WdfRequest.Detach(), Context->EndpointAddress, DataBuffer,
                                    [](WDFREQUEST Request, WDFIOTARGET, PWDF_REQUEST_COMPLETION_PARAMS Params, WDFCONTEXT)
                                    {
                                         auto status = Params->IoStatus.Status;
@@ -643,10 +699,11 @@ void CUsbDkRedirectorStrategy::ReadPipe(WDFREQUEST Request)
                                             TraceTransferError(WdfRequest, status, usbCompletionParams->UsbdStatus);
                                         }
 
+                                        RetireMaxPacketSizeShadowBuffer(WdfRequest.Context());
+
                                         CompleteTransferRequest(WdfRequest, status,
                                                                 usbCompletionParams->UsbdStatus,
                                                                 usbCompletionParams->Parameters.PipeRead.Length);
-
                                     });
         }
         break;

diff --git a/UsbDk/RedirectorStrategy.h b/UsbDk/RedirectorStrategy.h
@@ -57,6 +57,7 @@ class CUsbDkRedirectorQueueConfig : public CWdfSpecificQueue
 };
 
 class CRedirectorRequest;
+typedef struct tag_USBDK_REDIRECTOR_REQUEST_CONTEXT *PUSBDK_REDIRECTOR_REQUEST_CONTEXT;
 
 class CUsbDkRedirectorStrategy : public CUsbDkFilterStrategy
 {
@@ -116,6 +117,10 @@ class CUsbDkRedirectorStrategy : public CUsbDkFilterStrategy
                                    NTSTATUS Status,
                                    USBD_STATUS UsbdStatus);
 
+    bool SupplyMaxPacketSizeShadowBuffer(CRedirectorRequest &WdfRequest,
+                                         PUSBDK_REDIRECTOR_REQUEST_CONTEXT RequestContext);
+    static void RetireMaxPacketSizeShadowBuffer(PUSBDK_REDIRECTOR_REQUEST_CONTEXT RequestContext);
+
     CWdfUsbTarget m_Target;
 
     CUsbDkRedirectorQueueData m_IncomingDataQueue;

diff --git a/UsbDk/UsbTarget.cpp b/UsbDk/UsbTarget.cpp
@@ -110,7 +110,6 @@ void CWdfUsbPipe::Create(WDFUSBDEVICE Device, WDFUSBINTERFACE Interface, UCHAR P
 
     m_Pipe = WdfUsbInterfaceGetConfiguredPipe(m_Interface, PipeIndex, &m_Info);
     ASSERT(m_Pipe != nullptr);
-    WdfUsbTargetPipeSetNoMaximumPacketSizeCheck(m_Pipe);
 
     TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_USBTARGET,
                 "%!FUNC! Created pipe #%d, "
@@ -333,6 +332,19 @@ void CWdfUsbTarget::ReadPipeAsync(WDFREQUEST Request, ULONG64 EndpointAddress, W
     }
 }
 
+size_t CWdfUsbTarget::GetPipeMaxPacketSize(ULONG64 EndpointAddress)
+{
+    size_t res;
+
+    if (!DoPipeOperation<CWdfUsbInterface::SharedLock>(EndpointAddress, [&res](CWdfUsbPipe &Pipe) { res = Pipe.MaxPacketSize(); }))
+    {
+        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed because pipe was not found");
+        return 0;
+    }
+
+    return res;
+}
+
 void CWdfUsbTarget::ReadIsochronousPipeAsync(WDFREQUEST Request, ULONG64 EndpointAddress, WDFMEMORY Buffer,
                                              PULONG64 PacketSizes, size_t PacketNumber,
                                              PFN_WDF_REQUEST_COMPLETION_ROUTINE Completion)

diff --git a/UsbDk/UsbTarget.h b/UsbDk/UsbTarget.h
@@ -64,6 +64,11 @@ class CWdfUsbPipe : public CAllocatable<NonPagedPool, 'PUHR'>
         return m_Info.EndpointAddress;
     }
 
+    size_t MaxPacketSize() const
+    {
+        return m_Info.MaximumPacketSize;
+    }
+
 private:
     WDFUSBINTERFACE m_Interface = WDF_NO_HANDLE;
     WDFUSBDEVICE m_Device = WDF_NO_HANDLE;
@@ -76,6 +81,7 @@ class CWdfUsbPipe : public CAllocatable<NonPagedPool, 'PUHR'>
         PULONG64 PacketSizes,
         size_t PacketNumber,
         PFN_WDF_REQUEST_COMPLETION_ROUTINE Completion);
+
     CWdfUsbPipe(const CWdfUsbPipe&) = delete;
     CWdfUsbPipe& operator= (const CWdfUsbPipe&) = delete;
 };
@@ -158,7 +164,7 @@ class CWdfUsbTarget
     NTSTATUS AbortPipe(WDFREQUEST Request, ULONG64 EndpointAddress);
     NTSTATUS ResetPipe(WDFREQUEST Request, ULONG64 EndpointAddress);
     NTSTATUS ResetDevice(WDFREQUEST Request);
-
+    size_t GetPipeMaxPacketSize(ULONG64 EndpointAddress);
 private:
     void TracePipeNotFoundError(ULONG64 EndpointAddress);