Kill processes belonging to us on install and uninstall

src/Setup/winmain.cpp

@@ -10,13 +10,19 @@
 
 CAppModule _Module;
 
+typedef BOOL(WINAPI *SetDefaultDllDirectoriesFunction)(DWORD DirectoryFlags);
+
 int APIENTRY wWinMain(_In_ HINSTANCE hInstance,
                       _In_opt_ HINSTANCE hPrevInstance,
                       _In_ LPWSTR lpCmdLine,
                       _In_ int nCmdShow)
 {
 	// Attempt to mitigate http://textslashplain.com/2015/12/18/dll-hijacking-just-wont-die
-	SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_SYSTEM32);
+	HMODULE hKernel32 = LoadLibrary(L"kernel32.dll");
+	ATLASSERT(hKernel32 != NULL);
+
+	SetDefaultDllDirectoriesFunction pfn = (SetDefaultDllDirectoriesFunction) GetProcAddress(hKernel32, "SetDefaultDllDirectories");
+	if (pfn) { (*pfn)(LOAD_LIBRARY_SEARCH_SYSTEM32); }
 
 	int exitCode = -1;
 	CString cmdLine(lpCmdLine);

src/Squirrel/NativeMethods.cs

@@ -1,14 +1,37 @@
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Runtime.InteropServices;
 using System.Text;
 using System.Threading.Tasks;
 
 namespace Squirrel
 {
-    static class NativeMethods
+    public static class NativeMethods
     {
+        public static int GetParentProcessId()
+        {
+            var pbi = new PROCESS_BASIC_INFORMATION();
+
+            //Get a handle to our own process
+            IntPtr hProc = OpenProcess((ProcessAccess)0x001F0FFF, false, Process.GetCurrentProcess().Id);
+
+            try {
+                int sizeInfoReturned;
+                int queryStatus = NtQueryInformationProcess(hProc, (PROCESSINFOCLASS)0, ref pbi, pbi.Size, out sizeInfoReturned);
+            } finally {
+                if (!hProc.Equals(IntPtr.Zero)) {
+                    //Close handle and free allocated memory
+                    CloseHandle(hProc);
+                    hProc = IntPtr.Zero;
+                }
+            }
+
+            return (int)pbi.InheritedFromUniqueProcessId;
+        }
+
+
         [DllImport("version.dll", SetLastError = true)]
         [return:MarshalAs(UnmanagedType.Bool)] public static extern bool GetFileVersionInfo(
             string lpszFileName, 
@@ -27,5 +50,150 @@ public static extern int GetFileVersionInfoSize(
             string pSubBlock, 
             out IntPtr pValue, 
             out int len);
+
+        [DllImport("psapi.dll", SetLastError=true)]
+        public static extern bool EnumProcesses(
+            IntPtr pProcessIds, // pointer to allocated DWORD array
+            int cb,
+            out int pBytesReturned);
+
+        [DllImport("kernel32.dll", SetLastError=true)]
+        public static extern bool QueryFullProcessImageName(
+            IntPtr hProcess, 
+            [In] int justPassZeroHere,
+            [Out] StringBuilder lpImageFileName, 
+            [In] [MarshalAs(UnmanagedType.U4)] ref int nSize);
+
+        [DllImport("kernel32.dll", SetLastError=true)]
+        public static extern IntPtr OpenProcess(
+            ProcessAccess processAccess,
+            bool bInheritHandle,
+            int processId);
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        public static extern bool CloseHandle(IntPtr hHandle);
+
+        [DllImport("NTDLL.DLL", SetLastError=true)]
+        public static extern int NtQueryInformationProcess(IntPtr hProcess, PROCESSINFOCLASS pic, ref PROCESS_BASIC_INFORMATION pbi, int cb, out int pSize);
+
+        [DllImport("kernel32.dll", SetLastError=true)]
+        public static extern UInt32 WaitForSingleObject(IntPtr hHandle, UInt32 dwMilliseconds);
+
+        [DllImport("kernel32.dll", EntryPoint = "GetStdHandle")]
+        public static extern IntPtr GetStdHandle(StandardHandles nStdHandle);
+
+        [DllImport("kernel32.dll", EntryPoint = "AllocConsole")]
+        [return: MarshalAs(UnmanagedType.Bool)] 
+        public static extern bool AllocConsole();
+
+        [DllImport("kernel32.dll")]
+        public static extern bool AttachConsole(int pid);
+
+        [DllImport("Kernel32.dll", SetLastError=true)]
+        public static extern IntPtr BeginUpdateResource(string pFileName, bool bDeleteExistingResources);
+
+        [DllImport("Kernel32.dll", SetLastError=true)]
+        public static extern bool UpdateResource(IntPtr handle, string pType, IntPtr pName, short language, [MarshalAs(UnmanagedType.LPArray)] byte[] pData, int dwSize);
+
+        [DllImport("Kernel32.dll", SetLastError=true)]
+        public static extern bool EndUpdateResource(IntPtr handle, bool discard);
+    }
+
+    [Flags]
+    public enum ProcessAccess : uint {
+        All = 0x001F0FFF,
+        Terminate = 0x00000001,
+        CreateThread = 0x00000002,
+        VirtualMemoryOperation = 0x00000008,
+        VirtualMemoryRead = 0x00000010,
+        VirtualMemoryWrite = 0x00000020,
+        DuplicateHandle = 0x00000040,
+        CreateProcess = 0x000000080,
+        SetQuota = 0x00000100,
+        SetInformation = 0x00000200,
+        QueryInformation = 0x00000400,
+        QueryLimitedInformation = 0x00001000,
+        Synchronize = 0x00100000
+    }
+
+    public enum PROCESSINFOCLASS : int {
+        ProcessBasicInformation = 0, // 0, q: PROCESS_BASIC_INFORMATION, PROCESS_EXTENDED_BASIC_INFORMATION
+        ProcessQuotaLimits, // qs: QUOTA_LIMITS, QUOTA_LIMITS_EX
+        ProcessIoCounters, // q: IO_COUNTERS
+        ProcessVmCounters, // q: VM_COUNTERS, VM_COUNTERS_EX
+        ProcessTimes, // q: KERNEL_USER_TIMES
+        ProcessBasePriority, // s: KPRIORITY
+        ProcessRaisePriority, // s: ULONG
+        ProcessDebugPort, // q: HANDLE
+        ProcessExceptionPort, // s: HANDLE
+        ProcessAccessToken, // s: PROCESS_ACCESS_TOKEN
+        ProcessLdtInformation, // 10
+        ProcessLdtSize,
+        ProcessDefaultHardErrorMode, // qs: ULONG
+        ProcessIoPortHandlers, // (kernel-mode only)
+        ProcessPooledUsageAndLimits, // q: POOLED_USAGE_AND_LIMITS
+        ProcessWorkingSetWatch, // q: PROCESS_WS_WATCH_INFORMATION[]; s: void
+        ProcessUserModeIOPL,
+        ProcessEnableAlignmentFaultFixup, // s: BOOLEAN
+        ProcessPriorityClass, // qs: PROCESS_PRIORITY_CLASS
+        ProcessWx86Information,
+        ProcessHandleCount, // 20, q: ULONG, PROCESS_HANDLE_INFORMATION
+        ProcessAffinityMask, // s: KAFFINITY
+        ProcessPriorityBoost, // qs: ULONG
+        ProcessDeviceMap, // qs: PROCESS_DEVICEMAP_INFORMATION, PROCESS_DEVICEMAP_INFORMATION_EX
+        ProcessSessionInformation, // q: PROCESS_SESSION_INFORMATION
+        ProcessForegroundInformation, // s: PROCESS_FOREGROUND_BACKGROUND
+        ProcessWow64Information, // q: ULONG_PTR
+        ProcessImageFileName, // q: UNICODE_STRING
+        ProcessLUIDDeviceMapsEnabled, // q: ULONG
+        ProcessBreakOnTermination, // qs: ULONG
+        ProcessDebugObjectHandle, // 30, q: HANDLE
+        ProcessDebugFlags, // qs: ULONG
+        ProcessHandleTracing, // q: PROCESS_HANDLE_TRACING_QUERY; s: size 0 disables, otherwise enables
+        ProcessIoPriority, // qs: ULONG
+        ProcessExecuteFlags, // qs: ULONG
+        ProcessResourceManagement,
+        ProcessCookie, // q: ULONG
+        ProcessImageInformation, // q: SECTION_IMAGE_INFORMATION
+        ProcessCycleTime, // q: PROCESS_CYCLE_TIME_INFORMATION
+        ProcessPagePriority, // q: ULONG
+        ProcessInstrumentationCallback, // 40
+        ProcessThreadStackAllocation, // s: PROCESS_STACK_ALLOCATION_INFORMATION, PROCESS_STACK_ALLOCATION_INFORMATION_EX
+        ProcessWorkingSetWatchEx, // q: PROCESS_WS_WATCH_INFORMATION_EX[]
+        ProcessImageFileNameWin32, // q: UNICODE_STRING
+        ProcessImageFileMapping, // q: HANDLE (input)
+        ProcessAffinityUpdateMode, // qs: PROCESS_AFFINITY_UPDATE_MODE
+        ProcessMemoryAllocationMode, // qs: PROCESS_MEMORY_ALLOCATION_MODE
+        ProcessGroupInformation, // q: USHORT[]
+        ProcessTokenVirtualizationEnabled, // s: ULONG
+        ProcessConsoleHostProcess, // q: ULONG_PTR
+        ProcessWindowInformation, // 50, q: PROCESS_WINDOW_INFORMATION
+        ProcessHandleInformation, // q: PROCESS_HANDLE_SNAPSHOT_INFORMATION // since WIN8
+        ProcessMitigationPolicy, // s: PROCESS_MITIGATION_POLICY_INFORMATION
+        ProcessDynamicFunctionTableInformation,
+        ProcessHandleCheckingMode,
+        ProcessKeepAliveCount, // q: PROCESS_KEEPALIVE_COUNT_INFORMATION
+        ProcessRevokeFileHandles, // s: PROCESS_REVOKE_FILE_HANDLES_INFORMATION
+        MaxProcessInfoClass
+    };
+
+    [StructLayout(LayoutKind.Sequential, Pack = 1)]
+    public struct PROCESS_BASIC_INFORMATION {
+        public IntPtr ExitStatus;
+        public IntPtr PebBaseAddress;
+        public IntPtr AffinityMask;
+        public IntPtr BasePriority;
+        public UIntPtr UniqueProcessId;
+        public IntPtr InheritedFromUniqueProcessId;
+
+        public int Size {
+            get { return (int)Marshal.SizeOf(typeof(PROCESS_BASIC_INFORMATION)); }
+        }
+    }
+
+    public enum StandardHandles : int {
+        STD_INPUT_HANDLE = -10,
+        STD_OUTPUT_HANDLE = -11,
+        STD_ERROR_HANDLE = -12,
     }
 }

src/Squirrel/UpdateManager.InstallHelpers.cs

@@ -10,6 +10,7 @@
 using Microsoft.Win32;
 using NuGet;
 using Splat;
+using System.Reflection;
 
 namespace Squirrel
 {
@@ -102,6 +103,35 @@ public async Task<RegistryKey> CreateUninstallerRegistryEntry(string uninstallCm
                 return key;
             }
 
+            public void KillAllProcessesBelongingToPackage()
+            {
+                var ourExe = Assembly.GetEntryAssembly();
+                var ourExePath = ourExe != null ? ourExe.Location : null;
+
+                UnsafeUtility.EnumerateProcesses()
+                    .Where(x => {
+                        // Processes we can't query will have an empty process name, we can't kill them
+                        // anyways
+                        if (String.IsNullOrWhiteSpace(x.Item1)) return false;
+
+                        // Files that aren't in our root app directory are untouched
+                        if (!x.Item1.StartsWith(rootAppDirectory, StringComparison.OrdinalIgnoreCase)) return false;
+
+                        // Never kill our own EXE
+                        if (ourExePath != null && x.Item1.Equals(ourExePath, StringComparison.OrdinalIgnoreCase)) return false;
+
+                        var name = Path.GetFileName(x.Item1).ToLowerInvariant();
+                        if (name == "squirrel.exe" || name == "update.exe") return false;
+
+                        return true;
+                    })
+                    .ForEach(x => {
+                        try {
+                            this.WarnIfThrows(() => Process.GetProcessById(x.Item2).Kill());
+                        } catch {}
+                    });
+            }
+
             public Task<RegistryKey> CreateUninstallerRegistryEntry()
             {
                 var updateDotExe = Path.Combine(rootAppDirectory, "Update.exe");

src/Squirrel/UpdateManager.cs

@@ -89,6 +89,7 @@ public async Task FullUninstall()
             var applyReleases = new ApplyReleasesImpl(rootAppDirectory);
             await acquireUpdateLock();
 
+            this.KillAllExecutablesBelongingToPackage();
             await applyReleases.FullUninstall();
         }
 
@@ -145,6 +146,12 @@ public SemanticVersion CurrentlyInstalledVersion(string executable = null)
             return appDirName.ToSemanticVersion();
         }
 
+        public void KillAllExecutablesBelongingToPackage()
+        {
+            var installHelpers = new InstallHelperImpl(applicationName, rootAppDirectory);
+            installHelpers.KillAllProcessesBelongingToPackage();
+        }
+
         public string ApplicationName {
             get { return applicationName; }
         }

src/Squirrel/Utility.cs

@@ -586,6 +586,44 @@ enum MoveFileFlags
         }
     }
 
