From 9171a9216763b2e0fac79068e7b3b8d7ad97c118 Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Mon, 12 Aug 2019 09:55:48 +0200
Subject: installer: kill leftover processes forcibly

---
 installer/build.bat       |  4 +-
 installer/customactions.c | 97 +++++++++++++++++++++++++++++++++++++++++++++++
 installer/wireguard.wxs   |  8 ++++
 3 files changed, 107 insertions(+), 2 deletions(-)

diff --git a/installer/build.bat b/installer/build.bat
index 1f9fa729..b4730e1b 100644
--- a/installer/build.bat
+++ b/installer/build.bat
@@ -57,9 +57,9 @@ if exist .deps\prepared goto :build
 :msi
 	set PATH=%BUILDDIR%..\.deps\%~2-w64-mingw32-native\bin;%PATH%
 	set CC=%~2-w64-mingw32-gcc
-	set CFLAGS=-O3 -Wall -std=gnu11 -DWINVER=0x0601 -municode -DUNICODE -D_UNICODE -DNDEBUG
+	set CFLAGS=-O3 -Wall -std=gnu11 -DWINVER=0x0601 -D_WIN32_WINNT=0x0601 -municode -DUNICODE -D_UNICODE -DNDEBUG
 	set LDFLAGS=-shared -s -Wl,--kill-at -Wl,--major-os-version=6 -Wl,--minor-os-version=1 -Wl,--major-subsystem-version=6 -Wl,--minor-subsystem-version=1
-	set LDLIBS=-lmsi -lole32 -lshlwapi
+	set LDLIBS=-lmsi -lole32 -lshlwapi -lshell32 -luuid
 	if not exist "%~1" mkdir "%~1"
 	echo [+] Compiling %1
 	%CC% %CFLAGS% %LDFLAGS% -o "%~1\customactions.dll" customactions.c %LDLIBS% || exit /b 1
diff --git a/installer/customactions.c b/installer/customactions.c
index ae9e7d7f..496d80b1 100644
--- a/installer/customactions.c
+++ b/installer/customactions.c
@@ -4,10 +4,13 @@
  */
 
 #include <windows.h>
+#include <ntstatus.h>
+#include <tlhelp32.h>
 #include <msi.h>
 #include <msidefs.h>
 #include <msiquery.h>
 #include <shlwapi.h>
+#include <shlobj.h>
 #include <stdbool.h>
 #include <tchar.h>
 
@@ -274,3 +277,97 @@ out:
 		CoUninitialize();
 	return ERROR_SUCCESS;
 }
+
+struct file_id { DWORD volume, index_high, index_low; };
+
+static bool calculate_file_id(const TCHAR *path, struct file_id *id)
+{
+	BY_HANDLE_FILE_INFORMATION file_info = { 0 };
+	HANDLE file;
+	bool ret;
+
+	file = CreateFile(path, 0, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
+	if (file == INVALID_HANDLE_VALUE)
+		return false;
+	ret = GetFileInformationByHandle(file, &file_info);
+	CloseHandle(file);
+	if (!ret)
+		return false;
+	id->volume = file_info.dwVolumeSerialNumber;
+	id->index_high = file_info.nFileIndexHigh;
+	id->index_low = file_info.nFileIndexLow;
+	return true;
+}
+
+static bool calculate_known_file_id(const KNOWNFOLDERID *known_folder, const TCHAR *file, struct file_id *id)
+{
+	TCHAR *folder_path, process_path[MAX_PATH + 1];
+	bool ret = false;
+
+	if (SHGetKnownFolderPath(known_folder, KF_FLAG_DEFAULT, NULL, &folder_path) == S_OK) {
+		if (PathCombine(process_path, folder_path, file)) {
+			if (calculate_file_id(process_path, id))
+				ret = true;
+		}
+		CoTaskMemFree(folder_path);
+	}
+	return ret;
+}
+
+__declspec(dllexport) UINT __stdcall KillWireGuardProcesses(MSIHANDLE installer)
+{
+	HANDLE snapshot, process;
+	PROCESSENTRY32 entry = { .dwSize = sizeof(PROCESSENTRY32) };
+	TCHAR process_path[MAX_PATH + 1];
+	DWORD process_path_len;
+	struct file_id file_ids[3], file_id;
+	size_t file_ids_len = 0;
+	bool is_com_initialized = SUCCEEDED(CoInitialize(NULL));
+
+	if (calculate_known_file_id(&FOLDERID_System, TEXT("wg.exe"), &file_ids[file_ids_len]))
+		++file_ids_len;
+	if (calculate_known_file_id(&FOLDERID_SystemX86, TEXT("wg.exe"), &file_ids[file_ids_len]))
+		++file_ids_len;
+	if (calculate_known_file_id(&FOLDERID_ProgramFiles, TEXT("WireGuard\\wireguard.exe"), &file_ids[file_ids_len]))
+		++file_ids_len;
+	if (!file_ids_len)
+		goto out;
+
+	snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
+	if (snapshot == INVALID_HANDLE_VALUE)
+		goto out;
+
+	for (bool ret = Process32First(snapshot, &entry); ret; ret = Process32Next(snapshot, &entry)) {
+		if (_tcsicmp(entry.szExeFile, TEXT("wireguard.exe")) && _tcsicmp(entry.szExeFile, TEXT("wg.exe")))
+			continue;
+		process = OpenProcess(PROCESS_TERMINATE | PROCESS_QUERY_LIMITED_INFORMATION, false, entry.th32ProcessID);
+		if (!process)
+			continue;
+		process_path_len = _countof(process_path);
+		if (!QueryFullProcessImageName(process, 0, process_path, &process_path_len))
+			goto next;
+		if (!calculate_file_id(process_path, &file_id))
+			goto next;
+		ret = false;
+		for (size_t i = 0; i < file_ids_len; ++i) {
+			if (!memcmp(&file_id, &file_ids[i], sizeof(file_id))) {
+				ret = true;
+				break;
+			}
+		}
+		if (!ret)
+			goto next;
+		if (TerminateProcess(process, STATUS_DLL_INIT_FAILED_LOGOFF)) {
+			WaitForSingleObject(process, INFINITE);
+			log_messagef(installer, LOG_LEVEL_INFO, TEXT("Killed \"%1\" (pid %2!d!)"), process_path, entry.th32ProcessID);
+		}
+	next:
+		CloseHandle(process);
+	}
+	CloseHandle(snapshot);
+
+out:
+	if (is_com_initialized)
+		CoUninitialize();
+	return ERROR_SUCCESS;
+}
diff --git a/installer/wireguard.wxs b/installer/wireguard.wxs
index 7d7a7b5a..1918f849 100644
--- a/installer/wireguard.wxs
+++ b/installer/wireguard.wxs
@@ -120,6 +120,14 @@
 			<Custom Action="RemoveConfigFolder" After="DeleteServices">(NOT UPGRADINGPRODUCTCODE) AND (REMOVE="ALL")</Custom>
 		</InstallExecuteSequence>
 
+		<!--
+			Kill lingering processes
+		-->
+		<CustomAction Id="KillWireGuardProcesses" BinaryKey="customactions.dll" DllEntry="KillWireGuardProcesses" Execute="deferred" Impersonate="no" />
+		<InstallExecuteSequence>
+			<Custom Action="KillWireGuardProcesses" After="StopServices" />
+		</InstallExecuteSequence>
+
 		<!--
 			Launch wireguard.exe after setup complete
 		-->
-- 
cgit v1.2.3-59-g8ed1b