めいくりぷとのブログ

技術的なことをまったりと。

メインスレッドにメッセージウィンドウを仕込む

#ifndef MAKEULONGLONG
#define MAKEULONGLONG(ldw, hdw) ((ULONGLONG(hdw) << 32) | ((ldw) & 0xFFFFFFFF))
#endif

#ifndef MAXULONGLONG
#define MAXULONGLONG ((ULONGLONG)~((ULONGLONG)0))
#endif

// utils.cpp
DWORD GetMainThreadId(__in DWORD dwProcessId)
{
	ULONGLONG ullMinCreateTime;
	HANDLE hThread, hSnapshot;
	THREADENTRY32 th32;
	DWORD dwThreadId;

	hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	if (hSnapshot == INVALID_HANDLE_VALUE)
		return 0;

	th32.dwSize = sizeof(THREADENTRY32);

	if (!Thread32First(hSnapshot, &th32))
		return 0;

	ullMinCreateTime = MAXULONGLONG;
	dwThreadId = 0;

	do
	{
		if (th32.th32OwnerProcessID == dwProcessId)
		{
			hThread = OpenThread(THREAD_QUERY_INFORMATION, TRUE, th32.th32ThreadID);
			if (hThread != NULL)
			{
				FILETIME afTimes[4] = { 0 };
				ULONGLONG ullCreationTime;

				if (GetThreadTimes(hThread, &afTimes[0], &afTimes[1], &afTimes[2], &afTimes[3]))
				{
					ullCreationTime = MAKEULONGLONG(afTimes[0].dwLowDateTime, afTimes[0].dwHighDateTime);
					if (ullCreationTime < ullMinCreateTime)
					{
						ullMinCreateTime = ullCreationTime;
						dwThreadId = th32.th32ThreadID;
					}
				}
			}
		}
	} while (Thread32Next(hSnapshot, &th32));

	return dwThreadId;
}

// dllmain.cpp : DLL アプリケーションのエントリ ポイントを定義します。
#include "stdafx.h"

#include "CMsgHandler.hpp"
#include "utils.hpp"

#include <intrin.h>
#pragma intrinsic(_ReturnAddress) 

HINSTANCE g_hInstance;

VOID Detour_PostMessage(__in BOOL fEnable);
VOID Detour_SendMessage(__in BOOL fEnable);

VOID Detour_SendMessage(__in BOOL fEnable)
{
	static decltype(&SendMessageA) _SendMessageA = SendMessageA;
	static decltype(&SendMessageW) _SendMessageW = SendMessageW;

	decltype(&SendMessageA) SendMessageA_Hook = [](
		_In_ HWND   hWnd,
		_In_ UINT   Msg,
		_In_ WPARAM wParam,
		_In_ LPARAM lParam) -> LRESULT
	{
		HMODULE hModule;

		if (CMsgHandler::GetInstance()->GetHandle() == NULL)
		{
			if (GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<TCHAR*>(_ReturnAddress()), &hModule))
			{
				if (hModule == GetModuleHandle(NULL))
				{
					if (GetMainThreadId(GetCurrentProcessId()) == GetCurrentThreadId())
					{
						if (CMsgHandler::GetInstance()->Init(g_hInstance))
						{
							Detour_SendMessage(FALSE);
							Detour_PostMessage(FALSE);
						}
					}
				}
			}
		}

		return _SendMessageA(hWnd, Msg, wParam, lParam);
	};
	decltype(&SendMessageW) SendMessageW_Hook = [](
		_In_ HWND   hWnd,
		_In_ UINT   Msg,
		_In_ WPARAM wParam,
		_In_ LPARAM lParam) -> LRESULT
	{
		HMODULE hModule;

		if (CMsgHandler::GetInstance()->GetHandle() == NULL)
		{
			if (GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<TCHAR*>(_ReturnAddress()), &hModule))
			{
				if (hModule == GetModuleHandle(NULL))
				{
					if (GetMainThreadId(GetCurrentProcessId()) == GetCurrentThreadId())
					{
						if (CMsgHandler::GetInstance()->Init(g_hInstance))
						{
							Detour_SendMessage(FALSE);
							Detour_PostMessage(FALSE);
						}
					}
				}
			}
		}

		return _SendMessageW(hWnd, Msg, wParam, lParam);
	};

	DetourFunction(fEnable, reinterpret_cast<LPVOID*>(&_SendMessageA), SendMessageA_Hook);
	DetourFunction(fEnable, reinterpret_cast<LPVOID*>(&_SendMessageW), SendMessageW_Hook);
}

VOID Detour_PostMessage(__in BOOL fEnable)
{
	static decltype(&PostMessageA) _PostMessageA = PostMessageA;
	static decltype(&PostMessageW) _PostMessageW = PostMessageW;

	decltype(&PostMessageA) PostMessageA_Hook = [](
		_In_ HWND   hWnd,
		_In_ UINT   Msg,
		_In_ WPARAM wParam,
		_In_ LPARAM lParam) -> BOOL
	{
		HMODULE hModule;

		if (CMsgHandler::GetInstance()->GetHandle() == NULL)
		{
			if (GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<TCHAR*>(_ReturnAddress()), &hModule))
			{
				if (hModule == GetModuleHandle(NULL))
				{
					if (GetMainThreadId(GetCurrentProcessId()) == GetCurrentThreadId())
					{
						if (CMsgHandler::GetInstance()->Init(g_hInstance))
						{
							Detour_SendMessage(FALSE);
							Detour_PostMessage(FALSE);
						}
					}
				}
			}
		}

		return _PostMessageA(hWnd, Msg, wParam, lParam);
	};
	decltype(&PostMessageW) PostMessageW_Hook = [](
		_In_ HWND   hWnd,
		_In_ UINT   Msg,
		_In_ WPARAM wParam,
		_In_ LPARAM lParam) -> BOOL
	{
		HMODULE hModule;

		if (CMsgHandler::GetInstance()->GetHandle() == NULL)
		{
			if (GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, reinterpret_cast<TCHAR*>(_ReturnAddress()), &hModule))
			{
				if (hModule == GetModuleHandle(NULL))
				{
					if (GetMainThreadId(GetCurrentProcessId()) == GetCurrentThreadId())
					{
						if (CMsgHandler::GetInstance()->Init(g_hInstance))
						{
							Detour_SendMessage(FALSE);
							Detour_PostMessage(FALSE);
						}
					}
				}
			}
		}

		return _PostMessageW(hWnd, Msg, wParam, lParam);
	};

	DetourFunction(fEnable, reinterpret_cast<LPVOID*>(&_PostMessageA), PostMessageA_Hook);
	DetourFunction(fEnable, reinterpret_cast<LPVOID*>(&_PostMessageW), PostMessageW_Hook);
}

BOOL WINAPI OnAttachProcess(__in HINSTANCE hInstance)
{
	Detour_SendMessage(TRUE);
	Detour_PostMessage(TRUE);

	return TRUE;
}

BOOL APIENTRY DllMain(__in HINSTANCE hInstance, __in DWORD fdwReason, __in __reserved LPVOID lpvReserved)
{
	switch (fdwReason)
	{
		case DLL_PROCESS_ATTACH:
		{
			g_hInstance = hInstance;

			return OnAttachProcess(hInstance);
		}
		case DLL_PROCESS_DETACH:
		{
			FreeLibrary(hInstance);
			break;
		}
	}
	return TRUE;
}

GetMainThreadIdで、スレッドの作成時間を比較することで、一番最初に作成されたスレッドがメインスレッドと断定する。
メインスレッドから呼ばれるAPIであれば、これらのAPIに限らず何でもok