+    static unsafe class UnsafeUtility
+    {
+        public static List<Tuple<string, int>> EnumerateProcesses()
+        {
+            int length = 0;
+            var pids = new int[2048];
+
+            fixed(int* p = pids) {
+                if (!NativeMethods.EnumProcesses((IntPtr)p, sizeof(int) * pids.Length, out length)) {
+                    throw new Win32Exception("Failed to enumerate processes");
+                }
+
+                if (length < 1) throw new Exception("Failed to enumerate processes");
+            }
+
+            return Enumerable.Range(0, length)
+                .Where(i => pids[i] > 0)
+                .Select(i => {
+                    try {
+                        var hProcess = NativeMethods.OpenProcess(ProcessAccess.QueryLimitedInformation, false, pids[i]);
+                        if (hProcess == IntPtr.Zero) throw new Win32Exception();
+
+                        var sb = new StringBuilder(256);
+                        var capacity = sb.Capacity;
+                        if (!NativeMethods.QueryFullProcessImageName(hProcess, 0, sb, ref capacity)) {
+                            throw new Win32Exception();
+                        }
+
+                        NativeMethods.CloseHandle(hProcess);
+                        return Tuple.Create(sb.ToString(), pids[i]);
+                    } catch (Exception) {
+                        return Tuple.Create(default(string), pids[i]);
+                    }
+                })
+                .ToList();
+        }
+    }
+
     sealed class SingleGlobalInstance : IDisposable, IEnableLogger
     {
         IDisposable handle = null;