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

src/cli/proxy_mode.rs

@@ -14,7 +14,7 @@ pub async fn main(
     current_dir: PathBuf,
     process: &Process,
 ) -> anyhow::Result<ExitStatus> {
-    self_update::cleanup_self_updater(&process.cargo_home()?.join("bin"))?;
+    self_update::cleanup_self_updater(process, &process.cargo_home()?.join("bin"))?;
 
     let _setup = job::setup();
     let mut args = process.args_os().skip(1);

src/cli/rustup_mode.rs

@@ -703,7 +703,7 @@ pub async fn main(
         .bin("rustup")
         .complete();
 
-    self_update::cleanup_self_updater(&process.cargo_home()?.join("bin"))?;
+    self_update::cleanup_self_updater(process, &process.cargo_home()?.join("bin"))?;
 
     use clap::error::ErrorKind::*;
     let matches = match Rustup::try_parse_from(process.args_os()) {

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`):
 //!
@@ -78,28 +77,34 @@ use crate::{
 #[macro_use]
 mod msg;
 
+mod stage;
+#[cfg(feature = "test")]
+pub use stage::{Marker, SELF_UPDATE_DIRECTORY, updater_path};
+use stage::{PreparedUpdater, SelfUpdateLock};
+
 #[cfg(unix)]
 mod shell;
 
 #[cfg(unix)]
 mod unix;
 #[cfg(unix)]
-use unix::{add_to_path, remove_from_path};
+pub(crate) use unix::self_replace;
 #[cfg(unix)]
-pub(crate) use unix::{run_update, self_replace};
+use unix::{add_to_path, remove_from_path, 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_to_path, add_uninstall_registry_entry, remove_from_path, remove_uninstall_registry_entry,
+    run_update,
 };
-#[cfg(windows)]
-pub(crate) use windows::{run_update, self_replace};
 
 pub(crate) struct InstallOpts<'a> {
     pub default_host_tuple: Option<String>,
@@ -239,7 +244,7 @@ impl InstallOpts<'_> {
         process: &Process,
     ) -> anyhow::Result<()> {
         let cargo_bin = process.cargo_home()?.join("bin");
-        install_bins(&cargo_bin, force_hard_links(process))?;
+        install_bins(process, &cargo_bin, force_hard_links(process))?;
 
         #[cfg(unix)]
         unix::write_env_files(process)?;
@@ -529,10 +534,10 @@ impl SelfUpdateMode {
             SelfUpdatePermission::Permit => {}
         }
 
-        let setup_path = prepare_update(dl_cfg).await?;
+        let prepared_updater = prepare_update(dl_cfg).await?;
 
-        if let Some(setup_path) = &setup_path {
-            return run_update(setup_path, dl_cfg.process);
+        if let Some(prepared_updater) = prepared_updater {
+            return run_update(prepared_updater, dl_cfg.process);
         } else {
             // Try again in case we emitted "tool `{}` is already installed" last time.
             install_proxies(dl_cfg.process)?;
@@ -579,8 +584,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")
@@ -785,19 +788,8 @@ fn warn_if_default_linker_missing(process: &Process) {
     }
 }
 
-fn install_bins(bin_path: &Path, force_hard_links: bool) -> anyhow::Result<()> {
-    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_with_opts(bin_path, force_hard_links)
+fn install_bins(process: &Process, bin_path: &Path, force_hard_links: bool) -> anyhow::Result<()> {
+    SelfUpdateLock::lock(process)?.install_bins(bin_path, force_hard_links)
 }
 
 pub(crate) fn install_proxies(process: &Process) -> anyhow::Result<()> {
@@ -1134,21 +1126,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.
-///
-/// 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.
+/// Downloads the managed updater and runs it in replacement mode.
 ///
-/// 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);
 
@@ -1174,8 +1155,8 @@ 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_updater) => {
+            let Some(version) = get_and_parse_new_rustup_version(&prepared_updater) else {
                 error!("failed to get rustup version");
                 return Ok(ExitCode::FAILURE);
             };
@@ -1185,7 +1166,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_updater, cfg.process);
         }
         None => {
             let _ = common::show_channel_update(
@@ -1226,18 +1207,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<PreparedUpdater>> {
     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::lock(dl_cfg.process)?;
 
     // Get build tuple
     let tuple = TargetTuple::from_build();
@@ -1277,18 +1254,23 @@ pub(crate) async fn prepare_update(dl_cfg: &DownloadCfg<'_>) -> anyhow::Result<O
 
     // Get download path
     let download_url = utils::parse_url(&url)?;
+    let prepared_updater = self_update_lock.prepare_updater()?;
+    let setup_path: &Path = &prepared_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, setup_path)
         .download()
         .await?;
 
     // Mark as executable
-    utils::make_executable(&setup_path)?;
+    utils::make_executable(setup_path)?;
+
+    #[cfg(feature = "test")]
+    dl_cfg.process.checkpoint(CHECKPOINT_SELF_UPDATE_PREPARED);
 
-    Ok(Some(setup_path))
+    Ok(Some(prepared_updater))
 }
 
 async fn get_available_rustup_version(dl_cfg: &DownloadCfg<'_>) -> anyhow::Result<String> {
@@ -1382,16 +1364,16 @@ pub(crate) async fn check_rustup_update(dl_cfg: &DownloadCfg<'_>) -> anyhow::Res
 }
 
 #[tracing::instrument(level = "trace")]
-pub(crate) fn cleanup_self_updater(bin_path: &Path) -> anyhow::Result<()> {
-    let setup = bin_path.join(format!("rustup-init{EXE_SUFFIX}"));
-
-    if setup.exists() {
-        utils::remove_file("setup", &setup)?;
-    }
-
-    Ok(())
+pub(crate) fn cleanup_self_updater(process: &Process, bin_path: &Path) -> anyhow::Result<()> {
+    stage::cleanup(process, bin_path)
 }
 
+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, path::Path};
@@ -1403,7 +1385,7 @@ mod tests {
         dist::{PartialToolchainDesc, Profile},
         for_host,
         process::TestProcess,
-        test::{test_dir, with_rustup_home},
+        test::{Env, test_dir, with_rustup_home},
     };
 
     #[test]
@@ -1490,7 +1472,12 @@ 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");
-        super::install_bins(&cargo_home.join("bin"), false).unwrap();
+        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, &cargo_home.join("bin"), false).unwrap();
         assert!(cargo_home.exists());
     }
 }