From 3c9a88e78bbdc55a8912c09dbd033df5ae068543 Mon Sep 17 00:00:00 2001
From: Boring <1079299053@qq.com>
Date: Fri, 9 Dec 2022 11:12:58 +0800
Subject: [PATCH] Add support for SetUnhandledExceptionFilter

---
 MemoryModule/MemoryModule.cpp |  2 +-
 MemoryModule/MmpTls.cpp       |  8 +++++--
 a/dllmain.cpp                 | 42 +++++++++++++++++++++++++++++++++++
 a/m.def                       |  3 ++-
 test/test.cpp                 | 19 +++++++++++++---
 5 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/MemoryModule/MemoryModule.cpp b/MemoryModule/MemoryModule.cpp
index fbcd85d..8a0228c 100644
--- a/MemoryModule/MemoryModule.cpp
+++ b/MemoryModule/MemoryModule.cpp
@@ -239,7 +239,7 @@ NTSTATUS MemoryLoadLibrary(
 		LPVOID(old_header->OptionalHeader.ImageBase),
 		old_header->OptionalHeader.SizeOfImage,
 		MEM_RESERVE,
-		PAGE_READWRITE
+		PAGE_EXECUTE_READWRITE
 	);
 	if (!base) {
 		if (old_header->OptionalHeader.DllCharacteristics & IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE) {
diff --git a/MemoryModule/MmpTls.cpp b/MemoryModule/MmpTls.cpp
index 201a1d5..6757ae6 100644
--- a/MemoryModule/MmpTls.cpp
+++ b/MemoryModule/MmpTls.cpp
@@ -347,7 +347,11 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
         entry = entry->Flink;
     }
 
-    --MmpGlobalDataPtr->MmpTls->MmpActiveThreadCount;
+    if (record) {
+        --MmpGlobalDataPtr->MmpTls->MmpActiveThreadCount;
+    }
+
+    assert(0 < (int)MmpGlobalDataPtr->MmpTls->MmpActiveThreadCount);
 
     LeaveCriticalSection(&MmpGlobalDataPtr->MmpTls->MmpTlspLock);
 
@@ -371,7 +375,7 @@ VOID NTAPI HookLdrShutdownThread(VOID) {
     }
     else {
         if (MmpGlobalDataPtr->MmpTls->MmpTlsList.Flink != &MmpGlobalDataPtr->MmpTls->MmpTlsList) {
-            assert(false);
+            assert(NtCurrentTeb()->ThreadLocalStoragePointer == nullptr);
         }
     }
 
diff --git a/a/dllmain.cpp b/a/dllmain.cpp
index d02e6ea..d40eb91 100644
--- a/a/dllmain.cpp
+++ b/a/dllmain.cpp
@@ -146,3 +146,45 @@ int thread() {
     return -1;
 }
 
+DWORD Value;
+volatile LPDWORD lpAddr;
+
+LONG WINAPI Filter(_In_ struct _EXCEPTION_POINTERS* ExceptionInfo) {
+
+    if (ExceptionInfo->ExceptionRecord->ExceptionCode == STATUS_ACCESS_VIOLATION) {
+
+        lpAddr = &Value;
+
+        //  +++++++
+        // begin compiler specific
+        //  +++++++
+        
+        //ExceptionInfo->ContextRecord->Rip -= 7;
+        ExceptionInfo->ContextRecord->Rax = (ULONG_PTR)lpAddr;
+
+        //  +++++++
+        // end compiler specific
+        //  +++++++
+
+        return EXCEPTION_CONTINUE_EXECUTION;
+    }
+
+    return EXCEPTION_CONTINUE_SEARCH;
+}
+
+int unhandled_exception() {
+    auto filter = SetUnhandledExceptionFilter(Filter);
+    auto ff = SetUnhandledExceptionFilter(filter);
+
+    if (ff != Filter) {
+        printf("%p\t%p\t%p\nfailed\n", filter, ff, Filter);
+        return 0;
+    }
+
+    filter = SetUnhandledExceptionFilter(Filter);
+    lpAddr = nullptr;
+    *lpAddr = 1;
+    SetUnhandledExceptionFilter(filter);
+
+    return 1234;
+}
diff --git a/a/m.def b/a/m.def
index 47f3a56..8959306 100644
--- a/a/m.def
+++ b/a/m.def
@@ -4,4 +4,5 @@ test = __test__
 thread
 Socket = ws2_32.WSASocketW
 VerifyTruse = wintrust.WinVerifyTrust
-test_user32
\ No newline at end of file
+test_user32
+unhandled_exception
\ No newline at end of file
diff --git a/test/test.cpp b/test/test.cpp
index 68c4bf8..16f085b 100644
--- a/test/test.cpp
+++ b/test/test.cpp
@@ -1,4 +1,5 @@
 #include "../MemoryModule/stdafx.h"
+#include "../MemoryModule/LoadDllMemoryApi.h"
 #include <cstdio>
 
 static PVOID ReadDllFile(LPCSTR FileName) {
@@ -115,10 +116,22 @@ end:
     return 0;
 }
 
-int main() {
-    DisplayStatus();
+void test_uef() {
+    auto buffer = ReadDllFile("a.dll");
 
-    test();
+    HMODULE hm = LoadLibraryMemory(buffer);
+    auto pfn = GetProcAddress(hm, "unhandled_exception");
+
+    auto result = pfn();
+    if (result == 1234) {
+        printf("mmpp success\n");
+    }
+
+    return;
+}
+
+int main() {
+    test_uef();
 
     return 0;
 }