fix(self-update): make replacement atomic and race-safe

src/cli/self_update.rs

@@ -13,13 +13,12 @@
 //! * update the PATH in a system-specific way
 //! * run the equivalent of `rustup default stable`
 //!
-//! During upgrade (`rustup self upgrade`):
+//! During upgrade (`rustup self update`):
 //!
-//! * download rustup-init to $CARGO_HOME/bin/rustup-init
-//! * run rustup-init with appropriate flags to indicate
-//!   this is a self-upgrade
-//! * rustup-init copies bins and hardlinks into place. On windows
-//!   this happens *after* the upgrade command exits successfully.
+//! * download rustup-init to a managed path under `$RUSTUP_HOME`
+//! * run the downloaded binary in replacement mode
+//! * atomically replace rustup and update its proxy links. On Windows
+//!   this happens after the update command exits.
 //!
 //! During uninstall (`rustup self uninstall`):
 //!
@@ -77,29 +76,32 @@ use crate::{
 #[macro_use]
 mod msg;
 
+mod stage;
+use stage::{PreparedUpdate, SelfUpdateLock};
+
 #[cfg(unix)]
 mod shell;
 
 #[cfg(unix)]
 mod unix;
 #[cfg(unix)]
-use unix::{do_add_to_path, do_remove_from_path};
+pub(crate) use unix::self_replace;
 #[cfg(unix)]
-pub(crate) use unix::{run_update, self_replace};
+use unix::{do_add_to_path, do_remove_from_path, replace_rustup_binary, run_update};
 
 #[cfg(windows)]
 mod windows;
 #[cfg(windows)]
 pub use windows::complete_windows_uninstall;
+#[cfg(windows)]
+pub(crate) use windows::self_replace;
 #[cfg(all(windows, feature = "test"))]
 pub use windows::{RUSTUP_REGISTRY_TEST_ID, RegistryValueId, USER_PATH, get_path};
 #[cfg(windows)]
 use windows::{
     add_uninstall_registry_entry, do_add_to_path, do_remove_from_path,
-    remove_uninstall_registry_entry,
+    remove_uninstall_registry_entry, replace_rustup_binary, run_update,
 };
-#[cfg(windows)]
-pub(crate) use windows::{run_update, self_replace};
 
 pub(crate) struct InstallOpts<'a> {
     pub default_host_tuple: Option<String>,
@@ -533,10 +535,10 @@ impl SelfUpdateMode {
             SelfUpdatePermission::Permit => {}
         }
 
-        let setup_path = prepare_update(dl_cfg).await?;
+        let prepared_update = prepare_update(dl_cfg).await?;
 
-        if let Some(setup_path) = &setup_path {
-            return run_update(setup_path, dl_cfg.process);
+        if let Some(prepared_update) = prepared_update {
+            return run_update(prepared_update, dl_cfg.process);
         } else {
             // Try again in case we emitted "tool `{}` is already installed" last time.
             install_proxies(dl_cfg.process)?;
@@ -583,8 +585,6 @@ impl fmt::Display for SelfUpdateMode {
     }
 }
 
-static DEFAULT_UPDATE_ROOT: &str = "https://static.rust-lang.org/rustup";
-
 fn update_root(process: &Process) -> String {
     process
         .var("RUSTUP_UPDATE_ROOT")
@@ -773,19 +773,7 @@ fn warn_if_default_linker_missing(process: &Process) {
 }
 
 fn install_bins(process: &Process) -> anyhow::Result<()> {
-    let bin_path = process.cargo_home()?.join("bin");
-    let this_exe_path = utils::current_exe()?;
-    let rustup_path = bin_path.join(format!("rustup{EXE_SUFFIX}"));
-
-    utils::ensure_dir_exists("bin", &bin_path)?;
-    // NB: Even on Linux we can't just copy the new binary over the (running)
-    // old binary; we must unlink it first.
-    if rustup_path.exists() {
-        utils::remove_file("rustup-bin", &rustup_path)?;
-    }
-    utils::copy_file_symlink_to_source(&this_exe_path, &rustup_path)?;
-    utils::make_executable(&rustup_path)?;
-    install_proxies(process)
+    SelfUpdateLock::acquire(process)?.install_bins(process)
 }
 
 pub(crate) fn install_proxies(process: &Process) -> anyhow::Result<()> {
@@ -1122,21 +1110,10 @@ pub(crate) fn self_update_permitted(explicit: bool) -> anyhow::Result<SelfUpdate
     Ok(SelfUpdatePermission::Permit)
 }
 
-/// Self update downloads rustup-init to `$CARGO_HOME/bin/rustup-init`
-/// and runs it.
+/// Downloads the managed updater and runs it in replacement mode.
 ///
-/// It does a few things to accommodate self-delete problems on windows:
-///
-/// rustup-init is run in two stages, first with `--self-upgrade`,
-/// which displays update messages and asks for confirmations, etc;
-/// then with `--self-replace`, which replaces the rustup binary and
-/// hardlinks. The last step is done without waiting for confirmation
-/// on windows so that the running exe can be deleted.
-///
-/// Because it's again difficult for rustup-init to delete itself
-/// (and on windows this process will not be running to do it),
-/// rustup-init is stored in `$CARGO_HOME/bin`, and then deleted next
-/// time rustup runs.
+/// The updater is removed by a later rustup invocation because Windows
+/// cannot delete the updater while its process is still running.
 pub(crate) async fn update(cfg: &Cfg<'_>) -> anyhow::Result<ExitCode> {
     common::warn_if_host_is_emulated(cfg.process);
 
@@ -1162,8 +1139,9 @@ pub(crate) async fn update(cfg: &Cfg<'_>) -> anyhow::Result<ExitCode> {
     }
 
     match prepare_update(&DownloadCfg::new(cfg)).await? {
-        Some(setup_path) => {
-            let Some(version) = get_and_parse_new_rustup_version(&setup_path) else {
+        Some(prepared_update) => {
+            let Some(version) = get_and_parse_new_rustup_version(prepared_update.updater_path())
+            else {
                 error!("failed to get rustup version");
                 return Ok(ExitCode::FAILURE);
             };
@@ -1173,7 +1151,7 @@ pub(crate) async fn update(cfg: &Cfg<'_>) -> anyhow::Result<ExitCode> {
                 PackageUpdate::Rustup,
                 Ok(UpdateStatus::Updated(version)),
             );
-            return run_update(&setup_path, cfg.process);
+            return run_update(prepared_update, cfg.process);
         }
         None => {
             let _ = common::show_channel_update(
@@ -1214,18 +1192,14 @@ fn parse_new_rustup_version(version: String) -> String {
     String::from(matched_version)
 }
 
-pub(crate) async fn prepare_update(dl_cfg: &DownloadCfg<'_>) -> anyhow::Result<Option<PathBuf>> {
+async fn prepare_update(dl_cfg: &DownloadCfg<'_>) -> anyhow::Result<Option<PreparedUpdate>> {
     let cargo_home = dl_cfg.process.cargo_home()?;
     let rustup_path = cargo_home.join(format!("bin{MAIN_SEPARATOR}rustup{EXE_SUFFIX}"));
-    let setup_path = cargo_home.join(format!("bin{MAIN_SEPARATOR}rustup-init{EXE_SUFFIX}"));
 
     if !rustup_path.exists() {
         return Err(CliError::NotSelfInstalled { p: cargo_home }.into());
     }
-
-    if setup_path.exists() {
-        utils::remove_file("setup", &setup_path)?;
-    }
+    let self_update_lock = SelfUpdateLock::acquire(dl_cfg.process)?;
 
     // Get build tuple
     let tuple = TargetTuple::from_build();
@@ -1265,18 +1239,22 @@ pub(crate) async fn prepare_update(dl_cfg: &DownloadCfg<'_>) -> anyhow::Result<O
 
     // Get download path
     let download_url = utils::parse_url(&url)?;
+    let prepared_update = self_update_lock.prepare_updater()?;
 
     // Download new version
     info!("downloading self-update (new version: {available_version})");
     DownloadOptions::try_from(dl_cfg.process)?
-        .start(&download_url, &setup_path)
+        .start(&download_url, prepared_update.updater_path())
         .download()
         .await?;
 
     // Mark as executable
-    utils::make_executable(&setup_path)?;
+    utils::make_executable(prepared_update.updater_path())?;
 
-    Ok(Some(setup_path))
+    #[cfg(feature = "test")]
+    dl_cfg.process.checkpoint(CHECKPOINT_SELF_UPDATE_PREPARED);
+
+    Ok(Some(prepared_update))
 }
 
 async fn get_available_rustup_version(dl_cfg: &DownloadCfg<'_>) -> anyhow::Result<String> {
@@ -1371,19 +1349,18 @@ pub(crate) async fn check_rustup_update(dl_cfg: &DownloadCfg<'_>) -> anyhow::Res
 
 #[tracing::instrument(level = "trace")]
 pub(crate) fn cleanup_self_updater(process: &Process) -> anyhow::Result<()> {
-    let cargo_home = process.cargo_home()?;
-    let setup = cargo_home.join(format!("bin/rustup-init{EXE_SUFFIX}"));
-
-    if setup.exists() {
-        utils::remove_file("setup", &setup)?;
-    }
-
-    Ok(())
+    stage::cleanup(process)
 }
 
+static DEFAULT_UPDATE_ROOT: &str = "https://static.rust-lang.org/rustup";
+#[cfg(feature = "test")]
+pub const CHECKPOINT_SELF_UPDATE_PREPARED: &str = "self-update-prepared";
+#[cfg(feature = "test")]
+pub const CHECKPOINT_SELF_REPLACE_READY: &str = "self-replace-ready";
+
 #[cfg(test)]
 mod tests {
-    use std::collections::HashMap;
+    use std::{collections::HashMap, env::consts::EXE_SUFFIX, fs};
 
     use crate::{
         cli::self_update::InstallOpts,
@@ -1436,10 +1413,37 @@ info: default host tuple is {0}
     fn install_bins_creates_cargo_home() {
         let root_dir = test_dir().unwrap();
         let cargo_home = root_dir.path().join("cargo");
+        let rustup_home = root_dir.path().join("rustup");
         let mut vars = HashMap::new();
         vars.env("CARGO_HOME", cargo_home.to_string_lossy().to_string());
+        vars.env("RUSTUP_HOME", rustup_home);
         let tp = TestProcess::with_vars(vars);
         super::install_bins(&tp.process).unwrap();
         assert!(cargo_home.exists());
     }
+
+    #[test]
+    fn failed_atomic_replace_preserves_existing_rustup() {
+        let root_dir = test_dir().unwrap();
+        let rustup = root_dir.path().join(format!("rustup{EXE_SUFFIX}"));
+        fs::write(&rustup, "old rustup").unwrap();
+
+        super::replace_rustup_binary(&root_dir.path().join("missing"), &rustup).unwrap_err();
+
+        assert_eq!(fs::read_to_string(rustup).unwrap(), "old rustup");
+    }
+
+    #[test]
+    fn atomic_replace_publishes_pending_rustup() {
+        let root_dir = test_dir().unwrap();
+        let rustup = root_dir.path().join(format!("rustup{EXE_SUFFIX}"));
+        let pending = root_dir.path().join("pending");
+        fs::write(&rustup, "old rustup").unwrap();
+        fs::write(&pending, "new rustup").unwrap();
+
+        super::replace_rustup_binary(&pending, &rustup).unwrap();
+
+        assert_eq!(fs::read_to_string(rustup).unwrap(), "new rustup");
+        assert!(!pending.exists());
+    }
 }