From 139457073b2117a457828f1a7cd96aeca8ce53dc Mon Sep 17 00:00:00 2001 From: Pavel Gurvich Date: Thu, 27 Feb 2014 14:15:44 +0200 Subject: [PATCH] UsbDk: add enumeration of USB devices to ControlDevice Add functionality to ControlDevice that allows counting and enumerating USB devices Signed-off-by: Pavel Gurvich Signed-off-by: Dmitry Fleytman --- UsbDk/ControlDevice.cpp | 135 +++++++++++++++++++++++++++++++++++++++- UsbDk/ControlDevice.h | 10 +++ UsbDk/Public.h | 4 ++ 3 files changed, 146 insertions(+), 3 deletions(-) diff --git a/UsbDk/ControlDevice.cpp b/UsbDk/ControlDevice.cpp index abf0d6a..1c3a17b 100644 --- a/UsbDk/ControlDevice.cpp +++ b/UsbDk/ControlDevice.cpp @@ -3,7 +3,6 @@ #include "ControlDevice.tmh" #define MAX_DEVICE_ID_LEN (200) - #include "Public.h" class CUsbDkControlDeviceInit : public CDeviceInit @@ -47,7 +46,7 @@ void CUsbDkControlDeviceQueue::DeviceControl(WDFQUEUE Queue, UNREFERENCED_PARAMETER(OutputBufferLength); UNREFERENCED_PARAMETER(InputBufferLength); - NTSTATUS status = STATUS_INVALID_DEVICE_REQUEST; + NTSTATUS status; TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "%!FUNC! Request arrived"); @@ -73,12 +72,89 @@ void CUsbDkControlDeviceQueue::DeviceControl(WDFQUEUE Queue, wcsncpy(outBuff, TEXT("Pong!"), outBuffLen/sizeof(TCHAR)); WdfRequestSetInformation(Request, outBuffLen); status = STATUS_SUCCESS; + break; + } + case IOCTL_USBDK_COUNT_DEVICES: + { + TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "Called IOCTL_USBDK_COUNT_DEVICES\n"); + status = CountDevices(Request, Queue); + break; + } + case IOCTL_USBDK_ENUM_DEVICES: + { + TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "Called IOCTL_USBDK_ENUM_DEVICES\n"); + status = EnumerateDevices(Request, Queue); + break; + } + default: + { + TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "Wrong IoControlCode\n"); + status = STATUS_INVALID_DEVICE_REQUEST; + break; } - break; } WdfRequestComplete(Request, status); } +//------------------------------------------------------------------------------------------------------------ + +NTSTATUS CUsbDkControlDeviceQueue::CountDevices(WDFREQUEST Request, WDFQUEUE Queue) +{ + ULONG *numberDevices; + size_t outBuffLen; + auto status = WdfRequestRetrieveOutputBuffer(Request, 0, (PVOID *)&numberDevices, &outBuffLen); + if (!NT_SUCCESS(status)) + { + return status; + } + + if (outBuffLen < sizeof(*numberDevices)) + { + return STATUS_BUFFER_TOO_SMALL; + } + + auto devExt = UsbDkControlGetContext(WdfIoQueueGetDevice(Queue)); + + *numberDevices = devExt->UsbDkControl->CountDevices(); + + TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "%!FUNC! CountDevices returned %d", *numberDevices); + + WdfRequestSetInformation(Request, sizeof(*numberDevices)); + + return STATUS_SUCCESS; +} +//------------------------------------------------------------------------------------------------------------ + +NTSTATUS CUsbDkControlDeviceQueue::EnumerateDevices(WDFREQUEST Request, WDFQUEUE Queue) +{ + USB_DK_DEVICE_ID *outBuff; + size_t outBuffLen; + auto status = WdfRequestRetrieveOutputBuffer(Request, 0, (PVOID *)&outBuff, &outBuffLen); + if (!NT_SUCCESS(status)) + { + return status; + } + + auto devExt = UsbDkControlGetContext(WdfIoQueueGetDevice(Queue)); + + auto numberAllocatedDevices = outBuffLen / sizeof(USB_DK_DEVICE_ID); + size_t numberExistingDevices = 0; + auto res = devExt->UsbDkControl->EnumerateDevices(outBuff, numberAllocatedDevices, numberExistingDevices); + + TraceEvents(TRACE_LEVEL_INFORMATION, TRACE_CONTROLDEVICE, "%!FUNC! EnumerateDevices returned %llu devices", numberExistingDevices); + + if (res) + { + WdfRequestSetInformation(Request, outBuffLen); + return STATUS_SUCCESS; + } + else + { + WdfRequestSetInformation(Request, 0); + return STATUS_BUFFER_TOO_SMALL; + } +} +//------------------------------------------------------------------------------------------------------------ void CUsbDkControlDevice::DumpAllChildren() { @@ -88,6 +164,59 @@ void CUsbDkControlDevice::DumpAllChildren() return true; }); } +//------------------------------------------------------------------------------------------------------------ + +ULONG CUsbDkControlDevice::CountDevices() +{ + ULONG numberDevices = 0; + + m_FilterDevices.ForEach([&numberDevices](CUsbDkFilterDevice *Filter) + { + numberDevices += Filter->GetChildrenCount(); + return true; + }); + + return numberDevices; +} +//------------------------------------------------------------------------------------------------------------ + +bool CUsbDkControlDevice::EnumerateDevices(USB_DK_DEVICE_ID *outBuff, size_t numberAllocatedDevices, size_t &numberExistingDevices) +{ + numberExistingDevices = 0; + bool hasEnoughPlace; + m_FilterDevices.ForEach([this, &outBuff, numberAllocatedDevices, &numberExistingDevices, &hasEnoughPlace](CUsbDkFilterDevice *Filter) -> bool + { + hasEnoughPlace = EnumerateFilterChildren(Filter, outBuff, numberAllocatedDevices, numberExistingDevices); + return hasEnoughPlace; + }); + return hasEnoughPlace; +} +//------------------------------------------------------------------------------------------------------------ + +bool CUsbDkControlDevice::EnumerateFilterChildren(CUsbDkFilterDevice *Filter, USB_DK_DEVICE_ID *outBuff, size_t numberAllocatedDevices, size_t &numberExistingDevices) +{ + bool hasEnoughPlace = true; + Filter->EnumerateChildren + ( + [this, &outBuff, numberAllocatedDevices, &numberExistingDevices, &hasEnoughPlace](CUsbDkChildDevice *Child) -> bool + { + if (numberExistingDevices == numberAllocatedDevices) + { + TraceEvents(TRACE_LEVEL_ERROR, TRACE_CONTROLDEVICE, "%!FUNC! FAILED! Number existing devices is more than allocated buffer!"); + hasEnoughPlace = false; + return false; + } + + wcsncpy(outBuff->DeviceID, Child->DeviceID(), MAX_DEVICE_ID_LEN); + wcsncpy(outBuff->InstanceID, Child->InstanceID(), MAX_DEVICE_ID_LEN); + outBuff++; + numberExistingDevices++; + return true; + } + ); + return hasEnoughPlace; +} +//------------------------------------------------------------------------------------------------------------ NTSTATUS CUsbDkControlDevice::Create(WDFDRIVER Driver) { diff --git a/UsbDk/ControlDevice.h b/UsbDk/ControlDevice.h index 7950417..49cf253 100644 --- a/UsbDk/ControlDevice.h +++ b/UsbDk/ControlDevice.h @@ -7,6 +7,9 @@ #include "UsbDkUtil.h" #include "FilterDevice.h" +typedef struct tag_USB_DK_DEVICE_ID USB_DK_DEVICE_ID; +class CUsbDkFilterDevice; + class CUsbDkControlDeviceQueue : public CWdfDefaultQueue, public CAllocatable { public: @@ -22,6 +25,9 @@ private: size_t InputBufferLength, ULONG IoControlCode); + static NTSTATUS CountDevices(WDFREQUEST Request, WDFQUEUE Queue); + static NTSTATUS EnumerateDevices(WDFREQUEST Request, WDFQUEUE Queue); + CUsbDkControlDeviceQueue(const CUsbDkControlDeviceQueue&) = delete; CUsbDkControlDeviceQueue& operator= (const CUsbDkControlDeviceQueue&) = delete; }; @@ -39,6 +45,8 @@ public: //TODO: Temporary, until enumeration implementation void DumpAllChildren(); + ULONG CountDevices(); + bool EnumerateDevices(USB_DK_DEVICE_ID *outBuff, size_t numberAllocatedDevices, size_t &numberExistingDevices); static bool Allocate(); static void Deallocate() @@ -52,6 +60,8 @@ private: static CRefCountingHolder *m_UsbDkControlDevice; CWdmList m_FilterDevices; + + bool EnumerateFilterChildren(CUsbDkFilterDevice *Filter, USB_DK_DEVICE_ID *outBuff, size_t numberAllocatedDevices, size_t &numberExistingDevices); }; typedef struct _USBDK_CONTROL_DEVICE_EXTENSION { diff --git a/UsbDk/Public.h b/UsbDk/Public.h index 0ec098d..63f15af 100644 --- a/UsbDk/Public.h +++ b/UsbDk/Public.h @@ -32,3 +32,7 @@ DEFINE_GUID (GUID_DEVINTERFACE_UsbDk, // UsbDk Control Device IOCTLs #define IOCTL_USBDK_PING \ CTL_CODE( USBDK_DEVICE_TYPE, 0x850, METHOD_BUFFERED, FILE_READ_ACCESS ) +#define IOCTL_USBDK_COUNT_DEVICES \ + CTL_CODE( USBDK_DEVICE_TYPE, 0x851, METHOD_BUFFERED, FILE_READ_ACCESS ) +#define IOCTL_USBDK_ENUM_DEVICES \ + CTL_CODE( USBDK_DEVICE_TYPE, 0x852, METHOD_BUFFERED, FILE_READ_ACCESS )