use anyhow::{anyhow, Context, Result};
use futures_util::StreamExt;
use sha2::{Digest, Sha256};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
use crate::downloader::events::{DownloaderError, FailureKind, ProgressEvent};
use crate::downloader::resolve::CandidateUrl;
#[derive(Debug)]
pub struct FetchedArchive {
pub temp_dir: tempfile::TempDir,
pub archive_path: PathBuf,
pub winning_url: String,
}
pub async fn fetch(
client: &reqwest::Client,
name: &str,
candidates: &[CandidateUrl],
events: &tokio::sync::mpsc::UnboundedSender<ProgressEvent>,
cancel: Arc<AtomicBool>,
) -> Result<FetchedArchive, DownloaderError> {
if cancel.load(Ordering::SeqCst) {
return Err(DownloaderError::Cancelled);
}
let winner = head_probe_concurrent(client, candidates, name, events, &cancel).await?;
let temp_dir = tempfile::tempdir().map_err(|e| DownloaderError::Failed {
kind: FailureKind::DownloadInterrupted,
source: anyhow!(e).context("mkdir tempdir for download"),
})?;
let filename = winner
.url
.rsplit('/')
.next()
.unwrap_or("download.bin")
.to_string();
let archive_path = temp_dir.path().join(&filename);
let resp = client.get(&winner.url).send().await.map_err(|e| {
DownloaderError::Failed {
kind: FailureKind::DownloadInterrupted,
source: anyhow!(e).context("GET archive"),
}
})?;
let total = resp.content_length();
let mut stream = resp.bytes_stream();
let mut file = tokio::fs::File::create(&archive_path)
.await
.map_err(|e| DownloaderError::Failed {
kind: FailureKind::DownloadInterrupted,
source: anyhow!(e).context("create archive file"),
})?;
let mut got: u64 = 0;
while let Some(chunk) = stream.next().await {
if cancel.load(Ordering::SeqCst) {
return Err(DownloaderError::Cancelled);
}
let bytes = chunk.map_err(|e| DownloaderError::Failed {
kind: FailureKind::DownloadInterrupted,
source: anyhow!(e).context("read chunk"),
})?;
file.write_all(&bytes)
.await
.map_err(|e| DownloaderError::Failed {
kind: FailureKind::DownloadInterrupted,
source: anyhow!(e).context("write chunk"),
})?;
got += bytes.len() as u64;
let _ = events.send(ProgressEvent::Downloading {
name: name.to_string(),
got,
total,
});
}
file.flush()
.await
.map_err(|e| DownloaderError::Failed {
kind: FailureKind::DownloadInterrupted,
source: anyhow!(e).context("flush"),
})?;
drop(file);
let _ = events.send(ProgressEvent::Verifying {
name: name.to_string(),
});
if let Ok(resp) = client.get(format!("{}.sha256", winner.url)).send().await {
if resp.status().is_success() {
let expected_hex = resp.text().await.ok().and_then(parse_sha256_hex);
if let Some(expected) = expected_hex {
let actual = compute_sha256(&archive_path)
.await
.map_err(|e| DownloaderError::Failed {
kind: FailureKind::ChecksumMismatch,
source: e.context("compute sha256"),
})?;
if actual != expected {
return Err(DownloaderError::Failed {
kind: FailureKind::ChecksumMismatch,
source: anyhow!(
"expected sha256={expected}, got sha256={actual}"
),
});
}
}
}
}
Ok(FetchedArchive {
temp_dir,
archive_path,
winning_url: winner.url.clone(),
})
}
async fn head_probe_concurrent<'a>(
client: &reqwest::Client,
candidates: &'a [CandidateUrl],
name: &str,
events: &tokio::sync::mpsc::UnboundedSender<ProgressEvent>,
cancel: &Arc<AtomicBool>,
) -> Result<&'a CandidateUrl, DownloaderError> {
use futures_util::stream::{FuturesUnordered, StreamExt};
use std::time::Duration;
const PROBE_CONCURRENCY: usize = 16;
const PROBE_TIMEOUT: Duration = Duration::from_secs(5);
let sem = Arc::new(tokio::sync::Semaphore::new(PROBE_CONCURRENCY));
let mut tasks = FuturesUnordered::new();
for (idx, c) in candidates.iter().enumerate() {
let _ = events.send(ProgressEvent::UrlCandidate {
name: name.to_string(),
url: c.url.clone(),
});
let sem = sem.clone();
let client = client.clone();
let url = c.url.clone();
let cancel = cancel.clone();
tasks.push(async move {
let _permit = sem.acquire_owned().await.ok()?;
if cancel.load(Ordering::SeqCst) {
return None;
}
let fut = client.head(&url).send();
match tokio::time::timeout(PROBE_TIMEOUT, fut).await {
Ok(Ok(r)) if r.status().is_success() => Some(idx),
_ => None,
}
});
}
let mut best: Option<usize> = None;
while let Some(res) = tasks.next().await {
if cancel.load(Ordering::SeqCst) {
return Err(DownloaderError::Cancelled);
}
if let Some(idx) = res {
best = Some(idx);
break;
}
}
match best {
Some(idx) => Ok(&candidates[idx]),
None => Err(DownloaderError::Failed {
kind: FailureKind::AllUrlsFailed,
source: anyhow!("no candidate URL returned 2xx within {}s × {} probes", PROBE_TIMEOUT.as_secs(), candidates.len()),
}),
}
}
fn parse_sha256_hex(body: String) -> Option<String> {
let token = body.split_whitespace().next()?.to_lowercase();
if token.len() == 64 && token.chars().all(|c| c.is_ascii_hexdigit()) {
Some(token)
} else {
None
}
}
async fn compute_sha256(path: &std::path::Path) -> Result<String> {
use tokio::io::AsyncReadExt;
let mut f = tokio::fs::File::open(path).await.context("open for sha256")?;
let mut hasher = Sha256::new();
let mut buf = vec![0u8; 8192];
loop {
let n = f.read(&mut buf).await.context("read for sha256")?;
if n == 0 {
break;
}
hasher.update(&buf[..n]);
}
Ok(format!("{:x}", hasher.finalize()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_sha256_strict_hex_64() {
assert_eq!(
parse_sha256_hex("a".repeat(64) + " file.tar.gz\n"),
Some("a".repeat(64))
);
}
#[test]
fn parse_sha256_rejects_wrong_length() {
assert_eq!(parse_sha256_hex("deadbeef file.tar.gz".into()), None);
}
#[test]
fn parse_sha256_rejects_non_hex() {
assert_eq!(parse_sha256_hex("z".repeat(64)), None);
}
}