Only threads suspended by polyhook are now resumed. Added a thread ma…

…nager object for safety and ease of use.
stevemk14ebr · Sep 13, 2016 · 9bda6e0 · 9bda6e0
1 parent 60cdf06
commit 9bda6e0
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 39 deletions.
diff --git a/PolyHook/PolyHook.cpp b/PolyHook/PolyHook.cpp
@@ -62,37 +62,6 @@ void PLH::IHook::PrintError(const RuntimeError& Err) const
 		Err.GetString().c_str()); 
 }
 
-void PLH::IHook::ToggleThreadSuspension(DWORD CallingThreadId, bool Suspend)
-{
-	HANDLE h = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
-	if (h != INVALID_HANDLE_VALUE)
-		return;
-
-	THREADENTRY32 te;
-	te.dwSize = sizeof(te);
-	for (Thread32First(h, &te); Thread32Next(h, &te); )
-	{
-		if (te.dwSize >= FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) + sizeof(te.th32OwnerProcessID))
-		{
-			// Suspend all threads EXCEPT the one we want to keep running
-			if (te.th32ThreadID != CallingThreadId && te.th32OwnerProcessID == GetCurrentProcessId())
-			{
-				HANDLE thread = OpenThread(THREAD_SUSPEND_RESUME, FALSE, te.th32ThreadID);
-				if (thread == NULL)
-					continue;
-
-				if (Suspend)
-					SuspendThread(thread);
-				else
-					ResumeThread(thread);
-				CloseHandle(thread);
-			}
-		}
-		te.dwSize = sizeof(te);
-	}
-	CloseHandle(h);
-}
-
 PLH::RuntimeError PLH::IHook::GetLastError() const
 {
 	return m_LastError;
@@ -321,7 +290,11 @@ bool PLH::X86Detour::Hook()
 		XTrace("Function to small to hook\n");
 		return false;
 	}
-	this->ToggleThreadSuspension(GetCurrentThreadId(), true);
+
+	//TODO: Add single step support in case processes EIP is on/in the section we write to
+	Tools::ThreadManager ThreadMngr;
+	ThreadMngr.SuspendThreads();
+
 	m_Trampoline = new BYTE[m_hkLength + 30];   //Allocate Space for original plus extra to jump back and for jmp table
 	m_NeedFree = true;
 	VirtualProtect(m_Trampoline, m_hkLength + 30, PAGE_EXECUTE_READWRITE, &OldProtection); //Allow Execution
@@ -341,7 +314,7 @@ bool PLH::X86Detour::Hook()
 	for (int i = 5; i < m_OriginalLength; i++)
 		m_hkSrc[i] = 0x90;
 
-	this->ToggleThreadSuspension(GetCurrentThreadId(), false);
+	ThreadMngr.ResumeThreads();
 	FlushSrcInsCache();
 	m_Hooked = true;
 	PostError(RuntimeError(RuntimeError::Severity::Warning, "PolyHook x86Detour: Some opcodes may not be relocated properly"));
@@ -458,8 +431,10 @@ bool PLH::X64Detour::Hook()
 			return false;
 		}
 	}
+	//TODO: Add single step support in case processes RIP is on/in the section we write to
+	Tools::ThreadManager ThreadMngr;
+	ThreadMngr.SuspendThreads();
 
-	this->ToggleThreadSuspension(GetCurrentThreadId(), true);
 	memcpy(m_OriginalCode, m_hkSrc, m_hkLength);
 	memcpy(m_Trampoline, m_hkSrc, m_hkLength);
 	WriteAbsoluteJMP((DWORD64)&m_Trampoline[m_hkLength], (DWORD64)m_hkSrc + m_hkLength);
@@ -486,7 +461,7 @@ bool PLH::X64Detour::Hook()
 	for (int i = HookSize; i < m_OriginalLength; i++)
 		m_hkSrc[i] = 0x90;
 
-	this->ToggleThreadSuspension(GetCurrentThreadId(), false);
+	ThreadMngr.ResumeThreads();
 	FlushInstructionCache(GetCurrentProcess(), m_hkSrc, m_hkLength);
 	m_Hooked = true;
 	PostError(RuntimeError(RuntimeError::Severity::Warning, "PolyHook x64Detour: Relocation can be out of range"));

diff --git a/PolyHook/PolyHook.h b/PolyHook/PolyHook.h
@@ -13,6 +13,88 @@
 #pragma comment(lib,"capstone.lib")
 
 namespace PLH {
+	namespace Tools
+	{
+		class ThreadHandle
+		{
+		public:
+			//Thread ID, OpenThread's AccessFlag 
+			ThreadHandle(DWORD ThreadId, DWORD  DesiredAccessFlags) : m_ThreadId(ThreadId), m_IsSuspended(false)
+			{
+				m_hThread = OpenThread(DesiredAccessFlags, FALSE, ThreadId);
+			}
+
+			void ToggleSuspend(bool Suspend)
+			{
+				if (Suspend && !m_IsSuspended)
+				{
+					SuspendThread(m_hThread);
+					m_IsSuspended = true;
+				}else if (!Suspend && m_IsSuspended){
+					ResumeThread(m_hThread);
+					m_IsSuspended = false;
+				}
+			}
+
+			~ThreadHandle()
+			{
+				if (m_IsSuspended)
+					ToggleSuspend(false);
+
+				if (!m_hThread)
+					CloseHandle(m_hThread);
+			}
+		private:
+			bool m_IsSuspended;
+			HANDLE m_hThread;
+			DWORD m_ThreadId;
+		};
+
+		class ThreadManager
+		{
+		public:
+			void SuspendThreads()
+			{
+				UpdateThreadList(GetCurrentThreadId());
+				for (ThreadHandle& ThreadInstance : m_SuspendedThreads)
+				{
+					ThreadInstance.ToggleSuspend(true);
+				}
+			}
+			void ResumeThreads()
+			{
+				for (ThreadHandle& ThreadInstance : m_SuspendedThreads)
+				{
+					ThreadInstance.ToggleSuspend(false);
+				}
+			}
+		private:
+			void UpdateThreadList(DWORD CallingThreadId)
+			{
+				m_SuspendedThreads.clear();
+				HANDLE h = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
+				if (h != INVALID_HANDLE_VALUE)
+					return;
+
+				THREADENTRY32 te;
+				te.dwSize = sizeof(te);
+				for (Thread32First(h, &te); Thread32Next(h, &te); )
+				{
+					if (te.dwSize >= FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) + sizeof(te.th32OwnerProcessID))
+					{
+						if (te.th32ThreadID != CallingThreadId && te.th32OwnerProcessID == GetCurrentProcessId())
+						{
+							m_SuspendedThreads.emplace_back(te.th32ThreadID, THREAD_SUSPEND_RESUME);
+						}
+					}
+					te.dwSize = sizeof(te);
+				}
+				CloseHandle(h);
+			}
+			std::vector<Tools::ThreadHandle> m_SuspendedThreads;
+		};
+	}
+
 	class ASMHelper
 	{
 	public:
@@ -118,13 +200,9 @@ namespace PLH {
 
 		virtual RuntimeError GetLastError() const;
 		virtual void PrintError(const RuntimeError& Err) const;
-
 	protected:
 		virtual void PostError(const RuntimeError& Err);
 
-		//Does not suspend the thread with the ID passed in, Suspend == false calls resume thread
-		virtual void ToggleThreadSuspension(DWORD CallingThreadId, bool Suspend); 
-
 		RuntimeError m_LastError;
 	};