Overlay: Move DllMain logic into functions

1 files changed, 202 insertions, 188 deletions

diff --git a/overlay/lib.cpp b/overlay/lib.cpp
index f43d598a3..19efa359f 100644
--- a/overlay/lib.cpp
+++ b/overlay/lib.cpp
@@ -505,223 +505,237 @@ extern "C" __declspec(dllexport) unsigned int __cdecl GetOverlayMagicVersion() {
 	return OVERLAY_MAGIC_NUMBER;
 }
 
-extern "C" BOOL WINAPI DllMain(HINSTANCE, DWORD fdwReason, LPVOID) {
+void dllmainProcAttach() {
 	char procname[1024+64];
 	GetModuleFileNameA(NULL, procname, 1024);
 
-	switch (fdwReason) {
-		case DLL_PROCESS_ATTACH: {
-				Mutex::init();
-				char *p = strrchr(procname, '\\');
-				if (!p) {
-					// No blacklisting if the file has no path
-				} else if (GetProcAddress(NULL, "mumbleSelfDetection") != NULL) {
-					ods("Lib: Attached to self (own process). Blacklisted - no overlay injection.");
-					bBlackListed = TRUE;
-					bMumble = TRUE;
-				} else {
-					DWORD buffsize = MAX_PATH * 20; // Initial buffer size for registry operation
-
-					bool usewhitelist = false;
-					HKEY key = NULL;
-
-					char *buffer = new char[buffsize];
-
-					// check if we're using a whitelist or a blacklist
-					DWORD tmpsize = buffsize - 1;
-					bool success = (RegOpenKeyExA(HKEY_CURRENT_USER, "Software\\Mumble\\Mumble\\overlay", NULL, KEY_READ, &key) == ERROR_SUCCESS) &&
-					          (RegQueryValueExA(key, "usewhitelist", NULL, NULL, (LPBYTE)buffer, &tmpsize) == ERROR_SUCCESS);
-
-					if (success) {
-						buffer[tmpsize] = '\0';
-						usewhitelist = (_stricmp(buffer, "true") == 0);
-						// reset tmpsize to the buffers size (minus 1 char for str-termination), as it was changed by RegQuery
-						tmpsize = buffsize - 1;
-
-						// read the whitelist or blacklist (depending on which one we use)
-						DWORD ret;
-						while ((ret = RegQueryValueExA(key, usewhitelist ? "whitelist" : "blacklist", NULL, NULL, (LPBYTE)buffer, &tmpsize)) == ERROR_MORE_DATA) {
-							// Increase the buffsize according to the required size RegQuery wrote into tmpsize, so we can read the whole value
-							delete []buffer;
-							buffsize = tmpsize + 1;
-							buffer = new char[buffsize];
-						}
+	Mutex::init();
 
-						success = (ret == ERROR_SUCCESS);
-					}
+	char *p = strrchr(procname, '\\');
+	if (!p) {
+		// No blacklisting if the file has no path
+	} else if (GetProcAddress(NULL, "mumbleSelfDetection") != NULL) {
+		ods("Lib: Attached to self (own process). Blacklisted - no overlay injection.");
+		bBlackListed = TRUE;
+		bMumble = TRUE;
+	} else {
+		DWORD buffsize = MAX_PATH * 20; // Initial buffer size for registry operation
+
+		bool usewhitelist = false;
+		HKEY key = NULL;
+
+		char *buffer = new char[buffsize];
+
+		// check if we're using a whitelist or a blacklist
+		DWORD tmpsize = buffsize - 1;
+		bool success = (RegOpenKeyExA(HKEY_CURRENT_USER, "Software\\Mumble\\Mumble\\overlay", NULL, KEY_READ, &key) == ERROR_SUCCESS) &&
+		          (RegQueryValueExA(key, "usewhitelist", NULL, NULL, (LPBYTE)buffer, &tmpsize) == ERROR_SUCCESS);
+
+		if (success) {
+			buffer[tmpsize] = '\0';
+			usewhitelist = (_stricmp(buffer, "true") == 0);
+			// reset tmpsize to the buffers size (minus 1 char for str-termination), as it was changed by RegQuery
+			tmpsize = buffsize - 1;
+
+			// read the whitelist or blacklist (depending on which one we use)
+			DWORD ret;
+			while ((ret = RegQueryValueExA(key, usewhitelist ? "whitelist" : "blacklist", NULL, NULL, (LPBYTE)buffer, &tmpsize)) == ERROR_MORE_DATA) {
+				// Increase the buffsize according to the required size RegQuery wrote into tmpsize, so we can read the whole value
+				delete []buffer;
+				buffsize = tmpsize + 1;
+				buffer = new char[buffsize];
+			}
 
-					if (key)
-						RegCloseKey(key);
-
-					if (success) {
-						buffer[tmpsize] = '\0';
-						unsigned int pos = 0;
-
-						if (usewhitelist) {
-							// check if process is whitelisted
-							bool onwhitelist = false;
-							while (pos < buffsize && buffer[pos] != 0) {
-								if (_stricmp(procname, buffer + pos) == 0 || _stricmp(p+1, buffer + pos) == 0) {
-									ods("Lib: Overlay enabled for whitelisted '%s'", buffer + pos);
-									onwhitelist = true;
-									break;
-								}
-								pos += strlen(buffer + pos) + 1;
-							}
-
-							if (!onwhitelist) {
-								ods("Lib: No whitelist entry found for '%s', auto-blacklisted", procname);
-								bBlackListed = TRUE;
-								break;
-							}
-						} else {
-							// check if process is blacklisted
-							while (pos < buffsize && buffer[pos] != 0) {
-								if (_stricmp(procname, buffer + pos) == 0 || _stricmp(p+1, buffer + pos) == 0) {
-									ods("Lib: Overlay blacklist entry found for '%s'", buffer + pos);
-									bBlackListed = TRUE;
-									break;
-								}
-								pos += strlen(buffer + pos) + 1;
-							}
-						}
-					} else {
-						// If there is no list in the registry, fallback to using the default blacklist
-						ods("Lib: Overlay fallback to default blacklist");
-						int i = 0;
-						while (overlayBlacklist[i]) {
-							if (_stricmp(procname, overlayBlacklist[i]) == 0 || _stricmp(p+1, overlayBlacklist[i])==0) {
-								ods("Lib: Overlay default blacklist entry found for '%s'", overlayBlacklist[i]);
-								bBlackListed = TRUE;
-								break;
-							}
-							i++;
-						}
-					}
+			success = (ret == ERROR_SUCCESS);
+		}
 
-					// Make sure to always free/destroy buffer & heap
-					delete []buffer;
+		if (key)
+			RegCloseKey(key);
 
-					// if the processname is already found to be blacklisted, we can stop here
-					if (bBlackListed)
-						return TRUE;
+		if (success) {
+			buffer[tmpsize] = '\0';
+			unsigned int pos = 0;
 
-					// check if there is a "nooverlay" file in the executables folder, which would disable/blacklist the overlay
-					char fname[sizeof(procname)];
-					p = fname + (p - procname);
-					strncpy_s(fname, sizeof(fname), procname, p - procname + 1);
+			if (usewhitelist) {
+				// check if process is whitelisted
+				bool onwhitelist = false;
+				while (pos < buffsize && buffer[pos] != 0) {
+					if (_stricmp(procname, buffer + pos) == 0 || _stricmp(p+1, buffer + pos) == 0) {
+						ods("Lib: Overlay enabled for whitelisted '%s'", buffer + pos);
+						onwhitelist = true;
+						break;
+					}
+					pos += strlen(buffer + pos) + 1;
+				}
 
-					strcpy_s(p+1, 64, "nooverlay");
-					HANDLE h = CreateFile(fname, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
-					if (h != INVALID_HANDLE_VALUE) {
-						CloseHandle(h);
-						ods("Lib: Overlay disable %s found", fname);
+				if (!onwhitelist) {
+					ods("Lib: No whitelist entry found for '%s', auto-blacklisted", procname);
+					bBlackListed = TRUE;
+					break;
+				}
+			} else {
+				// check if process is blacklisted
+				while (pos < buffsize && buffer[pos] != 0) {
+					if (_stricmp(procname, buffer + pos) == 0 || _stricmp(p+1, buffer + pos) == 0) {
+						ods("Lib: Overlay blacklist entry found for '%s'", buffer + pos);
 						bBlackListed = TRUE;
-						return TRUE;
+						break;
 					}
+					pos += strlen(buffer + pos) + 1;
+				}
+			}
+		} else {
+			// If there is no list in the registry, fallback to using the default blacklist
+			ods("Lib: Overlay fallback to default blacklist");
+			int i = 0;
+			while (overlayBlacklist[i]) {
+				if (_stricmp(procname, overlayBlacklist[i]) == 0 || _stricmp(p+1, overlayBlacklist[i])==0) {
+					ods("Lib: Overlay default blacklist entry found for '%s'", overlayBlacklist[i]);
+					bBlackListed = TRUE;
+					break;
+				}
+				i++;
+			}
+		}
 
-					// check for "debugoverlay" file, which would enable overlay debugging
-					strcpy_s(p+1, 64, "debugoverlay");
-					h = CreateFile(fname, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
-					if (h != INVALID_HANDLE_VALUE) {
-						CloseHandle(h);
-						ods("Lib: Overlay debug %s found", fname);
-						bDebug = TRUE;
-					}
+		// Make sure to always free/destroy buffer & heap
+		delete []buffer;
 
-					// check for blacklisting for loading WPF library
-					checkForWPF();
+		// if the processname is already found to be blacklisted, we can stop here
+		if (bBlackListed)
+			return;
 
-					if (bBlackListed)
-						return TRUE;
-				}
-				ods("Lib: ProcAttach: %s", procname);
+		// check if there is a "nooverlay" file in the executables folder, which would disable/blacklist the overlay
+		char fname[sizeof(procname)];
+		p = fname + (p - procname);
+		strncpy_s(fname, sizeof(fname), procname, p - procname + 1);
+
+		strcpy_s(p+1, 64, "nooverlay");
+		HANDLE h = CreateFile(fname, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
+		if (h != INVALID_HANDLE_VALUE) {
+			CloseHandle(h);
+			ods("Lib: Overlay disable %s found", fname);
+			bBlackListed = TRUE;
+			return;
+		}
 
-				OSVERSIONINFOEX ovi;
-				memset(&ovi, 0, sizeof(ovi));
-				ovi.dwOSVersionInfoSize = sizeof(ovi);
-				GetVersionEx(reinterpret_cast<OSVERSIONINFO *>(&ovi));
-				bIsWin8 = (ovi.dwMajorVersion >= 7) || ((ovi.dwMajorVersion == 6) &&(ovi.dwBuildNumber >= 9200));
+		// check for "debugoverlay" file, which would enable overlay debugging
+		strcpy_s(p+1, 64, "debugoverlay");
+		h = CreateFile(fname, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
+		if (h != INVALID_HANDLE_VALUE) {
+			CloseHandle(h);
+			ods("Lib: Overlay debug %s found", fname);
+			bDebug = TRUE;
+		}
 
-				ods("Lib: bIsWin8: %i", bIsWin8);
+		// check for blacklisting for loading WPF library
+		checkForWPF();
 
-				hHookMutex = CreateMutex(NULL, false, "MumbleHookMutex");
-				if (hHookMutex == NULL) {
-					ods("Lib: CreateMutex failed");
-					return TRUE;
-				}
+		if (bBlackListed)
+			return;
+	}
+	ods("Lib: ProcAttach: %s", procname);
 
-				DWORD dwSharedSize = sizeof(SharedData) + sizeof(Direct3D9Data) + sizeof(DXGIData);
+	OSVERSIONINFOEX ovi;
+	memset(&ovi, 0, sizeof(ovi));
+	ovi.dwOSVersionInfoSize = sizeof(ovi);
+	GetVersionEx(reinterpret_cast<OSVERSIONINFO *>(&ovi));
+	bIsWin8 = (ovi.dwMajorVersion >= 7) || ((ovi.dwMajorVersion == 6) &&(ovi.dwBuildNumber >= 9200));
 
-				hMapObject = CreateFileMapping(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, dwSharedSize, "MumbleOverlayPrivate");
-				if (hMapObject == NULL) {
-					ods("Lib: CreateFileMapping failed");
-					return TRUE;
-				}
+	ods("Lib: bIsWin8: %i", bIsWin8);
 
-				bool bInit = (GetLastError() != ERROR_ALREADY_EXISTS);
+	hHookMutex = CreateMutex(NULL, false, "MumbleHookMutex");
+	if (hHookMutex == NULL) {
+		ods("Lib: CreateMutex failed");
+		return;
+	}
 
-				sd = static_cast<SharedData *>(MapViewOfFile(hMapObject, FILE_MAP_ALL_ACCESS, 0, 0, dwSharedSize));
+	DWORD dwSharedSize = sizeof(SharedData) + sizeof(Direct3D9Data) + sizeof(DXGIData);
 
-				if (sd == NULL) {
-					ods("Lib: MapViewOfFile Failed");
-					return TRUE;
-				}
+	hMapObject = CreateFileMapping(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, dwSharedSize, "MumbleOverlayPrivate");
+	if (hMapObject == NULL) {
+		ods("Lib: CreateFileMapping failed");
+		return;
+	}
 
-				if (bInit)
-					memset(sd, 0, dwSharedSize);
+	bool bInit = (GetLastError() != ERROR_ALREADY_EXISTS);
 
-				unsigned char *raw = (unsigned char *) sd;
-				d3dd = reinterpret_cast<Direct3D9Data *>(raw + sizeof(SharedData));
-				dxgi = reinterpret_cast<DXGIData *>(raw + sizeof(SharedData) + sizeof(Direct3D9Data));
+	sd = static_cast<SharedData *>(MapViewOfFile(hMapObject, FILE_MAP_ALL_ACCESS, 0, 0, dwSharedSize));
 
+	if (sd == NULL) {
+		ods("Lib: MapViewOfFile Failed");
+		return;
+	}
 
-				if (! bMumble) {
-					// Hook our own LoadLibrary functions so we notice when a new library (like the d3d ones) is loaded.
-					hhLoad.setup(reinterpret_cast<voidFunc>(LoadLibraryA), reinterpret_cast<voidFunc>(MyLoadLibrary));
-					hhLoadW.setup(reinterpret_cast<voidFunc>(LoadLibraryW), reinterpret_cast<voidFunc>(MyLoadLibraryW));
-					hhFree.setup(reinterpret_cast<voidFunc>(FreeLibrary), reinterpret_cast<voidFunc>(MyFreeLibrary));
+	if (bInit)
+		memset(sd, 0, dwSharedSize);
 
-					checkD3D9Hook(true);
-					checkDXGIHook(true);
-					checkOpenGLHook();
-					ods("Lib: Injected into %s", procname);
-				}
-			}
+	unsigned char *raw = (unsigned char *) sd;
+	d3dd = reinterpret_cast<Direct3D9Data *>(raw + sizeof(SharedData));
+	dxgi = reinterpret_cast<DXGIData *>(raw + sizeof(SharedData) + sizeof(Direct3D9Data));
+
+
+	if (! bMumble) {
+		// Hook our own LoadLibrary functions so we notice when a new library (like the d3d ones) is loaded.
+		hhLoad.setup(reinterpret_cast<voidFunc>(LoadLibraryA), reinterpret_cast<voidFunc>(MyLoadLibrary));
+		hhLoadW.setup(reinterpret_cast<voidFunc>(LoadLibraryW), reinterpret_cast<voidFunc>(MyLoadLibraryW));
+		hhFree.setup(reinterpret_cast<voidFunc>(FreeLibrary), reinterpret_cast<voidFunc>(MyFreeLibrary));
+
+		checkD3D9Hook(true);
+		checkDXGIHook(true);
+		checkOpenGLHook();
+		ods("Lib: Injected into %s", procname);
+	}
+}
+
+void dllmainProcDetach() {
+
+	hhLoad.restore(true);
+	hhLoad.reset();
+	hhLoadW.restore(true);
+	hhLoadW.reset();
+	hhFree.restore(true);
+	hhFree.reset();
+
+	if (sd)
+		UnmapViewOfFile(sd);
+	if (hMapObject)
+		CloseHandle(hMapObject);
+	if (hHookMutex)
+		CloseHandle(hHookMutex);
+}
+
+void dllmainThreadAttach() {
+	static bool bTriedHook = false;
+	if (!bBlackListed && sd && ! bTriedHook) {
+		bTriedHook = true;
+		checkForWPF();
+
+		if (!bBlackListed) {
+			checkD3D9Hook();
+			checkDXGIHook();
+			checkOpenGLHook();
+			ods("Lib: Injected to thread of %s", procname);
+		}
+	}
+}
+
+extern "C" BOOL WINAPI DllMain(HINSTANCE, DWORD fdwReason, LPVOID) {
+	char procname[1024+64];
+	GetModuleFileNameA(NULL, procname, 1024);
+
+	switch (fdwReason) {
+		case DLL_PROCESS_ATTACH:
+			ods("Lib: ProcAttach: %s", procname);
+			dllmainProcAttach();
 			break;
-		case DLL_PROCESS_DETACH: {
-				ods("Lib: ProcDetach: %s", procname);
-
-				hhLoad.restore(true);
-				hhLoad.reset();
-				hhLoadW.restore(true);
-				hhLoadW.reset();
-				hhFree.restore(true);
-				hhFree.reset();
-
-				if (sd)
-					UnmapViewOfFile(sd);
-				if (hMapObject)
-					CloseHandle(hMapObject);
-				if (hHookMutex)
-					CloseHandle(hHookMutex);
-			}
+		case DLL_PROCESS_DETACH:
+			ods("Lib: ProcDetach: %s", procname);
+			dllmainProcDetach();
 			break;
-		case DLL_THREAD_ATTACH: {
-				ods("Lib: ThreadAttach: %s", procname);
-				static bool bTriedHook = false;
-				if (!bBlackListed && sd && ! bTriedHook) {
-					bTriedHook = true;
-					checkForWPF();
-
-					if (!bBlackListed) {
-						checkD3D9Hook();
-						checkDXGIHook();
-						checkOpenGLHook();
-						ods("Lib: Injected to thread of %s", procname);
-					}
-				}
-			}
+		case DLL_THREAD_ATTACH:
+			ods("Lib: ThreadAttach: %s", procname);
+			dllmainThreadAttach();
 			break;
 		default:
 			break;