cargo_source/
lib.rs

1use std::{fs::{self, File}, io::Write, error::Error};
2use toml_edit::{Document, Item, Table, value};
3use ansi_term::Colour::{Green, Red};
4use clap::{arg, Command};
5use std::time::Instant;
6use reqwest;
7use std::thread;
8
9pub type R<T> = Result<T, Box<dyn Error>>;
10
11pub const REGISTRIES: [(&str, &str, &str); 5] = [
12    ("rsproxy", "https://rsproxy.cn/crates.io-index", "字节"),
13    ("ustc", "git://mirrors.ustc.edu.cn/crates.io-index", "中国科学技术大学"),
14    ("sjtu", "https://mirrors.sjtug.sjtu.edu.cn/git/crates.io-index/", "上海交通大学"),
15    ("tuna", "https://mirrors.tuna.tsinghua.edu.cn/git/crates.io-index.git", "清华大学"),
16    ("rustcc", "https://code.aliyun.com/rustcc/crates.io-index.git", "rustcc社区"),
17];
18
19const TMPL: &str = r#"
20[source.crates-io]
21registry = "https://github.com/rust-lang/crates.io-index"
22"#;
23pub fn pad_end(input: &str, total_length: usize, padding_char: char) -> String {
24    let input_length = input.chars().count();
25    if input_length >= total_length {
26        input.to_string()
27    } else {
28        let padding_length = total_length - input_length;
29        let padding = padding_char.to_string().repeat(padding_length);
30        input.to_string() + &padding
31    }
32}
33
34pub struct CargoConfig {
35    pub document: Document,
36    pub registries: Vec<(String, String)>,
37    path: String
38}
39
40impl CargoConfig {
41    pub fn new() -> R<Self> {
42        let mut path = CargoConfig::get_config().unwrap_or("".to_string());
43        let mut toml_str = match path.len() {
44            0 => {
45                path = CargoConfig::gen_config_dir()?;
46                TMPL.to_string()
47            },
48            _ => fs::read_to_string(&path)?
49        };
50
51        if toml_str.len() == 0 { toml_str = TMPL.to_string() }
52        let doc = toml_str.parse::<Document>()?;
53        let mut registries: Vec<(String, String)> = Vec::new();
54        if let Some(source) = doc.as_table().get("source") {
55            source.as_table().unwrap().iter().for_each(|(key, val)| {
56                if let Some(v) = val.get("registry") {
57                    let v = v.as_value().unwrap().as_str().unwrap().to_string();
58                    registries.push((key.to_string(),  v));
59                }
60            });
61        }
62        Ok(
63            CargoConfig {
64                document: doc,
65                registries,
66                path
67            }
68        )
69    }
70
71    pub fn check_registry(&mut self, registry: &str) {
72        if registry == "crates-io" {
73            let doc = &mut self.document;
74            let crates_io = doc["source"]["crates-io"].as_table_mut().unwrap();
75            crates_io.remove("replace-with");
76            self.write_to_file().unwrap();
77            println!("Registry has been replaced with {}", registry);
78            return
79        }
80        let in_local_config = self.registries.iter().any(|item| item.0 == registry);
81        if in_local_config {
82            let doc = &mut self.document;
83            doc["source"]["crates-io"]["replace-with"] = value(registry);
84            self.write_to_file().unwrap();
85            println!("Registry has been replaced with {}", registry);
86            return
87        } 
88
89        let mut url = String::from("");
90        let in_recommendation_list = REGISTRIES.into_iter().any(|item| {
91            if item.0 == registry {
92                url.push_str(item.1);
93                true
94            } else {
95                false
96            }
97        });
98        if in_recommendation_list {
99            self.insert_registry(registry, &url);
100            let doc = &mut self.document;
101            doc["source"]["crates-io"]["replace-with"] = value(registry);
102            self.write_to_file().unwrap();
103            println!("Registry has been replaced with {}", registry);
104            return
105        } 
106        println!("there is no any registry named {} in recommendation list.", registry);
107    }
108
109    pub fn write_to_file(&self) -> R<()>{
110        let updated_toml = self.document.to_string();
111        let mut file = File::create(&self.path)?;
112        file.write_all(updated_toml.as_bytes())?;
113        Ok(())
114    }
115
116    pub fn insert_registry(&mut self, name: &str, url: &str) {
117        let mut new_table = Table::new();
118        new_table["registry"] = value(url);
119        self.document["source"][name] = Item::Table(new_table);
120    }
121    
122    pub fn get_config() -> R<String> {
123        let dir = dirs::home_dir().ok_or("No home directory")?;
124        let mut dir = dir.to_str().unwrap().to_string();
125        dir.push_str("/.cargo/");
126        let mut entries = fs::read_dir(&dir)?;
127        let exist = entries.any(|entry| {
128            let file_name = entry.unwrap().file_name();
129            if file_name.to_str().unwrap().contains("config") {
130                dir.push_str(file_name.to_str().unwrap());
131                true
132            } else {
133                false
134            }
135        });
136        if exist { Ok(dir) } else { Ok("".to_string())}
137    }
138
139    pub fn gen_config_dir() -> R<String> {
140        let dir = dirs::home_dir().ok_or("No home directory")?;
141        let mut dir = dir.to_str().unwrap().to_string();
142        dir.push_str("/.cargo/config");
143        Ok(dir)
144    }
145
146}
147
148pub struct Cli {
149    pub command: Command,
150}
151
152impl Default for Cli {
153    fn default() -> Self {
154        let command = Command::new("cargo-source")
155            .version(env!("CARGO_PKG_VERSION"))
156            .about("Crates's registry manager")
157            .arg_required_else_help(true)
158            .subcommand_required(true)
159            .subcommands([
160                Command::new("list").about("List all the registries").alias("ls"),
161                Command::new("use")
162                    .about("Change registry to registry")
163                    .arg(arg!(<registry> "registry").required(true)),
164                Command::new("add")
165                    .about("Add one custom registry")
166                    .arg(arg!(<name> "Name of registry").required(true))
167                    .arg(arg!(<url> "Url of registry").required(true)),
168                Command::new("test").about("Test the spead of all the registries").alias("t"),
169            ]);
170            // .after_help(
171            //     "If you find 【cargo-source】 is useful, or you are a experienced Rust developer, or you have the interest in the project, then welcome to submit PRs and help maintain 【cargo-source】. \n \
172            //     ",
173            // );
174        Self { command }
175    }
176}
177
178impl Cli {
179    pub fn run_command(&mut self, args: Vec<String>) -> R<()> {
180        match self.command.try_get_matches_from_mut(args)?.subcommand() {
181            Some(("list", _)) => self.ls()?,
182            Some(("use", sub_m)) => {
183                if let Some(c) = sub_m.get_one::<String>("registry") {
184                    self.switch(c);
185                }
186            }
187            Some(("add", sub_m)) => {
188                let registry_name = sub_m.get_one::<String>("name").unwrap();
189                let registry_url = sub_m.get_one::<String>("url").unwrap();
190                self.insert_registry(registry_name, registry_url);
191            }
192            Some(("test", _)) => self.test(),
193            _ => (),
194        }
195        Ok(())
196    }
197    fn ls(&self) -> R<()> {
198        let cargo_config = CargoConfig::new().unwrap();
199        println!("Recommended registries:");
200        for (tag, url, desc) in REGISTRIES {
201            println!(
202                "  {} | {} | {} ",
203                Green.paint(pad_end(tag, 10, ' ')),
204                url,
205                desc
206            )
207        }
208
209        let default_registry = value("crates-io");
210        let replace_with = cargo_config
211            .document
212            .as_table()
213            .get("source")
214            .ok_or("No source config")?
215            .get("crates-io")
216            .ok_or("No crates-io config")?
217            .get("replace-with")
218            .unwrap_or(&default_registry)
219            .as_str()
220            .unwrap();
221        if cargo_config.registries.len() > 0 {
222            println!("\n-------------------------------------------\nLocal config registries:");
223        }
224        cargo_config.registries.iter()
225            .for_each(|(name, registry)| {
226                let tag = if name == replace_with {
227                    format!("* {}", name)
228                } else {
229                    format!("  {}", name)
230                };
231               
232                println!(
233                    "{}\t| {}",
234                    pad_end(&tag, 12, ' '),
235                    registry.trim_matches('"')
236                )
237            }
238        );
239        println!("\n 说明:*表示当前使用的源,可以通过cargo source use xxxx 来切换源");
240        Ok(())
241    }
242
243    fn switch(&self, registry: &str) {
244        let cargo_config = CargoConfig::new();
245        match cargo_config {
246            Ok(mut result) => result.check_registry(registry),
247            Err(err) => println!("{}", err),
248        }
249    }
250
251    fn insert_registry(&self, name: &str, url: &str) {
252        let cargo_config = CargoConfig::new();
253        match cargo_config {
254            Ok(mut result) => {
255                result.insert_registry(name, url);
256                if let Err(err) = result.write_to_file() {
257                    println!("{err}")
258                }
259            }
260            Err(err) => println!("{}", err),
261        }
262        println!("{name}, {url}")
263    }
264
265    fn test(&self) {
266        REGISTRIES.into_iter().map(|item| {
267            thread::spawn(move || -> String {
268                let start_time = Instant::now();
269                let url = if item.1.starts_with("git") {
270                    let addr = item.1.split("://").collect::<Vec<&str>>()[1];
271                    format!("http://{}", addr)
272                } else {
273                    item.1.to_string()
274                };
275                let response = reqwest::blocking::get(url);
276                let duration = start_time.elapsed();
277                // Ok(format!("{}\t| {} | {:?}", item.0, item.1, duration))
278                match response {
279                    Ok(_) => format!("{}\t| {} | {:?}", item.0, item.1, duration),
280                    Err(_) => format!("{}\t| {} | {:?}", item.0, item.1, Red.paint("Failed"))
281                }
282            })
283        })
284        .for_each(|handle| { 
285            let result = handle.join().unwrap();
286            println!("{}", result)
287        })
288    }
289}
290
291