PingPlugin.dll: Removed use of TerminateThread

Instead of termination in Finalize(), the thread is left to run and perform required cleanup before exit. This also fixes a memory leak caused by termination in some cases.

Plugins/PluginPing/Ping.cpp

@@ -21,8 +21,6 @@
 #include <Icmpapi.h>
 #include "../../Library/Export.h"	// Rainmeter's exported functions
 
-#include "../../Library/DisableThreadLibraryCalls.h"	// contains DllMain entry point
-
 struct MeasureData
 {
 	IPAddr destAddr;
@@ -30,7 +28,7 @@ struct MeasureData
 	double timeoutValue;
 	DWORD updateRate;
 	DWORD updateCounter;
-	HANDLE threadHandle;
+	bool threadActive;
 	double value;
 
 	MeasureData() :
@@ -39,26 +37,37 @@ struct MeasureData
 		timeoutValue(),
 		updateRate(),
 		updateCounter(),
-		threadHandle(),
+		threadActive(false),
 		value()
 	{
 	}
 };
 
 static CRITICAL_SECTION g_CriticalSection;
-static UINT g_Instances = 0;
+
+BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved)
+{
+	switch (fdwReason)
+	{
+	case DLL_PROCESS_ATTACH:
+		InitializeCriticalSection(&g_CriticalSection);
+
+		// Disable DLL_THREAD_ATTACH and DLL_THREAD_DETACH notification calls.
+		DisableThreadLibraryCalls(hinstDLL);
+		break;
+
+	case DLL_PROCESS_DETACH:
+		DeleteCriticalSection(&g_CriticalSection);
+		break;
+	}
+
+	return TRUE;
+}
 
 PLUGIN_EXPORT void Initialize(void** data, void* rm)
 {
 	MeasureData* measure = new MeasureData;
 	*data = measure;
-
-	if (g_Instances == 0)
-	{
-		InitializeCriticalSection(&g_CriticalSection);
-	}
-
-	++g_Instances;
 }
 
 PLUGIN_EXPORT void Reload(void* data, void* rm, double* maxValue)
@@ -108,59 +117,71 @@ PLUGIN_EXPORT void Reload(void* data, void* rm, double* maxValue)
 	measure->timeoutValue = RmReadDouble(rm, L"TimeoutValue", 30000.0);
 }
 
-DWORD WINAPI NetworkThreadProc(LPVOID pParam)
+DWORD WINAPI NetworkThreadProc(void* pParam)
 {
+	// NOTE: Do not use CRT functions (since thread was created by CreateThread())!
+
 	MeasureData* measure = (MeasureData*)pParam;
+	const DWORD bufferSize = sizeof(ICMP_ECHO_REPLY) + 32;
+	BYTE buffer[bufferSize];
 
-	const DWORD replySize = sizeof(ICMP_ECHO_REPLY) + 32;
-	BYTE* reply = new BYTE[replySize];
-
+	double value = 0.0;
 	HANDLE hIcmpFile = IcmpCreateFile();
-
 	if (hIcmpFile != INVALID_HANDLE_VALUE)
 	{
-		IcmpSendEcho(hIcmpFile, measure->destAddr, NULL, 0, NULL, reply, replySize, measure->timeout);
+		IcmpSendEcho(hIcmpFile, measure->destAddr, NULL, 0, NULL, buffer, bufferSize, measure->timeout);
 		IcmpCloseHandle(hIcmpFile);
+
+		ICMP_ECHO_REPLY* reply = (ICMP_ECHO_REPLY*)buffer;
+		value = (reply->Status == IP_REQ_TIMED_OUT) ? measure->timeoutValue : reply->RoundTripTime;
 	}
 
-	EnterCriticalSection(&g_CriticalSection);
+	HMODULE module = NULL;
 
-	ICMP_ECHO_REPLY* pReply = (ICMP_ECHO_REPLY*)reply;
-	if (pReply->Status == IP_REQ_TIMED_OUT)
+	EnterCriticalSection(&g_CriticalSection);
+	if (measure->threadActive)
 	{
-		measure->value = measure->timeoutValue;
+		measure->value = value;
+		measure->threadActive = false;
 	}
 	else
 	{
-		measure->value = pReply->RoundTripTime;
+		// Thread is not attached to an existing measure any longer, so delete
+		// unreferenced data.
+		delete measure;
+
+		DWORD flags = GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT;
+		GetModuleHandleEx(flags, (LPCWSTR)DllMain, &module);
 	}
-
-	delete [] reply;
-
-	CloseHandle(measure->threadHandle);
-	measure->threadHandle = NULL;
-
 	LeaveCriticalSection(&g_CriticalSection);
 
+	if (module)
+	{
+		// Decrement the ref count and possibly unload the module if this is
+		// the last instance.
+		FreeLibraryAndExitThread(module, 0);
+	}
+
 	return 0;
 }
 
 PLUGIN_EXPORT double Update(void* data)
 {
 	MeasureData* measure = (MeasureData*)data;
-	double value = 0.0;
 
 	EnterCriticalSection(&g_CriticalSection);
-	value = measure->value;
-	LeaveCriticalSection(&g_CriticalSection);
-
-	if (measure->threadHandle == NULL)
+	if (!measure->threadActive)
 	{
 		if (measure->updateCounter == 0)
 		{
 			// Launch a new thread to fetch the web data
 			DWORD id;
-			measure->threadHandle = CreateThread(NULL, 0, NetworkThreadProc, measure, 0, &id);
+			HANDLE thread = CreateThread(NULL, 0, NetworkThreadProc, measure, 0, &id);
+			if (thread)
+			{
+				CloseHandle(thread);
+				measure->threadActive = true;
+			}
 		}
 
 		measure->updateCounter++;
@@ -170,6 +191,9 @@ PLUGIN_EXPORT double Update(void* data)
 		}
 	}
 
+	double value = measure->value;
+	LeaveCriticalSection(&g_CriticalSection);
+
 	return value;
 }
 
@@ -178,18 +202,20 @@ PLUGIN_EXPORT void Finalize(void* data)
 	MeasureData* measure = (MeasureData*)data;
 
 	EnterCriticalSection(&g_CriticalSection);
-	if (measure->threadHandle)
+	if (measure->threadActive)
 	{
-		// Should really wait until the thread finishes instead terminating it...
-		TerminateThread(measure->threadHandle, 0);
+		// Increment ref count of this module so that it will not be unloaded prior to
+		// thread completion.
+		DWORD flags = GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS;
+		HMODULE module;
+		GetModuleHandleEx(flags, (LPCWSTR)DllMain, &module);
+
+		// Thread will perform cleanup.
+		measure->threadActive = false;
+	}
+	else
+	{
+		delete measure;
 	}
 	LeaveCriticalSection(&g_CriticalSection);
-
-	--g_Instances;
-	if (g_Instances == 0)
-	{
-		DeleteCriticalSection(&g_CriticalSection);
-	}
-
-	delete measure;
-}
+}