Fixed heap corruption issues + fixed a bug where a device could get s…

…tuck in an invalid state when disconnected mid-mode switch

Keystone/Apple.c

@@ -156,6 +156,22 @@ APPLE_CONNECTION_MODE GetAppleMode(
 	return out;
 }
 
+UCHAR BestConfigurationForMode(APPLE_CONNECTION_MODE Mode) {
+	switch (Mode) {
+	case APPLE_MODE_BASE_NETWORK_TETHER_VALERIA:
+	case APPLE_MODE_BASE_NETWORK_TETHER:
+		return 6;
+	case APPLE_MODE_BASE_NETWORK_VALERIA:
+	case APPLE_MODE_BASE_NETWORK:
+	case APPLE_MODE_BASE_VALERIA:
+		return 5;
+	case APPLE_MODE_BASE:
+		return 4;
+	default: 
+		return 1;
+	}
+}
+
 const unsigned char APPLE_MODE_CAPABILITIES[IU_NUMBER_OF_APPLE_MODES][IU_MAX_NUMBER_OF_CONFIGURATIONS + 1][IU_NUMBER_OF_FEATURES] = {
 	{ // APPLE_MODE_UNKNOWN
 		{0, 0, 0, 0, 0, 0}, 

Keystone/Apple.h

@@ -11,3 +11,4 @@ APPLE_CONNECTION_MODE GetAppleMode(
 	IN PIU_DEVICE Dev
 );
 
+UCHAR BestConfigurationForMode(APPLE_CONNECTION_MODE Mode);

Keystone/Child.c

@@ -272,7 +272,7 @@ VOID ActivateChildren(
 ) {
 	WDFCHILDLIST ChildList = WdfFdoGetDefaultChildList(Dev->Self);
 	PIU_DEVICE_STORE deviceStore = DriverGetContext(Dev->Driver);
-	if (InterlockedAdd(&deviceStore->Devices[Dev->DeviceNum].DeviceState, 0) != IU_DEVICE_OPERATIONAL) {
+	if (Dev->AppleMode != deviceStore->Devices[Dev->DeviceNum].DesiredMode) {
 		LOG_INFO("Device not operational yet, cancelling scan");
 		return;
 	}

Keystone/Device.c

@@ -79,6 +79,34 @@ NTSTATUS KeystoneCreateDevice(
 		return status;
 	}
 
+	//find and initialize device in device store
+	PIU_DEVICE_STORE deviceStore = DriverGetContext(dev->Driver);
+	LONG deviceNum = -1;
+	for (LONG i = 0; i < IU_MAX_NUMBER_OF_DEVICES; i++) {
+		if (wcsncmp(deviceStore->Devices[i].Udid, dev->Udid, IU_MAX_UDID_LENGTH) == 0) {
+			LOG_INFO("Device was reconnected");
+			deviceNum = i;
+			break;
+		}
+	}
+	if (deviceNum == -1) {
+		LOG_INFO("Device is new connection");
+		for (LONG i = 0; i < IU_MAX_NUMBER_OF_DEVICES; i++) {
+			if (wcsnlen(deviceStore->Devices[i].Udid, IU_MAX_UDID_LENGTH) == 0) {
+				LOG_INFO("Found spot for device");
+				deviceNum = i;
+				wcscpy_s(deviceStore->Devices[i].Udid, IU_MAX_UDID_LENGTH, dev->Udid);
+				deviceStore->Devices[i].DesiredMode = APPLE_MODE_BASE_NETWORK_TETHER_VALERIA;
+				break;
+			}
+		}
+	}
+	if (deviceNum == -1) { 
+		LOG_ERROR("Too many devices attached to this driver!");
+		return STATUS_UNSUCCESSFUL;
+	}
+	dev->DeviceNum = (UCHAR)deviceNum;
+
 	//TODO: WMI registration?
 
 	return status;
@@ -144,87 +172,64 @@ NTSTATUS KeystoneEvtDeviceD0Entry(
 	PIU_DEVICE dev = DeviceGetContext(Device);
 	LOG_INFO("%ws Entering D0", dev->Udid);
 
-	//find and initialize device in device store
-	PIU_DEVICE_STORE deviceStore = DriverGetContext(dev->Driver);
-	LONG deviceNum = -1;
-	for (LONG i = 0; i < IU_MAX_NUMBER_OF_DEVICES; i++) {
-		if (wcscmp(deviceStore->Devices[i].Udid, dev->Udid) == 0) {
-			LOG_INFO("Device was reconnected");
-			deviceNum = i;
-			break;
-		}
-	}
-	if (deviceNum == -1) {
-		LOG_INFO("Device is new connection");
-		for (LONG i = 0; i < IU_MAX_NUMBER_OF_DEVICES; i++) {
-			if (InterlockedAdd(&deviceStore->Devices[i].DeviceState, 0) == IU_DEVICE_DISCONNECTED) {
-				LOG_INFO("Found spot for device");
-				deviceNum = i;
-				wcscpy(deviceStore->Devices[i].Udid, dev->Udid);
-				break;
-			}
-		}
-	}
-	if (deviceNum == -1) {
-		LOG_ERROR("Too many devices attached to this driver!");
-		return STATUS_UNSUCCESSFUL;
-	}
-	dev->DeviceNum = (UCHAR)deviceNum;
-
 	//get the device descriptor
 	status = FillDeviceDescriptor(dev);
 	if (!NT_SUCCESS(status)) {
 		LOG_ERROR("Failed to get device descriptor");
 		return status;
 	}
 
+	// Get current apple mode
+	APPLE_CONNECTION_MODE curAppleMode = GetAppleMode(dev);
+	if (!NT_SUCCESS(status)) {
+		LOG_ERROR("Failed to get apple mode");
+		return status;
+	}
+	LOG_INFO("cur apple mode %d", curAppleMode);
+
 	//set the device into the best possible apple mode
-	IU_DEVICE_STATE prevState = InterlockedAdd(&deviceStore->Devices[deviceNum].DeviceState, 0);
-	if (prevState == IU_DEVICE_DISCONNECTED) {
+	PIU_DEVICE_STORE deviceStore = DriverGetContext(dev->Driver);
+	if (curAppleMode == deviceStore->Devices[dev->DeviceNum].DesiredMode) {
+		LOG_INFO("Device is set to correct mode, moving to configure device");
+		goto ConfigureDevice;
+	}
+
+	if (curAppleMode == 1) { //if haven't attempted to set first mode
 		LOG_INFO("First enabling super network mode (4)");
-		status = SetAppleMode(dev, 4);
+		status = SetAppleMode(dev, APPLE_MODE_BASE_NETWORK_TETHER);
 		if (NT_SUCCESS(status)) {
-			InterlockedExchange(&deviceStore->Devices[deviceNum].DeviceState, IU_DEVICE_AWAITING_RECONNECT_PENDING_VALERIA);
 			return STATUS_SUCCESS;
 		} 
+
 		LOG_INFO("Couldn't set super network mode, enabling regular network mode");
-		status = SetAppleMode(dev, 3);
+		deviceStore->Devices[dev->DeviceNum].DesiredMode = APPLE_MODE_BASE_NETWORK_VALERIA;
+		status = SetAppleMode(dev, APPLE_MODE_BASE_NETWORK);
 		if (NT_SUCCESS(status)) {
-			InterlockedExchange(&deviceStore->Devices[deviceNum].DeviceState, IU_DEVICE_AWAITING_RECONNECT_PENDING_VALERIA);
 			return STATUS_SUCCESS;
 		} 
+
 		LOG_INFO("Couldn't set regular network mode, confinuing in initial mode");
-		InterlockedExchange(&deviceStore->Devices[deviceNum].DeviceState, IU_DEVICE_AWAITING_RECONNECT_PENDING_VALERIA);
+		deviceStore->Devices[dev->DeviceNum].DesiredMode = APPLE_MODE_BASE_VALERIA;
 	}
-	// refresh prev state in case above continues in regular mode, try to get valeria still
-	prevState = InterlockedAdd(&deviceStore->Devices[deviceNum].DeviceState, 0);
-	if (prevState == IU_DEVICE_AWAITING_RECONNECT_PENDING_VALERIA) {
+
+	if (curAppleMode == 3 || curAppleMode == 4 || curAppleMode == 1) {
 		LOG_INFO("Adding valeria functionality");
 		status = SetAppleMode(dev, 2);
 		if (NT_SUCCESS(status)) {
-			InterlockedExchange(&deviceStore->Devices[deviceNum].DeviceState, IU_DEVICE_AWAITING_RECONNECT);
 			return STATUS_SUCCESS;
 		}
 		LOG_INFO("Couldn't add valeria function");
+		deviceStore->Devices[dev->DeviceNum].DesiredMode = curAppleMode;
 	}
 
-	// Get current apple mode
-	APPLE_CONNECTION_MODE curAppleMode = GetAppleMode(dev);
-	if (!NT_SUCCESS(status)) {
-		LOG_ERROR("Failed to get apple mode");
-		return status;
-	}
-	LOG_INFO("cur apple mode %d", curAppleMode);
-
+ConfigureDevice:
 	//set configuration mode
-	status = SetConfigurationByValue(dev, 6);
+	status = SetConfigurationByValue(dev, BestConfigurationForMode(dev->AppleMode));
 	if (!NT_SUCCESS(status)) {
 		LOG_ERROR("Failed to set configuration");
 		return status;
 	}
 
-	InterlockedExchange(&deviceStore->Devices[deviceNum].DeviceState, IU_DEVICE_OPERATIONAL);
-
 	//create interface for talking with user app
 	status = WdfDeviceCreateDeviceInterface(
 		Device,
@@ -257,7 +262,6 @@ NTSTATUS KeystoneEvtDeviceD0Entry(
 	//	DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, "\n");
 	//}
 
-
 	return status;
 }
 
@@ -272,15 +276,6 @@ NTSTATUS KeystoneEvtDeviceD0Exit(
 
 	WdfDeviceSetDeviceInterfaceState(Device, &GUID_DEVINTERFACE_Keystone, NULL, FALSE);
 
-	IU_DEVICE_STATE curState = InterlockedAdd(&deviceStore->Devices[Dev->DeviceNum].DeviceState, 0);
-	//if was awaiting reconnect, then do not delete from store
-	if (curState == IU_DEVICE_AWAITING_RECONNECT || curState == IU_DEVICE_AWAITING_RECONNECT_PENDING_VALERIA) {
-		LOG_INFO("Device is reconnecting, not marking disc");
-		return STATUS_SUCCESS;
-	}
-
-	RtlZeroMemory(&deviceStore->Devices[Dev->DeviceNum], sizeof(deviceStore->Devices[Dev->DeviceNum]));
-	LOG_INFO("Marking device as disconnected");
 	return STATUS_SUCCESS;
 }
 

Keystone/Driver.h

@@ -23,7 +23,7 @@ typedef enum _IU_DEVICE_STATE {
 typedef struct _IU_DEVICE_STORE {
 	struct {
 		WCHAR Udid[IU_MAX_UDID_LENGTH];
-		/* IU_DEVICE_STATE */ volatile LONG DeviceState;
+		APPLE_CONNECTION_MODE DesiredMode;
 	} Devices[IU_MAX_NUMBER_OF_DEVICES];
 } IU_DEVICE_STORE, *PIU_DEVICE_STORE, DRIVER_CONTEXT, *PDRIVER_CONTEXT;
 

Messenger/Api.c

@@ -72,7 +72,7 @@ LONG WINAPI MSGGetDevices(
 	}
 	LONG deviceFlags = 0;
 	for (LONG i = 0; i < IU_MAX_NUMBER_OF_DEVICES; i++) {
-		if (MSGContext->Devices[i].SymbolicLink != NULL) {
+		if (wcsnlen(MSGContext->Devices[i].SymbolicLink, MSG_MAX_SYMLINK_LENGTH) != 0) {
 			deviceFlags |= (1 << i);
 		}
 	}

Messenger/Messenger.c

@@ -23,11 +23,15 @@ CONFIGRET AddDevice(
 	}
 
 	PMESSENGER_DEVICE_CONTEXT MSGDeviceContext = &MSGContext->Devices[ind];
-	SIZE_T linkLen = wcslen(SymbolicLink);
+	SIZE_T linkLen = wcsnlen(SymbolicLink, MSG_MAX_SYMLINK_LENGTH);
 	MSGDeviceContext->DeviceInd = (UCHAR)ind;
 	MSGDeviceContext->ParentContext = MSGContext;
 	// fill Symbolic Link
-	MSGDeviceContext->SymbolicLink = SymbolicLink;
+	if (linkLen >= MSG_MAX_SYMLINK_LENGTH - 1) {
+		MSG_DEBUG(L"[IUMSG] Device Symlink too large");
+		return CR_BUFFER_SMALL;
+	}
+	wcscpy_s(MSGDeviceContext->SymbolicLink, MSG_MAX_SYMLINK_LENGTH, SymbolicLink);
 
 	//init lock
 	InitializeCriticalSection(&MSGDeviceContext->Lock);
@@ -139,12 +143,6 @@ VOID RemoveDevice( //CALLED OUTSIDE OF USB HANDLE CALLBACK
 		MSGDeviceContext->LockInitialized = FALSE;
 	}
 
-	// free symlink string
-	if (MSGDeviceContext->SymbolicLink) {
-		MSG_DEBUG(L"[IUMSG] Freeing device symlink\n");
-		MSGDeviceContext->SymbolicLink = NULL;
-	}
-
 	// reset ShouldRegister
 	memset(MSGDeviceContext, 0, sizeof(MESSENGER_DEVICE_CONTEXT));
 }
@@ -253,7 +251,7 @@ DWORD WINAPI USBInterfaceCallback( // when devices added or removed
 		AddDevice(MSGContext, symbolicLink, TRUE);
 		break;
 	case CM_NOTIFY_ACTION_DEVICEINTERFACEREMOVAL:
-		MSG_DEBUG(L"[IUMSG] Removing device...\n");
+		MSG_DEBUG(L"[IUMSG] Removing device %ws\n", symbolicLink);
 		if (deviceInd == -1) {
 			MSG_DEBUG(L"[IUMSG] Device to be removed not found!\n");
 			goto Cleanup;
@@ -283,34 +281,38 @@ CONFIGRET RetrieveExistingDevices(
 	CONFIGRET ret;
 
 	LONG reqLen = 0;
-	ret = CM_Get_Device_Interface_List_Size(
-		&reqLen,
-		&GUID_DEVINTERFACE_Keystone,
-		NULL,
-		CM_GET_DEVICE_INTERFACE_LIST_PRESENT
-	);
-	if (ret != CR_SUCCESS) {
-		MSG_DEBUG(L"[IUMSG] Failed to get interface list size\n");
-		return ret;
-	}
-	MSG_DEBUG(L"[IUMSG] Interface list size is %d\n", reqLen);
+	PWSTR buf = NULL;
+	do {
+		ret = CM_Get_Device_Interface_List_Size(
+			&reqLen,
+			&GUID_DEVINTERFACE_Keystone,
+			NULL,
+			CM_GET_DEVICE_INTERFACE_LIST_PRESENT
+		);
+		if (ret != CR_SUCCESS) {
+			MSG_DEBUG(L"[IUMSG] Failed to get interface list size\n");
+			break;
+		}
+		MSG_DEBUG(L"[IUMSG] Interface list size is %d\n", reqLen);
 
-	PWCHAR buf = HeapAlloc(MSGContext->Heap, HEAP_ZERO_MEMORY, reqLen);
-	if (buf == NULL) {
-		MSG_DEBUG(L"[IUMSG] Failed to allocate memory for interface list\n");
-		return CR_OUT_OF_MEMORY;
-	}
+		buf = HeapAlloc(MSGContext->Heap, HEAP_ZERO_MEMORY, reqLen * sizeof(WCHAR));
+		if (buf == NULL) {
+			MSG_DEBUG(L"[IUMSG] Failed to allocate memory for interface list\n");
+			ret = CR_OUT_OF_MEMORY;
+			break;
+		}
+
+		ret = CM_Get_Device_Interface_List(
+			&GUID_DEVINTERFACE_Keystone,
+			NULL,
+			buf,
+			reqLen,
+			CM_GET_DEVICE_INTERFACE_LIST_PRESENT
+		);
+	} while (ret == CR_BUFFER_SMALL);
 
-	ret = CM_Get_Device_Interface_ListW(
-		&GUID_DEVINTERFACE_Keystone,
-		NULL,
-		buf,
-		reqLen,
-		CM_GET_DEVICE_INTERFACE_LIST_PRESENT
-	);
 	if (ret != CR_SUCCESS) {
-		MSG_DEBUG(L"[IUMSG] Failed to get full interface list\n");
-		return ret;
+		goto Cleanup;
 	}
 
 	PWCHAR curChar = buf;
@@ -326,6 +328,9 @@ CONFIGRET RetrieveExistingDevices(
 		reqLen -= strLen + 1;
 	}
 
+Cleanup:
+	if (buf)
+		HeapFree(MSGContext->Heap, 0, buf);
 	return ret;
 }
 

Messenger/Messenger.vcxproj

@@ -112,6 +112,7 @@
       <ConformanceMode>true</ConformanceMode>
       <PrecompiledHeader>Use</PrecompiledHeader>
       <PrecompiledHeaderFile>pch.h</PrecompiledHeaderFile>
+      <RuntimeLibrary>MultiThreadedDebug</RuntimeLibrary>
     </ClCompile>
     <Link>
       <SubSystem>Windows</SubSystem>
@@ -133,6 +134,7 @@
       <ConformanceMode>true</ConformanceMode>
       <PrecompiledHeader>Use</PrecompiledHeader>
       <PrecompiledHeaderFile>pch.h</PrecompiledHeaderFile>
+      <RuntimeLibrary>MultiThreadedDebug</RuntimeLibrary>
     </ClCompile>
     <Link>
       <SubSystem>Windows</SubSystem>

Messenger/Utils.c

@@ -7,7 +7,7 @@ LONG GetDeviceIndex(
 ) {
 	//check if device already exists
 	for (LONG i = 0; i < IU_MAX_NUMBER_OF_DEVICES; i++) {
-		if (MSGContext->Devices[i].SymbolicLink == NULL)
+		if (wcsnlen(MSGContext->Devices[i].SymbolicLink, MSG_MAX_SYMLINK_LENGTH) == 0)
 			continue;
 		if (wcscmp(SymbolicLink, MSGContext->Devices[i].SymbolicLink) == 0) 
 			return i;
@@ -20,7 +20,7 @@ LONG FindEmptyDeviceIndex(
 ) {
 	//find spot for device
 	for (LONG i = 0; i < IU_MAX_NUMBER_OF_DEVICES; i++) {
-		if (MSGContext->Devices[i].SymbolicLink == NULL) {
+		if (wcsnlen(MSGContext->Devices[i].SymbolicLink, MSG_MAX_SYMLINK_LENGTH) == 0) {
 			return i;
 		}
 	}

Messenger/Utils.h

@@ -6,6 +6,8 @@
 
 #define MSG_DEBUG(...) {WCHAR cad[512]; swprintf_s(cad, 511, __VA_ARGS__); OutputDebugString(cad);}
 
+#define MSG_MAX_SYMLINK_LENGTH 100
+
 struct _MESSENGER_CONTEXT;
 
 typedef struct _MESSENGER_DEVICE_CONTEXT {
@@ -14,7 +16,7 @@ typedef struct _MESSENGER_DEVICE_CONTEXT {
 	CRITICAL_SECTION Lock;
 	PTP_WORK pWork;
 
-	PWCHAR SymbolicLink;
+	WCHAR SymbolicLink[MSG_MAX_SYMLINK_LENGTH];
 	struct _MESSENGER_CONTEXT* ParentContext;
 	BOOLEAN ShouldUnregister; //set to true from within callback context
 	BOOLEAN LockInitialized;