diff --git a/src/detection_win.cpp b/src/detection_win.cpp index cb6fb1c..9802feb 100644 --- a/src/detection_win.cpp +++ b/src/detection_win.cpp @@ -2,8 +2,9 @@ #include #include #include -#include - +#include +#include +#include // Include Windows headers #include @@ -19,6 +20,8 @@ using namespace std; +typedef std::basic_string tstring; + /********************************** * Local defines **********************************/ @@ -132,7 +135,7 @@ void cbAfter(uv_work_t* req) { } // Delete Item in case of removal - if(isAdded == false) { + if(!isAdded) { delete currentDevice; } @@ -371,7 +374,7 @@ DWORD WINAPI ListenerThread( LPVOID lpParam ) { } -void BuildInitialDeviceList() { +void BuildInitialDeviceList() { DWORD dwFlag = (DIGCF_ALLCLASSES | DIGCF_PRESENT); HDEVINFO hDevInfo = DllSetupDiGetClassDevs(NULL, "USB", NULL, dwFlag); @@ -380,28 +383,28 @@ void BuildInitialDeviceList() { } SP_DEVINFO_DATA* pspDevInfoData = (SP_DEVINFO_DATA*) HeapAlloc(GetProcessHeap(), 0, sizeof(SP_DEVINFO_DATA)); - pspDevInfoData->cbSize = sizeof(SP_DEVINFO_DATA); - for(int i=0; DllSetupDiEnumDeviceInfo(hDevInfo, i, pspDevInfoData); i++) { - DWORD nSize=0 ; - TCHAR buf[MAX_PATH]; + if (pspDevInfoData) { + pspDevInfoData->cbSize = sizeof(SP_DEVINFO_DATA); + for(int i=0; DllSetupDiEnumDeviceInfo(hDevInfo, i, pspDevInfoData); i++) { + DWORD nSize=0 ; + TCHAR buf[MAX_PATH]; - if (!DllSetupDiGetDeviceInstanceId(hDevInfo, pspDevInfoData, buf, sizeof(buf), &nSize)) { - break; - } - NormalizeSlashes(buf); + if (!DllSetupDiGetDeviceInstanceId(hDevInfo, pspDevInfoData, buf, sizeof(buf), &nSize)) { + break; + } + NormalizeSlashes(buf); - DeviceItem_t* item = new DeviceItem_t(); - item->deviceState = DeviceState_Connect; + DeviceItem_t* item = new DeviceItem_t(); + item->deviceState = DeviceState_Connect; - DWORD DataT; - DllSetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData, SPDRP_LOCATION_INFORMATION, &DataT, (PBYTE)buf, MAX_PATH, &nSize); - DllSetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData, SPDRP_HARDWAREID, &DataT, (PBYTE)(buf + nSize - 1), MAX_PATH - nSize, &nSize); + DWORD DataT; + DllSetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData, SPDRP_LOCATION_INFORMATION, &DataT, (PBYTE)buf, MAX_PATH, &nSize); + DllSetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData, SPDRP_HARDWAREID, &DataT, (PBYTE)(buf + nSize - 1), MAX_PATH - nSize, &nSize); - AddItemToList(buf, item); - ExtractDeviceInfo(hDevInfo, pspDevInfoData, buf, MAX_PATH, &item->deviceParams); - } + AddItemToList(buf, item); + ExtractDeviceInfo(hDevInfo, pspDevInfoData, buf, MAX_PATH, &item->deviceParams); + } - if(pspDevInfoData) { HeapFree(GetProcessHeap(), 0, pspDevInfoData); } @@ -469,79 +472,77 @@ void UpdateDevice(PDEV_BROADCAST_DEVICEINTERFACE pDevInf, WPARAM wParam, DeviceS // \\?\USB#Vid_04e8&Pid_503b#0002F9A9828E0F06#{a5dcbf10-6530-11d2-901f-00c04fb951ed} // convert to // USB\Vid_04e8&Pid_503b\0002F9A9828E0F06 - CString szDevId = pDevInf->dbcc_name+4; - int idx = szDevId.ReverseFind(_T('#')); - - szDevId.Truncate(idx); - szDevId.Replace(_T('#'), _T('\\')); - szDevId.MakeUpper(); + tstring szDevId = pDevInf->dbcc_name+4; + auto idx = szDevId.rfind(_T('#')); - CString szClass; - idx = szDevId.Find(_T('\\')); - szClass = szDevId.Left(idx); + if (idx != tstring::npos) szDevId.resize(idx); + std::replace(begin(szDevId), end(szDevId), _T('#'), _T('\\')); + auto to_upper = [] (TCHAR ch) { return std::use_facet>(std::locale()).toupper(ch); }; + transform(begin(szDevId), end(szDevId), begin(szDevId), to_upper); + tstring szClass; + idx = szDevId.find(_T('\\')); + if (idx != tstring::npos) szClass = szDevId.substr(0, idx); // if we are adding device, we only need present devices // otherwise, we need all devices DWORD dwFlag = DBT_DEVICEARRIVAL != wParam ? DIGCF_ALLCLASSES : (DIGCF_ALLCLASSES | DIGCF_PRESENT); - HDEVINFO hDevInfo = DllSetupDiGetClassDevs(NULL, szClass, NULL, dwFlag); + HDEVINFO hDevInfo = DllSetupDiGetClassDevs(NULL, szClass.c_str(), NULL, dwFlag); if(INVALID_HANDLE_VALUE == hDevInfo) { return; } SP_DEVINFO_DATA* pspDevInfoData = (SP_DEVINFO_DATA*) HeapAlloc(GetProcessHeap(), 0, sizeof(SP_DEVINFO_DATA)); - pspDevInfoData->cbSize = sizeof(SP_DEVINFO_DATA); - for(int i=0; DllSetupDiEnumDeviceInfo(hDevInfo, i, pspDevInfoData); i++) { - DWORD nSize=0 ; - TCHAR buf[MAX_PATH]; - - if (!DllSetupDiGetDeviceInstanceId(hDevInfo, pspDevInfoData, buf, sizeof(buf), &nSize)) { - break; - } - NormalizeSlashes(buf); - - if(szDevId == buf) { - WaitForSingleObject(deviceChangedSentEvent, INFINITE); - - DWORD DataT; - DWORD nSize; - DllSetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData, SPDRP_LOCATION_INFORMATION, &DataT, (PBYTE)buf, MAX_PATH, &nSize); - DllSetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData, SPDRP_HARDWAREID, &DataT, (PBYTE)(buf + nSize - 1), MAX_PATH - nSize, &nSize); - - if(state == DeviceState_Connect) { - DeviceItem_t* device = new DeviceItem_t(); - - AddItemToList(buf, device); - ExtractDeviceInfo(hDevInfo, pspDevInfoData, buf, MAX_PATH, &device->deviceParams); + if (pspDevInfoData) { + pspDevInfoData->cbSize = sizeof(SP_DEVINFO_DATA); + for (int i = 0; DllSetupDiEnumDeviceInfo(hDevInfo, i, pspDevInfoData); i++) { + DWORD nSize = 0; + TCHAR buf[MAX_PATH]; - currentDevice = &device->deviceParams; - isAdded = true; + if (!DllSetupDiGetDeviceInstanceId(hDevInfo, pspDevInfoData, buf, sizeof(buf), &nSize)) { + break; } - else { - - ListResultItem_t* item = NULL; - if(IsItemAlreadyStored(buf)) { - DeviceItem_t* deviceItem = GetItemFromList(buf); - if(deviceItem) - { - item = CopyElement(&deviceItem->deviceParams); + NormalizeSlashes(buf); + + if (szDevId == buf) { + WaitForSingleObject(deviceChangedSentEvent, INFINITE); + + DWORD DataT; + DWORD nSize; + DllSetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData, SPDRP_LOCATION_INFORMATION, &DataT, (PBYTE) buf, MAX_PATH, &nSize); + DllSetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData, SPDRP_HARDWAREID, &DataT, (PBYTE)(buf + nSize - 1), MAX_PATH - nSize, &nSize); + + if (state == DeviceState_Connect) { + DeviceItem_t *device = new DeviceItem_t(); + + AddItemToList(buf, device); + ExtractDeviceInfo(hDevInfo, pspDevInfoData, buf, MAX_PATH, &device->deviceParams); + + currentDevice = &device->deviceParams; + isAdded = true; + } else { + + ListResultItem_t *item = NULL; + if (IsItemAlreadyStored(buf)) { + DeviceItem_t *deviceItem = GetItemFromList(buf); + if (deviceItem) { + item = CopyElement(&deviceItem->deviceParams); + } + RemoveItemFromList(deviceItem); + delete deviceItem; } - RemoveItemFromList(deviceItem); - delete deviceItem; - } - if(item == NULL) { - item = new ListResultItem_t(); - ExtractDeviceInfo(hDevInfo, pspDevInfoData, buf, MAX_PATH, item); + if (item == NULL) { + item = new ListResultItem_t(); + ExtractDeviceInfo(hDevInfo, pspDevInfoData, buf, MAX_PATH, item); + } + currentDevice = item; + isAdded = false; } - currentDevice = item; - isAdded = false; - } - break; + break; + } } - } - if (pspDevInfoData) { HeapFree(GetProcessHeap(), 0, pspDevInfoData); }