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 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 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