try to catch relaunches through the steam client

Author: devinacker

DebugLoop.cpp

@@ -52,13 +52,18 @@ static void HandleNewProcess(DWORD pid, LPCWSTR filePath)
 
 	PID_DEBUG("starting %ws\n", pid, filePath);
 
+	GetCurrentDirectoryW(MAX_PATH, g_pathw);
+	std::wstring currentDir(g_pathw);
+
 	GetModuleFileNameW(NULL, g_pathw, MAX_PATH);
 	wcsrchr(g_pathw, L'\\')[1] = 0;
 	SetCurrentDirectoryW(g_pathw);
 
 	WIN32_FIND_DATAW findData;
 	HANDLE hFile = FindFirstFileW(L".\\UnrealKey64.dll", &findData);
 
+	SetCurrentDirectoryW(currentDir.c_str());
+
 	if (hFile != INVALID_HANDLE_VALUE)
 	{
 		// got DLL path, load it in the remote process
@@ -137,6 +142,21 @@ static void HandleIndexDataMessage(const PipeMessage& message)
 	}
 }
 
+// ----------------------------------------------------------------------------
+static bool HandleSteamAppIDMessage(const PipeMessage& message)
+{
+	snprintf(g_path, sizeof(g_path), "%u", message.msgUInt);
+
+	if (SetEnvironmentVariableA("SteamAppID", g_path))
+	{
+		PID_DEBUG("Set SteamAppID successfully, app will restart...\n", message.pid);
+		return true;
+	}
+
+	PID_DEBUG("Couldn't set SteamAppID (error 0x%x)\n", message.pid, GetLastError());
+	return false;
+}
+
 // ----------------------------------------------------------------------------
 static bool BreakpointEvent(const DEBUG_EVENT& event, bool &detach)
 {
@@ -255,32 +275,44 @@ static BOOL WINAPI CtrlHandler(DWORD type)
 }
 
 // ----------------------------------------------------------------------------
-int DebugLoop(LPCWSTR appPath)
+static bool StartProcess(LPCWSTR appPath, HANDLE hJob, PROCESS_INFORMATION& pi)
 {
-	JOBOBJECT_EXTENDED_LIMIT_INFORMATION jobInfo = { 0 };
-	HANDLE hJob = CreateJobObjectW(NULL, NULL);
-	jobInfo.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
-	SetInformationJobObject(hJob, JobObjectExtendedLimitInformation, &jobInfo, sizeof(jobInfo));
-
-	PROCESS_INFORMATION pi = { 0 };
 	STARTUPINFO si = { 0 };
 	si.cb = sizeof(si);
 
 	if (!CreateProcessW(appPath, NULL, NULL, NULL, FALSE,
 		CREATE_SUSPENDED | CREATE_BREAKAWAY_FROM_JOB, NULL, NULL, &si, &pi))
 	{
 		printf("couldn't launch %ws (error 0x%x)\n", appPath, GetLastError());
-		return 1;
+		return false;
 	}
 
-	printf("Starting the game now.\nClose the game or press Ctrl-C to stop.\n\n");
-
 	HandleNewProcess(pi.dwProcessId, appPath);
 	AssignProcessToJobObject(hJob, pi.hProcess);
 	ResumeThread(pi.hThread);
 	CloseHandle(pi.hThread);
 
+	return true;
+}
+
+// ----------------------------------------------------------------------------
+int DebugLoop(LPCWSTR appPath)
+{
+	JOBOBJECT_EXTENDED_LIMIT_INFORMATION jobInfo = { 0 };
+	HANDLE hJob = CreateJobObjectW(NULL, NULL);
+	jobInfo.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+	SetInformationJobObject(hJob, JobObjectExtendedLimitInformation, &jobInfo, sizeof(jobInfo));
+
+	PROCESS_INFORMATION pi = { 0 };
+
+	printf("Starting the game now.\nClose the game or press Ctrl-C to stop.\n\n");
+
+	if (!StartProcess(appPath, hJob, pi))
+		return 1;
+
 	g_running = true;
+	int rc = 0;
+	bool restart = false;
 	DEBUG_EVENT event;
 	PipeMessage message;
 	DWORD dwTemp;
@@ -294,6 +326,17 @@ int DebugLoop(LPCWSTR appPath)
 		{
 			PID_DEBUG("process exited with code 0x%x\n", pi.dwProcessId, dwTemp);
 
+			// if we want to try to restart with a Steam App ID set (or something), do that now
+			if (restart)
+			{
+				restart = false;
+
+				if (StartProcess(appPath, hJob, pi))
+					continue;
+				else
+					rc = 1;
+			}
+
 			g_running = false;
 			break;
 		}
@@ -315,6 +358,10 @@ int DebugLoop(LPCWSTR appPath)
 				case IndexDataMessage:
 					HandleIndexDataMessage(message);
 					break;
+
+				case SteamAppIDMessage:
+					restart = HandleSteamAppIDMessage(message);
+					break;
 				}
 
 				HANDLE hThread = OpenThread(THREAD_SUSPEND_RESUME, FALSE, message.tid);
