From 923a982d1f31ff5715205597637d023133ff0d53 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Misty=20De=20M=C3=A9o?= <misty@axo.dev>
Date: Mon, 16 Sep 2024 14:34:51 -0700
Subject: [PATCH] feat: use self_replace to delete self

---
 Cargo.lock               | 12 ++++++++++++
 axoupdater/Cargo.toml    |  3 +++
 axoupdater/src/errors.rs |  7 +++++++
 axoupdater/src/lib.rs    | 40 +++++++++++++++++-----------------------
 4 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index a0cb74c..1cd08da 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -155,6 +155,7 @@ dependencies = [
  "gazenot",
  "homedir",
  "miette",
+ "self-replace",
  "serde",
  "tempfile",
  "thiserror",
@@ -1427,6 +1428,17 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "self-replace"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "03ec815b5eab420ab893f63393878d89c90fdd94c0bcc44c07abb8ad95552fb7"
+dependencies = [
+ "fastrand 2.1.0",
+ "tempfile",
+ "windows-sys 0.52.0",
+]
+
 [[package]]
 name = "semver"
 version = "1.0.22"
diff --git a/axoupdater/Cargo.toml b/axoupdater/Cargo.toml
index 2f15205..d816fd1 100644
--- a/axoupdater/Cargo.toml
+++ b/axoupdater/Cargo.toml
@@ -35,3 +35,6 @@ tokio = { version = "1.36.0", features = ["full"], optional = true }
 # errors
 miette = "7.2.0"
 thiserror = "1.0.58"
+
+[target.'cfg(windows)'.dependencies]
+self-replace = "1.5.0"
diff --git a/axoupdater/src/errors.rs b/axoupdater/src/errors.rs
index 81f1359..7251a7a 100644
--- a/axoupdater/src/errors.rs
+++ b/axoupdater/src/errors.rs
@@ -159,4 +159,11 @@ pub enum AxoupdateError {
         /// to the terminal when running the installer.
         stderr: Option<String>,
     },
+
+    /// self_replace/self_delete failed
+    #[error(
+        "Cleaning up the previous version failed; a copy of the old version has been left behind."
+    )]
+    #[diagnostic(help("This probably isn't your fault; please open an issue at https://github.com/axodotdev/axoupdater!"))]
+    CleanupFailed {},
 }
diff --git a/axoupdater/src/lib.rs b/axoupdater/src/lib.rs
index a8d55fb..41a31aa 100644
--- a/axoupdater/src/lib.rs
+++ b/axoupdater/src/lib.rs
@@ -13,12 +13,16 @@ pub use release::*;
 
 use std::{
     env::{self, args},
+    ffi::OsStr,
     process::Stdio,
 };
 
 #[cfg(unix)]
 use std::{fs::File, os::unix::fs::PermissionsExt};
 
+#[cfg(windows)]
+use self_replace;
+
 use axoasset::LocalAsset;
 use axoprocess::Cmd;
 pub use axotag::Version;
@@ -453,34 +457,21 @@ impl AxoUpdater {
             installer_path
         };
 
-        // Before we update, move ourselves to a temporary directory.
+        // Before we update, rename ourselves to a temporary name.
         // This is necessary because Windows won't let an actively-running
         // executable be overwritten.
         // If the update fails, we'll move it back to where it was before
         // we began the update process.
-        // NOTE: this TempDir needs to be held alive for the whole function.
-        let temp_root;
-        let to_restore = if cfg!(target_family = "windows") {
-            // If we know the install prefix, place the temporary directory for
-            // the executable there. This is because the `rename` syscall can't
-            // rename an executable across filesystem bounds, and there's a
-            // chance the system temporary directory could be on a different
-            // drive than the executable is.
-            // (A copy-and-delete move won't work because Windows won't let us
-            // delete a running executable - the same reason we're moving it out
-            // of the way to begin with!)
-            // See https://github.com/axodotdev/axoupdater/issues/120.
-            temp_root = if let Ok(prefix) = self.install_prefix_root() {
-                TempDir::new_in(prefix)?
-            } else {
-                TempDir::new()?
-            };
-            let old_path = std::env::current_exe()?;
-            let old_filename = old_path.file_name().expect("current binary has no name!?");
-            let ourselves = temp_root.path().join(old_filename);
-            std::fs::rename(&old_path, &ourselves)?;
+        let old_path;
+        let to_restore = if cfg!(windows) {
+            old_path = std::env::current_exe()?;
+            let filename = old_path.file_name();
+            let old_filename = filename.expect("current binary has no name!?");
 
-            Some((ourselves, old_path))
+            let mut new_filename = old_filename.to_os_string();
+            new_filename.push(OsStr::new(".__selfdelete__.exe"));
+
+            Some((new_filename, old_filename))
         } else {
             None
         };
@@ -550,6 +541,9 @@ impl AxoUpdater {
             if let Some((ourselves, old_path)) = to_restore {
                 std::fs::rename(ourselves, old_path)?;
             }
+        } else {
+            #[cfg(windows)]
+            self_replace::self_delete().map_err(|_| AxoupdateError::CleanupFailed {})?;
         }
 
         // Return the original AxoprocessError if we failed to launch