Prevent crashes from partial in-place updates

UniGetUI.iss

@@ -29,7 +29,10 @@ DefaultDirName="{autopf64}\UniGetUI"
 DisableProgramGroupPage=yes
 DisableDirPage=no
 DirExistsWarning=no
-CloseApplications=no
+; Force-close any process holding files we overwrite (backstop for the kill in PrepareToInstall).
+CloseApplications=force
+CloseApplicationsFilter=*.exe,*.dll
+RestartApplications=no
 ; Default to per-user install mode and let the dialog opt into all-users installs when needed.
 PrivilegesRequired=lowest
 PrivilegesRequiredOverridesAllowed=dialog
@@ -97,19 +100,70 @@ begin
   WizardForm.Bevel1.Visible := True;
 end;
 
-procedure TaskKill(FileName: String);
+// Kills all instances of an image and loops until none remain (taskkill returns 0 while killing, 128 when none left).
+procedure TaskKillWait(FileName: String);
 var
-  ResultCode: Integer;
+  ResultCode, Attempts: Integer;
 begin
-    Exec('taskkill.exe', '/f /im ' + '"' + FileName + '"', '', SW_HIDE,
-     ewWaitUntilTerminated, ResultCode);
+    Attempts := 0;
+    repeat
+        if not Exec('taskkill.exe', '/f /im "' + FileName + '"', '', SW_HIDE,
+            ewWaitUntilTerminated, ResultCode) then
+            Break;
+        if ResultCode <> 0 then
+            Break;
+        Sleep(500);
+        Attempts := Attempts + 1;
+    until Attempts >= 10;
 end;
 
 procedure KillRunningApps;
 begin
-    TaskKill('WingetUI.exe');
-    TaskKill('UniGetUI.exe');
-    TaskKill('UniGetUI.Avalonia.exe');
+    TaskKillWait('WingetUI.exe');
+    TaskKillWait('UniGetUI.exe');
+    TaskKillWait('UniGetUI.Avalonia.exe');
+    // Elevator (gsudo cache) and pinget live in {app} and lock their own files.
+    TaskKillWait('UniGetUI Elevator.exe');
+    TaskKillWait('pinget.exe');
+    Sleep(1000); // let the OS release file handles before copying
+
+end;
+
+function GetCurrentProcessId: Cardinal; external 'GetCurrentProcessId@kernel32.dll stdcall';
+
+function UpdateMarkerPath(): String;
+begin
+    Result := ExpandConstant('{app}\.unigetui-update-in-progress');
+end;
+
+// Marker holds our PID; the app blocks only while this installer runs. Name MUST match UpdateInProgressGuard.MarkerFileName.
+procedure WriteUpdateMarker;
+var
+    Pid: Int64;
+begin
+    ForceDirectories(ExpandConstant('{app}'));
+    Pid := GetCurrentProcessId;
+    SaveStringToFile(UpdateMarkerPath(), IntToStr(Pid), False);
+end;
+
+procedure RemoveUpdateMarker;
+begin
+    DeleteFile(UpdateMarkerPath());
+end;
+
+// Runs before any file is copied: shut everything down, then mark the copy window.
+function PrepareToInstall(var NeedsRestart: Boolean): String;
+begin
+    KillRunningApps;
+    WriteUpdateMarker;
+    Result := '';
+end;
+
+// Clear the marker once the copy is done, before the post-install launch.
+procedure CurStepChanged(CurStep: TSetupStep);
+begin
+    if CurStep = ssPostInstall then
+        RemoveUpdateMarker;
 end;
 
 function CmdLineParamExists(const Value: string): Boolean;
@@ -132,6 +186,8 @@ procedure ExitProcess(exitCode:integer);
 
 procedure DeinitializeSetup();
 begin
+    RemoveUpdateMarker; // also clear on abort, before ssPostInstall
+
     if (CustomExitCode <> 0) then
     begin
         DelTree(ExpandConstant('{tmp}'), True, True, True);
@@ -215,8 +271,8 @@ Root: HKA; Subkey: "Software\Classes\UniGetUI.PackageBundle\shell\open\command";
 Source: "{srcexe}"; DestDir: "{app}"; DestName: "UniGetUI.Installer.exe"; Flags: external ignoreversion; Tasks: regularinstall; Check: not CmdLineParamExists('/NoDeployInstaller');
 ; Deploy integrity tree
 Source: "unigetui_bin\IntegrityTree.json"; DestDir: "{app}"; Flags: createallsubdirs ignoreversion recursesubdirs;
