HookHelper.cpp

#include "pch.h"
#include "HookHelper.hpp"
#include "detours.h"

using namespace AcrylicEverywhere;
namespace AcrylicEverywhere::HookHelper::Detours
{
	// Begin to install hooks
	HRESULT Begin();
	// End attaching
	HRESULT End(bool commit = true);
}

bool HookHelper::IsBadReadPtr(const void* ptr)
{
	bool result{ false };
	MEMORY_BASIC_INFORMATION mbi{};
	if (VirtualQuery(ptr, &mbi, sizeof(mbi)))
	{
		DWORD mask{ PAGE_READONLY | PAGE_READWRITE | PAGE_WRITECOPY | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY };
		result = !(mbi.Protect & mask);

		// Check the page is not a guard page
		if (mbi.Protect & (PAGE_GUARD | PAGE_NOACCESS))
		{
			result = true;
		}
	}
	else
	{
		result = true;
	}

	return result;
}

PVOID HookHelper::WritePointer(PVOID* memoryAddress, PVOID value) try
{
	THROW_HR_IF_NULL(E_INVALIDARG, memoryAddress);

	MEMORY_BASIC_INFORMATION mbi{};
	THROW_LAST_ERROR_IF(
		VirtualQuery(
			memoryAddress,
			&mbi,
			sizeof(MEMORY_BASIC_INFORMATION)
		) == 0
	);
	THROW_IF_WIN32_BOOL_FALSE(
		VirtualProtect(
			mbi.BaseAddress,
			mbi.RegionSize,
			PAGE_EXECUTE_READWRITE,
			&mbi.Protect
		)
	);
	PVOID oldValue{ InterlockedExchangePointer(memoryAddress, value) };
	THROW_IF_WIN32_BOOL_FALSE(
		VirtualProtect(
			mbi.BaseAddress,
			mbi.RegionSize,
			mbi.Protect,
			&mbi.Protect
		)
	);

	return oldValue;
}
catch (...) { return nullptr; }
void HookHelper::WriteMemory(PVOID memoryAddress, const std::function<void()>&& callback) try
{
	THROW_HR_IF_NULL(E_INVALIDARG, memoryAddress);

	MEMORY_BASIC_INFORMATION mbi{};
	THROW_LAST_ERROR_IF(
		VirtualQuery(
			memoryAddress,
			&mbi,
			sizeof(MEMORY_BASIC_INFORMATION)
		) == 0
	);
	THROW_IF_WIN32_BOOL_FALSE(
		VirtualProtect(
			mbi.BaseAddress,
			mbi.RegionSize,
			PAGE_EXECUTE_READWRITE,
			&mbi.Protect
		)
	);
	callback();
	THROW_IF_WIN32_BOOL_FALSE(
		VirtualProtect(
			mbi.BaseAddress,
			mbi.RegionSize,
			mbi.Protect,
			&mbi.Protect
		)
	);
}
catch(...) {}

