Fix autoupdate nuking the app on Windows (#41571)

Closes #41477

Release Notes:

- N/A

Cargo.lock

@@ -1351,6 +1351,7 @@ dependencies = [
  "anyhow",
  "log",
  "simplelog",
+ "tempfile",
  "windows 0.61.3",
  "winresource",
 ]

crates/auto_update/src/auto_update.rs

@@ -331,6 +331,16 @@ impl AutoUpdater {
 
     pub fn start_polling(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
         cx.spawn(async move |this, cx| {
+            #[cfg(target_os = "windows")]
+            {
+                use util::ResultExt;
+
+                cleanup_windows()
+                    .await
+                    .context("failed to cleanup old directories")
+                    .log_err();
+            }
+
             loop {
                 this.update(cx, |this, cx| this.poll(UpdateCheckType::Automatic, cx))?;
                 cx.background_executor().timer(POLL_INTERVAL).await;
@@ -923,6 +933,32 @@ async fn install_release_macos(
     Ok(None)
 }
 
+#[cfg(target_os = "windows")]
+async fn cleanup_windows() -> Result<()> {
+    use util::ResultExt;
+
+    let parent = std::env::current_exe()?
+        .parent()
+        .context("No parent dir for Zed.exe")?
+        .to_owned();
+
+    // keep in sync with crates/auto_update_helper/src/updater.rs
+    smol::fs::remove_dir(parent.join("updates"))
+        .await
+        .context("failed to remove updates dir")
+        .log_err();
+    smol::fs::remove_dir(parent.join("install"))
+        .await
+        .context("failed to remove install dir")
+        .log_err();
+    smol::fs::remove_dir(parent.join("old"))
+        .await
+        .context("failed to remove old version dir")
+        .log_err();
+
+    Ok(())
+}
+
 async fn install_release_windows(downloaded_installer: PathBuf) -> Result<Option<PathBuf>> {
     let output = Command::new(downloaded_installer)
         .arg("/verysilent")

crates/auto_update_helper/Cargo.toml

@@ -21,6 +21,9 @@ simplelog.workspace = true
 [target.'cfg(target_os = "windows")'.dependencies]
 windows.workspace = true
 
+[target.'cfg(target_os = "windows")'.dev-dependencies]
+tempfile.workspace = true
+
 [target.'cfg(target_os = "windows")'.build-dependencies]
 winresource = "0.1"
 

crates/auto_update_helper/src/updater.rs

@@ -1,4 +1,5 @@
 use std::{
+    cell::LazyCell,
     path::Path,
     time::{Duration, Instant},
 };
@@ -11,210 +12,274 @@ use windows::Win32::{
 
 use crate::windows_impl::WM_JOB_UPDATED;
 
-type Job = fn(&Path) -> Result<()>;
+pub(crate) struct Job {
+    pub apply: Box<dyn Fn(&Path) -> Result<()>>,
+    pub rollback: Box<dyn Fn(&Path) -> Result<()>>,
+}
 
-#[cfg(not(test))]
-pub(crate) const JOBS: &[Job] = &[
-    // Delete old files
-    |app_dir| {
-        let zed_executable = app_dir.join("Zed.exe");
-        log::info!("Removing old file: {}", zed_executable.display());
-        std::fs::remove_file(&zed_executable).context(format!(
-            "Failed to remove old file {}",
-            zed_executable.display()
-        ))
-    },
-    |app_dir| {
-        let zed_cli = app_dir.join("bin\\zed.exe");
-        log::info!("Removing old file: {}", zed_cli.display());
-        std::fs::remove_file(&zed_cli)
-            .context(format!("Failed to remove old file {}", zed_cli.display()))
-    },
-    |app_dir| {
-        let zed_wsl = app_dir.join("bin\\zed");
-        log::info!("Removing old file: {}", zed_wsl.display());
-        std::fs::remove_file(&zed_wsl)
-            .context(format!("Failed to remove old file {}", zed_wsl.display()))
-    },
-    // TODO: remove after a few weeks once everyone is on the new version and this file never exists
-    |app_dir| {
-        let open_console = app_dir.join("OpenConsole.exe");
-        if open_console.exists() {
-            log::info!("Removing old file: {}", open_console.display());
-            std::fs::remove_file(&open_console).context(format!(
-                "Failed to remove old file {}",
-                open_console.display()
-            ))?
+impl Job {
+    pub fn mkdir(name: &'static Path) -> Self {
+        Job {
+            apply: Box::new(move |app_dir| {
+                let dir = app_dir.join(name);
+                std::fs::create_dir_all(&dir)
+                    .context(format!("Failed to create directory {}", dir.display()))
+            }),
+            rollback: Box::new(move |app_dir| {
+                let dir = app_dir.join(name);
+                std::fs::remove_dir_all(&dir)
+                    .context(format!("Failed to remove directory {}", dir.display()))
+            }),
         }
-        Ok(())
-    },
-    |app_dir| {
-        let archs = ["x64", "arm64"];
-        for arch in archs {
-            let open_console = app_dir.join(format!("{arch}\\OpenConsole.exe"));
-            if open_console.exists() {
-                log::info!("Removing old file: {}", open_console.display());
-                std::fs::remove_file(&open_console).context(format!(
-                    "Failed to remove old file {}",
-                    open_console.display()
-                ))?
-            }
+    }
+
+    pub fn mkdir_if_exists(name: &'static Path, check: &'static Path) -> Self {
+        Job {
+            apply: Box::new(move |app_dir| {
+                let dir = app_dir.join(name);
+                let check = app_dir.join(check);
+
+                if check.exists() {
+                    std::fs::create_dir_all(&dir)
+                        .context(format!("Failed to create directory {}", dir.display()))?
+                }
+                Ok(())
+            }),
+            rollback: Box::new(move |app_dir| {
+                let dir = app_dir.join(name);
+
+                if dir.exists() {
+                    std::fs::remove_dir_all(&dir)
+                        .context(format!("Failed to remove directory {}", dir.display()))?
+                }
+
+                Ok(())
+            }),
         }
-        Ok(())
-    },
-    |app_dir| {
-        let conpty = app_dir.join("conpty.dll");
-        log::info!("Removing old file: {}", conpty.display());
-        std::fs::remove_file(&conpty)
-            .context(format!("Failed to remove old file {}", conpty.display()))
-    },
-    // Copy new files
-    |app_dir| {
-        let zed_executable_source = app_dir.join("install\\Zed.exe");
-        let zed_executable_dest = app_dir.join("Zed.exe");
-        log::info!(
-            "Copying new file {} to {}",
-            zed_executable_source.display(),
-            zed_executable_dest.display()
-        );
-        std::fs::copy(&zed_executable_source, &zed_executable_dest)
-            .map(|_| ())
-            .context(format!(
-                "Failed to copy new file {} to {}",
-                zed_executable_source.display(),
-                zed_executable_dest.display()
-            ))
-    },
-    |app_dir| {
-        let zed_cli_source = app_dir.join("install\\bin\\zed.exe");
-        let zed_cli_dest = app_dir.join("bin\\zed.exe");
-        log::info!(
-            "Copying new file {} to {}",
-            zed_cli_source.display(),
-            zed_cli_dest.display()
-        );
-        std::fs::copy(&zed_cli_source, &zed_cli_dest)
-            .map(|_| ())
-            .context(format!(
-                "Failed to copy new file {} to {}",
-                zed_cli_source.display(),
-                zed_cli_dest.display()
-            ))
-    },
-    |app_dir| {
-        let zed_wsl_source = app_dir.join("install\\bin\\zed");
-        let zed_wsl_dest = app_dir.join("bin\\zed");
-        log::info!(
-            "Copying new file {} to {}",
-            zed_wsl_source.display(),
-            zed_wsl_dest.display()
-        );
-        std::fs::copy(&zed_wsl_source, &zed_wsl_dest)
-            .map(|_| ())
-            .context(format!(
-                "Failed to copy new file {} to {}",
-                zed_wsl_source.display(),
-                zed_wsl_dest.display()
-            ))
-    },
-    |app_dir| {
-        let archs = ["x64", "arm64"];
-        for arch in archs {
-            let open_console_source = app_dir.join(format!("install\\{arch}\\OpenConsole.exe"));
-            let open_console_dest = app_dir.join(format!("{arch}\\OpenConsole.exe"));
-            if open_console_source.exists() {
+    }
+
+    pub fn move_file(filename: &'static Path, new_filename: &'static Path) -> Self {
+        Job {
+            apply: Box::new(move |app_dir| {
+                let old_file = app_dir.join(filename);
+                let new_file = app_dir.join(new_filename);
                 log::info!(
-                    "Copying new file {} to {}",
-                    open_console_source.display(),
-                    open_console_dest.display()
+                    "Moving file: {}->{}",
+                    old_file.display(),
+                    new_file.display()
                 );
-                let parent = open_console_dest.parent().context(format!(
-                    "Failed to get parent directory of {}",
-                    open_console_dest.display()
-                ))?;
-                std::fs::create_dir_all(parent)
-                    .context(format!("Failed to create directory {}", parent.display()))?;
-                std::fs::copy(&open_console_source, &open_console_dest)
-                    .map(|_| ())
-                    .context(format!(
-                        "Failed to copy new file {} to {}",
-                        open_console_source.display(),
-                        open_console_dest.display()
-                    ))?
-            }
-        }
-        Ok(())
-    },
-    |app_dir| {
-        let conpty_source = app_dir.join("install\\conpty.dll");
-        let conpty_dest = app_dir.join("conpty.dll");
-        log::info!(
-            "Copying new file {} to {}",
-            conpty_source.display(),
-            conpty_dest.display()
-        );
-        std::fs::copy(&conpty_source, &conpty_dest)
-            .map(|_| ())
-            .context(format!(
-                "Failed to copy new file {} to {}",
-                conpty_source.display(),
-                conpty_dest.display()
-            ))
-    },
-    // Clean up installer folder and updates folder
-    |app_dir| {
-        let updates_folder = app_dir.join("updates");
-        log::info!("Cleaning up: {}", updates_folder.display());
-        std::fs::remove_dir_all(&updates_folder).context(format!(
-            "Failed to remove updates folder {}",
-            updates_folder.display()
-        ))
-    },
-    |app_dir| {
-        let installer_folder = app_dir.join("install");
-        log::info!("Cleaning up: {}", installer_folder.display());
-        std::fs::remove_dir_all(&installer_folder).context(format!(
-            "Failed to remove installer folder {}",
-            installer_folder.display()
-        ))
-    },
-];
 
+                std::fs::rename(&old_file, new_file)
+                    .context(format!("Failed to move file {}", old_file.display()))
+            }),
+            rollback: Box::new(move |app_dir| {
+                let old_file = app_dir.join(filename);
+                let new_file = app_dir.join(new_filename);
+                log::info!(
+                    "Rolling back file move: {}->{}",
+                    old_file.display(),
+                    new_file.display()
+                );
+
+                std::fs::rename(&new_file, &old_file).context(format!(
+                    "Failed to rollback file move {}->{}",
+                    new_file.display(),
+                    old_file.display()
+                ))
+            }),
+        }
+    }
+
+    pub fn move_if_exists(filename: &'static Path, new_filename: &'static Path) -> Self {
+        Job {
+            apply: Box::new(move |app_dir| {
+                let old_file = app_dir.join(filename);
+                let new_file = app_dir.join(new_filename);
+
+                if old_file.exists() {
+                    log::info!(
+                        "Moving file: {}->{}",
+                        old_file.display(),
+                        new_file.display()
+                    );
+
+                    std::fs::rename(&old_file, new_file)
+                        .context(format!("Failed to move file {}", old_file.display()))?;
+                }
+
+                Ok(())
+            }),
+            rollback: Box::new(move |app_dir| {
+                let old_file = app_dir.join(filename);
+                let new_file = app_dir.join(new_filename);
+
+                if new_file.exists() {
+                    log::info!(
+                        "Rolling back file move: {}->{}",
+                        old_file.display(),
+                        new_file.display()
+                    );
+
+                    std::fs::rename(&new_file, &old_file).context(format!(
+                        "Failed to rollback file move {}->{}",
+                        new_file.display(),
+                        old_file.display()
+                    ))?
+                }
+
+                Ok(())
+            }),
+        }
+    }
+
+    pub fn rmdir_nofail(filename: &'static Path) -> Self {
+        Job {
+            apply: Box::new(move |app_dir| {
+                let filename = app_dir.join(filename);
+                log::info!("Removing file: {}", filename.display());
+                if let Err(e) = std::fs::remove_dir_all(&filename) {
+                    log::warn!("Failed to remove directory: {}", e);
+                }
+
+                Ok(())
+            }),
+            rollback: Box::new(move |app_dir| {
+                let filename = app_dir.join(filename);
+                anyhow::bail!(
+                    "Delete operations cannot be rolled back, file: {}",
+                    filename.display()
+                )
+            }),
+        }
+    }
+}
+
+// app is single threaded
+#[cfg(not(test))]
+#[allow(clippy::declare_interior_mutable_const)]
+pub(crate) const JOBS: LazyCell<[Job; 22]> = LazyCell::new(|| {
+    fn p(value: &str) -> &Path {
+        Path::new(value)
+    }
+    [
+        // Move old files
+        // Not deleting because installing new files can fail
+        Job::mkdir(p("old")),
+        Job::move_file(p("Zed.exe"), p("old\\Zed.exe")),
+        Job::mkdir(p("old\\bin")),
+        Job::move_file(p("bin\\Zed.exe"), p("old\\bin\\Zed.exe")),
+        Job::move_file(p("bin\\zed"), p("old\\bin\\zed")),
+        //
+        // TODO: remove after a few weeks once everyone is on the new version and this file never exists
+        Job::move_if_exists(p("OpenConsole.exe"), p("old\\OpenConsole.exe")),
+        Job::mkdir(p("old\\x64")),
+        Job::mkdir(p("old\\arm64")),
+        Job::move_if_exists(p("x64\\OpenConsole.exe"), p("old\\x64\\OpenConsole.exe")),
+        Job::move_if_exists(
+            p("arm64\\OpenConsole.exe"),
+            p("old\\arm64\\OpenConsole.exe"),
+        ),
+        //
+        Job::move_file(p("conpty.dll"), p("old\\conpty.dll")),
+        // Copy new files
+        Job::move_file(p("install\\Zed.exe"), p("Zed.exe")),
+        Job::move_file(p("install\\bin\\Zed.exe"), p("bin\\Zed.exe")),
+        Job::move_file(p("install\\bin\\zed"), p("bin\\zed")),
+        //
+        Job::mkdir_if_exists(p("x64"), p("install\\x64")),
+        Job::mkdir_if_exists(p("arm64"), p("install\\arm64")),
+        Job::move_if_exists(
+            p("install\\x64\\OpenConsole.exe"),
+            p("x64\\OpenConsole.exe"),
+        ),
+        Job::move_if_exists(
+            p("install\\arm64\\OpenConsole.exe"),
+            p("arm64\\OpenConsole.exe"),
+        ),
+        //
+        Job::move_file(p("install\\conpty.dll"), p("conpty.dll")),
+        // Cleanup installer and updates folder
+        Job::rmdir_nofail(p("updates")),
+        Job::rmdir_nofail(p("install")),
+        // Cleanup old installation
+        Job::rmdir_nofail(p("old")),
+    ]
+});
+
+// app is single threaded
 #[cfg(test)]
-pub(crate) const JOBS: &[Job] = &[
-    |_| {
-        std::thread::sleep(Duration::from_millis(1000));
-        if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
-            match config.as_str() {
-                "err" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
-                _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
-            }
-        } else {
-            Ok(())
-        }
-    },
-    |_| {
-        std::thread::sleep(Duration::from_millis(1000));
-        if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
-            match config.as_str() {
-                "err" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
-                _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
-            }
-        } else {
-            Ok(())
-        }
-    },
-];
+#[allow(clippy::declare_interior_mutable_const)]
+pub(crate) const JOBS: LazyCell<[Job; 9]> = LazyCell::new(|| {
+    fn p(value: &str) -> &Path {
+        Path::new(value)
+    }
+    [
+        Job {
+            apply: Box::new(|_| {
+                std::thread::sleep(Duration::from_millis(1000));
+                if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
+                    match config.as_str() {
+                        "err1" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
+                        "err2" => Ok(()),
+                        _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
+                    }
+                } else {
+                    Ok(())
+                }
+            }),
+            rollback: Box::new(|_| {
+                unsafe { std::env::set_var("ZED_AUTO_UPDATE_RB", "rollback1") };
+                Ok(())
+            }),
+        },
+        Job::mkdir(p("test1")),
+        Job::mkdir_if_exists(p("test_exists"), p("test1")),
+        Job::mkdir_if_exists(p("test_missing"), p("dont")),
+        Job {
+            apply: Box::new(|folder| {
+                std::fs::write(folder.join("test1/test"), "test")?;
+                Ok(())
+            }),
+            rollback: Box::new(|folder| {
+                std::fs::remove_file(folder.join("test1/test"))?;
+                Ok(())
+            }),
+        },
+        Job::move_file(p("test1/test"), p("test1/moved")),
+        Job::move_if_exists(p("test1/test"), p("test1/noop")),
+        Job {
+            apply: Box::new(|_| {
+                std::thread::sleep(Duration::from_millis(1000));
+                if let Ok(config) = std::env::var("ZED_AUTO_UPDATE") {
+                    match config.as_str() {
+                        "err1" => Ok(()),
+                        "err2" => Err(std::io::Error::other("Simulated error")).context("Anyhow!"),
+                        _ => panic!("Unknown ZED_AUTO_UPDATE value: {}", config),
+                    }
+                } else {
+                    Ok(())
+                }
+            }),
+            rollback: Box::new(|_| Ok(())),
+        },
+        Job::rmdir_nofail(p("test1/nofolder")),
+    ]
+});
 
 pub(crate) fn perform_update(app_dir: &Path, hwnd: Option<isize>, launch: bool) -> Result<()> {
     let hwnd = hwnd.map(|ptr| HWND(ptr as _));
 
-    for job in JOBS.iter() {
+    let mut last_successful_job = None;
+    'outer: for (i, job) in JOBS.iter().enumerate() {
         let start = Instant::now();
         loop {
-            anyhow::ensure!(start.elapsed().as_secs() <= 2, "Timed out");
-            match (*job)(app_dir) {
+            if start.elapsed().as_secs() > 2 {
+                log::error!("Timed out, rolling back");
+                break 'outer;
+            }
+            match (job.apply)(app_dir) {
                 Ok(_) => {
+                    last_successful_job = Some(i);
                     unsafe { PostMessageW(hwnd, WM_JOB_UPDATED, WPARAM(0), LPARAM(0))? };
                     break;
                 }
@@ -223,6 +288,7 @@ pub(crate) fn perform_update(app_dir: &Path, hwnd: Option<isize>, launch: bool)
                     let io_err = err.downcast_ref::<std::io::Error>().unwrap();
                     if io_err.kind() == std::io::ErrorKind::NotFound {
                         log::warn!("File or folder not found.");
+                        last_successful_job = Some(i);
                         unsafe { PostMessageW(hwnd, WM_JOB_UPDATED, WPARAM(0), LPARAM(0))? };
                         break;
                     }
@@ -233,6 +299,28 @@ pub(crate) fn perform_update(app_dir: &Path, hwnd: Option<isize>, launch: bool)
             }
         }
     }
+
+    if last_successful_job
+        .map(|job| job != JOBS.len() - 1)
+        .unwrap_or(true)
+    {
+        let Some(last_successful_job) = last_successful_job else {
+            anyhow::bail!("Autoupdate failed, nothing to rollback");
+        };
+
+        for job in (0..=last_successful_job).rev() {
+            let job = &JOBS[job];
+            if let Err(e) = (job.rollback)(app_dir) {
+                anyhow::bail!(
+                    "Job rollback failed, the app might be left in an inconsistent state: ({:?})",
+                    e
+                );
+            }
+        }
+
+        anyhow::bail!("Autoupdate failed, rollback successful");
+    }
+
     if launch {
         #[allow(clippy::disallowed_methods, reason = "doesn't run in the main binary")]
         let _ = std::process::Command::new(app_dir.join("Zed.exe")).spawn();
@@ -247,12 +335,27 @@ mod test {
 
     #[test]
     fn test_perform_update() {
-        let app_dir = std::path::Path::new("C:/");
+        let app_dir = tempfile::tempdir().unwrap();
+        let app_dir = app_dir.path();
         assert!(perform_update(app_dir, None, false).is_ok());
 
+        let app_dir = tempfile::tempdir().unwrap();
+        let app_dir = app_dir.path();
         // Simulate a timeout
-        unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err") };
+        unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err1") };
         let ret = perform_update(app_dir, None, false);
-        assert!(ret.is_err_and(|e| e.to_string().as_str() == "Timed out"));
+        assert!(
+            ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, nothing to rollback")
+        );
+
+        let app_dir = tempfile::tempdir().unwrap();
+        let app_dir = app_dir.path();
+        // Simulate a timeout
+        unsafe { std::env::set_var("ZED_AUTO_UPDATE", "err2") };
+        let ret = perform_update(app_dir, None, false);
+        assert!(
+            ret.is_err_and(|e| e.to_string().as_str() == "Autoupdate failed, rollback successful")
+        );
+        assert!(std::env::var("ZED_AUTO_UPDATE_RB").is_ok_and(|e| e == "rollback1"));
     }
 }