385 lines
9.7 KiB
C
385 lines
9.7 KiB
C
|
#include "stdafx.h"
|
||
|
|
|
||
|
|
#include "systray.h"
|
||
|
|
|
||
|
|
#include <stdio.h>
|
||
|
|
#include <initguid.h>
|
||
|
|
#include <usbioctl.h>
|
||
|
|
#include <wmium.h>
|
||
|
|
#include <tchar.h>
|
||
|
|
#include <setupapi.h>
|
||
|
|
|
||
|
|
#define USBUIMENU 100
|
||
|
|
|
||
|
|
#define NUM_HCS_TO_CHECK 10
|
||
|
|
|
||
|
|
typedef int (CALLBACK *USBERRORMESSAGESCALLBACK)
|
||
|
|
(PUSB_CONNECTION_NOTIFICATION,LPTSTR);
|
||
|
|
|
||
|
|
extern HINSTANCE g_hInstance;
|
||
|
|
|
||
|
|
static BOOL g_bUSBUIEnabled = FALSE;
|
||
|
|
static BOOL g_bUSBUIIconShown = FALSE;
|
||
|
|
static HINSTANCE g_hUsbWatch = NULL;
|
||
|
|
static USBERRORMESSAGESCALLBACK g_UsbHandler = NULL;
|
||
|
|
static BOOL g_bSubstituteDll = FALSE;
|
||
|
|
static TCHAR g_strSubstituteDll[MAX_PATH];
|
||
|
|
static HANDLE g_hWait = NULL;
|
||
|
|
|
||
|
|
int _cdecl main(){
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
#define USBUI_OffsetToPtr(Base, Offset) ((PBYTE)((PBYTE)Base + Offset))
|
||
|
|
|
||
|
|
LPTSTR USBUI_CountedStringToSz(LPTSTR lpString)
|
||
|
|
{
|
||
|
|
SHORT usNameLength;
|
||
|
|
LPTSTR lpStringPlusNull;
|
||
|
|
|
||
|
|
usNameLength = * (USHORT *) lpString;
|
||
|
|
|
||
|
|
lpStringPlusNull = (LPTSTR) LocalAlloc(LMEM_ZEROINIT,
|
||
|
|
sizeof(TCHAR) * (usNameLength+1));
|
||
|
|
|
||
|
|
if (lpStringPlusNull != NULL) {
|
||
|
|
lpString = (LPTSTR) USBUI_OffsetToPtr(lpString, sizeof(USHORT));
|
||
|
|
|
||
|
|
wcsncpy( lpStringPlusNull, lpString, usNameLength );
|
||
|
|
|
||
|
|
lpStringPlusNull[usNameLength] = TEXT('0');
|
||
|
|
// _tcscpy( lpStringPlusNull + usNameLength, _TEXT("") );
|
||
|
|
}
|
||
|
|
|
||
|
|
return lpStringPlusNull;
|
||
|
|
}
|
||
|
|
|
||
|
|
void USBUI_EventCallbackRoutine(PWNODE_HEADER WnodeHeader, UINT_PTR NotificationContext)
|
||
|
|
{
|
||
|
|
PWNODE_SINGLE_INSTANCE wNode = (PWNODE_SINGLE_INSTANCE)WnodeHeader;
|
||
|
|
PUSB_CONNECTION_NOTIFICATION usbConnectionNotification;
|
||
|
|
LPGUID eventGuid = &WnodeHeader->Guid;
|
||
|
|
LPTSTR strInstanceName;
|
||
|
|
|
||
|
|
if (memcmp(&GUID_USB_WMI_STD_DATA, eventGuid, sizeof(GUID)) == 0) {
|
||
|
|
usbConnectionNotification = (PUSB_CONNECTION_NOTIFICATION)
|
||
|
|
USBUI_OffsetToPtr(wNode,
|
||
|
|
wNode->DataBlockOffset);
|
||
|
|
|
||
|
|
//
|
||
|
|
// Get the instance name
|
||
|
|
//
|
||
|
|
strInstanceName =
|
||
|
|
USBUI_CountedStringToSz((LPTSTR)
|
||
|
|
USBUI_OffsetToPtr(wNode,
|
||
|
|
wNode->OffsetInstanceName));
|
||
|
|
if (strInstanceName) {
|
||
|
|
if (g_hUsbWatch && g_UsbHandler) {
|
||
|
|
USBUIEngageHandler:
|
||
|
|
g_UsbHandler(usbConnectionNotification, strInstanceName);
|
||
|
|
} else {
|
||
|
|
if (g_bSubstituteDll) {
|
||
|
|
g_hUsbWatch = LoadLibrary(g_strSubstituteDll);
|
||
|
|
} else {
|
||
|
|
g_hUsbWatch = LoadLibrary(TEXT("usbui.dll"));
|
||
|
|
}
|
||
|
|
g_UsbHandler = (USBERRORMESSAGESCALLBACK)
|
||
|
|
GetProcAddress(g_hUsbWatch, "USBErrorHandler");
|
||
|
|
goto USBUIEngageHandler;
|
||
|
|
}
|
||
|
|
LocalFree(strInstanceName);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
VOID USBUI_WaitRoutineCallback(WMIHANDLE Handle, BOOLEAN Unused) {
|
||
|
|
ASSERT(!Unused);
|
||
|
|
UnregisterWaitEx(g_hWait, NULL);
|
||
|
|
g_hWait = NULL;
|
||
|
|
WmiReceiveNotifications(1, &Handle, USBUI_EventCallbackRoutine, (ULONG_PTR)NULL);
|
||
|
|
RegisterWaitForSingleObject(&g_hWait,
|
||
|
|
Handle,
|
||
|
|
USBUI_WaitRoutineCallback,
|
||
|
|
Handle, // context
|
||
|
|
INFINITE,
|
||
|
|
WT_EXECUTELONGFUNCTION | WT_EXECUTEONLYONCE);
|
||
|
|
}
|
||
|
|
|
||
|
|
int USBUI_ErrorMessagesEnable(BOOL fEnable)
|
||
|
|
{
|
||
|
|
ULONG status = ERROR_SUCCESS;
|
||
|
|
BOOL result;
|
||
|
|
static WMIHANDLE hWmi = NULL;
|
||
|
|
|
||
|
|
if (fEnable) {
|
||
|
|
ASSERT(!g_hWait);
|
||
|
|
ASSERT(!hWmi);
|
||
|
|
status = WmiOpenBlock((LPGUID) &GUID_USB_WMI_STD_DATA,
|
||
|
|
WMIGUID_NOTIFICATION | SYNCHRONIZE,
|
||
|
|
&hWmi);
|
||
|
|
|
||
|
|
if (!status) {
|
||
|
|
result = RegisterWaitForSingleObject(&g_hWait,
|
||
|
|
hWmi,
|
||
|
|
USBUI_WaitRoutineCallback,
|
||
|
|
hWmi, // context
|
||
|
|
INFINITE,
|
||
|
|
WT_EXECUTELONGFUNCTION | WT_EXECUTEONLYONCE);
|
||
|
|
status = result ? 0 : ERROR_INVALID_FUNCTION;
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
ASSERT(hWmi);
|
||
|
|
if (g_hWait) {
|
||
|
|
result = UnregisterWait(g_hWait);
|
||
|
|
}
|
||
|
|
if (hWmi) {
|
||
|
|
status = WmiCloseBlock(hWmi);
|
||
|
|
}
|
||
|
|
hWmi = NULL;
|
||
|
|
g_hWait = NULL;
|
||
|
|
if (g_hUsbWatch) {
|
||
|
|
|
||
|
|
// This allows us to replace the library
|
||
|
|
|
||
|
|
FreeLibrary(g_hUsbWatch);
|
||
|
|
g_hUsbWatch = NULL;
|
||
|
|
g_UsbHandler = NULL;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return status;
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
void USBUI_Notify(HWND hwnd, WPARAM wParam, LPARAM lParam)
|
||
|
|
{
|
||
|
|
|
||
|
|
switch (lParam)
|
||
|
|
{
|
||
|
|
case WM_RBUTTONUP:
|
||
|
|
{
|
||
|
|
USBUI_Menu(hwnd, 1, TPM_RIGHTBUTTON);
|
||
|
|
}
|
||
|
|
break;
|
||
|
|
|
||
|
|
case WM_LBUTTONDOWN:
|
||
|
|
{
|
||
|
|
SetTimer(hwnd, USBUI_TIMER_ID, GetDoubleClickTime()+100, NULL);
|
||
|
|
}
|
||
|
|
break;
|
||
|
|
|
||
|
|
case WM_LBUTTONDBLCLK:
|
||
|
|
{
|
||
|
|
KillTimer(hwnd, USBUI_TIMER_ID);
|
||
|
|
USBUI_Toggle();
|
||
|
|
}
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void USBUI_Toggle()
|
||
|
|
{
|
||
|
|
USBUI_SetState(!g_bUSBUIEnabled);
|
||
|
|
}
|
||
|
|
|
||
|
|
void USBUI_Timer(HWND hwnd)
|
||
|
|
{
|
||
|
|
KillTimer(hwnd, USBUI_TIMER_ID);
|
||
|
|
USBUI_Menu(hwnd, 0, TPM_LEFTBUTTON);
|
||
|
|
}
|
||
|
|
/*
|
||
|
|
HMENU USBUI_CreateMenu()
|
||
|
|
{
|
||
|
|
HMENU hmenu;
|
||
|
|
LPSTR lpszMenu1;
|
||
|
|
|
||
|
|
hmenu = CreatePopupMenu();
|
||
|
|
|
||
|
|
if (!hmenu)
|
||
|
|
{
|
||
|
|
return NULL;
|
||
|
|
}
|
||
|
|
|
||
|
|
lpszMenu1 = LoadDynamicString(g_bUSBUIEnabled?IDS_USBUIDISABLE:IDS_USBUIENABLE);
|
||
|
|
|
||
|
|
// AppendMenu(hmenu,MF_STRING,USBUIMENU,lpszMenu1);
|
||
|
|
SysTray_AppendMenuString (hmenu,USBUIMENU,lpszMenu1);
|
||
|
|
|
||
|
|
SetMenuDefaultItem(hmenu,USBUIMENU,FALSE);
|
||
|
|
|
||
|
|
DeleteDynamicString(lpszMenu1);
|
||
|
|
|
||
|
|
return hmenu;
|
||
|
|
}
|
||
|
|
*/
|
||
|
|
void USBUI_Menu(HWND hwnd, UINT uMenuNum, UINT uButton)
|
||
|
|
{
|
||
|
|
POINT pt;
|
||
|
|
UINT iCmd;
|
||
|
|
HMENU hmenu = 0;
|
||
|
|
|
||
|
|
GetCursorPos(&pt);
|
||
|
|
|
||
|
|
// hmenu = USBUI_CreateMenu();
|
||
|
|
|
||
|
|
if (!hmenu)
|
||
|
|
{
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
SetForegroundWindow(hwnd);
|
||
|
|
|
||
|
|
iCmd = TrackPopupMenu(hmenu, uButton | TPM_RETURNCMD | TPM_NONOTIFY, pt.x, pt.y, 0, hwnd, NULL);
|
||
|
|
|
||
|
|
DestroyMenu(hmenu);
|
||
|
|
|
||
|
|
switch (iCmd)
|
||
|
|
{
|
||
|
|
case USBUIMENU:
|
||
|
|
{
|
||
|
|
USBUI_Toggle();
|
||
|
|
}
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
BOOL USBUI_SetState(BOOL NewState)
|
||
|
|
{
|
||
|
|
int retValue;
|
||
|
|
|
||
|
|
if (g_bUSBUIEnabled != NewState) {
|
||
|
|
//
|
||
|
|
// Only enable it if not already enabled
|
||
|
|
//
|
||
|
|
retValue = (int) USBUI_ErrorMessagesEnable (NewState);
|
||
|
|
g_bUSBUIEnabled = retValue ? g_bUSBUIEnabled : NewState;
|
||
|
|
}
|
||
|
|
return g_bUSBUIEnabled;
|
||
|
|
}
|
||
|
|
|
||
|
|
BOOL
|
||
|
|
IsErrorCheckingEnabled()
|
||
|
|
{
|
||
|
|
DWORD ErrorCheckingEnabled = TRUE, size;
|
||
|
|
HKEY hKey;
|
||
|
|
|
||
|
|
//
|
||
|
|
// Check the registry value ErrorCheckingEnabled to make sure that we should
|
||
|
|
// be enabling this.
|
||
|
|
//
|
||
|
|
if (ERROR_SUCCESS ==
|
||
|
|
RegOpenKeyEx(HKEY_LOCAL_MACHINE,
|
||
|
|
TEXT("SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Usb"),
|
||
|
|
0,
|
||
|
|
KEY_READ,
|
||
|
|
&hKey)) {
|
||
|
|
|
||
|
|
// Get the ErrorCheckingEnabled value
|
||
|
|
|
||
|
|
size = sizeof(DWORD);
|
||
|
|
RegQueryValueEx(hKey,
|
||
|
|
TEXT("ErrorCheckingEnabled"),
|
||
|
|
0,
|
||
|
|
NULL,
|
||
|
|
(LPBYTE) &ErrorCheckingEnabled,
|
||
|
|
&size);
|
||
|
|
|
||
|
|
if (ErrorCheckingEnabled) {
|
||
|
|
|
||
|
|
// Look for a substitute dll for usbui.dll
|
||
|
|
|
||
|
|
size = MAX_PATH*sizeof(TCHAR);
|
||
|
|
|
||
|
|
if (ERROR_SUCCESS ==
|
||
|
|
RegQueryValueEx(hKey,
|
||
|
|
TEXT("SubstituteDll"),
|
||
|
|
0,
|
||
|
|
NULL,
|
||
|
|
(LPBYTE) g_strSubstituteDll,
|
||
|
|
&size)) {
|
||
|
|
g_bSubstituteDll = TRUE;
|
||
|
|
} else {
|
||
|
|
g_bSubstituteDll = FALSE;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return (BOOL) ErrorCheckingEnabled;
|
||
|
|
}
|
||
|
|
|
||
|
|
BOOL USBUI_Init(HWND hWnd)
|
||
|
|
{
|
||
|
|
TCHAR HCName[16];
|
||
|
|
BOOL ControllerFound = FALSE;
|
||
|
|
int HCNum;
|
||
|
|
HDEVINFO hHCDev;
|
||
|
|
|
||
|
|
//
|
||
|
|
// Check the registry to make sure that it is turned on
|
||
|
|
//
|
||
|
|
if (!IsErrorCheckingEnabled()) {
|
||
|
|
return FALSE;
|
||
|
|
}
|
||
|
|
|
||
|
|
//
|
||
|
|
// Check for the existence of a USB controller.
|
||
|
|
// If there is one, load and initialize USBUI.dll which will check for
|
||
|
|
// usb error messages. If we can't open a controller, than we shouldn't
|
||
|
|
// load a USB watch dll.
|
||
|
|
//
|
||
|
|
for (HCNum = 0; HCNum < NUM_HCS_TO_CHECK; HCNum++)
|
||
|
|
{
|
||
|
|
wsprintf(HCName, TEXT("\\\\.\\HCD%d"), HCNum);
|
||
|
|
|
||
|
|
hHCDev = CreateFile(HCName,
|
||
|
|
GENERIC_READ,
|
||
|
|
FILE_SHARE_READ,
|
||
|
|
NULL,
|
||
|
|
OPEN_EXISTING,
|
||
|
|
0,
|
||
|
|
NULL);
|
||
|
|
//
|
||
|
|
// If the handle is valid, then we've successfully opened a Host
|
||
|
|
// Controller.
|
||
|
|
//
|
||
|
|
|
||
|
|
if (hHCDev != INVALID_HANDLE_VALUE) {
|
||
|
|
CloseHandle(hHCDev);
|
||
|
|
return TRUE;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
hHCDev = SetupDiGetClassDevs(&GUID_CLASS_USB_HOST_CONTROLLER,
|
||
|
|
NULL,
|
||
|
|
NULL,
|
||
|
|
(DIGCF_DEVICEINTERFACE | DIGCF_PRESENT));
|
||
|
|
if(hHCDev == INVALID_HANDLE_VALUE) {
|
||
|
|
return FALSE;
|
||
|
|
}
|
||
|
|
|
||
|
|
SetupDiDestroyDeviceInfoList(hHCDev);
|
||
|
|
return TRUE;
|
||
|
|
}
|
||
|
|
|
||
|
|
//
|
||
|
|
// Called at init time and whenever services are enabled/disabled.
|
||
|
|
//
|
||
|
|
BOOL USBUI_CheckEnable(HWND hWnd, BOOL bSvcEnabled)
|
||
|
|
{
|
||
|
|
BOOL bEnable = bSvcEnabled && USBUI_Init(hWnd);
|
||
|
|
|
||
|
|
if (bEnable != g_bUSBUIEnabled)
|
||
|
|
{
|
||
|
|
//
|
||
|
|
// state change
|
||
|
|
//
|
||
|
|
USBUI_SetState(bEnable);
|
||
|
|
}
|
||
|
|
|
||
|
|
return(bEnable);
|
||
|
|
}
|
||
|
|
|
||
|
|
|