From e4cceb851259120726eefcb58baa8fdb1540d3b9 Mon Sep 17 00:00:00 2001 From: Birunthan Mohanathas Date: Thu, 28 Jun 2012 12:57:41 +0300 Subject: [PATCH] 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 | 120 ++++++++++++++++++++++-------------- 1 file changed, 73 insertions(+), 47 deletions(-) diff --git a/Plugins/PluginPing/Ping.cpp b/Plugins/PluginPing/Ping.cpp index e5e56c5c..5000dc3a 100644 --- a/Plugins/PluginPing/Ping.cpp +++ b/Plugins/PluginPing/Ping.cpp @@ -21,8 +21,6 @@ #include #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; -} \ No newline at end of file +}