use std::error::Error; use std::fmt::{Debug, Display, Formatter}; use std::path::{Path, PathBuf}; use futures::{stream, StreamExt, TryStream, TryStreamExt}; use log::debug; use reqwest::{Client, Method, RequestBuilder}; use sha1_smol::{Digest, Sha1}; use tokio::fs; use tokio::fs::File; use tokio::io::{self, AsyncWriteExt}; use crate::util; use crate::util::{FileVerifyError, IntegrityError, USER_AGENT}; pub trait Download: Debug + Display { // return Ok(None) to skip downloading this file async fn prepare(&mut self, client: &Client) -> Result, Box>; async fn handle_chunk(&mut self, chunk: &[u8]) -> Result<(), Box>; async fn finish(&mut self) -> Result<(), Box>; } pub trait FileDownload: Download { fn get_path(&self) -> &Path; } pub struct MultiDownloader<'j, T: Download + 'j, I: Iterator> { jobs: I, nconcurrent: usize } #[derive(Debug, Clone, Copy)] pub enum Phase { Prepare, Send, Receive, HandleChunk, Finish } impl Display for Phase { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { /* an error occurred while (present participle) ... */ Self::Prepare => f.write_str("preparing the request"), Self::Send => f.write_str("sending the request"), Self::Receive => f.write_str("receiving response data"), Self::HandleChunk => f.write_str("handling response data"), Self::Finish => f.write_str("finishing the request"), } } } pub struct PhaseDownloadError<'j, T: Download> { phase: Phase, inner: Box, job: &'j T } impl Debug for PhaseDownloadError<'_, T> { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("PhaseDownloadError") .field("phase", &self.phase) .field("inner", &self.inner) .field("job", &self.job) .finish() } } impl Display for PhaseDownloadError<'_, T> { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "error while {} ({}): {}", self.phase, self.job, self.inner) } } impl Error for PhaseDownloadError<'_, T> { fn source(&self) -> Option<&(dyn Error + 'static)> { Some(&*self.inner) } } impl<'j, T: Download> PhaseDownloadError<'j, T> { fn new(phase: Phase, inner: Box, job: &'j T) -> Self { PhaseDownloadError { phase, inner, job } } } impl<'j, T: Download + 'j, I: Iterator> MultiDownloader<'j, T, I> { pub fn new(jobs: I) -> MultiDownloader<'j, T, I> { Self::with_concurrent(jobs, 24) } pub fn with_concurrent(jobs: I, n: usize) -> MultiDownloader<'j, T, I> { assert!(n > 0); MultiDownloader { jobs, nconcurrent: n } } pub async fn perform(self, client: &'j Client) -> impl TryStream> { stream::iter(self.jobs).map(move |job| Ok(async move { macro_rules! map_err { ($result:expr, $phase:expr, $job:expr) => { match $result { Ok(v) => v, Err(e) => return Err(PhaseDownloadError::new($phase, e.into(), $job)) } } } let Some(rq) = map_err!(job.prepare(client).await, Phase::Prepare, job) else { return Ok(()) }; let rq = rq.header(reqwest::header::USER_AGENT, USER_AGENT); let mut data = map_err!(map_err!(rq.send().await, Phase::Send, job).error_for_status(), Phase::Send, job).bytes_stream(); while let Some(bytes) = data.next().await { let bytes = map_err!(bytes, Phase::Receive, job); map_err!(job.handle_chunk(bytes.as_ref()).await, Phase::HandleChunk, job); } job.finish().await.map_err(|e| PhaseDownloadError::new(Phase::Finish, e, job))?; Ok(()) })).try_buffer_unordered(self.nconcurrent) } } pub struct VerifiedDownload { url: String, expect_size: Option, expect_sha1: Option, path: PathBuf, file: Option, sha1: Sha1, tally: usize } impl Debug for VerifiedDownload { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("VerifiedDownload") .field("url", &self.url) .field("expect_size", &self.expect_size) .field("expect_sha1", &self.expect_sha1) .field("path", &self.path).finish() } } impl Display for VerifiedDownload { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "downloading {} to {}", self.url, self.path.display()) } } impl VerifiedDownload { pub fn new(url: &str, path: &Path, expect_size: Option, expect_sha1: Option) -> VerifiedDownload { VerifiedDownload { url: url.to_owned(), path: path.to_owned(), expect_size, expect_sha1, file: None, sha1: Sha1::new(), tally: 0 } } pub fn with_size(mut self, expect: usize) -> VerifiedDownload { self.expect_size = Some(expect); self } pub fn with_sha1(mut self, expect: Digest) -> VerifiedDownload { self.expect_sha1.replace(expect); self } pub fn get_url(&self) -> &str { &self.url } pub fn get_expect_size(&self) -> Option { self.expect_size } pub fn get_expect_sha1(&self) -> Option { self.expect_sha1 } pub async fn make_dirs(&self) -> Result<(), io::Error> { fs::create_dir_all(self.path.parent().expect("download created with no containing directory (?)")).await } async fn open_output(&mut self) -> Result<(), io::Error> { self.file.replace(File::create(&self.path).await?); Ok(()) } } impl Download for VerifiedDownload { async fn prepare(&mut self, client: &Client) -> Result, Box> { if !util::should_download(&self.path, self.expect_size, self.expect_sha1).await? { return Ok(None) } // potentially racy to close the file and reopen it... :/ self.open_output().await?; Ok(Some(client.request(Method::GET, &self.url))) } async fn handle_chunk(&mut self, chunk: &[u8]) -> Result<(), Box> { self.file.as_mut().unwrap().write_all(chunk).await?; self.tally += chunk.len(); self.sha1.update(chunk); Ok(()) } async fn finish(&mut self) -> Result<(), Box> { let digest = self.sha1.digest(); if let Some(d) = self.expect_sha1 { if d != digest { debug!("Could not download {}: sha1 mismatch (exp {}, got {}).", self.path.display(), d, digest); return Err(IntegrityError::Sha1Mismatch { expect: d, actual: digest }.into()); } } else if let Some(s) = self.expect_size { if s != self.tally { debug!("Could not download {}: size mismatch (exp {}, got {}).", self.path.display(), s, self.tally); return Err(IntegrityError::SizeMismatch { expect: s, actual: self.tally }.into()); } } debug!("Successfully downloaded {} ({} bytes).", self.path.display(), self.tally); // release the file descriptor (don't want to wait until it's dropped automatically because idk when that would be) drop(self.file.take().unwrap()); Ok(()) } } impl FileDownload for VerifiedDownload { fn get_path(&self) -> &Path { &self.path } } pub async fn verify_files(files: impl Iterator) -> Result<(), FileVerifyError> { stream::iter(files) .map(|dl| Ok(async move { debug!("Verifying library {}", dl.get_path().display()); util::verify_file(dl.get_path(), dl.get_expect_size(), dl.get_expect_sha1()).await })) .try_buffer_unordered(32) .try_fold((), |_, _| async {Ok(())}) .await }