PVOID HookHelper::InjectCallbackToThread(DWORD threadId, const std::function<void()>& callback) try
{
#ifdef _WIN64
wil::unique_handle threadHandle{ OpenThread(THREAD_SET_CONTEXT | THREAD_GET_CONTEXT | THREAD_SUSPEND_RESUME, FALSE, threadId) };
	THROW_LAST_ERROR_IF_NULL(threadHandle);

	THROW_LAST_ERROR_IF(SuspendThread(threadHandle.get()) == -1);
	auto cleanUp = wil::scope_exit([&]
	{
		THROW_LAST_ERROR_IF(ResumeThread(threadHandle.get()) == -1);
	});

	CONTEXT context{};
	context.ContextFlags = CONTEXT_FULL;
	THROW_LAST_ERROR_IF(!GetThreadContext(threadHandle.get(), &context));

	auto proxyCallback = [](ULONG_PTR parameter) -> void
	{
		auto callback_ptr{ reinterpret_cast<std::function<void()>*>(parameter) };
		auto callback{ *callback_ptr };
		delete callback_ptr;

		callback();
	};

	static BYTE s_shellCode[]
	{
		// sub rsp, 28h
		0x48, 0x83, 0xec, 0x28,
		// mov [rsp + 18], rax
		0x48, 0x89, 0x44, 0x24, 0x18,
		// mov [rsp + 10h], rcx
		0x48, 0x89, 0x4c, 0x24, 0x10,
		// mov rcx, function_parameter
		0x48, 0xb9, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
		// mov rax, function_pointer
		0x48, 0xb8, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
		// call rax
		0xff, 0xd0,
		// mov rcx, [rsp + 10h]
		0x48, 0x8b, 0x4c, 0x24, 0x10,
		// mov rax, [rsp + 18h]
		0x48, 0x8b, 0x44, 0x24, 0x18,
		// add rsp, 28h
		0x48, 0x83, 0xc4, 0x28,
		// mov r11, context.Rip
		0x49, 0xbb, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33, 0x33,
		// jmp r11
		0x41, 0xff, 0xe3
	};
	*reinterpret_cast<PVOID*>(&s_shellCode[16]) = static_cast<void*>(new std::function<void()>{ callback });
	*reinterpret_cast<PVOID*>(&s_shellCode[26]) = static_cast<void*>(static_cast<void(WINAPI*)(ULONG_PTR)>(proxyCallback));
	*reinterpret_cast<ULONG_PTR*>(&s_shellCode[52]) = context.Rip;

	auto buffer
	{ 
		VirtualAlloc(
			nullptr,
			sizeof(s_shellCode),
			MEM_COMMIT | MEM_RESERVE,
			PAGE_EXECUTE_READWRITE
		) 
	};
	memcpy_s(buffer, sizeof(s_shellCode), s_shellCode, sizeof(s_shellCode));
	context.Rip = reinterpret_cast<ULONG_PTR>(buffer);
	THROW_LAST_ERROR_IF(!SetThreadContext(threadHandle.get(), &context));

	return buffer;
#else
	return nullptr;
#endif
}
catch(...) { return nullptr; }

HMODULE HookHelper::GetProcessModule(HANDLE processHandle, std::wstring_view dllPath)
{
	HMODULE targetModule{ nullptr };
	DWORD bytesNeeded{ 0 };
	if (!EnumProcessModules(processHandle, nullptr, 0, &bytesNeeded))
	{
		return targetModule;
	}
	DWORD moduleCount{ bytesNeeded / sizeof(HMODULE) };
	auto moduleList{ std::make_unique<HMODULE[]>(moduleCount) };
	if (!EnumProcessModules(processHandle, moduleList.get(), bytesNeeded, &bytesNeeded))
	{
		return targetModule;
	}

	for (DWORD i = 0; i < moduleCount; i++)
	{
		HMODULE moduleHandle{ moduleList[i] };
		WCHAR modulePath[MAX_PATH + 1];
		GetModuleFileNameExW(processHandle, moduleHandle, modulePath, MAX_PATH);

		if (!_wcsicmp(modulePath, dllPath.data()))
		{
			targetModule = moduleHandle;
			break;
		}
	}

	return targetModule;
}

