From 9e99fb7f85c264609289da4fff0358d43c523cb8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Benjamin=20H=C3=B6glinger-Stelzer?= <nefarius@dhmx.at>
Date: Wed, 17 Jul 2024 23:16:24 +0200
Subject: [PATCH] Working on #374

---
 sys/Device.h       | 10 +++++++
 sys/DsHidMiniDrv.c | 72 ++++++++++++++++++++++++++++++++++++++++++++--
 sys/DsUsb.c        | 13 +++++----
 3 files changed, 87 insertions(+), 8 deletions(-)

diff --git a/sys/Device.h b/sys/Device.h
index f45b4c4e..216c06df 100644
--- a/sys/Device.h
+++ b/sys/Device.h
@@ -357,6 +357,16 @@ typedef struct _DEVICE_CONTEXT
 
 	} RumbleControlState;
 
+	//
+	// Instance ID of this device
+	// 
+	WDFMEMORY InstanceId;
+
+	//
+	// TRUE if the mode property didn't match the configuration provided one
+	// 
+	BOOLEAN HidModeMismatch;
+
 } DEVICE_CONTEXT, * PDEVICE_CONTEXT;
 
 //
diff --git a/sys/DsHidMiniDrv.c b/sys/DsHidMiniDrv.c
index bda3ac1e..fe98a77c 100644
--- a/sys/DsHidMiniDrv.c
+++ b/sys/DsHidMiniDrv.c
@@ -354,12 +354,81 @@ DMF_DsHidMini_Open(
 	const WDFDEVICE device = DMF_ParentDeviceGet(DmfModule);
 	const PDEVICE_CONTEXT pDevCtx = DeviceGetContext(device);
 	DMF_CONFIG_VirtualHidMini* pHidCfg = DMF_ModuleConfigGet(moduleContext->DmfModuleVirtualHidMini);
+	const DS_HID_DEVICE_MODE startupHidMode = pDevCtx->Configuration.HidDeviceMode;
+
+	WDF_DEVICE_PROPERTY_DATA propertyData;
+	ULONG requiredSize = 0;
+	DEVPROPTYPE propType = DEVPROP_TYPE_EMPTY;
 
 	//
 	// Load settings
 	// 
 	ConfigLoadForDevice(pDevCtx, FALSE);
 
+	const DS_HID_DEVICE_MODE configHidMode = pDevCtx->Configuration.HidDeviceMode;
+
+	//
+	// The property value on device creation differs from what is specified in the configuration
+	// This will cause issues if igfilter is enabled on non-XI-modes so we restart ourself here
+	// 
+	if (startupHidMode != configHidMode)
+	{
+		TraceVerbose(
+			TRACE_DSHIDMINIDRV,
+			"HID Device Mode mismatch; got %hs but should be %hs",
+			G_HID_DEVICE_MODE_NAMES[startupHidMode],
+			G_HID_DEVICE_MODE_NAMES[configHidMode]
+		);
+
+		WDF_DEVICE_PROPERTY_DATA_INIT(&propertyData, &DEVPKEY_DsHidMini_RW_HidDeviceMode);
+		propertyData.Flags |= PLUGPLAY_PROPERTY_PERSISTENT;
+		propertyData.Lcid = LOCALE_NEUTRAL;
+
+		//
+		// Sync mode up with value from configuration
+		// 
+		status = WdfDeviceAssignProperty(
+			device,
+			&propertyData,
+			DEVPROP_TYPE_BYTE,
+			sizeof(BYTE),
+			&pDevCtx->Configuration.HidDeviceMode
+		);
+
+		if (!NT_SUCCESS(status))
+			goto exit;
+
+		WDF_DEVICE_PROPERTY_DATA_INIT(&propertyData, &DEVPKEY_Device_InstanceId);
+		propertyData.Flags |= PLUGPLAY_PROPERTY_PERSISTENT;
+		propertyData.Lcid = LOCALE_NEUTRAL;
+
+		//
+		// Grab our instance ID
+		// 
+		if (!NT_SUCCESS(status = WdfDeviceAllocAndQueryPropertyEx(
+			device,
+			&propertyData,
+			DS3_POOL_TAG,
+			WDF_NO_OBJECT_ATTRIBUTES,
+			&pDevCtx->InstanceId,
+			&propType
+		)))
+		{
+			TraceError(
+				TRACE_DSHIDMINIDRV,
+				"WdfDeviceAllocAndQueryPropertyEx failed with status %!STATUS!",
+				status
+			);
+			EventWriteFailedWithNTStatus(__FUNCTION__, L"WdfDeviceAllocAndQueryPropertyEx", status);
+
+			// TODO: what else to do here?
+		}
+		else
+		{
+			pDevCtx->HidModeMismatch = TRUE;
+		}
+	}
+
 	pHidCfg->VendorId = pDevCtx->VendorId;
 	pHidCfg->ProductId = pDevCtx->ProductId;
 	pHidCfg->VersionNumber = pDevCtx->VersionNumber;
@@ -457,8 +526,7 @@ DMF_DsHidMini_Open(
 	//
 	// Set currently used HID mode
 	// 
-
-	WDF_DEVICE_PROPERTY_DATA propertyData;
+	
 	WDF_DEVICE_PROPERTY_DATA_INIT(&propertyData, &DEVPKEY_DsHidMini_RW_HidDeviceMode);
 	propertyData.Flags |= PLUGPLAY_PROPERTY_PERSISTENT;
 	propertyData.Lcid = LOCALE_NEUTRAL;
diff --git a/sys/DsUsb.c b/sys/DsUsb.c
index ab5f8f0a..e1c59719 100644
--- a/sys/DsUsb.c
+++ b/sys/DsUsb.c
@@ -881,9 +881,9 @@ USB_CyclePort(
 
 	if (hubFound)
 	{
-		ULONG interfaceListSize = 0;
+		ULONG interfaceListCchSize = 0;
 		ret = CM_Get_Device_Interface_List_SizeW(
-			&interfaceListSize,
+			&interfaceListCchSize,
 			(LPGUID)&GUID_DEVINTERFACE_USB_HUB,
 			hubInstanceId,
 			CM_GET_DEVICE_INTERFACE_LIST_PRESENT
@@ -895,13 +895,14 @@ USB_CyclePort(
 			goto exit;
 		}
 
-		interfaceList = LocalAlloc(LPTR, interfaceListSize);
+		const ULONG interfaceListByteSize = interfaceListCchSize * sizeof(WCHAR);
+		interfaceList = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, interfaceListByteSize);
 
 		ret = CM_Get_Device_Interface_ListW(
 			(LPGUID)&GUID_DEVINTERFACE_USB_HUB,
 			hubInstanceId,
 			interfaceList,
-			interfaceListSize,
+			interfaceListCchSize,
 			CM_GET_DEVICE_INTERFACE_LIST_PRESENT
 		);
 
@@ -928,7 +929,7 @@ USB_CyclePort(
 		}
 
 		USB_CYCLE_PORT_PARAMS cycleParams = {
-			connectionIndex, 0
+			connectionIndex, STATUS_SUCCESS
 		};
 
 		const BOOL success = DeviceIoControl(
@@ -951,7 +952,7 @@ USB_CyclePort(
 		CloseHandle(hubHandle);
 
 	if (interfaceList)
-		LocalFree(interfaceList);
+		HeapFree(GetProcessHeap(), 0, interfaceList);
 
 	return result;
 }