-; Deploy executable files
-Source: "unigetui_bin\{#MyAppExeName}"; DestDir: "{app}"; Flags: ignoreversion; BeforeInstall: KillRunningApps;
+; Deploy executable files (running instances already killed in PrepareToInstall).
+Source: "unigetui_bin\{#MyAppExeName}"; DestDir: "{app}"; Flags: ignoreversion;
 Source: "unigetui_bin\*"; DestDir: "{app}"; Flags: createallsubdirs ignoreversion recursesubdirs;
 ; Make installation portable (if required)
 Source: "InstallerExtras\ForceUniGetUIPortable"; DestDir: "{app}"; Tasks: portableinstall

src/UniGetUI.Avalonia/Program.cs

@@ -1,6 +1,7 @@
 using System;
 using Avalonia;
 using UniGetUI.Avalonia.Infrastructure;
+using UniGetUI.Core.Data;
 
 namespace UniGetUI.Avalonia;
 
@@ -12,6 +13,17 @@ sealed class Program
     [STAThread]
     public static void Main(string[] args)
     {
+        // Bail out if the installer is mid-swap (try/catch so the guard never blocks a normal launch).
+        try
+        {
+            if (UpdateInProgressGuard.IsUpdateInProgress())
+            {
+                Environment.ExitCode = 0;
+                return;
+            }
+        }
+        catch { }
+
         AvaloniaAppHost.Run(args);
     }
 

src/UniGetUI.Core.Data.Tests/UpdateInProgressGuardTests.cs

@@ -0,0 +1,89 @@
+namespace UniGetUI.Core.Data.Tests
+{
+    public class UpdateInProgressGuardTests : IDisposable
+    {
+        // {root}/app stands in for {app}; {root} is its always-empty parent.
+        private readonly string _root;
+        private readonly string _appDir;
+
+        private static readonly Func<int, bool> ProcessAlive = _ => true;
+        private static readonly Func<int, bool> ProcessDead = _ => false;
+
+        public UpdateInProgressGuardTests()
+        {
+            _root = Path.Combine(Path.GetTempPath(), "ugui-guard-" + Guid.NewGuid().ToString("N"));
+            _appDir = Path.Combine(_root, "app");
+            Directory.CreateDirectory(_appDir);
+        }
+
+        public void Dispose()
+        {
+            try { Directory.Delete(_root, recursive: true); }
+            catch { }
+        }
+
+        private static string WriteMarker(string directory, string content = "1234")
+        {
+            Directory.CreateDirectory(directory);
+            string path = Path.Combine(directory, UpdateInProgressGuard.MarkerFileName);
+            File.WriteAllText(path, content);
+            return path;
+        }
+
+        [Fact]
+        public void NoMarker_ReturnsFalse()
+        {
+            Assert.False(UpdateInProgressGuard.IsUpdateInProgress(_appDir, ProcessAlive));
+        }
+
+        [Fact]
+        public void MarkerWithRunningInstaller_ReturnsTrue()
+        {
+            WriteMarker(_appDir);
+            Assert.True(UpdateInProgressGuard.IsUpdateInProgress(_appDir, ProcessAlive));
+        }
+
+        [Fact]
+        public void MarkerInParentWithRunningInstaller_ReturnsTrue()
+        {
+            // Avalonia runs from {app}\Avalonia; marker is in {app}.
+            WriteMarker(_appDir);
+            string child = Path.Combine(_appDir, "Avalonia");
+            Directory.CreateDirectory(child);
+
+            Assert.True(UpdateInProgressGuard.IsUpdateInProgress(child, ProcessAlive));
+        }
+
+        [Fact]
+        public void MarkerWithDeadInstaller_ReturnsFalseAndIsDeleted()
+        {
+            string marker = WriteMarker(_appDir);
+
+            Assert.False(UpdateInProgressGuard.IsUpdateInProgress(_appDir, ProcessDead));
+            Assert.False(File.Exists(marker)); // stale marker is cleaned up
+        }
+
+        [Fact]
+        public void MarkerWithUnreadableContent_ReturnsFalseAndIsKept()
+        {
+            string marker = WriteMarker(_appDir, "not-a-pid");
+
+            Assert.False(UpdateInProgressGuard.IsUpdateInProgress(_appDir, ProcessAlive));
+            Assert.True(File.Exists(marker)); // not deleted: could be a partial write
+        }
+
+        [Fact]
+        public void RealProcessCheck_TreatsCurrentProcessAsRunning()
+        {
+            // Exercises the real IsProcessRunning via the parameterless overload.
+            WriteMarker(_appDir, Environment.ProcessId.ToString());
+            Assert.True(UpdateInProgressGuard.IsUpdateInProgress(_appDir));
+        }
+
+        [Fact]
+        public void MarkerFileName_MatchesInstallerContract()
+        {
+            Assert.Equal(".unigetui-update-in-progress", UpdateInProgressGuard.MarkerFileName);
+        }
+    }
+}