void HookHelper::WalkIAT(PVOID baseAddress, std::string_view dllName, std::function<bool(PVOID* functionAddress, LPCSTR functionNameOrOrdinal, BOOL importedByName)> callback) try
{
	THROW_HR_IF(E_INVALIDARG, IsBadReadPtr(baseAddress));
	THROW_HR_IF(E_INVALIDARG, dllName.empty());
	THROW_HR_IF_NULL(E_INVALIDARG, baseAddress);
	THROW_HR_IF_NULL(E_INVALIDARG, callback);

	ULONG size{ 0ul };
	auto importDescriptor
	{
		static_cast<PIMAGE_IMPORT_DESCRIPTOR>(
			ImageDirectoryEntryToData(
				baseAddress,
				TRUE,
				IMAGE_DIRECTORY_ENTRY_IMPORT,
				&size
			)
		)
	};

	THROW_HR_IF_NULL(E_INVALIDARG, importDescriptor);

	bool found = false;
	while (importDescriptor->Name)
	{
		auto moduleName = reinterpret_cast<LPCSTR>(reinterpret_cast<UINT_PTR>(baseAddress) + importDescriptor->Name);

		if (!_stricmp(moduleName, dllName.data()))
		{
			found = true;
			break;
		}

		importDescriptor++;
	}

	THROW_HR_IF(HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND), !found);

	auto thunk{ reinterpret_cast<PIMAGE_THUNK_DATA>(reinterpret_cast<UINT_PTR>(baseAddress) + importDescriptor->FirstThunk) };
	auto nameThunk = reinterpret_cast<PIMAGE_THUNK_DATA>(reinterpret_cast<UINT_PTR>(baseAddress) + importDescriptor->OriginalFirstThunk);

	bool result{ true };
	while (thunk->u1.Function)
	{
		LPCSTR functionName{ nullptr };
		auto functionAddress = reinterpret_cast<PVOID*>(&thunk->u1.Function);

		BOOL importedByName{ !IMAGE_SNAP_BY_ORDINAL(nameThunk->u1.Ordinal) };
		if (importedByName)
		{
			functionName = reinterpret_cast<PIMAGE_IMPORT_BY_NAME>(
				RVA_TO_ADDR(baseAddress, static_cast<RVA>(nameThunk->u1.AddressOfData))
				)->Name;
		}
		else
		{
			functionName = MAKEINTRESOURCEA(IMAGE_ORDINAL(nameThunk->u1.Ordinal));
		}

		result = callback(functionAddress, functionName, importedByName);
		if (!result)
		{
			break;
		}

		thunk++;
		nameThunk++;
	}
}
catch (...) {}

void HookHelper::WalkDelayloadIAT(PVOID baseAddress, std::string_view dllName, std::function<bool(HMODULE* moduleHandle, PVOID* functionAddress, LPCSTR functionNameOrOrdinal, BOOL importedByName)> callback) try
{
	THROW_HR_IF(E_INVALIDARG, IsBadReadPtr(baseAddress));
	THROW_HR_IF(E_INVALIDARG, dllName.empty());
	THROW_HR_IF_NULL(E_INVALIDARG, baseAddress);
	THROW_HR_IF_NULL(E_INVALIDARG, callback);

	ULONG size{ 0ul };
	auto importDescriptor
	{
		static_cast<PIMAGE_DELAYLOAD_DESCRIPTOR>(
			ImageDirectoryEntryToData(
				baseAddress,
				TRUE,
				IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT,
				&size
			)
		)
	};

	THROW_HR_IF_NULL(E_INVALIDARG, importDescriptor);

	bool found = false;
	while (importDescriptor->DllNameRVA)
	{
		auto moduleName = reinterpret_cast<LPCSTR>(
			RVA_TO_ADDR(baseAddress, importDescriptor->DllNameRVA)
			);

		if (!_stricmp(moduleName, dllName.data()))
		{
			found = true;
			break;
		}

		importDescriptor++;
	}

	THROW_HR_IF(HRESULT_FROM_WIN32(ERROR_FILE_NOT_FOUND), !found);

	auto attributes{ importDescriptor->Attributes.RvaBased };
	THROW_WIN32_IF_MSG(ERROR_FILE_NOT_FOUND, attributes != 1, "Unsupported delay loaded dll![%hs]", dllName.data());

	auto moduleHandle = reinterpret_cast<HMODULE*>(
		RVA_TO_ADDR(baseAddress, importDescriptor->ModuleHandleRVA)
		);
	auto thunk = reinterpret_cast<PIMAGE_THUNK_DATA>(
		RVA_TO_ADDR(baseAddress, importDescriptor->ImportAddressTableRVA)
		);
	auto nameThunk = reinterpret_cast<PIMAGE_THUNK_DATA>(
		RVA_TO_ADDR(baseAddress, importDescriptor->ImportNameTableRVA)
		);

	bool result{ true };
	while (thunk->u1.Function)
	{
		LPCSTR functionName{ nullptr };
		auto functionAddress = reinterpret_cast<PVOID*>(&thunk->u1.Function);

		BOOL importedByName{ !IMAGE_SNAP_BY_ORDINAL(nameThunk->u1.Ordinal) };
		if (importedByName)
		{
			functionName = reinterpret_cast<PIMAGE_IMPORT_BY_NAME>(
				RVA_TO_ADDR(baseAddress, nameThunk->u1.AddressOfData)
				)->Name;
		}
		else
		{
			functionName = MAKEINTRESOURCEA(IMAGE_ORDINAL(nameThunk->u1.Ordinal));
		}

		result = callback(moduleHandle, functionAddress, functionName, importedByName);
		if (!result)
		{
			break;
		}

		thunk++;
		nameThunk++;
	}
}
catch (...) {}

