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::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        Command::new("tar")
138            .current_dir(&root_dir)
139            .args(["-xzf", &toolchain_asset_name, "-C", &toolchain_dir.to_string_lossy()])
140            .status()?;
141
142        // Move the toolchain to a randomly named directory in the 'toolchains' folder
143        let toolchains_dir = root_dir.join("toolchains");
144        fs::create_dir_all(&toolchains_dir)?;
145        let random_string: String =
146            rand::thread_rng().sample_iter(&Alphanumeric).take(10).map(char::from).collect();
147        let new_toolchain_dir = toolchains_dir.join(random_string);
148        fs::rename(&toolchain_dir, &new_toolchain_dir)?;
149
150        // Link the new toolchain directory to rustup
151        Command::new("rustup")
152            .current_dir(&root_dir)
153            .args([
154                "toolchain",
155                "link",
156                RUSTUP_TOOLCHAIN_NAME,
157                &new_toolchain_dir.to_string_lossy(),
158            ])
159            .status()?;
160        println!("Successfully linked toolchain to rustup.");
161
162        // Ensure permissions.
163        let bin_dir = new_toolchain_dir.join("bin");
164        let rustlib_bin_dir = new_toolchain_dir.join(format!("lib/rustlib/{target}/bin"));
165        for entry in fs::read_dir(bin_dir)?.chain(fs::read_dir(rustlib_bin_dir)?) {
166            let entry = entry?;
167            if entry.path().is_file() {
168                let mut perms = entry.metadata()?.permissions();
169                perms.set_mode(0o755);
170                fs::set_permissions(entry.path(), perms)?;
171            }
172        }
173
174        Ok(())
175    }
176}
177
178pub async fn download_file(
179    client: &Client,
180    url: &str,
181    file: &mut (impl tokio::io::AsyncWrite + Unpin),
182) -> std::result::Result<(), String> {
183    use futures::StreamExt;
184    use tokio::io::AsyncWriteExt;
185
186    let mut headers = reqwest::header::HeaderMap::new();
187    headers.insert(reqwest::header::ACCEPT, "application/octet-stream".parse().unwrap());
188    let res = send_with_retry(client, reqwest::Method::GET, url, Some(headers), "Download")
189        .await
190        .map_err(|e| e.to_string())?;
191
192    let total_size =
193        res.content_length().ok_or(format!("Failed to get content length from '{}'", &url))?;
194
195    let pb = ProgressBar::new(total_size);
196    pb.set_style(ProgressStyle::default_bar()
197        .template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})").unwrap()
198        .progress_chars("#>-"));
199
200    let mut downloaded: u64 = 0;
201    let mut stream = res.bytes_stream();
202    while let Some(item) = stream.next().await {
203        let chunk = item.or(Err("Error while downloading file"))?;
204        file.write_all(&chunk).await.or(Err("Error while writing to file"))?;
205        let new = (downloaded + (chunk.len() as u64)).min(total_size);
206        downloaded = new;
207        pb.set_position(new);
208    }
209    pb.finish();
210
211    Ok(())
212}