From 05ecb770c9995e7acf6d0a4136d5ba3324f12613 Mon Sep 17 00:00:00 2001
From: zeffy <zeffy@users.noreply.github.com>
Date: Fri, 2 Mar 2018 19:46:35 -0800
Subject: [PATCH] crash detection fixes

---
 src/wufuc/callbacks.c | 11 ++++++++---
 src/wufuc/context.c   | 25 ++++++++++++++++++++++++-
 src/wufuc/context.h   |  6 +++++-
 src/wufuc/rundll32.c  |  8 ++++----
 4 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/wufuc/callbacks.c b/src/wufuc/callbacks.c
index 0792171..d956838 100644
--- a/src/wufuc/callbacks.c
+++ b/src/wufuc/callbacks.c
@@ -112,9 +112,14 @@ abort_hook:
         }
         MH_EnableHook(MH_ALL_HOOKS);
 
-        // wait for unload event or parent mutex to be abandoned.
+        // wait for unload event or the main mutex to be released or abandoned,
         // for example if the user killed rundll32.exe with task manager.
-        result = WaitForMultipleObjects(ctx->count, ctx->handles, FALSE, INFINITE);
+
+        // we use ctx_wait_any_unsafe here because contexts created by
+        // ctx_duplicate_context are not initialized by ctx_create,
+        // and have no critical section to lock, so they are only used to
+        // hold static values to send to another process.
+        ctx_wait_any_unsafe(ctx, false);
         trace(L"Unload condition has been met.");
 
         switch ( result ) {
@@ -133,7 +138,7 @@ close_handles:
         CloseHandle(ctx->uevent);
         VirtualFree(ctx, 0, MEM_RELEASE);
 unload:
-        trace(L"Freeing library and exiting thread.");
+        trace(L"Unloading wufuc and exiting thread.");
         FreeLibraryAndExitThread(PIMAGEBASE, 0);
 }
 
diff --git a/src/wufuc/context.c b/src/wufuc/context.c
index 69910d0..cf73e33 100644
--- a/src/wufuc/context.c
+++ b/src/wufuc/context.c
@@ -265,6 +265,16 @@ DWORD ctx_add_event(context *ctx,
         return result;
 }
 
+DWORD ctx_wait_all_unsafe(context *ctx,
+        bool Alertable)
+{
+        return WaitForMultipleObjectsEx(ctx->count,
+                ctx->handles,
+                TRUE,
+                INFINITE,
+                Alertable);
+}
+
 bool ctx_wait_all(context *ctx,
         bool Alertable,
         DWORD *pResult)
@@ -289,6 +299,16 @@ bool ctx_wait_all(context *ctx,
         return true;
 }
 
+DWORD ctx_wait_any_unsafe(context *ctx,
+        bool Alertable)
+{
+        return WaitForMultipleObjectsEx(ctx->count,
+                ctx->handles,
+                FALSE,
+                INFINITE,
+                Alertable);
+}
+
 DWORD ctx_wait_any(context *ctx,
         bool Alertable,
         DWORD *pResult,
@@ -361,13 +381,16 @@ bool ctx_duplicate_context(const context *pSrc,
 
         if ( !DuplicateHandle(hSrcProcess, pSrc->mutex, hProcess, &pDst->mutex, SYNCHRONIZE, FALSE, 0) )
                 return false;
+        pDst->count++;
 
         if ( !DuplicateHandle(hSrcProcess, pSrc->uevent, hProcess, &pDst->uevent, SYNCHRONIZE, FALSE, 0) ) {
 close_mutex:
                 CloseHandle(pDst->mutex);
                 pDst->mutex = INVALID_HANDLE_VALUE;
+                pDst->count = 0;
                 return false;
         }
+        pDst->count++;
         if ( !DuplicateHandle(hSrcProcess, Handle, hProcess, &hTarget, 0, FALSE, DesiredAccess)
                 || (index = ctx_add_handle(pDst, hTarget, Tag)) == -1 ) {
 
@@ -376,6 +399,6 @@ close_mutex:
                 goto close_mutex;
 
         }
-        pDst->tags[0] = index;
+        pDst->mutex_tag = index;
         return true;
 }
diff --git a/src/wufuc/context.h b/src/wufuc/context.h
index 6bbdb5e..a9b0baa 100644
--- a/src/wufuc/context.h
+++ b/src/wufuc/context.h
@@ -51,9 +51,13 @@ DWORD ctx_add_event(context *ctx,
         const wchar_t *StringSecurityDescriptor,
         DWORD Tag,
         HANDLE *pEventHandle);
+DWORD ctx_wait_all_unsafe(context *ctx,
+        bool Alertable);
 bool ctx_wait_all(context *ctx,
         bool Alertable,
         DWORD *pResult);
+DWORD ctx_wait_any_unsafe(context *ctx,
+        bool Alertable);
 DWORD ctx_wait_any(context *ctx,
         bool Alertable,
         DWORD *pResult,
@@ -65,4 +69,4 @@ bool ctx_duplicate_context(const context *pSrc,
         context *pDst,
         HANDLE Handle,
         DWORD DesiredAccess,
-        DWORD Tag);
+        DWORD Tag);
\ No newline at end of file
diff --git a/src/wufuc/rundll32.c b/src/wufuc/rundll32.c
index 722c789..13496bf 100644
--- a/src/wufuc/rundll32.c
+++ b/src/wufuc/rundll32.c
@@ -58,20 +58,20 @@ void CALLBACK RUNDLL32_StartW(HWND hwnd, HINSTANCE hinst, LPWSTR lpszCmdLine, in
                                         Unloading = true;
                                         break;
                                 } else if ( ret == WAIT_OBJECT_0 + 1 ) {
-                                        trace(L"Unload event was set!");
+                                        trace(L"Unload event was signaled!");
                                         Unloading = true;
                                         break;
                                 } else if ( (ret >= ERROR_ABANDONED_WAIT_0 || ret < WAIT_ABANDONED_0 + count) ) {
                                         g_ServiceHostCrashCount++;
-                                        trace(L"A process that wufuc injected into has crashed %Iu times!!! PID=%lu ",
+                                        trace(L"A process that wufuc injected into has crashed %Iu times!!! PID=%lu",
                                                 g_ServiceHostCrashCount, tag);
 
+                                        ReleaseMutex(handle); // release the abandoned mutex
+                                        ctx_close_and_remove_handle(&ctx, handle);
                                         if ( g_ServiceHostCrashCount >= WUFUC_CRASH_THRESHOLD ) {
                                                 trace(L"Crash threshold has been reached, disabling wufuc until next reboot!");
                                                 Unloading = true;
                                         }
-                                        
-                                        Unloading = true;
                                 } else if ( ret == WAIT_FAILED ) {
                                         trace(L"Wait failed! GLE=%lu", GetLastError());
                                         Unloading = true;