PVOID* HookHelper::GetIAT(PVOID baseAddress, std::string_view dllName, LPCSTR targetFunctionNameOrOrdinal)
{
	PVOID* originalFunction{ nullptr };
	WalkIAT(baseAddress, dllName, [&](PVOID* functionAddress, LPCSTR functionNameOrOrdinal, BOOL importedByName) -> bool
	{
		if (
			(importedByName == TRUE && targetFunctionNameOrOrdinal && !strcmp(functionNameOrOrdinal, targetFunctionNameOrOrdinal)) ||
			(importedByName == FALSE && functionNameOrOrdinal == targetFunctionNameOrOrdinal)
		)
		{
			originalFunction = functionAddress;

			return false;
		}

		return true;
	});

	return originalFunction;
}
std::pair<HMODULE*, PVOID*> HookHelper::GetDelayloadIAT(PVOID baseAddress, std::string_view dllName, LPCSTR targetFunctionNameOrOrdinal, bool resolveAPI)
{
	std::pair<HMODULE*, PVOID*> originalInfo{nullptr, nullptr};
	WalkDelayloadIAT(baseAddress, dllName, [&](HMODULE* moduleHandle, PVOID* functionAddress, LPCSTR functionNameOrOrdinal, BOOL importedByName) -> bool
	{
		if (
			(importedByName == TRUE && targetFunctionNameOrOrdinal && (reinterpret_cast<DWORD64>(targetFunctionNameOrOrdinal) & 0xFFFF0000) != 0 && !strcmp(functionNameOrOrdinal, targetFunctionNameOrOrdinal)) ||
			(importedByName == FALSE && functionNameOrOrdinal == targetFunctionNameOrOrdinal)
		)
		{
			originalInfo.first = moduleHandle;
			originalInfo.second = functionAddress;

			if (resolveAPI)
			{
				ResolveDelayloadIAT(originalInfo, baseAddress, dllName, targetFunctionNameOrOrdinal);
			}
			return false;
		}

		return true;
	});

	return originalInfo;
}

PVOID HookHelper::WriteIAT(PVOID baseAddress, std::string_view dllName, LPCSTR targetFunctionNameOrOrdinal, PVOID detourFunction)
{
	PVOID originalFunction{ nullptr };

	auto functionAddress{ GetIAT(baseAddress, dllName, targetFunctionNameOrOrdinal) };
	if (functionAddress)
	{
		originalFunction = *functionAddress;

		WriteMemory(functionAddress, [&]()
		{
			InterlockedExchangePointer(functionAddress, detourFunction);
		});
	}

	return originalFunction;
}
std::pair<HMODULE, PVOID> HookHelper::WriteDelayloadIAT(PVOID baseAddress, std::string_view dllName, LPCSTR targetFunctionNameOrOrdinal, PVOID detourFunction, std::optional<HMODULE> newModuleHandle)
{
	HMODULE originalModule{ nullptr };
	PVOID originalFunction{ nullptr };

	auto [moduleHandle, functionAddress]{GetDelayloadIAT(baseAddress, dllName, targetFunctionNameOrOrdinal, true)};
	if (functionAddress)
	{
		originalModule = *moduleHandle;
		originalFunction = *functionAddress;

		if (newModuleHandle)
		{
			WriteMemory(moduleHandle, [&]()
			{
				InterlockedExchangePointer(reinterpret_cast<PVOID*>(moduleHandle), newModuleHandle.value());
			});
		}
		WriteMemory(functionAddress, [&]()
		{
			InterlockedExchangePointer(functionAddress, detourFunction);
		});
	}

	return std::make_pair(originalModule, originalFunction);
}