@@ -405,5 +452,5 @@ int DebugLoop(LPCWSTR appPath)
 		printf("\n\n");
 	}
 
-	return 0;
+	return rc;
 }

common.h

@@ -9,6 +9,7 @@ enum PipeMessageType
 	StringMessage,
 	NewProcessMessage,
 	IndexDataMessage,
+	SteamAppIDMessage,
 };
 
 struct PipeProcessData
@@ -36,5 +37,6 @@ struct PipeMessage
 		WCHAR msgString[1024];
 		PipeProcessData msgProcess;
 		PipeIndexData msgIndex;
+		UINT msgUInt;
 	};
 };

dllmain.cpp

@@ -36,18 +36,20 @@ std::vector<PakIndex> g_pakIndexes;
 static HANDLE g_hPipe;
 
 // WinAPI function hooks
-auto orig_ReadFile = ReadFile;
 auto orig_CreateProcessW = CreateProcessW;
+auto orig_ShellExecuteA = ShellExecuteA;
+auto orig_ShellExecuteW = ShellExecuteW;
+auto orig_ReadFile = ReadFile;
 NTSTATUS (*orig_ZwSetInformationThread)(HANDLE, ULONG, PVOID, ULONG);
 
 // ----------------------------------------------------------------------------
-static void SendPipeMessage(PipeMessage *msg)
+static void SendPipeMessage(PipeMessage& msg)
 {
-	msg->pid = GetCurrentProcessId();
-	msg->tid = GetCurrentThreadId();
+	msg.pid = GetCurrentProcessId();
+	msg.tid = GetCurrentThreadId();
 
 	DWORD dummy;
-	if (WriteFile(g_hPipe, msg, sizeof(*msg), &dummy, NULL))
+	if (WriteFile(g_hPipe, &msg, sizeof(msg), &dummy, NULL))
 	{
 		// wait for launcher to handle message
 		SuspendThread(GetCurrentThread());
@@ -65,7 +67,7 @@ static void SendStringMessage(const WCHAR *str, ...)
 	_vsnwprintf_s(msg.msgString, sizeof(msg.msgString) / sizeof(WCHAR), str, va);
 	va_end(va);
 
-	SendPipeMessage(&msg);
+	SendPipeMessage(msg);
 }
 
 // ----------------------------------------------------------------------------
@@ -116,7 +118,7 @@ static BOOL WINAPI hook_CreateProcessW(LPCWSTR name, LPWSTR cmdLine, LPSECURITY_
 
 		msg.msgProcess.pid = pi->dwProcessId;
 		msg.msgProcess.flags = flags;
-		SendPipeMessage(&msg);
+		SendPipeMessage(msg);
 
 		// if the process wasn't already supposed to be suspended, unsuspend it now
 		if (!(flags & CREATE_SUSPENDED))
@@ -130,6 +132,49 @@ static BOOL WINAPI hook_CreateProcessW(LPCWSTR name, LPWSTR cmdLine, LPSECURITY_
 	return FALSE;
 }
 
+// ----------------------------------------------------------------------------
+static void RunWithSteamAppID(UINT appID)
+{
+	SendStringMessage(L"Tried to launch Steam w/ app ID %u", appID);
+	
+	PipeMessage msg;
+	msg.msgType = SteamAppIDMessage;
+	msg.msgUInt = appID;
+	SendPipeMessage(msg);
+}
+
+// ----------------------------------------------------------------------------
+static HINSTANCE WINAPI hook_ShellExecuteA(HWND hwnd, LPCSTR lpOperation, LPCSTR lpFile, LPCSTR lpParams, LPCSTR lpDir, INT nShowCmd)
+{
+	UINT appID;
+
+	if (lpOperation && lpFile
+		&& !strcmp(lpOperation, "open")
+		&& sscanf_s(lpFile, "steam://run/%u", &appID) == 1)
+	{
+		RunWithSteamAppID(appID);
+		return 0;
+	}
+
+	return orig_ShellExecuteA(hwnd, lpOperation, lpFile, lpParams, lpDir, nShowCmd);
+}
+
+// ----------------------------------------------------------------------------
+static HINSTANCE WINAPI hook_ShellExecuteW(HWND hwnd, LPCWSTR lpOperation, LPCWSTR lpFile, LPCWSTR lpParams, LPCWSTR lpDir, INT nShowCmd)
+{
+	UINT appID;
+
+	if (lpFile && lpParams
+		&& wcsstr(lpFile, L"steam.exe")
+		&& swscanf_s(lpParams, L"steam://run/%u", &appID) == 1)
+	{
+		RunWithSteamAppID(appID);
+		return 0;
+	}
+
+	return orig_ShellExecuteW(hwnd, lpOperation, lpFile, lpParams, lpDir, nShowCmd);
+}
+
 // ----------------------------------------------------------------------------
 static BOOL WINAPI hook_ReadFile(HANDLE hFile, LPVOID lpBuffer, DWORD nSize, LPDWORD lpSize, LPOVERLAPPED lpOverlapped)
 {
@@ -174,7 +219,7 @@ static BOOL WINAPI hook_ReadFile(HANDLE hFile, LPVOID lpBuffer, DWORD nSize, LPD
 			//	SendStringMessage(L"Reading encrypted pak index to 0x%p for %ws", lpBuffer, msg.msgIndex.filePath);
 				SendStringMessage(L"Reading encrypted pak index for %ws", wcsrchr(msg.msgIndex.filePath, L'\\') + 1);
 			//	SendStringMessage(L"Buffer size is 0x%x", nSize);
-				SendPipeMessage(&msg);
+				SendPipeMessage(msg);
 			}
 		}
 	}
@@ -214,6 +259,8 @@ static void Init()
 	if (!hookStatus)
 	{
 		MH_CreateHook(CreateProcessW, hook_CreateProcessW, (LPVOID*)&orig_CreateProcessW);
+		MH_CreateHook(ShellExecuteA, hook_ShellExecuteA, (LPVOID*)&orig_ShellExecuteA);
+		MH_CreateHook(ShellExecuteW, hook_ShellExecuteW, (LPVOID*)&orig_ShellExecuteW);
 		MH_CreateHook(ReadFile, hook_ReadFile, (LPVOID*)&orig_ReadFile);
 		MH_CreateHook(GetProcAddress(GetModuleHandleA("ntdll"), "ZwSetInformationThread"),
 			hook_ZwSetInformationThread, (LPVOID*)&orig_ZwSetInformationThread);
@@ -225,12 +272,12 @@ static void Init()
 		}
 		else
 		{
-			SendStringMessage(L"MinHook enable failed: %s", MH_StatusToString(hookStatus));
+			SendStringMessage(L"MinHook enable failed: %hs", MH_StatusToString(hookStatus));
 		}
 	}
 	else
 	{
-		SendStringMessage(L"MinHook init failed: %s", MH_StatusToString(hookStatus));
+		SendStringMessage(L"MinHook init failed: %hs", MH_StatusToString(hookStatus));
 	}
 }