From 9e20396e17c7e38ff2d42b270fdae1788356d3b8 Mon Sep 17 00:00:00 2001
From: Martijn Gribnau <garm@ilumeo.com>
Date: Tue, 4 May 2021 18:27:01 +0200
Subject: [PATCH] Allow rustup-init to install the default toolchain from a
 toolchain file

---
 rustup-init.sh        |  1 +
 src/cli/setup_mode.rs | 24 ++++++++++++++++++++++--
 src/config.rs         | 40 ++++++++++++++++++++++++++++++++++++++++
 tests/cli-self-upd.rs | 25 +++++++++++++++++++++++++
 4 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/rustup-init.sh b/rustup-init.sh
index c7622bb2ecb..a012542e791 100755
--- a/rustup-init.sh
+++ b/rustup-init.sh
@@ -43,6 +43,7 @@ OPTIONS:
         --default-host <default-host>              Choose a default host triple
         --default-toolchain <default-toolchain>    Choose a default toolchain to install
         --default-toolchain none                   Do not install any toolchains
+        --from-file <rust-toolchain-file>          Use the default toolchain specified by the rust-toolchain file
         --profile [minimal|default|complete]       Choose a profile
     -c, --component <components>...                Component name to also install
     -t, --target <targets>...                      Target name to also install
diff --git a/src/cli/setup_mode.rs b/src/cli/setup_mode.rs
index f54accbd806..2927372657b 100644
--- a/src/cli/setup_mode.rs
+++ b/src/cli/setup_mode.rs
@@ -4,8 +4,8 @@ use clap::{App, AppSettings, Arg};
 use super::common;
 use super::self_update::{self, InstallOpts};
 use crate::dist::dist::Profile;
-use crate::process;
 use crate::utils::utils;
+use crate::{process, InstallFromToolchainFileCfg};
 
 pub fn main() -> Result<utils::ExitCode> {
     let args: Vec<_> = process().args().collect();
@@ -57,6 +57,12 @@ pub fn main() -> Result<utils::ExitCode> {
                 .takes_value(true)
                 .help("Choose a default toolchain to install"),
         )
+        .arg(
+            Arg::with_name("from-file")
+                .long("from-file")
+                .takes_value(true)
+                .help("Use the default toolchain specified by the rust-toolchain file")
+        )
         .arg(
             Arg::with_name("profile")
                 .long("profile")
@@ -103,11 +109,25 @@ pub fn main() -> Result<utils::ExitCode> {
         }
         Err(e) => return Err(e.into()),
     };
+
+    let toolchain_cfg = match matches
+        .value_of("from-file")
+        .map(InstallFromToolchainFileCfg::try_from_file)
+    {
+        Some(Ok(cfg)) => Some(cfg),
+        Some(Err(e)) => return Err(e.into()),
+        None => None,
+    };
+
     let no_prompt = matches.is_present("no-prompt");
     let verbose = matches.is_present("verbose");
     let quiet = matches.is_present("quiet");
     let default_host = matches.value_of("default-host").map(ToOwned::to_owned);
-    let default_toolchain = matches.value_of("default-toolchain").map(ToOwned::to_owned);
+    let default_toolchain = matches
+        .value_of("default-toolchain")
+        .map(ToOwned::to_owned)
+        .or_else(|| toolchain_cfg.and_then(|cfg| cfg.channel()));
+
     let profile = matches
         .value_of("profile")
         .expect("Unreachable: Clap should supply a default");
diff --git a/src/config.rs b/src/config.rs
index da64657e05d..5570a683826 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::ffi::OsStr;
 use std::fmt::{self, Display};
 use std::io;
 use std::path::{Path, PathBuf};
@@ -997,6 +998,45 @@ enum ParseMode {
     Both,
 }
 
+/// Configuration used with `rustup-init`, to install a default toolchain from a toolchain file.
+#[derive(Debug, Default, PartialEq, Eq)]
+pub struct InstallFromToolchainFileCfg {
+    toolchain_file: OverrideFile,
+}
+
+impl InstallFromToolchainFileCfg {
+    /// Try to create a new `InstallFromToolchainFileCfg`.
+    ///
+    /// If the given toolchain file path (given through the `path` argument) ends in a `toml` file
+    /// extension, only the TOML parse mode will be used.
+    ///
+    /// Fails if the toolchain file does not exist, is not a file, can not be read
+    /// or if the toolchain file contents cannot be parsed.
+    pub fn try_from_file(path: impl AsRef<Path>) -> Result<Self> {
+        let path = path.as_ref();
+        let parse_mode = Self::parse_mode(path);
+        let contents = utils::read_file("toolchain file", path)?;
+
+        let toolchain_file = Cfg::parse_override_file(&contents, parse_mode)?;
+
+        Ok(Self { toolchain_file })
+    }
+
+    /// The `channel` defined by a toolchain file corresponds to the `toolchain` field when installing
+    /// a toolchain with `rustup-init`.
+    pub fn channel(&self) -> Option<String> {
+        self.toolchain_file.toolchain.channel.to_owned()
+    }
+
+    /// Allow only the legacy format, if the file doesn't end in a toml file extension.
+    fn parse_mode(path: &Path) -> ParseMode {
+        path.extension()
+            .filter(|ext| ext == &OsStr::new("toml"))
+            .map(|_| ParseMode::OnlyToml)
+            .unwrap_or(ParseMode::Both)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/tests/cli-self-upd.rs b/tests/cli-self-upd.rs
index 916a37fc8ee..9f323e81d4b 100644
--- a/tests/cli-self-upd.rs
+++ b/tests/cli-self-upd.rs
@@ -914,3 +914,28 @@ fn install_minimal_profile() {
         expect_component_not_executable(config, "cargo");
     });
 }
+
+#[test]
+fn install_from_rust_toolchain_file() {
+    clitools::setup(Scenario::SimpleV2, &|config| {
+        let cwd = config.current_dir();
+        let toolchain_file = cwd.join("rust-toolchain.toml");
+
+        raw::write_file(
+            &toolchain_file,
+            r#"[toolchain]
+channel = "1.1.0""#,
+        )
+        .unwrap();
+
+        expect_ok(
+            config,
+            &["rustup-init", "-y", "--from-file", "rust-toolchain.toml"],
+        );
+        expect_stdout_ok(
+            config,
+            &["rustup", "toolchain", "list"],
+            &format!("1.1.0-{} (default)", this_host_triple()),
+        );
+    })
+}