1use futures::stream::{self, StreamExt};
2use sha1::{Digest, Sha1};
3use sha2::Sha256;
4use std::path::{Path, PathBuf};
5use tokio::fs;
6use tokio::io::AsyncWriteExt;
7
8use crate::error::{HexoError, Result};
9
10const MAX_RETRIES: u32 = 3;
11const ASSETS_URL: &str = "https://resources.download.minecraft.net";
12
13#[derive(Debug, Clone)]
14pub struct DownloadTask {
15 pub url: String,
16 pub path: PathBuf,
17 pub sha1: Option<String>,
18 pub sha256: Option<String>,
19}
20
21impl DownloadTask {
22 pub fn new(url: impl Into<String>, path: impl Into<PathBuf>) -> Self {
23 Self {
24 url: url.into(),
25 path: path.into(),
26 sha1: None,
27 sha256: None,
28 }
29 }
30
31 pub fn with_sha1(mut self, sha1: impl Into<String>) -> Self {
32 self.sha1 = Some(sha1.into());
33 self
34 }
35
36 pub fn with_sha256(mut self, sha256: impl Into<String>) -> Self {
38 self.sha256 = Some(sha256.into());
39 self
40 }
41
42 pub fn asset(hash: &str, assets_dir: &Path) -> Self {
44 let prefix = &hash[..2];
45 let url = format!("{}/{}/{}", ASSETS_URL, prefix, hash);
46 let path = assets_dir.join("objects").join(prefix).join(hash);
47 Self::new(url, path).with_sha1(hash.to_string())
48 }
49}
50
51pub async fn verify_sha1(path: &Path, expected: &str) -> bool {
53 let Ok(data) = fs::read(path).await else {
54 return false;
55 };
56 let mut hasher = Sha1::new();
57 hasher.update(&data);
58 let result = hex::encode(hasher.finalize());
59 result.eq_ignore_ascii_case(expected)
60}
61
62pub async fn verify_sha256(path: &Path, expected: &str) -> bool {
64 let Ok(data) = fs::read(path).await else {
65 return false;
66 };
67 hex::encode(Sha256::digest(&data)).eq_ignore_ascii_case(expected)
68}
69
70async fn verify_task(task: &DownloadTask) -> bool {
72 if let Some(expected) = &task.sha1 {
73 if !verify_sha1(&task.path, expected).await {
74 return false;
75 }
76 }
77 if let Some(expected) = &task.sha256 {
78 if !verify_sha256(&task.path, expected).await {
79 return false;
80 }
81 }
82 true
83}
84
85pub async fn download_file(task: &DownloadTask) -> Result<()> {
87 download_file_with_progress(task, |_, _| {}).await
88}
89
90pub async fn download_file_with_progress<F>(task: &DownloadTask, progress: F) -> Result<()>
94where
95 F: Fn(usize, usize) + Send + Sync,
96{
97 if task.path.exists() && verify_task(task).await {
98 let size = fs::metadata(&task.path).await?.len() as usize;
99 progress(size, size);
100 return Ok(());
101 }
102
103 if let Some(parent) = task.path.parent() {
104 fs::create_dir_all(parent).await?;
105 }
106
107 let client = reqwest::Client::new();
108 let mut last_err = None;
109
110 for _ in 0..MAX_RETRIES {
111 match try_download(&client, &task.url, &task.path, &progress).await {
112 Ok(()) => {
113 if !verify_task(task).await {
114 last_err = Some(HexoError::ChecksumMismatch {
115 path: task.path.display().to_string(),
116 });
117 continue;
118 }
119 return Ok(());
120 }
121 Err(e) => {
122 last_err = Some(e);
123 }
124 }
125 }
126
127 Err(last_err.unwrap_or(HexoError::DownloadFailed {
128 url: task.url.clone(),
129 }))
130}
131
132async fn try_download<F>(
133 client: &reqwest::Client,
134 url: &str,
135 path: &Path,
136 progress: &F,
137) -> Result<()>
138where
139 F: Fn(usize, usize) + Send + Sync,
140{
141 let response = client.get(url).send().await?.error_for_status()?;
142 let total = response.content_length().unwrap_or(0) as usize;
143 let mut downloaded = 0usize;
144 progress(0, total);
145 let mut file = fs::File::create(path).await?;
146 let mut stream = response.bytes_stream();
147
148 use futures::StreamExt as _;
149 while let Some(chunk) = stream.next().await {
150 let chunk = chunk?;
151 file.write_all(&chunk).await?;
152 downloaded = downloaded.saturating_add(chunk.len());
153 progress(downloaded, total);
154 }
155 file.flush().await?;
156 Ok(())
157}
158
159pub async fn download_batch<F>(
164 tasks: Vec<DownloadTask>,
165 concurrency: usize,
166 progress: F,
167) -> Result<()>
168where
169 F: Fn(usize, usize) + Send + Sync + 'static,
170{
171 let total = tasks.len();
172 let progress = std::sync::Arc::new(progress);
173 let counter = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
174
175 let results: Vec<Result<()>> = stream::iter(tasks)
176 .map(|task| {
177 let progress = progress.clone();
178 let counter = counter.clone();
179 async move {
180 let result = download_file(&task).await;
181 let done = counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1;
182 progress(done, total);
183 result
184 }
185 })
186 .buffer_unordered(concurrency)
187 .collect()
188 .await;
189
190 for r in results {
191 r?;
192 }
193 Ok(())
194}
195
196#[cfg(test)]
197mod tests {
198 use super::*;
199 use sha1::{Digest, Sha1};
200 use std::io::Write;
201 use tempfile::NamedTempFile;
202
203 const ABC_SHA256: &str = "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad";
204
205 #[tokio::test]
206 async fn byte_progress_reports_retries_and_cached_completion() {
207 let temp = tempfile::tempdir().unwrap();
208 let path = temp.path().join("archive");
209 let (url, server) = serve(vec![(200, "bad"), (200, "abc")]).await;
210 let task = DownloadTask::new(url, &path).with_sha256(ABC_SHA256);
211 let events = std::sync::Mutex::new(Vec::new());
212 tokio::time::timeout(
213 std::time::Duration::from_secs(5),
214 download_file_with_progress(&task, |done, total| {
215 events.lock().unwrap().push((done, total))
216 }),
217 )
218 .await
219 .unwrap()
220 .unwrap();
221 tokio::time::timeout(std::time::Duration::from_secs(5), server)
222 .await
223 .unwrap()
224 .unwrap();
225 {
226 let events = events.lock().unwrap();
227 assert_eq!(events.iter().filter(|event| **event == (0, 3)).count(), 2);
228 assert_eq!(events.last(), Some(&(3, 3)));
229 assert!(events.iter().all(|(done, total)| *total == 3 && *done <= 3));
230 }
231 events.lock().unwrap().clear();
232 download_file_with_progress(&task, |done, total| {
233 events.lock().unwrap().push((done, total))
234 })
235 .await
236 .unwrap();
237 assert_eq!(*events.lock().unwrap(), vec![(3, 3)]);
238 }
239
240 async fn serve(responses: Vec<(u16, &'static str)>) -> (String, tokio::task::JoinHandle<()>) {
241 use tokio::io::AsyncReadExt;
242 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
243 let url = format!("http://{}/archive", listener.local_addr().unwrap());
244 let handle = tokio::spawn(async move {
245 for (status, body) in responses {
246 let (mut socket, _) = listener.accept().await.unwrap();
247 let mut request = Vec::new();
248 while !request.ends_with(b"\r\n\r\n") {
249 request.push(socket.read_u8().await.unwrap());
250 }
251 let response = format!("HTTP/1.1 {status} Test\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len());
252 socket.write_all(response.as_bytes()).await.unwrap();
253 }
254 });
255 (url, handle)
256 }
257
258 #[tokio::test]
259 async fn sha256_checks_missing_corrupt_and_uppercase_hashes() {
260 let temp = tempfile::tempdir().unwrap();
261 let path = temp.path().join("archive");
262 assert!(!verify_sha256(&path, ABC_SHA256).await);
263 fs::write(&path, b"abc").await.unwrap();
264 assert!(verify_sha256(&path, ABC_SHA256).await);
265 assert!(verify_sha256(&path, &ABC_SHA256.to_uppercase()).await);
266 fs::write(&path, b"corrupt").await.unwrap();
267 assert!(!verify_sha256(&path, ABC_SHA256).await);
268 }
269
270 #[tokio::test]
271 async fn sha256_download_retries_and_replaces_invalid_cache() {
272 for cached in [false, true] {
273 let temp = tempfile::tempdir().unwrap();
274 let path = temp.path().join("archive");
275 if cached {
276 fs::write(&path, b"stale").await.unwrap();
277 }
278 let (url, server) = serve(vec![(200, "corrupt"), (200, "abc")]).await;
279 let task = DownloadTask::new(url, &path).with_sha256(ABC_SHA256);
280 tokio::time::timeout(std::time::Duration::from_secs(5), download_file(&task))
281 .await
282 .unwrap()
283 .unwrap();
284 tokio::time::timeout(std::time::Duration::from_secs(5), server)
285 .await
286 .unwrap()
287 .unwrap();
288 assert_eq!(fs::read(&path).await.unwrap(), b"abc");
289 }
290 }
291
292 #[tokio::test]
293 async fn sha256_download_rejects_corruption_after_retries() {
294 let temp = tempfile::tempdir().unwrap();
295 let (url, server) = serve(vec![(200, "corrupt"); MAX_RETRIES as usize]).await;
296 let task = DownloadTask::new(url, temp.path().join("archive")).with_sha256(ABC_SHA256);
297 let result = tokio::time::timeout(std::time::Duration::from_secs(5), download_file(&task))
298 .await
299 .unwrap();
300 assert!(matches!(result, Err(HexoError::ChecksumMismatch { .. })));
301 tokio::time::timeout(std::time::Duration::from_secs(5), server)
302 .await
303 .unwrap()
304 .unwrap();
305 }
306
307 #[tokio::test]
308 async fn all_configured_checksums_must_match() {
309 let temp = tempfile::tempdir().unwrap();
310 let path = temp.path().join("archive");
311 fs::write(&path, b"abc").await.unwrap();
312 let task = DownloadTask::new("http://127.0.0.1:0/archive", &path)
313 .with_sha1(sha1_of(b"abc"))
314 .with_sha256(ABC_SHA256);
315 download_file(&task).await.unwrap();
316 assert!(!verify_task(&task.clone().with_sha1(sha1_of(b"wrong"))).await);
317 assert!(!verify_task(&task.with_sha256("0".repeat(64))).await);
318 }
319
320 #[tokio::test]
321 async fn http_errors_do_not_overwrite_cached_file() {
322 let temp = tempfile::tempdir().unwrap();
323 let path = temp.path().join("archive");
324 fs::write(&path, b"existing").await.unwrap();
325 let (url, server) = serve(vec![(404, "not found"); MAX_RETRIES as usize]).await;
326 let task = DownloadTask::new(url, &path).with_sha256(ABC_SHA256);
327 let result = tokio::time::timeout(std::time::Duration::from_secs(5), download_file(&task))
328 .await
329 .unwrap();
330 assert!(result.is_err());
331 tokio::time::timeout(std::time::Duration::from_secs(5), server)
332 .await
333 .unwrap()
334 .unwrap();
335 assert_eq!(fs::read(&path).await.unwrap(), b"existing");
336 }
337
338 fn sha1_of(data: &[u8]) -> String {
339 let mut hasher = Sha1::new();
340 hasher.update(data);
341 hex::encode(hasher.finalize())
342 }
343
344 #[tokio::test]
345 async fn verify_sha1_correct() {
346 let data = b"hello hexomc-lib";
347 let mut f = NamedTempFile::new().unwrap();
348 f.write_all(data).unwrap();
349 let expected = sha1_of(data);
350 assert!(verify_sha1(f.path(), &expected).await);
351 }
352
353 #[tokio::test]
354 async fn verify_sha1_wrong_hash() {
355 let data = b"hello hexomc-lib";
356 let mut f = NamedTempFile::new().unwrap();
357 f.write_all(data).unwrap();
358 assert!(!verify_sha1(f.path(), "0000000000000000000000000000000000000000").await);
359 }
360
361 #[tokio::test]
362 async fn verify_sha1_missing_file() {
363 assert!(!verify_sha1(std::path::Path::new("/nonexistent/file.bin"), "abc").await);
364 }
365
366 #[tokio::test]
367 async fn download_skips_existing_valid_file() {
368 let data = b"cached content";
369 let mut f = NamedTempFile::new().unwrap();
370 f.write_all(data).unwrap();
371 let sha1 = sha1_of(data);
372
373 let task = DownloadTask::new("http://127.0.0.1:0/nonexistent", f.path()).with_sha1(sha1);
375
376 assert!(download_file(&task).await.is_ok());
377 }
378}