thcrap_tsa: properly load game icon in RegisterClass

Most (maybe all) games don't specify their icon in RegisterClassA, which makes Windows assign a default icon to their window. We detour the RegisterClassA call and assign the correct icon in there.
thpatch · Nov 22, 2023 · 37aef68 · 37aef68
1 parent efcd838
commit 37aef68
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 6 deletions.
diff --git a/thcrap/src/shelllink.cpp b/thcrap/src/shelllink.cpp
@@ -123,7 +123,7 @@ typedef struct GRPICONDIRENTRY
 } GRPICONDIRENTRY;
 #pragma pack(pop)
 
-BOOL CALLBACK CopyIconGroupCallback(HMODULE hModule, LPCWSTR lpType, LPWSTR lpName, LONG_PTR lParam)
+static BOOL CALLBACK CopyIconGroupCallback(HMODULE hModule, LPCWSTR lpType, LPWSTR lpName, LONG_PTR lParam)
 {
 	LPWSTR *iconGroupId = (LPWSTR*)lParam;
 	if (IS_INTRESOURCE(lpName)) {
@@ -135,6 +135,13 @@ BOOL CALLBACK CopyIconGroupCallback(HMODULE hModule, LPCWSTR lpType, LPWSTR lpNa
 	return FALSE;
 }
 
+LPWSTR GetIconGroupResourceId(HMODULE hModule)
+{
+	LPWSTR iconGroupId = nullptr;
+	EnumResourceNamesW(hModule, RT_GROUP_ICON, CopyIconGroupCallback, (intptr_t)&iconGroupId);
+	return iconGroupId;
+}
+
 void CopyIconGroup(HANDLE hUpdate, const std::filesystem::path& icon_path)
 {
 	HMODULE hIconExe = LoadLibraryExW(icon_path.wstring().c_str(), nullptr, LOAD_LIBRARY_AS_DATAFILE);
@@ -143,14 +150,13 @@ void CopyIconGroup(HANDLE hUpdate, const std::filesystem::path& icon_path)
 	}
 	defer(FreeLibrary(hIconExe));
 
-	LPWSTR iconGroupId = nullptr;
-	EnumResourceNamesW(hIconExe, RT_GROUP_ICON, CopyIconGroupCallback, (intptr_t)&iconGroupId);
-	defer(if (!IS_INTRESOURCE(iconGroupId)) {
-		free(iconGroupId);
-	});
+	LPWSTR iconGroupId = GetIconGroupResourceId(hIconExe);
 	if (!iconGroupId) {
 		return;
 	}
+	defer(if (!IS_INTRESOURCE(iconGroupId)) {
+		free(iconGroupId);
+	});
 
 	auto [iconGroupData, iconGroupSize] = GetResource(hIconExe, iconGroupId, RT_GROUP_ICON);
 	if (!iconGroupData || iconGroupSize < sizeof(GRPICONDIR)) {

diff --git a/thcrap/src/shelllink.h b/thcrap/src/shelllink.h
@@ -33,3 +33,8 @@ HRESULT CreateLink(
 
 // Create shortcuts for the given games
 int CreateShortcuts(const char *run_cfg_fn, games_js_entry *games, enum ShortcutsDestination destination, enum ShortcutsType shortcut_type);
+
+// Get the resource ID of the first icon group for the given module.
+// If the return value is a resource name (IS_INTRESOURCE(iconGroupId) is false),
+// it must be freed with thcrap_free();
+LPWSTR GetIconGroupResourceId(HMODULE hModule);
diff --git a/thcrap/thcrap.def b/thcrap/thcrap.def
@@ -320,6 +320,7 @@ EXPORTS
 	; ---------
 	CreateLink
 	CreateShortcuts
+	GetIconGroupResourceId
 
 	; Stack
 	; -----

diff --git a/thcrap_tsa/src/win32_tsa.cpp b/thcrap_tsa/src/win32_tsa.cpp
@@ -17,8 +17,14 @@ typedef BOOL WINAPI GetWindowRect_type(
 	LPRECT lpRect
 );
 
+typedef ATOM WINAPI RegisterClassW_type(
+	const WNDCLASSW *lpWndClass
+);
+
 W32U8_DETOUR_CHAIN_DEF(CreateWindowEx);
+W32U8_DETOUR_CHAIN_DEF(RegisterClass);
 DETOUR_CHAIN_DEF(GetWindowRect);
+DETOUR_CHAIN_DEF(RegisterClassW);
 /// -------------
 
 /// Window position hacks
@@ -136,13 +142,44 @@ HWND WINAPI tsa_CreateWindowExA(
 	VLA_FREE(custom_title);
 	return ret;
 }
+
+ATOM WINAPI tsa_RegisterClassA(const WNDCLASSA *lpWndClass)
+{
+	if (lpWndClass->hIcon == NULL) {
+		LPWSTR iconGroupId = GetIconGroupResourceId(GetModuleHandle(NULL));
+		if (iconGroupId) {
+			const_cast<WNDCLASSA*>(lpWndClass)->hIcon = LoadIcon(GetModuleHandle(NULL), iconGroupId);
+			if (!IS_INTRESOURCE(iconGroupId)) {
+				thcrap_free(iconGroupId);
+			}
+		}
+	}
+	return chain_RegisterClassU(lpWndClass);
+}
+
+ATOM WINAPI tsa_RegisterClassW(const WNDCLASSW *lpWndClass)
+{
+	if (lpWndClass->hIcon == NULL) {
+		LPWSTR iconGroupId = GetIconGroupResourceId(GetModuleHandle(NULL));
+		if (iconGroupId) {
+			const_cast<WNDCLASSW*>(lpWndClass)->hIcon = LoadIcon(GetModuleHandle(NULL), iconGroupId);
+			if (!IS_INTRESOURCE(iconGroupId)) {
+				thcrap_free(iconGroupId);
+			}
+		}
+	}
+	return chain_RegisterClassW(lpWndClass);
+}
+
 /// ---------------------
 
 void tsa_mod_detour(void)
 {
 	detour_chain("user32.dll", 1,
 		"CreateWindowExA", tsa_CreateWindowExA, &chain_CreateWindowExU,
 		"GetWindowRect", tsa_GetWindowRect, &chain_GetWindowRect,
+		"RegisterClassA", tsa_RegisterClassA, &chain_RegisterClassU,
+		"RegisterClassW", tsa_RegisterClassW, &chain_RegisterClassU,
 		NULL
 	);
 }