UsbTarget: Protect pipes from concurrent accesses

Signed-off-by: Dmitry Fleytman <dfleytma@redhat.com>

2 files changed, 115 insertions, 72 deletions

diff --git a/UsbDk/UsbTarget.cpp b/UsbDk/UsbTarget.cpp
index 71286ff..0e56761 100644
--- a/UsbDk/UsbTarget.cpp
+++ b/UsbDk/UsbTarget.cpp
@@ -40,6 +40,8 @@ NTSTATUS CWdfUsbInterface::SetAltSetting(ULONG64 AltSettingIdx)
         return status;
     }
 
+    ExclusiveLock m_LockedContext(m_PipesLock);
+
     m_Pipes.reset();
 
     m_NumPipes = WdfUsbInterfaceGetNumConfiguredPipes(m_Interface);
@@ -67,19 +69,6 @@ NTSTATUS CWdfUsbInterface::SetAltSetting(ULONG64 AltSettingIdx)
     return STATUS_SUCCESS;
 }
 
-CWdfUsbPipe *CWdfUsbInterface::FindPipeByEndpointAddress(ULONG64 EndpointAddress)
-{
-    for (UCHAR i = 0; i < m_NumPipes; i++)
-    {
-        if (m_Pipes[i].EndpointAddress() == EndpointAddress)
-        {
-            return &m_Pipes[i];
-        }
-    }
-
-    return nullptr;
-}
-
 NTSTATUS CWdfUsbInterface::Reset(WDFREQUEST Request)
 {
     NTSTATUS status = STATUS_SUCCESS;
@@ -311,38 +300,20 @@ NTSTATUS CWdfUsbTarget::SetInterfaceAltSetting(ULONG64 InterfaceIdx, ULONG64 Alt
     TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_USBTARGET, "%!FUNC! setting #%d for interface #%d",
                 static_cast<UCHAR>(AltSettingIdx), static_cast<UCHAR>(InterfaceIdx));
 
-    //TODO: Stop read/write queue before interface reconfiguration
     return m_Interfaces[InterfaceIdx].SetAltSetting(AltSettingIdx);
 }
 
-CWdfUsbPipe *CWdfUsbTarget::FindPipeByEndpointAddress(ULONG64 EndpointAddress)
-{
-    CWdfUsbPipe *Pipe = nullptr;
-
-    for (UCHAR i = 0; i < m_NumInterfaces; i++)
-    {
-        Pipe = m_Interfaces[i].FindPipeByEndpointAddress(EndpointAddress);
-        if (Pipe != nullptr)
-        {
-            break;
-        }
-    }
-
-    return Pipe;
-}
-
 void CWdfUsbTarget::WritePipeAsync(WDFREQUEST Request, ULONG64 EndpointAddress, WDFMEMORY Buffer, PFN_WDF_REQUEST_COMPLETION_ROUTINE Completion)
 {
     CWdfRequest WdfRequest(Request);
 
-    CWdfUsbPipe *Pipe = FindPipeByEndpointAddress(EndpointAddress);
-    if (Pipe != nullptr)
+    if (!DoPipeOperation<CWdfUsbInterface::SharedLock>(EndpointAddress,
+                                                       [&WdfRequest, Buffer, Completion](CWdfUsbPipe &Pipe)
+                                                       {
+                                                           Pipe.WriteAsync(WdfRequest, Buffer, Completion);
+                                                       }))
     {
-        Pipe->WriteAsync(WdfRequest, Buffer, Completion);
-    }
-    else
-    {
-        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed: Pipe 0x%llu not found", EndpointAddress);
+        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed because pipe was not found");
         WdfRequest.SetStatus(STATUS_NOT_FOUND);
     }
 }