src/UniGetUI.Core.Data/InternalsVisibleTo.cs

@@ -0,0 +1,3 @@
+using System.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("UniGetUI.Core.Data.Tests")]

src/UniGetUI.Core.Data/UpdateInProgressGuard.cs

@@ -0,0 +1,76 @@
+using System.Diagnostics;
+
+namespace UniGetUI.Core.Data
+{
+    // Blocks UI startup while the Windows installer is replacing files in {app} (see UniGetUI.iss),
+    // so an instance launched mid-update doesn't load a half-written binary set and crash.
+    public static class UpdateInProgressGuard
+    {
+        // MUST match the marker name written by UniGetUI.iss. The file holds the installer's PID.
+        public const string MarkerFileName = ".unigetui-update-in-progress";
+
+        public static bool IsUpdateInProgress()
+        {
+            if (!OperatingSystem.IsWindows())
+                return false;
+
+            return IsUpdateInProgress(AppContext.BaseDirectory);
+        }
+
+        // Checks the running dir and its parent (the Avalonia UI runs from {app}\Avalonia).
+        internal static bool IsUpdateInProgress(string baseDirectory)
+            => IsUpdateInProgress(baseDirectory, IsProcessRunning);
+
+        internal static bool IsUpdateInProgress(string baseDirectory, Func<int, bool> isProcessRunning)
+        {
+            try
+            {
+                if (MarkerIsActive(baseDirectory, isProcessRunning))
+                    return true;
+
+                string? parent = Directory
+                    .GetParent(baseDirectory.TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar))
+                    ?.FullName;
+                return parent is not null && MarkerIsActive(parent, isProcessRunning);
+            }
+            catch
+            {
+                return false;
+            }
+        }
+
+        // Active only while the installer that wrote the PID is still running, so the guard tracks
+        // the real copy window (any duration) and self-heals if the installer dies without cleanup.
+        private static bool MarkerIsActive(string directory, Func<int, bool> isProcessRunning)
+        {
+            string marker = Path.Combine(directory, MarkerFileName);
+            if (!File.Exists(marker))
+                return false;
+
+            if (!int.TryParse(File.ReadAllText(marker).Trim(), out int pid))
+                return false; // unreadable (e.g. racing the installer's write) — leave it alone
+
+            if (isProcessRunning(pid))
+                return true;
+
+            try { File.Delete(marker); } catch { /* stale: installer is gone */ }
+            return false;
+        }
+
+        private static bool IsProcessRunning(int pid)
+        {
+            if (pid <= 0)
+                return false;
+
+            try
+            {
+                using var process = Process.GetProcessById(pid);
+                return !process.HasExited;
+            }
+            catch (ArgumentException)
+            {
+                return false;
+            }
+        }
+    }
+}

src/UniGetUI/EntryPoint.cs

@@ -51,6 +51,13 @@ private static void Main(string[] args)
                     Environment.ExitCode = WinUiHeadlessHost.RunAsync(args).GetAwaiter().GetResult();
                     return;
                 }
+                else if (UpdateInProgressGuard.IsUpdateInProgress())
+                {
+                    // Update is replacing install files; the installer relaunches us when done.
+                    Logger.Warn("An update is replacing install files; aborting UI startup until it completes.");
+                    Environment.ExitCode = 0;
+                    return;
+                }
                 else if (!ModernAppLauncher.IsClassicModeEnabled())
                 {
                     ModernAppLauncher.Launch(args);