void HookHelper::ResolveDelayloadIAT(const std::pair<HMODULE*, PVOID*>& info, PVOID baseAddress, std::string_view dllName, LPCSTR targetFunctionNameOrOrdinal)
{
	auto& [moduleHandle, functionAddress]{info};
	if (DetourGetContainingModule(*functionAddress) == baseAddress || DetourGetContainingModule(DetourCodeFromPointer(*functionAddress, nullptr)) == baseAddress)
	{
		if (!(*moduleHandle))
		{
			HMODULE importModule{ LoadLibraryA(dllName.data()) };

			WriteMemory(moduleHandle, [&]()
			{
				InterlockedExchangePointer(reinterpret_cast<PVOID*>(moduleHandle), importModule);
			});
			WriteMemory(functionAddress, [&]()
			{
				InterlockedExchangePointer(reinterpret_cast<PVOID*>(functionAddress), GetProcAddress(importModule, targetFunctionNameOrOrdinal));
			});
		}
		else
		{
			WriteMemory(functionAddress, [&]()
			{
				InterlockedExchangePointer(reinterpret_cast<PVOID*>(functionAddress), GetProcAddress(*moduleHandle, targetFunctionNameOrOrdinal));
			});
		}
	}
}

PUCHAR HookHelper::MatchAsmCode(PUCHAR startAddress, std::pair<UCHAR*, size_t> asmCodeInfo, std::vector<std::pair<size_t, size_t>> asmSliceInfo, size_t maxSearchBytes)
{
	PUCHAR matchedAddress{ nullptr };
	if (!startAddress)
	{
		return matchedAddress;
	}

	size_t currentMaxSearchBytes{ 0 };
	size_t matchedBytesLength{ 0 };
	size_t sliceIndex{ 0 };

	auto [asmCode, asmCodeLength] { asmCodeInfo };
	while (matchedBytesLength < asmCodeLength)
	{
		auto [section_length, gap_length] { asmSliceInfo[sliceIndex] };
		
		if (section_length)
		{
			if (sliceIndex == 0)
			{
				while (memcmp(startAddress, &asmCode[matchedBytesLength], section_length) != 0)
				{
					startAddress++;
					maxSearchBytes--;
					if (!maxSearchBytes)
					{
						goto leave;
					}

					continue;
				}
			}
			else if (memcmp(startAddress, &asmCode[matchedBytesLength], section_length) != 0)
			{
				startAddress = matchedAddress + 1;
				maxSearchBytes--;
				matchedAddress = nullptr;
				matchedBytesLength = 0;
				sliceIndex = 0;
				maxSearchBytes = currentMaxSearchBytes;
				if (!maxSearchBytes)
				{
					goto leave;
				}

				continue;
			}

			if (sliceIndex == 0)
			{
				matchedAddress = startAddress;
				currentMaxSearchBytes = maxSearchBytes;
			}
		}
		startAddress += (section_length + gap_length);
		maxSearchBytes -= (section_length + gap_length);
		matchedBytesLength += (section_length + gap_length);
		sliceIndex++;

		if (!maxSearchBytes || sliceIndex >= asmSliceInfo.size())
		{
			goto leave;
		}
	}

leave:
	return matchedAddress;
}