@@ -351,14 +322,13 @@ void CWdfUsbTarget::ReadPipeAsync(WDFREQUEST Request, ULONG64 EndpointAddress, W
 {
     CWdfRequest WdfRequest(Request);
 
-    CWdfUsbPipe *Pipe = FindPipeByEndpointAddress(EndpointAddress);
-    if (Pipe != nullptr)
-    {
-        Pipe->ReadAsync(WdfRequest, Buffer, Completion);
-    }
-    else
+    if (!DoPipeOperation<CWdfUsbInterface::SharedLock>(EndpointAddress,
+                                                       [&WdfRequest, Buffer, Completion](CWdfUsbPipe &Pipe)
+                                                       {
+                                                           Pipe.ReadAsync(WdfRequest, Buffer, Completion);
+                                                       }))
     {
-        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed: Pipe 0x%llu not found", EndpointAddress);
+        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed because pipe was not found");
         WdfRequest.SetStatus(STATUS_NOT_FOUND);
     }
 }
@@ -369,14 +339,13 @@ void CWdfUsbTarget::ReadIsochronousPipeAsync(WDFREQUEST Request, ULONG64 Endpoin
 {
     CWdfRequest WdfRequest(Request);
 
-    CWdfUsbPipe *Pipe = FindPipeByEndpointAddress(EndpointAddress);
-    if (Pipe != nullptr)
+    if (!DoPipeOperation<CWdfUsbInterface::SharedLock>(EndpointAddress,
+                                                       [&WdfRequest, Buffer, PacketSizes, PacketNumber, Completion](CWdfUsbPipe &Pipe)
+                                                       {
+                                                           Pipe.ReadIsochronousAsync(WdfRequest, Buffer, PacketSizes, PacketNumber, Completion);
+                                                       }))
     {
-        Pipe->ReadIsochronousAsync(WdfRequest, Buffer, PacketSizes, PacketNumber, Completion);
-    }
-    else
-    {
-        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed: Pipe 0x%llu not found", EndpointAddress);
+        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed because pipe was not found");
         WdfRequest.SetStatus(STATUS_NOT_FOUND);
     }
 }
@@ -387,50 +356,72 @@ void CWdfUsbTarget::WriteIsochronousPipeAsync(WDFREQUEST Request, ULONG64 Endpoi
 {
     CWdfRequest WdfRequest(Request);
 
-    CWdfUsbPipe *Pipe = FindPipeByEndpointAddress(EndpointAddress);
-    if (Pipe != nullptr)
+    if (!DoPipeOperation<CWdfUsbInterface::SharedLock>(EndpointAddress,
+                                                       [&WdfRequest, Buffer, PacketSizes, PacketNumber, Completion](CWdfUsbPipe &Pipe)
+                                                       {
+                                                           Pipe.WriteIsochronousAsync(WdfRequest, Buffer, PacketSizes, PacketNumber, Completion);
+                                                       }))
     {
-        Pipe->WriteIsochronousAsync(WdfRequest, Buffer, PacketSizes, PacketNumber, Completion);
-    }
-    else
-    {
-        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed: Pipe 0x%llu not found", EndpointAddress);
+        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed because pipe was not found");
         WdfRequest.SetStatus(STATUS_NOT_FOUND);
     }
 }
 
 NTSTATUS CWdfUsbTarget::AbortPipe(WDFREQUEST Request, ULONG64 EndpointAddress)
 {
-    auto Pipe = FindPipeByEndpointAddress(EndpointAddress);
-    if (Pipe != nullptr)
+    NTSTATUS status;
+
+    //AbortPipe does not require locking because is scheduled sequentially
+    //with SetAltSettings which is only operation that changes pipes array
+
+    if (!DoPipeOperation<CWdfUsbInterface::NeitherLock>(EndpointAddress,
+                                                        [&status, &Request](CWdfUsbPipe &Pipe)
+                                                        {
+                                                            status = Pipe.Abort(Request);
+                                                        }))
     {
-        return Pipe->Abort(Request);
+        status = STATUS_NOT_FOUND;
     }
-    else
+
+    if (!NT_SUCCESS(status))
     {
-        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed: Pipe 0x%llu not found", EndpointAddress);
-        return STATUS_NOT_FOUND;
+        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed: %!STATUS!", status);
     }
+
+    return status;
 }
 
 NTSTATUS CWdfUsbTarget::ResetPipe(WDFREQUEST Request, ULONG64 EndpointAddress)
 {
-    auto Pipe = FindPipeByEndpointAddress(EndpointAddress);
-    if (Pipe != nullptr)
+    NTSTATUS status;
+
+    //AbortPipe does not require locking because is scheduled sequentially
+    //with SetAltSettings which is only operation that changes pipes array
+
+    if (!DoPipeOperation<CWdfUsbInterface::NeitherLock>(EndpointAddress,
+                                                        [&status, &Request](CWdfUsbPipe &Pipe)
+                                                        {
+                                                            status = Pipe.Reset(Request);
+                                                        }))
     {
-        return Pipe->Reset(Request);
+        status = STATUS_NOT_FOUND;
     }
-    else
+
+    if (!NT_SUCCESS(status))
     {
-        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed: Pipe 0x%llu not found", EndpointAddress);
-        return STATUS_NOT_FOUND;
+        TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "%!FUNC! Failed: %!STATUS!", status);
     }
+
+    return status;
 }
 
 NTSTATUS CWdfUsbTarget::ResetDevice(WDFREQUEST Request)
 {
     NTSTATUS status = STATUS_SUCCESS;
 
+    //ResetDevice does not require locking because is scheduled sequentially
+    //with SetAltSettings which is only operation that changes pipes array
+
     for (UCHAR i = 0; i < m_NumInterfaces; i++)
     {
         auto currentStatus = m_Interfaces[i].Reset(Request);
@@ -468,3 +459,8 @@ NTSTATUS CWdfUsbTarget::ControlTransferAsync(CWdfRequest &WdfRequest, PWDF_USB_C
 
     return status;
 }
+
+void CWdfUsbTarget::TracePipeNotFoundError(ULONG64 EndpointAddress)
+{
+    TraceEvents(TRACE_LEVEL_ERROR, TRACE_USBTARGET, "Pipe %llu not found", EndpointAddress);
+}
diff --git a/UsbDk/UsbTarget.h b/UsbDk/UsbTarget.h
index e663b96..516bcc8 100644
--- a/UsbDk/UsbTarget.h
+++ b/UsbDk/UsbTarget.h
@@ -24,6 +24,7 @@
 #pragma once
 
 #include "Alloc.h"
+#include "UsbDkUtil.h"
 #include "Urb.h"
 
 class CWdfRequest;
@@ -88,13 +89,44 @@ public:
     NTSTATUS Create(WDFUSBDEVICE Device, UCHAR InterfaceIdx);
     NTSTATUS SetAltSetting(ULONG64 AltSettingIdx);
 
-    CWdfUsbPipe *FindPipeByEndpointAddress(ULONG64 EndpointAddress);
+    template<typename TLockingStrategy, typename TFunctor>
+    bool DoPipeOperation(ULONG64 EndpointAddress, TFunctor Functor)
+    {
+        TLockingStrategy m_LockedContext(m_PipesLock);
+
+        for (UCHAR i = 0; i < m_NumPipes; i++)
+        {
+            if (m_Pipes[i].EndpointAddress() == EndpointAddress)
+            {
+                Functor(m_Pipes[i]);
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     NTSTATUS Reset(WDFREQUEST Request);
 
+    class Lock : public CWdmExSpinLock
+    {
+    public:
+        void NoLock() {};
+        void LockShared() { CWdmExSpinLock::LockShared(); }
+        void UnlockShared() { CWdmExSpinLock::UnlockShared(); }
+        void LockExclusive() { CWdmExSpinLock::LockExclusive(); }
+        void UnlockExclusive() { CWdmExSpinLock::UnlockExclusive(); }
+    };
+
+    using SharedLock = CBaseLockedContext < Lock, &Lock::LockShared, &Lock::UnlockShared >;
+    using ExclusiveLock = CBaseLockedContext < Lock, &Lock::LockExclusive, &Lock::UnlockExclusive > ;
+    using NeitherLock = CBaseLockedContext < Lock, &Lock::NoLock, &Lock::NoLock >;
+
 private:
     WDFUSBDEVICE m_UsbDevice;
     WDFUSBINTERFACE m_Interface;
 
+    Lock m_PipesLock;
     CObjHolder<CWdfUsbPipe, CVectorDeleter<CWdfUsbPipe> > m_Pipes;
     BYTE m_NumPipes = 0;
 
@@ -128,7 +160,22 @@ public:
     NTSTATUS ResetDevice(WDFREQUEST Request);
 
 private:
-    CWdfUsbPipe *FindPipeByEndpointAddress(ULONG64 EndpointAddress);
+    void TracePipeNotFoundError(ULONG64 EndpointAddress);
+
+    template<typename TLockingStrategy, typename TFunctor>
+    bool DoPipeOperation(ULONG64 EndpointAddress, TFunctor Functor)
+    {
+        for (UCHAR i = 0; i < m_NumInterfaces; i++)
+        {
+            if (m_Interfaces[i].DoPipeOperation<TLockingStrategy, TFunctor>(EndpointAddress, Functor))
+            {
+                return true;
+            }
+        }
+
+        TracePipeNotFoundError(EndpointAddress);
+        return false;
+    }
 
     WDFDEVICE m_Device = WDF_NO_HANDLE;
     WDFUSBDEVICE m_UsbDevice = WDF_NO_HANDLE;