oma_fetch/
lib.rs

1use std::{borrow::Cow, cmp::Ordering, fmt::Debug, path::PathBuf, sync::Arc, time::Duration};
2
3use ahash::AHashMap;
4use bon::Builder;
5use checksum::Checksum;
6use download::{BuilderError, SingleDownloader, SuccessSummary};
7
8use oma_logger::debug;
9use reqwest::{Method, Response, Url};
10use reqwest_middleware::{ClientWithMiddleware, RequestBuilder};
11use serde::{Deserialize, Serialize};
12use tokio::task::JoinSet;
13
14pub mod checksum;
15pub mod download;
16pub use crate::download::SingleDownloadError;
17
18pub use reqwest;
19
20#[derive(Clone, Default, Builder)]
21pub struct DownloadEntry {
22    pub source: Vec<DownloadSource>,
23    pub filename: String,
24    dir: PathBuf,
25    final_dir: Option<PathBuf>,
26    hash: Option<Checksum>,
27    allow_resume: bool,
28    msg: Option<Cow<'static, str>>,
29    /// Whether this entry downloads the by-hash variant of an index file with
30    /// a fallback to the traditional by-name path of the same file. When
31    /// `true`, a `NextUrl` event means a same-file fallback rather than a
32    /// different mirror.
33    #[builder(default)]
34    by_hash_fallback: bool,
35}
36
37impl Debug for DownloadEntry {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        f.debug_struct("DownloadEntry")
40            .field("source", &self.source)
41            .field("filename", &self.filename)
42            .field("dir", &self.dir)
43            .field("hash", &self.hash.as_ref().map(|c| c.to_string()))
44            .field("allow_resume", &self.allow_resume)
45            .field("msg", &self.msg)
46            .field("by_hash_fallback", &self.by_hash_fallback)
47            .finish()
48    }
49}
50
51#[derive(Debug, Clone, Default, PartialEq, Eq, Copy)]
52pub enum CompressType {
53    Bz2,
54    Gzip,
55    Xz,
56    Zstd,
57    Lzma,
58    Lz4,
59    #[default]
60    None,
61}
62
63impl From<&str> for CompressType {
64    fn from(s: &str) -> Self {
65        match s {
66            "xz" => CompressType::Xz,
67            "gz" => CompressType::Gzip,
68            "bz2" => CompressType::Bz2,
69            "zst" => CompressType::Zstd,
70            _ => CompressType::None,
71        }
72    }
73}
74
75#[derive(Debug, Clone)]
76pub struct DownloadSource {
77    pub url: String,
78    pub source_type: DownloadSourceType,
79    /// The compression format of the file this source points to. Used to
80    /// decompress the downloaded file. This is per-source so that a single
81    /// [`DownloadEntry`] can fall back between different compression formats
82    /// (e.g. `Packages.zst` → `Packages.xz` → `Packages.gz`), matching apt's
83    /// `Acquire::CompressionTypes::Order` behavior.
84    pub file_type: CompressType,
85}
86
87#[derive(Debug, PartialEq, Eq, Clone)]
88pub enum DownloadSourceType {
89    Http,
90    Local(bool),
91}
92
93impl PartialOrd for DownloadSourceType {
94    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
95        Some(self.cmp(other))
96    }
97}
98
99impl Ord for DownloadSourceType {
100    fn cmp(&self, other: &Self) -> Ordering {
101        match self {
102            DownloadSourceType::Http => match other {
103                DownloadSourceType::Http => Ordering::Equal,
104                DownloadSourceType::Local { .. } => Ordering::Less,
105            },
106            DownloadSourceType::Local { .. } => match other {
107                DownloadSourceType::Http => Ordering::Greater,
108                DownloadSourceType::Local { .. } => Ordering::Equal,
109            },
110        }
111    }
112}
113
114#[derive(Debug, Serialize, Deserialize)]
115pub enum Event {
116    ChecksumMismatch {
117        index: usize,
118        filename: String,
119        times: usize,
120    },
121    Timeout {
122        filename: String,
123        times: usize,
124    },
125    GlobalProgressAdd(u64),
126    GlobalProgressSub(u64),
127    ProgressDone(usize),
128    NewProgressSpinner {
129        index: usize,
130        total: usize,
131        msg: String,
132    },
133    NewProgressBar {
134        index: usize,
135        total: usize,
136        msg: String,
137        size: u64,
138    },
139    ProgressInc {
140        index: usize,
141        size: u64,
142    },
143    NextUrl {
144        index: usize,
145        file_name: String,
146        err: SingleDownloadError,
147        /// `true` when the failed source was the by-hash variant and the next
148        /// source is the traditional by-name path of the same file.
149        by_hash: bool,
150    },
151    DownloadDone {
152        index: usize,
153        msg: Box<str>,
154    },
155    Failed {
156        file_name: String,
157        error: SingleDownloadError,
158    },
159    AllDone,
160    NewGlobalProgressBar(u64),
161}
162
163#[derive(Serialize, Deserialize)]
164pub(crate) enum SingleDownloadErrorHelper {
165    SetPermission { source: String },
166    OpenAsWriteMode { source: String },
167    Open { source: String },
168    Create { source: String },
169    Seek { source: String },
170    Write { source: String },
171    Flush { source: String },
172    Remove { source: String },
173    CreateSymlink { source: String },
174    ReqwestMiddlewareError { source: String },
175    BrokenPipe { source: String },
176    SendRequestTimeout,
177    DownloadTimeout,
178    ChecksumMismatch,
179    AcquireError,
180}
181
182#[derive(Builder)]
183pub struct DownloadManager {
184    client: ClientWithMiddleware,
185    download_list: Box<[DownloadEntry]>,
186    #[builder(default = 4)]
187    threads: usize,
188    #[builder(default = 3)]
189    retry_times: usize,
190    #[builder(default)]
191    total_size: u64,
192    #[builder(default = Duration::from_secs(15))]
193    timeout: Duration,
194}
195
196#[derive(Debug)]
197pub struct Summary {
198    pub success: Vec<SuccessSummary>,
199    pub failed: Vec<String>,
200}
201
202impl Summary {
203    pub fn is_download_success(&self) -> bool {
204        self.failed.is_empty()
205    }
206
207    pub fn has_wrote(&self) -> bool {
208        self.success.iter().any(|x| x.wrote)
209    }
210}
211
212impl DownloadManager {
213    /// Start download
214    pub async fn start_download<F, Fut>(mut self, callback: F) -> Result<Summary, BuilderError>
215    where
216        F: Fn(Event) -> Fut + Send + Sync + 'static,
217        Fut: std::future::Future<Output = ()> + Send + 'static,
218    {
219        let mut list = vec![];
220        let len = self.download_list.len();
221
222        let mut source_locks = AHashMap::new();
223
224        for (i, c) in std::mem::take(&mut self.download_list)
225            .into_iter()
226            .enumerate()
227        {
228            let source_key = if let Some(src) = c.source.first() {
229                if let Ok(url) = Url::parse(&src.url) {
230                    format!("{}://{}", url.scheme(), url.host_str().unwrap_or("unknown"))
231                } else {
232                    "fallback_source".to_string()
233                }
234            } else {
235                "unknown_source".to_string()
236            };
237
238            let source_sem = source_locks
239                .entry(source_key)
240                .or_insert_with(|| Arc::new(tokio::sync::Semaphore::new(self.threads)))
241                .clone();
242
243            let single = SingleDownloader::builder()
244                .client(self.client.clone())
245                .download_list_index(i)
246                .entry(c)
247                .total(len)
248                .retry_times(self.retry_times)
249                .timeout(self.timeout)
250                .build()?;
251
252            list.push((single, source_sem));
253        }
254
255        if self.total_size != 0 {
256            callback(Event::NewGlobalProgressBar(self.total_size)).await;
257        }
258
259        let mut set = JoinSet::new();
260        let callback_arc = Arc::new(callback);
261
262        for (single, source_sem) in list {
263            let cb = callback_arc.clone();
264
265            set.spawn(async move {
266                let _permit = match source_sem.acquire_owned().await {
267                    Ok(p) => Some(p),
268                    Err(_) => {
269                        cb(Event::Failed {
270                            file_name: single.entry.filename.to_string(),
271                            error: SingleDownloadError::AcquireError,
272                        })
273                        .await;
274
275                        return download::DownloadResult::Failed {
276                            file_name: single.entry.filename,
277                        };
278                    }
279                };
280
281                single.try_download(&*cb).await
282            });
283        }
284
285        let mut success = vec![];
286        let mut failed = vec![];
287
288        while let Some(res) = set.join_next().await {
289            match res {
290                Ok(download::DownloadResult::Success(success_summary)) => {
291                    success.push(success_summary);
292                }
293                Ok(download::DownloadResult::Failed { file_name }) => {
294                    failed.push(file_name);
295                }
296                Err(_) => {
297                    failed.push("task_panicked".to_string());
298                }
299            }
300        }
301
302        callback_arc(Event::AllDone).await;
303
304        Ok(Summary { success, failed })
305    }
306}
307
308pub async fn send_request_with_url_and_method(
309    url: &str,
310    client: &ClientWithMiddleware,
311    method: Method,
312) -> Result<Response, reqwest_middleware::Error> {
313    let resp = client.request(method, url).send().await?;
314    let headers = resp.headers();
315
316    debug!(
317        "\nDownload URL: {url}\nStatus: {}\nHeaders: {headers:#?}",
318        resp.status()
319    );
320
321    let resp = resp.error_for_status()?;
322
323    Ok(resp)
324}
325
326pub async fn send_request(request: RequestBuilder) -> Result<Response, reqwest_middleware::Error> {
327    let resp = request.send().await?;
328    let headers = resp.headers();
329    let url = resp.url();
330
331    debug!(
332        "\nDownload URL: {url}\nStatus: {}\nHeaders: {headers:#?}",
333        resp.status()
334    );
335
336    let resp = resp.error_for_status()?;
337
338    Ok(resp)
339}
oma_fetch/lib.rs

oma_fetch/
lib.rs
