AbsoluteTouchEx patch to inject into existing process

absolutetouchex-injection.patch

diff --git "a/atdll/atdll.cpp" "b/atdll/atdll.cpp"
index 280349b..22e818c 100644
--- "a/atdll/atdll.cpp"
+++ "b/atdll/atdll.cpp"
@@ -9,6 +9,8 @@
 #include <hidsdi.h>
 #include <hidpi.h>
 #include <hidusage.h>
+#include <TlHelp32.h>
+#include <iostream>
 #include "detours.h"
 
 // Version number
@@ -173,6 +175,9 @@ static thread_local RAWINPUT t_injectedInput;
 // Holds the current primary touch point ID
 static thread_local ULONG t_primaryContactID;
 
+// Holds the hWnd to add hotkey for
+static HWND g_hWndToHook;
+
 // Allocates a malloc_ptr with the given size. The size must be
 // greater than or equal to sizeof(T).
 template<typename T>
@@ -883,19 +900,30 @@ AT_WndProcHook(
 	WPARAM wParam,
 	LPARAM lParam)
 {
+	if (g_hWndToHook == hWnd) {
+		std::cout << "Running deferred hotkey hooking" << std::endl;
+		RegisterHotKey(hWnd, HOTKEY_ENABLE_ID, HOTKEY_ENABLE_MOD, HOTKEY_ENABLE_VK);
+		RegisterHotKey(hWnd, HOTKEY_CALIBRATION_ID, HOTKEY_CALIBRATION_MOD, HOTKEY_CALIBRATION_VK);
+		RegisterHotKey(hWnd, HOTKEY_LOAD_ID, HOTKEY_LOAD_MOD, HOTKEY_LOAD_VK);
+		RegisterHotKey(hWnd, HOTKEY_SAVE_ID, HOTKEY_SAVE_MOD, HOTKEY_SAVE_VK);
+		g_hWndToHook = NULL;
+	}
 	if (message == WM_HOTKEY && wParam == HOTKEY_ENABLE_ID) {
 		g_enabled = !g_enabled;
 		debugf("Absolute touch mode -> %s", g_enabled ? "ON" : "OFF");
 		return 0;
@@ -935,7 +966,7 @@ static BOOL WINAPI
 AT_RegisterRawInputDevicesHook(
 	PCRAWINPUTDEVICE pRawInputDevices,
 	UINT uiNumDevices,
-    UINT cbSize)
+	UINT cbSize, bool isHookingExistingDevices = false)
 {
 	if (uiNumDevices == 0) {
 		debugf("RegisterRawInputDevices called with no devices");
@@ -953,16 +984,22 @@ AT_RegisterRawInputDevicesHook(
 			// chances are that only one window will be receiving raw input. Ignore errors.
 			// This is a bit ugly, but it works well enough for our purposes.
 			debugf("Registering global hotkeys with hWnd=%p", hWnd);
+			if (isHookingExistingDevices && g_hWndToHook == NULL) {
+				std::cout << "Deferring hotkey hooking to hooked WndProc" << std::endl;
+				g_hWndToHook = hWnd;
+			}
+			else if (!isHookingExistingDevices) {
 				RegisterHotKey(hWnd, HOTKEY_ENABLE_ID, HOTKEY_ENABLE_MOD, HOTKEY_ENABLE_VK);
 				RegisterHotKey(hWnd, HOTKEY_CALIBRATION_ID, HOTKEY_CALIBRATION_MOD, HOTKEY_CALIBRATION_VK);
 				RegisterHotKey(hWnd, HOTKEY_LOAD_ID, HOTKEY_LOAD_MOD, HOTKEY_LOAD_VK);
 				RegisterHotKey(hWnd, HOTKEY_SAVE_ID, HOTKEY_SAVE_MOD, HOTKEY_SAVE_VK);
-
+			}
 
 			debugf("Registering touchpad input with hWnd=%p", hWnd);
 			try {
@@ -1133,10 +1171,53 @@ AT_PrintSystemInfo()
 	}
 }
 
+// Callback for EnumThreadWindows.
+BOOL CALLBACK EnumWindowProc(HWND hwnd, LPARAM /*lParam*/) {
+#ifdef _WIN64
+	WNDPROC windowProc = (WNDPROC)g_originalGetWindowLongPtrW(hwnd, GWL_WNDPROC);
+#else
+	WNDPROC windowProc = (WNDPROC)g_originalGetWindowLongW(hwnd, GWL_WNDPROC);
+#endif
+	if (windowProc == NULL) {
+		std::cout << "Window proc is null, skipping this window" << std::endl;
+		return TRUE;
+	}
+	if (windowProc == AT_WndProcHook) {
+		std::cout << "Window seems to be already hooked, skipping" << std::endl;
+		return TRUE;
+}
+#ifdef _WIN64
+	WNDPROC prevWindowProc = (WNDPROC)g_originalSetWindowLongPtrW(hwnd, GWL_WNDPROC, (LONG_PTR)&AT_WndProcHook);
+#else
+	WNDPROC prevWindowProc = (WNDPROC)g_originalSetWindowLongW(hwnd, GWL_WNDPROC, (LONG)&AT_WndProcHook);
+#endif
+	g_originalWndProcs[hwnd] = prevWindowProc;
+	std::cout << "Patched existing window proc, hwnd = " << hwnd << std::endl;
+	return TRUE;
+}
+
 // Called at startup to patch all the WinAPI functions.
 static void
 AT_Initialize()
 {
+	AllocConsole();
 	DetourTransactionBegin();
 	DetourUpdateThread(GetCurrentThread());
 	DetourAttach((void**)&g_originalRegisterRawInputDevices, AT_RegisterRawInputDevicesHook);
@@ -1152,6 +1233,30 @@ AT_Initialize()
 	if (DetourTransactionCommit() != NO_ERROR) {
 		throw std::runtime_error("Failed to commit Detours transaction");
 	}
+
+	// Enumerate all current windows and replace their procs
+	HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
+	if (snapshot == INVALID_HANDLE_VALUE) {
+		std::cerr << "Invalid CreateToolhelp32Snapshot" << std::endl;
+		return;
+	}
+	THREADENTRY32 te;
+	te.dwSize = sizeof(THREADENTRY32);
+	if (Thread32First(snapshot, &te)) {
+		do {
+			EnumThreadWindows(te.th32ThreadID, EnumWindowProc, NULL);
+		} while (Thread32Next(snapshot, &te));
+	}
+	else {
+		std::cerr << "Cannot get first thread from snapshot" << std::endl;
+	}
+	// Enumerate raw input devices and re-register if needed
+	PRAWINPUTDEVICE devices;
+	UINT numDevices;
+	GetRegisteredRawInputDevices(NULL, &numDevices, sizeof(RAWINPUTDEVICE));
+	devices = (PRAWINPUTDEVICE)malloc(numDevices * sizeof(RAWINPUTDEVICE));
+	GetRegisteredRawInputDevices(devices, &numDevices, sizeof(RAWINPUTDEVICE));
+	AT_RegisterRawInputDevicesHook(devices, numDevices, sizeof(RAWINPUTDEVICE), true);
 }
 
 // Called at shutdown to unpatch all the WinAPI functions
@@ -1184,6 +1289,7 @@ AT_Uninitialize()
 	if (DetourTransactionCommit() != NO_ERROR) {
 		throw std::runtime_error("Failed to commit Detours transaction");
 	}
+	FreeConsole();
 }
 
 BOOL APIENTRY
diff --git "a/atloader/atloader.cpp" "b/atloader/atloader.cpp"
index 0fc3c84..448539c 100644
--- "a/atloader/atloader.cpp"
+++ "b/atloader/atloader.cpp"
@@ -2,6 +2,8 @@
 #include <iostream>
 #include <string>
 #include <Windows.h>
+#include <TlHelp32.h>
+#include <conio.h>
 #include "Detours.h"
 
 // C++ exception wrapping the Win32 GetLastError() status
@@ -40,28 +42,49 @@ main(int argc, char *argv[])
 	std::string exePath(path);
 	std::string exeDir = exePath.substr(0, exePath.find_last_of('\\'));
 	std::string atDllPath = exeDir + "\\atdll.dll";
-    std::string targetPath = exeDir + "\\attest.exe";
-    if (argc >= 2) {
-        targetPath = argv[1];
+	std::wstring targetProcess = L"osu!.exe";
+
+	PROCESSENTRY32 pe32;
+	HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
+	if (snapshot == INVALID_HANDLE_VALUE) {
+		std::cerr << "Invalid handle for CreateToolhelp32Snapshot" << std::endl;
+		return 1;
 	}
 
-    STARTUPINFOA startupInfo = { 0 };
-    startupInfo.cb = sizeof(startupInfo);
-    PROCESS_INFORMATION procInfo;
-    if (!DetourCreateProcessWithDllExA(
-        targetPath.c_str(),
-        nullptr,
-        nullptr,
-        nullptr,
-        TRUE,
-        0,
-        nullptr,
-        nullptr,
-        &startupInfo,
-        &procInfo,
-        atDllPath.c_str(),
-        nullptr))
-    {
-        std::cerr << "Failed to create process: 0x" << std::hex << GetLastError() << std::endl;
+	pe32.dwSize = sizeof(PROCESSENTRY32);
+
+	if (!Process32First(snapshot, &pe32)) {
+		CloseHandle(snapshot);
+		std::cerr << "Process32First failed" << std::endl;
+		return 1;
+	}
+
+	while (Process32Next(snapshot, &pe32)) {
+		std::wcout << "Name = " << pe32.szExeFile << ", pid = " << pe32.th32ProcessID << std::endl;;
+		if (targetProcess.compare(pe32.szExeFile) == 0) {
+			std::cout << "May be the one we want" << std::endl;
+			std::cout << "Injecting..." << std::endl;
+			/* try to open the process, this is very finicky */
+			HANDLE process = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pe32.th32ProcessID);
+			if (process) {
+				// try to inject ourself using CreateRemoteThread
+				LPVOID LoadLibraryAddr = (LPVOID)GetProcAddress(GetModuleHandle(L"kernel32.dll"), "LoadLibraryA");
+				LPVOID LLParam = (LPVOID)VirtualAllocEx(process, NULL, atDllPath.length(), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
+				WriteProcessMemory(process, LLParam, atDllPath.c_str(), atDllPath.length(), NULL);
+				HANDLE thread = CreateRemoteThread(process, NULL, NULL, (LPTHREAD_START_ROUTINE)LoadLibraryAddr, LLParam, NULL, NULL);
+				if (thread != NULL) {
+					std::cout << "Success!" << std::endl;
+				}
+				else {
+					std::cout << "Failure, code = " << GetLastError() << std::endl;
+				}
+			}
+			else {
+				std::cout << "Failed to open process, code = " << GetLastError() << std::endl;
+			}
+			break;
+		}
 	}
+	_getch();
+	return 0;
 }