std::pair<UCHAR*, size_t> HookHelper::GetTextSectionInfo(PVOID baseAddress)
{
	auto ntHeader{ ImageNtHeader(baseAddress) };
	auto sectionCount{ ntHeader->FileHeader.NumberOfSections };
	auto section{ IMAGE_FIRST_SECTION(ntHeader) };
	for (WORD i = 0; i < sectionCount; i++)
	{
		if (!strcmp(reinterpret_cast<const char*>(section->Name), ".text"))
		{
			return std::make_pair<UCHAR*, size_t>(
				reinterpret_cast<UCHAR*>(reinterpret_cast<ULONG_PTR>(baseAddress) + section->VirtualAddress),
				section->Misc.VirtualSize
			);
		}
		section++;
	}

	return std::make_pair<UCHAR*, size_t>(nullptr, 0);
}

HookHelper::ImageMapper::ImageMapper(std::wstring_view dllPath) : 
	m_fileHandle{
		CreateFileW(
			dllPath.data(),
			GENERIC_READ,
			FILE_SHARE_READ,
			nullptr,
			OPEN_EXISTING,
			FILE_ATTRIBUTE_NORMAL,
			nullptr
		)
	},
	m_fileMapping{
		CreateFileMappingFromApp(
			m_fileHandle.get(),
			nullptr,
			PAGE_READONLY | SEC_IMAGE_NO_EXECUTE,
			0,
			nullptr
		)
	},
	m_baseAddress{
		MapViewOfFileFromApp(
			m_fileMapping.get(),
			FILE_MAP_READ,
			0,
			0
		)
	}
{
}

HRESULT HookHelper::Detours::Begin()
{
	DetourSetIgnoreTooSmall(TRUE);
	RETURN_IF_WIN32_ERROR(DetourTransactionBegin());
	RETURN_IF_WIN32_ERROR(DetourUpdateThread(GetCurrentThread()));
	return S_OK;
}

HRESULT HookHelper::Detours::End(bool commit)
{
	return ( commit ? HRESULT_FROM_WIN32(DetourTransactionCommit()) : HRESULT_FROM_WIN32(DetourTransactionAbort()) );
}

HRESULT HookHelper::Detours::Write(const std::function<void()>&& callback) try
{
	HRESULT hr{ HookHelper::Detours::Begin() };
	if (FAILED(hr))
	{
		return hr;
	}

	callback();

	return HookHelper::Detours::End(true);
}
catch (...)
{
	LOG_CAUGHT_EXCEPTION();
	LOG_IF_FAILED(HookHelper::Detours::End(false));
	return wil::ResultFromCaughtException();
}

void HookHelper::Detours::Attach(std::string_view dllName, std::string_view funcName, PVOID* realFuncAddr, PVOID hookFuncAddr)
{
	THROW_HR_IF_NULL(E_INVALIDARG, realFuncAddr);
	THROW_HR_IF_NULL(E_INVALIDARG, *realFuncAddr);
	*realFuncAddr = DetourFindFunction(dllName.data(), funcName.data());
	THROW_LAST_ERROR_IF_NULL(*realFuncAddr);

	Attach(realFuncAddr, hookFuncAddr);
}

void HookHelper::Detours::Attach(PVOID* realFuncAddr, PVOID hookFuncAddr)
{
	THROW_HR_IF_NULL(E_INVALIDARG, realFuncAddr);
	THROW_HR_IF_NULL(E_INVALIDARG, *realFuncAddr);
	THROW_HR_IF_NULL(E_INVALIDARG, hookFuncAddr);

	THROW_IF_WIN32_ERROR(DetourAttach(realFuncAddr, hookFuncAddr));
}

void HookHelper::Detours::Detach(PVOID* realFuncAddr, PVOID hookFuncAddr)
{
	THROW_HR_IF_NULL(E_INVALIDARG, realFuncAddr);
	THROW_HR_IF_NULL(E_INVALIDARG, *realFuncAddr);
	THROW_HR_IF_NULL(E_INVALIDARG, hookFuncAddr);

	THROW_IF_WIN32_ERROR(DetourDetach(realFuncAddr, hookFuncAddr));
}