From d2fc2405c3b82df9a87b706fdb2ef67ffaa08bdb Mon Sep 17 00:00:00 2001
From: Kawe Mazidjatari <48657826+Mauler125@users.noreply.github.com>
Date: Tue, 27 Dec 2022 23:08:59 +0100
Subject: [PATCH] Add ability to whitelist return addresses from exception
 filter

Programmer could add and remove whitelisted addresses. If an exception occurs, the system checks the callstack up to 'MAX_IMI_SEARCH' frames (defined in CCrashHandler), and returns true if found (returning early with 'EXCEPTION_CONTINUE_SEARCH' so we use whatever exception handler is set for this particular case), false otherwise.
---
 r5dev/public/utility/crashhandler.cpp | 60 +++++++++++++++++++++++----
 r5dev/public/utility/crashhandler.h   | 13 +++++-
 2 files changed, 63 insertions(+), 10 deletions(-)

diff --git a/r5dev/public/utility/crashhandler.cpp b/r5dev/public/utility/crashhandler.cpp
index e66e0578..45b2c2ad 100644
--- a/r5dev/public/utility/crashhandler.cpp
+++ b/r5dev/public/utility/crashhandler.cpp
@@ -193,6 +193,14 @@ void CCrashHandler::FormatBuildInfo()
 	m_svBuffer.append(fmt::format("build_id: {:d}\n", g_SDKDll.m_pNTHeaders->FileHeader.TimeDateStamp));
 }
 
+//-----------------------------------------------------------------------------
+// Purpose: formats the module, address and exception
+//-----------------------------------------------------------------------------
+void CCrashHandler::FormatExceptionAddress()
+{
+	FormatExceptionAddress(static_cast<LPCSTR>(m_pExceptionPointers->ExceptionRecord->ExceptionAddress));
+}
+
 //-----------------------------------------------------------------------------
 // Purpose: formats the module, address and exception
 // Input  : pExceptionAddress - 
@@ -200,11 +208,6 @@ void CCrashHandler::FormatBuildInfo()
 void CCrashHandler::FormatExceptionAddress(LPCSTR pExceptionAddress)
 {
 	HMODULE hCrashedModule;
-	if (!pExceptionAddress)
-	{
-		pExceptionAddress = static_cast<LPCSTR>(m_pExceptionPointers->ExceptionRecord->ExceptionAddress);
-	}
-
 	if (!GetModuleHandleExA(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, pExceptionAddress, &hCrashedModule))
 	{
 		m_svBuffer.append(fmt::format("\t!!!unknown-module!!!: 0x{:016X}\n", reinterpret_cast<uintptr_t>(pExceptionAddress)));
@@ -418,6 +421,38 @@ void CCrashHandler::GetCallStack()
 	m_nCapturedFrames = RtlCaptureStackBackTrace(2, NUM_FRAMES_TO_CAPTURE, m_ppStackTrace, NULL);
 }
 
+//-----------------------------------------------------------------------------
+// Purpose: adds a whitelist exception address
+//-----------------------------------------------------------------------------
+void CCrashHandler::AddWhitelist(void* pWhitelist)
+{
+	m_WhiteList.insert(pWhitelist);
+}
+
+//-----------------------------------------------------------------------------
+// Purpose: removes a whitelist exception address
+//-----------------------------------------------------------------------------
+void CCrashHandler::RemoveWhitelist(void* pWhitelist)
+{
+	m_WhiteList.erase(pWhitelist);
+}
+
+//-----------------------------------------------------------------------------
+// Purpose: checks if callstack contains whitelisted addresses up to 'MAX_IMI_SEARCH' frames.
+// Output : true is exist, false otherwise
+//-----------------------------------------------------------------------------
+bool CCrashHandler::HasWhitelist()
+{
+	for (WORD i = 0; i < MAX_IMI_SEARCH; i++)
+	{
+		if (m_WhiteList.count(m_ppStackTrace[i]))
+		{
+			return true;
+		}
+	}
+	return false;
+}
+
 //-----------------------------------------------------------------------------
 // Purpose: writes the formatted exception buffer to a file on the disk
 //-----------------------------------------------------------------------------
@@ -481,19 +516,28 @@ long __stdcall ExceptionFilter(EXCEPTION_POINTERS* exceptionInfo)
 		return EXCEPTION_CONTINUE_SEARCH;
 	}
 
+	// Don't run when a debugger is present.
 	if (IsDebuggerPresent())
 	{
 		g_CrashHandler->Unlock();
 		return EXCEPTION_CONTINUE_SEARCH;
 	}
 
+	g_CrashHandler->GetCallStack();
+
+	// Don't run when exception return address is in whitelist.
+	// This is usefull for when we want to use a different exception handler instead.
+	if (g_CrashHandler->HasWhitelist())
+	{
+		g_CrashHandler->Unlock();
+		return EXCEPTION_CONTINUE_SEARCH;
+	}
+
 	// Kill on recursive call.
 	if (g_CrashHandler->GetState())
 		ExitProcess(1u);
 	g_CrashHandler->SetState(true);
 
-	g_CrashHandler->GetCallStack();
-
 	g_CrashHandler->FormatCrash();
 	g_CrashHandler->FormatCallstack();
 	g_CrashHandler->FormatRegisters();
@@ -502,7 +546,7 @@ long __stdcall ExceptionFilter(EXCEPTION_POINTERS* exceptionInfo)
 	g_CrashHandler->FormatBuildInfo();
 
 	g_CrashHandler->WriteFile();
-	g_CrashHandler->CreateMessageProcess();
+	g_CrashHandler->CreateMessageProcess(); // Display the message to the user.
 
 	g_CrashHandler->Unlock();
 
diff --git a/r5dev/public/utility/crashhandler.h b/r5dev/public/utility/crashhandler.h
index 38cb3407..7d9d955f 100644
--- a/r5dev/public/utility/crashhandler.h
+++ b/r5dev/public/utility/crashhandler.h
@@ -35,8 +35,13 @@ public:
 	const char* ExceptionToString(DWORD nExceptionCode) const;
 	void SetExceptionPointers(EXCEPTION_POINTERS* pExceptionPointers) { m_pExceptionPointers = pExceptionPointers; };
 
-	void WriteFile();
+	void AddWhitelist(void* pWhiteList);
+	void RemoveWhitelist(void* pWhiteList);
+	bool HasWhitelist();
+
 	void GetCallStack();
+	void WriteFile();
+
 	void CreateMessageProcess();
 
 private:
@@ -44,7 +49,8 @@ private:
 	//-------------------------------------------------------------------------
 	// Internals: 
 	//-------------------------------------------------------------------------
-	void FormatExceptionAddress(LPCSTR pExceptionAddress = nullptr);
+	void FormatExceptionAddress();
+	void FormatExceptionAddress(LPCSTR pExceptionAddress);
 	void FormatExceptionCode();
 
 	void FormatAPU(const char* pszRegister, DWORD64 nContent);
@@ -55,6 +61,7 @@ private:
 private:
 	enum
 	{
+		MAX_IMI_SEARCH = 7,
 		NUM_FRAMES_TO_CAPTURE = 128
 	};
 
@@ -69,6 +76,8 @@ private:
 
 	bool m_bCallState;       // Set when called to prevent recursive calls.
 	bool m_bCrashMsgCreated; // Set when crashmsg.exe is created to prevent recursive messages.
+
+	std::set<void*> m_WhiteList;
 	mutable std::mutex m_Mutex;
 };