Skip to main content

sp1_cli/commands/
install_toolchain.rs

1use std::{
2    fs::{self},
3    io::Read,
4    process::Command,
5};
6
7use anyhow::{Context, Result};
8use clap::Parser;
9use dirs::home_dir;
10use indicatif::{ProgressBar, ProgressStyle};
11use rand::{distributions::Alphanumeric, Rng};
12use reqwest::Client;
13
14#[cfg(target_family = "unix")]
15use std::os::unix::fs::PermissionsExt;
16
17use crate::{
18    get_target, get_toolchain_asset_url, is_supported_target, send_with_retry,
19    RUSTUP_TOOLCHAIN_NAME,
20};
21
22#[derive(Parser)]
23#[command(name = "install-toolchain", about = "Install the cargo-prove toolchain.")]
24pub struct InstallToolchainCmd {
25    #[arg(short, long, env = "GITHUB_TOKEN")]
26    pub token: Option<String>,
27}
28
29impl InstallToolchainCmd {
30    #[allow(clippy::uninlined_format_args)]
31    pub async fn run(&self) -> Result<()> {
32        // Check if rust is installed.
33        if Command::new("rustup")
34            .arg("--version")
35            .stdout(std::process::Stdio::null())
36            .stderr(std::process::Stdio::null())
37            .status()
38            .is_err()
39        {
40            return Err(anyhow::anyhow!(
41                "Rust is not installed. Please install Rust from https://rustup.rs/ and try again."
42            ));
43        }
44
45        // Setup client with optional token.
46        let client_builder = Client::builder().user_agent("Mozilla/5.0");
47        let client = if let Some(ref token) = self.token {
48            client_builder
49                .default_headers({
50                    let mut headers = reqwest::header::HeaderMap::new();
51                    headers.insert(
52                        reqwest::header::AUTHORIZATION,
53                        reqwest::header::HeaderValue::from_str(&format!("token {token}")).unwrap(),
54                    );
55                    headers
56                })
57                .build()?
58        } else {
59            client_builder.build()?
60        };
61
62        // Setup variables.
63        let root_dir = home_dir().unwrap().join(".sp1");
64        match fs::read_dir(&root_dir) {
65            Ok(entries) =>
66            {
67                #[allow(clippy::manual_flatten)]
68                for entry in entries {
69                    if let Ok(entry) = entry {
70                        let entry_path = entry.path();
71                        let entry_name = entry_path.file_name().unwrap();
72                        if entry_path.is_dir()
73                            && entry_name != "bin"
74                            && entry_name != "circuits"
75                            && entry_name != "toolchains"
76                        {
77                            if let Err(err) = fs::remove_dir_all(&entry_path) {
78                                println!("Failed to remove directory {entry_path:?}: {err}");
79                            }
80                        } else if entry_path.is_file() {
81                            if let Err(err) = fs::remove_file(&entry_path) {
82                                println!("Failed to remove file {entry_path:?}: {err}");
83                            }
84                        }
85                    }
86                }
87            }
88            Err(_) => println!("No existing ~/.sp1 directory to remove."),
89        }
90        println!("Successfully cleaned up ~/.sp1 directory.");
91        match fs::create_dir_all(&root_dir) {
92            Ok(_) => println!("Successfully created ~/.sp1 directory."),
93            Err(err) => println!("Failed to create ~/.sp1 directory: {err}"),
94        };
95
96        assert!(
97            is_supported_target(),
98            "Unsupported architecture. Please build the toolchain from source."
99        );
100        let target = get_target();
101        let toolchain_asset_name = format!("rust-toolchain-{target}.tar.gz");
102        let toolchain_archive_path = root_dir.join(toolchain_asset_name.clone());
103        let toolchain_dir = root_dir.join(&target);
104
105        let toolchain_asset_url = get_toolchain_asset_url(&client, target.to_string()).await?;
106
107        // Download the toolchain via the GitHub API. Using the API asset URL with
108        // Accept: application/octet-stream works correctly with authentication,
109        // unlike browser download URLs which redirect to a CDN that rejects the
110        // Authorization header.
111        let mut file = tokio::fs::File::create(toolchain_archive_path).await.unwrap();
112        download_file(&client, toolchain_asset_url.as_str(), &mut file)
113            .await
114            .map_err(|e| anyhow::anyhow!(e))?;
115
116        // Remove the existing toolchain from rustup, if it exists.
117        let mut child = Command::new("rustup")
118            .current_dir(&root_dir)
119            .args(["toolchain", "remove", RUSTUP_TOOLCHAIN_NAME])
120            .stdout(std::process::Stdio::piped())
121            .spawn()?;
122        let res = child.wait();
123        match res {
124            Ok(_) => {
125                let mut stdout = child.stdout.take().unwrap();
126                let mut content = String::new();
127                stdout.read_to_string(&mut content).unwrap();
128                if !content.contains("no toolchain installed") {
129                    println!("Successfully removed existing toolchain.");
130                }
131            }
132            Err(_) => println!("Failed to remove existing toolchain."),
133        }
134
135        // Unpack the toolchain.
136        fs::create_dir_all(toolchain_dir.clone())?;
137        run_checked(Command::new("tar").current_dir(&root_dir).args([
138            "-xzf",
139            &toolchain_asset_name,
140            "-C",
141            &toolchain_dir.to_string_lossy(),
142        ]))?;
143
144        // Move the toolchain to a randomly named directory in the 'toolchains' folder
145        let toolchains_dir = root_dir.join("toolchains");
146        fs::create_dir_all(&toolchains_dir)?;
147        let random_string: String =
148            rand::thread_rng().sample_iter(&Alphanumeric).take(10).map(char::from).collect();
149        let new_toolchain_dir = toolchains_dir.join(random_string);
150        fs::rename(&toolchain_dir, &new_toolchain_dir)?;
151
152        // Link the new toolchain directory to rustup
153        run_checked(Command::new("rustup").current_dir(&root_dir).args([
154            "toolchain",
155            "link",
156            RUSTUP_TOOLCHAIN_NAME,
157            &new_toolchain_dir.to_string_lossy(),
158        ]))?;
159        println!("Successfully linked toolchain to rustup.");
160
161        // Ensure permissions.
162        let bin_dir = new_toolchain_dir.join("bin");
163        let rustlib_bin_dir = new_toolchain_dir.join(format!("lib/rustlib/{target}/bin"));
164        for entry in fs::read_dir(bin_dir)?.chain(fs::read_dir(rustlib_bin_dir)?) {
165            let entry = entry?;
166            if entry.path().is_file() {
167                let mut perms = entry.metadata()?.permissions();
168                perms.set_mode(0o755);
169                fs::set_permissions(entry.path(), perms)?;
170            }
171        }
172
173        Ok(())
174    }
175}
176
177/// Run `command` to completion, returning an error if it could not be spawned
178/// or if it exited with a non-zero status. `Command::status` only reports a
179/// spawn/wait failure, so a process that runs and exits non-zero must be caught
180/// explicitly here.
181fn run_checked(command: &mut Command) -> Result<()> {
182    let status = command.status().with_context(|| format!("while executing `{command:?}`"))?;
183
184    status.success().then_some(()).with_context(|| format!("`{command:?}` failed: {status}"))
185}
186
187pub async fn download_file(
188    client: &Client,
189    url: &str,
190    file: &mut (impl tokio::io::AsyncWrite + Unpin),
191) -> std::result::Result<(), String> {
192    use futures::StreamExt;
193    use tokio::io::AsyncWriteExt;
194
195    let mut headers = reqwest::header::HeaderMap::new();
196    headers.insert(reqwest::header::ACCEPT, "application/octet-stream".parse().unwrap());
197    let res = send_with_retry(client, reqwest::Method::GET, url, Some(headers), "Download")
198        .await
199        .map_err(|e| e.to_string())?;
200
201    let total_size =
202        res.content_length().ok_or(format!("Failed to get content length from '{}'", url))?;
203
204    let pb = ProgressBar::new(total_size);
205    pb.set_style(ProgressStyle::default_bar()
206        .template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})").unwrap()
207        .progress_chars("#>-"));
208
209    let mut downloaded: u64 = 0;
210    let mut stream = res.bytes_stream();
211    while let Some(item) = stream.next().await {
212        let chunk = item.or(Err("Error while downloading file"))?;
213        file.write_all(&chunk).await.or(Err("Error while writing to file"))?;
214        let new = (downloaded + (chunk.len() as u64)).min(total_size);
215        downloaded = new;
216        pb.set_position(new);
217    }
218    pb.finish();
219
220    Ok(())
221}
222
223#[cfg(all(test, unix))]
224mod tests {
225    use super::*;
226
227    #[test]
228    fn run_checked_errors_on_nonzero_exit() {
229        let err = run_checked(Command::new("sh").args(["-c", "exit 7"])).unwrap_err();
230        assert!(err.to_string().contains("failed"), "unexpected error: {err}");
231    }
232
233    #[test]
234    fn run_checked_succeeds_on_zero_exit() {
235        run_checked(Command::new("sh").args(["-c", "exit 0"])).unwrap();
236    }
237}