use std::fs::{File, OpenOptions};
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};
use purecrypto::hash::{Digest, Sha256};
use crate::error::{Error, Result};
use crate::http::Request;
use crate::resume::{self, Kind, ResumeState};
use crate::tmpfile::{write_all_at, TempBlob};
pub type ProgressFn = Box<dyn FnMut(u64, Option<u64>) + Send>;
pub struct DownloadOptions {
pub max_retries: u32,
pub segment_size: Option<u64>,
pub segments: Option<usize>,
pub parallelism: usize,
pub prefer_http11: bool,
pub expected_sha256: Option<[u8; 32]>,
pub max_size: Option<u64>,
pub max_time: Option<Duration>,
pub limit_rate: Option<u64>,
pub low_speed: Option<(u64, u64)>,
pub initial_backoff: Duration,
pub max_backoff: Duration,
pub progress: Option<ProgressFn>,
pub tmp_spill_threshold: Option<u64>,
pub tmp_dir: Option<PathBuf>,
}
impl Default for DownloadOptions {
fn default() -> Self {
DownloadOptions {
max_retries: 5,
segment_size: None,
segments: None,
parallelism: 1,
prefer_http11: true,
expected_sha256: None,
max_size: None,
max_time: None,
limit_rate: None,
low_speed: None,
initial_backoff: Duration::from_millis(500),
max_backoff: Duration::from_secs(30),
progress: None,
tmp_spill_threshold: None,
tmp_dir: None,
}
}
}
impl std::fmt::Debug for DownloadOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DownloadOptions")
.field("max_retries", &self.max_retries)
.field("segment_size", &self.segment_size)
.field("segments", &self.segments)
.field("parallelism", &self.parallelism)
.field("prefer_http11", &self.prefer_http11)
.field("expected_sha256", &self.expected_sha256.is_some())
.field("max_size", &self.max_size)
.field("max_time", &self.max_time)
.field("limit_rate", &self.limit_rate)
.field("low_speed", &self.low_speed)
.field("initial_backoff", &self.initial_backoff)
.field("max_backoff", &self.max_backoff)
.field("progress", &self.progress.is_some())
.field("tmp_spill_threshold", &self.tmp_spill_threshold)
.field("tmp_dir", &self.tmp_dir)
.finish()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DownloadOutcome {
pub bytes_written: u64,
pub total: Option<u64>,
pub resumed_from: u64,
}
pub fn download(url: &str, path: &Path, mut opts: DownloadOptions) -> Result<DownloadOutcome> {
if let Some(decoded) = decode_data_uri(url) {
let bytes = checked_data_uri(decoded, &mut opts)?;
let n = bytes.len() as u64;
write_atomically(path, |f| {
f.write_all(&bytes).map_err(Error::Io)?;
Ok(n)
})?;
return Ok(DownloadOutcome {
bytes_written: n,
total: Some(n),
resumed_from: 0,
});
}
Request::get(url)?.download_resumable(path, opts)
}
fn checked_data_uri(decoded: Result<Vec<u8>>, opts: &mut DownloadOptions) -> Result<Vec<u8>> {
let bytes = decoded?;
let n = bytes.len() as u64;
if let Some(max) = opts.max_size {
if n > max {
return Err(Error::BadResponse("maximum file size exceeded".into()));
}
}
if let Some(want) = opts.expected_sha256 {
if Sha256::digest(&bytes) != want {
return Err(Error::BadResponse(
"data URI failed SHA-256 verification".into(),
));
}
}
if let Some(cb) = opts.progress.as_mut() {
cb(n, Some(n));
}
Ok(bytes)
}
pub fn download_to_tmp(url: &str, mut opts: DownloadOptions) -> Result<TempBlob> {
if let Some(decoded) = decode_data_uri(url) {
let bytes = checked_data_uri(decoded, &mut opts)?;
let blob = new_blob(&opts);
blob.write_at(0, &bytes).map_err(Error::Io)?;
return Ok(blob);
}
Request::get(url)?.download_to_tmp(opts)
}
fn new_blob(opts: &DownloadOptions) -> TempBlob {
let blob = TempBlob::with_threshold(
opts.tmp_spill_threshold
.unwrap_or(crate::tmpfile::DEFAULT_SPILL_THRESHOLD),
);
match &opts.tmp_dir {
Some(dir) => blob.in_dir(dir),
None => blob,
}
}
pub fn fetch_to_file(url: &str, path: &Path, mut opts: DownloadOptions) -> Result<DownloadOutcome> {
match url_scheme(url).as_deref() {
Some("data") | Some("http") | Some("https") => download(url, path, opts),
Some("magnet") => Err(Error::UnsupportedScheme(
"magnet: use the bittorrent module (front door has no peer discovery)".into(),
)),
Some(_) => {
let n = fetch_via_transfer(url, path, &mut opts)?;
Ok(DownloadOutcome {
bytes_written: n,
total: Some(n),
resumed_from: 0,
})
}
None => Err(Error::InvalidUrl(url.to_string())),
}
}
fn url_scheme(url: &str) -> Option<String> {
let colon = url.find(':')?;
let scheme = &url[..colon];
if scheme.is_empty()
|| !scheme
.bytes()
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'+' | b'-' | b'.'))
{
return None;
}
Some(scheme.to_ascii_lowercase())
}
pub fn fetch_to_tmp(url: &str, mut opts: DownloadOptions) -> Result<TempBlob> {
match url_scheme(url).as_deref() {
Some("data") | Some("http") | Some("https") => download_to_tmp(url, opts),
Some("magnet") => Err(Error::UnsupportedScheme(
"magnet: use the bittorrent module (front door has no peer discovery)".into(),
)),
Some(_) => {
let mut parsed = crate::url::Url::parse(url)?;
parsed.set_idn(true)?;
let mut blob = new_blob(&opts);
let n = crate::transfer::transfer_url_to_with(
&parsed,
&crate::net::NetConfig::default(),
&mut blob,
)?;
one_shot_checks(
n,
&mut opts,
|len| hash_blob_prefix(&blob, len),
|| {
let _ = blob.set_len(0);
},
)?;
Ok(blob)
}
None => Err(Error::InvalidUrl(url.to_string())),
}
}
fn fetch_via_transfer(url: &str, path: &Path, opts: &mut DownloadOptions) -> Result<u64> {
let mut parsed = crate::url::Url::parse(url)?;
parsed.set_idn(true)?;
write_atomically(path, |file| {
let n = crate::transfer::transfer_url_to_with(
&parsed,
&crate::net::NetConfig::default(),
file,
)?;
file.flush().map_err(Error::Io)?;
let hash_file = file.try_clone().map_err(Error::Io)?;
one_shot_checks(n, opts, |len| hash_file_prefix(hash_file, len), || {})?;
Ok(n)
})
}
fn write_atomically(path: &Path, fill: impl FnOnce(&mut File) -> Result<u64>) -> Result<u64> {
static SEQ: AtomicUsize = AtomicUsize::new(0);
let name = path
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_default();
let tmp = path.with_file_name(format!(
".{name}.rsurltmp.{}.{}",
std::process::id(),
SEQ.fetch_add(1, Ordering::Relaxed)
));
let mut file = std::fs::OpenOptions::new()
.write(true)
.read(true)
.create_new(true)
.open(&tmp)
.map_err(Error::Io)?;
let res = fill(&mut file).and_then(|n| {
drop(file);
std::fs::rename(&tmp, path).map_err(Error::Io)?;
Ok(n)
});
if res.is_err() {
let _ = std::fs::remove_file(&tmp);
}
res
}
fn one_shot_checks(
n: u64,
opts: &mut DownloadOptions,
hash: impl FnOnce(u64) -> io::Result<[u8; 32]>,
discard: impl FnOnce(),
) -> Result<()> {
if let Some(max) = opts.max_size {
if n > max {
discard();
return Err(Error::BadResponse("maximum file size exceeded".into()));
}
}
if let Some(want) = opts.expected_sha256 {
match hash(n) {
Ok(got) if got == want => {}
Ok(_) => {
discard();
return Err(Error::BadResponse(
"downloaded file failed SHA-256 verification".into(),
));
}
Err(e) => return Err(Error::Io(e)),
}
}
if let Some(cb) = opts.progress.as_mut() {
cb(n, Some(n));
}
Ok(())
}
fn decode_data_uri(url: &str) -> Option<Result<Vec<u8>>> {
let b = url.as_bytes();
if b.len() < 5 || !b[..5].eq_ignore_ascii_case(b"data:") {
return None;
}
let rest = &url[5..];
let Some((meta, data)) = rest.split_once(',') else {
return Some(Err(Error::InvalidUrl(
"data URI missing ',' separator".into(),
)));
};
let is_base64 = meta
.split(';')
.any(|token| token.trim().eq_ignore_ascii_case("base64"));
let bytes = if is_base64 {
match crate::tls::client_auth::base64_decode(data) {
Some(b) => b,
None => return Some(Err(Error::InvalidUrl("invalid base64 in data URI".into()))),
}
} else {
percent_decode_to_bytes(data)
};
Some(Ok(bytes))
}
fn percent_decode_to_bytes(s: &str) -> Vec<u8> {
let b = s.as_bytes();
let mut out = Vec::with_capacity(b.len());
let mut i = 0;
while i < b.len() {
if b[i] == b'%' && i + 3 <= b.len() {
if let (Some(h), Some(l)) = (hex_nibble(b[i + 1]), hex_nibble(b[i + 2])) {
out.push((h << 4) | l);
i += 3;
continue;
}
}
out.push(b[i]);
i += 1;
}
out
}
fn hex_nibble(c: u8) -> Option<u8> {
match c {
b'0'..=b'9' => Some(c - b'0'),
b'a'..=b'f' => Some(c - b'a' + 10),
b'A'..=b'F' => Some(c - b'A' + 10),
_ => None,
}
}
impl Request {
pub fn download_resumable(self, path: &Path, opts: DownloadOptions) -> Result<DownloadOutcome> {
Downloader::new(self, Arc::new(FileStore::new(path)), opts).run()
}
pub fn download_to_tmp(self, opts: DownloadOptions) -> Result<TempBlob> {
let blob = Arc::new(new_blob(&opts));
let store = Arc::new(TmpStore {
blob: Arc::clone(&blob),
});
Downloader::new(self, store, opts).run()?;
Arc::try_unwrap(blob)
.map_err(|_| Error::Io(io::Error::other("temp blob still shared after download")))
}
}
trait PartStore: Send + Sync {
fn write_at(&self, at: u64, buf: &[u8]) -> io::Result<()>;
fn set_len(&self, n: u64) -> io::Result<()>;
fn discard(&self);
fn hash_prefix(&self, len: u64) -> io::Result<[u8; 32]>;
fn finalize(&self, real_size: u64) -> io::Result<()>;
fn save_state(&self, real_size: u64, kind: Kind, meta: &[u8]) {
let _ = (real_size, kind, meta);
}
fn load_state(&self) -> Option<ResumeState> {
None
}
}
struct FileStore {
final_path: PathBuf,
part: PathBuf,
handle: RwLock<Option<File>>,
}
impl FileStore {
fn new(path: &Path) -> Self {
FileStore {
final_path: path.to_path_buf(),
part: resume::part_path(path),
handle: RwLock::new(None),
}
}
fn with_file<R>(&self, f: impl FnOnce(&File) -> io::Result<R>) -> io::Result<R> {
{
let g = self.handle.read().unwrap_or_else(|e| e.into_inner());
if let Some(fh) = g.as_ref() {
return f(fh);
}
}
let mut g = self.handle.write().unwrap_or_else(|e| e.into_inner());
if g.is_none() {
*g = Some(
OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&self.part)?,
);
}
match g.as_ref() {
Some(fh) => f(fh),
None => Err(io::Error::other("part file handle vanished")),
}
}
fn close_handle(&self) {
*self.handle.write().unwrap_or_else(|e| e.into_inner()) = None;
}
}
impl PartStore for FileStore {
fn write_at(&self, at: u64, buf: &[u8]) -> io::Result<()> {
self.with_file(|f| write_all_at(f, at, buf))
}
fn set_len(&self, n: u64) -> io::Result<()> {
self.with_file(|f| f.set_len(n))
}
fn discard(&self) {
self.close_handle();
let _ = std::fs::remove_file(&self.part);
}
fn hash_prefix(&self, len: u64) -> io::Result<[u8; 32]> {
hash_prefix(&self.part, len)
}
fn finalize(&self, real_size: u64) -> io::Result<()> {
self.close_handle();
resume::finalize(&self.part, &self.final_path, real_size)
}
fn save_state(&self, real_size: u64, kind: Kind, meta: &[u8]) {
let _ = resume::write_state(&self.part, real_size, kind, meta);
}
fn load_state(&self) -> Option<ResumeState> {
resume::read_state(&self.part).ok().flatten()
}
}
struct TmpStore {
blob: Arc<TempBlob>,
}
impl PartStore for TmpStore {
fn write_at(&self, at: u64, buf: &[u8]) -> io::Result<()> {
self.blob.write_at(at, buf)
}
fn set_len(&self, n: u64) -> io::Result<()> {
self.blob.set_len(n)
}
fn discard(&self) {
let _ = self.blob.set_len(0);
}
fn hash_prefix(&self, len: u64) -> io::Result<[u8; 32]> {
hash_blob_prefix(&self.blob, len)
}
fn finalize(&self, real_size: u64) -> io::Result<()> {
self.blob.set_len(real_size)
}
}
#[derive(Clone, Default, PartialEq, Eq)]
struct Validators {
url: String,
etag: String,
last_modified: String,
}
impl Validators {
fn if_range(&self) -> Option<&str> {
if !self.etag.is_empty() && !self.etag.starts_with("W/") {
Some(&self.etag)
} else if !self.last_modified.is_empty() {
Some(&self.last_modified)
} else {
None
}
}
}
struct Downloader {
base: Request,
store: Arc<dyn PartStore>,
url_key: String,
opts: DownloadOptions,
}
enum Attempt {
Done { written: u64, total: Option<u64> },
Transient {
written: u64,
resumable: bool,
err: Error,
},
Fatal(Error),
}
enum SegErr {
Fallback,
Fatal(Error),
}
enum Bootstrap {
Full(u64),
Ranged {
total: u64,
validators: Validators,
chunk_key: u32,
plan: Vec<(u64, u64)>,
bitmap: Vec<u8>,
},
}
impl Downloader {
fn new(req: Request, store: Arc<dyn PartStore>, opts: DownloadOptions) -> Self {
let url = req.url();
let url_key = format!("{}://{}:{}{}", url.scheme, url.host, url.port, url.path);
let mut base = req.follow_redirects(true).decompress(false);
if opts.prefer_http11 {
base = base.http11_only();
}
if let Some(t) = opts.max_time {
base = base.max_time(t);
}
if let Some((_, secs)) = opts.low_speed {
base = base.read_timeout(Some(Duration::from_secs(secs.max(1))));
}
Downloader {
base,
store,
url_key,
opts,
}
}
fn run(mut self) -> Result<DownloadOutcome> {
let segmented =
self.opts.segments.is_some() || self.opts.segment_size.is_some_and(|s| s > 0);
if segmented {
match self.run_segmented() {
Ok(outcome) => return Ok(outcome),
Err(SegErr::Fatal(e)) => return Err(e),
Err(SegErr::Fallback) => { }
}
}
self.run_single()
}
fn run_single(&mut self) -> Result<DownloadOutcome> {
let (mut have, mut validators, mut known_total) = self.load_stream_state();
let resumed_from = have;
let mut budget = self.opts.max_retries;
let mut attempt_no: u32 = 0;
loop {
match self.attempt_single(have, &mut validators, &mut known_total) {
Attempt::Done { written, total } => {
self.verify_and_finalize(written, total)?;
return Ok(DownloadOutcome {
bytes_written: written,
total,
resumed_from,
});
}
Attempt::Fatal(e) => return Err(e),
Attempt::Transient {
written,
resumable,
err,
} => {
let progressed = resumable && written > have;
if resumable || written < have {
have = written;
}
if progressed {
budget = self.opts.max_retries;
} else if budget == 0 {
return Err(err);
} else {
budget -= 1;
}
attempt_no += 1;
self.backoff(attempt_no);
}
}
}
}
fn attempt_single(
&mut self,
have: u64,
validators: &mut Validators,
known_total: &mut Option<u64>,
) -> Attempt {
let mut req = self.base.clone();
if have > 0 {
req = req.header("Range", &format!("bytes={have}-"));
if let Some(v) = validators.if_range() {
req = req.header("If-Range", v);
}
}
let reader = match req.send_reader() {
Ok(r) => r,
Err(e) => {
return classify_pre_body(e, have);
}
};
let status = reader.status();
if status == 416 {
let total = parse_unsatisfied_total(reader.header("content-range"));
if have > 0 && total.or(*known_total) == Some(have) {
return Attempt::Done {
written: have,
total: Some(have),
};
}
if have > 0 && total.is_some() {
self.store.discard();
*known_total = None;
return Attempt::Transient {
written: 0,
resumable: false,
err: Error::BadResponse("resource changed size during resume".into()),
};
}
return Attempt::Fatal(status_error(status, &reader));
}
if (300..400).contains(&status) {
return Attempt::Fatal(Error::BadResponse(format!(
"unexpected redirect status {status}"
)));
}
if (400..500).contains(&status) {
return Attempt::Fatal(status_error(status, &reader));
}
if status >= 500 {
return Attempt::Transient {
written: have,
resumable: false,
err: status_error(status, &reader),
};
}
let (offset, total, resumable) = if status == 206 {
match parse_content_range(reader.header("content-range")) {
Some(cr)
if cr.start == have && (known_total.is_none() || cr.total == *known_total) =>
{
*known_total = cr.total;
(have, cr.total, true)
}
_ => {
self.store.discard();
*known_total = None;
return Attempt::Transient {
written: 0,
resumable: false,
err: Error::BadResponse("range offset mismatch on resume".into()),
};
}
}
} else {
*validators = self.validators_from(&reader);
let total = reader
.header("content-length")
.and_then(|v| v.trim().parse::<u64>().ok());
let accepts_ranges = reader
.header("accept-ranges")
.is_some_and(|v| v.to_ascii_lowercase().contains("bytes"));
*known_total = total;
(0, total, have == 0 && accepts_ranges)
};
if let Some(max) = self.opts.max_size {
if total.is_some_and(|t| t > max) {
return Attempt::Fatal(Error::BadResponse("maximum file size exceeded".into()));
}
}
self.stream_to_store(reader, offset, total, validators, resumable)
}
fn stream_to_store(
&mut self,
mut reader: crate::http::BodyReader,
offset: u64,
total: Option<u64>,
validators: &Validators,
resumable: bool,
) -> Attempt {
if let Some(t) = total {
if let Err(e) = self.store.set_len(t) {
return Attempt::Fatal(Error::Io(e));
}
} else if offset == 0 {
let _ = self.store.set_len(0);
}
let started = Instant::now();
let mut last_save = started;
let mut written = offset;
let mut buf = [0u8; 64 * 1024];
loop {
let n = match reader.read(&mut buf) {
Ok(0) => break,
Ok(n) => n,
Err(e) => {
self.persist_stream(total, written, validators);
return Attempt::Transient {
written,
resumable,
err: Error::Io(e),
};
}
};
if let Some(max) = self.opts.max_size {
if written + n as u64 > max {
return Attempt::Fatal(Error::BadResponse("maximum file size exceeded".into()));
}
}
if let Some(rate) = self.opts.limit_rate.filter(|r| *r > 0) {
let target =
Duration::from_secs_f64((written + n as u64 - offset) as f64 / rate as f64);
let elapsed = started.elapsed();
if target > elapsed {
std::thread::sleep(target - elapsed);
}
}
if let Err(e) = self.store.write_at(written, &buf[..n]) {
self.persist_stream(total, written, validators);
return Attempt::Transient {
written,
resumable,
err: Error::Io(e),
};
}
written += n as u64;
if let Some(cb) = self.opts.progress.as_mut() {
cb(written, total);
}
if let Some((min, secs)) = self.opts.low_speed {
let el = started.elapsed().as_secs();
if el >= secs && (written - offset) / el.max(1) < min {
self.persist_stream(total, written, validators);
return Attempt::Transient {
written,
resumable,
err: low_speed_error(),
};
}
}
if last_save.elapsed() >= Duration::from_secs(1) {
self.persist_stream(total, written, validators);
last_save = Instant::now();
}
}
Attempt::Done {
written,
total: total.or(Some(written)),
}
}
fn persist_stream(&self, total: Option<u64>, done: u64, validators: &Validators) {
if let Some(total) = total {
let meta = stream_meta(total, done, validators);
self.store.save_state(total, Kind::HttpStream, &meta);
}
}
fn load_stream_state(&self) -> (u64, Validators, Option<u64>) {
if let Some(st) = self.store.load_state() {
if st.kind == Kind::HttpStream {
if let Some((done, v)) = parse_stream_meta(&st.meta) {
if v.url == self.url_key && done <= st.real_size {
return (done, v, Some(st.real_size));
}
}
}
}
(0, Validators::default(), None)
}
fn run_segmented(&mut self) -> std::result::Result<DownloadOutcome, SegErr> {
let meter = Arc::new(ByteMeter::new(
self.opts.progress.take(),
self.opts.limit_rate,
self.opts.low_speed,
));
let out = self.run_segmented_metered(&meter);
meter.report(true);
self.opts.progress = meter.take_cb();
out
}
fn run_segmented_metered(
&mut self,
meter: &Arc<ByteMeter>,
) -> std::result::Result<DownloadOutcome, SegErr> {
let (total, validators, chunk_key, plan, bitmap, resumed_from) =
if let Some((total, validators, chunk_key, plan, bitmap)) = self.resume_ranged() {
let done = plan_done_bytes(&plan, &bitmap);
meter.set_total(total);
meter.seed(done);
(total, validators, chunk_key, plan, bitmap, done)
} else {
match self.bootstrap(meter)? {
Bootstrap::Full(written) => {
self.verify_and_finalize(written, Some(written))
.map_err(SegErr::Fatal)?;
return Ok(DownloadOutcome {
bytes_written: written,
total: Some(written),
resumed_from: 0,
});
}
Bootstrap::Ranged {
total,
validators,
chunk_key,
plan,
bitmap,
} => (total, validators, chunk_key, plan, bitmap, 0),
}
};
self.run_chunks(&plan, total, chunk_key, &validators, bitmap, meter)?;
self.verify_and_finalize(total, Some(total))
.map_err(SegErr::Fatal)?;
Ok(DownloadOutcome {
bytes_written: total,
total: Some(total),
resumed_from,
})
}
#[allow(clippy::type_complexity)]
fn resume_ranged(&self) -> Option<(u64, Validators, u32, Vec<(u64, u64)>, Vec<u8>)> {
let st = self.store.load_state()?;
if st.kind != Kind::HttpRanged {
return None;
}
let total = st.real_size;
let (stored_key, validators, bitmap) = parse_ranged_full(&st.meta)?;
if validators.url != self.url_key {
return None;
}
let (chunk_key, plan) = self.chunk_plan(total).ok()?;
if chunk_key != stored_key || bitmap.len() != plan.len().div_ceil(8) {
return None;
}
Some((total, validators, chunk_key, plan, bitmap))
}
fn bootstrap(&mut self, meter: &ByteMeter) -> std::result::Result<Bootstrap, SegErr> {
let mut budget = self.opts.max_retries;
let mut attempt_no = 0u32;
loop {
let req = self.base.clone().header("Range", "bytes=0-");
let mut reader = match req.send_reader() {
Ok(r) => r,
Err(e) if is_transient(&e) && budget > 0 => {
budget -= 1;
attempt_no += 1;
self.backoff(attempt_no);
continue;
}
Err(e) => return Err(SegErr::Fatal(e)),
};
let status = reader.status();
if status == 200 {
let total = reader
.header("content-length")
.and_then(|v| v.trim().parse::<u64>().ok());
if let Some(max) = self.opts.max_size {
if total.is_some_and(|t| t > max) {
return Err(SegErr::Fatal(Error::BadResponse(
"maximum file size exceeded".into(),
)));
}
}
self.prepare_part(total)?;
if let Some(t) = total {
meter.set_total(t);
}
let p = pump_to_store(
&mut reader,
self.store.as_ref(),
0,
total.unwrap_or(u64::MAX),
Some(meter),
);
if let Some(e) = p.err {
if budget == 0 {
return Err(SegErr::Fatal(e));
}
meter.rewind(p.wrote);
budget -= 1;
attempt_no += 1;
self.backoff(attempt_no);
continue;
}
return Ok(Bootstrap::Full(p.wrote));
}
if status == 416 {
self.prepare_part(Some(0))?;
return Ok(Bootstrap::Full(0));
}
if (400..500).contains(&status) {
return Err(SegErr::Fatal(status_error(status, &reader)));
}
if status >= 500 {
if budget == 0 {
return Err(SegErr::Fatal(status_error(status, &reader)));
}
budget -= 1;
attempt_no += 1;
self.backoff(attempt_no);
continue;
}
if status != 206 {
return Err(SegErr::Fatal(Error::BadResponse(format!(
"unexpected status {status}"
))));
}
let total = match parse_content_range(reader.header("content-range")) {
Some(ContentRange {
start: 0,
total: Some(t),
..
}) => t,
_ => return Err(SegErr::Fallback), };
if total == 0 {
self.prepare_part(Some(0))?;
return Ok(Bootstrap::Full(0));
}
if let Some(max) = self.opts.max_size {
if total > max {
return Err(SegErr::Fatal(Error::BadResponse(
"maximum file size exceeded".into(),
)));
}
}
meter.set_total(total);
let validators = self.validators_from(&reader);
let (chunk_key, plan) = match self.chunk_plan(total) {
Ok(x) => x,
Err(SegErr::Fallback) => {
self.prepare_part(Some(total))?;
let p = pump_to_store(&mut reader, self.store.as_ref(), 0, total, Some(meter));
if let Some(e) = p.err {
if budget == 0 {
return Err(SegErr::Fatal(e));
}
meter.rewind(p.wrote);
budget -= 1;
attempt_no += 1;
self.backoff(attempt_no);
continue;
}
return Ok(Bootstrap::Full(p.wrote));
}
Err(e) => return Err(e),
};
self.prepare_part(Some(total))?;
let map_len = plan.len().div_ceil(8);
let mut bitmap = vec![0u8; map_len];
let (_, end0) = plan[0];
let want0 = end0 + 1;
let got0 = pump_to_store(&mut reader, self.store.as_ref(), 0, want0, Some(meter)).wrote;
drop(reader);
if got0 < want0 {
match fetch_chunk_streaming(
&self.base,
self.store.as_ref(),
got0,
end0,
total,
validators.if_range().unwrap_or(""),
self.retry(),
meter,
) {
ChunkResult::Ok => {}
ChunkResult::Fallback => return Err(SegErr::Fallback),
ChunkResult::Fatal(e) => return Err(SegErr::Fatal(e)),
}
}
bit_set(&mut bitmap, 0);
self.persist_ranged(chunk_key, total, &validators, &bitmap);
meter.report(true);
return Ok(Bootstrap::Ranged {
total,
validators,
chunk_key,
plan,
bitmap,
});
}
}
fn prepare_part(&self, total: Option<u64>) -> std::result::Result<(), SegErr> {
if let Some(t) = total {
self.store
.set_len(t)
.map_err(|e| SegErr::Fatal(Error::Io(e)))?;
}
Ok(())
}
fn chunk_plan(&self, total: u64) -> std::result::Result<(u32, Vec<(u64, u64)>), SegErr> {
if let Some(n_req) = self.opts.segments {
let by_size = total.div_ceil(MIN_SEGMENT_BYTES).max(1) as usize;
let upper = by_size.clamp(1, MAX_SEGMENT_CHUNKS);
let n = n_req.clamp(1, upper);
if n < 2 {
return Err(SegErr::Fallback);
}
let seg = total / n as u64;
let plan: Vec<(u64, u64)> = (0..n)
.map(|i| {
let start = i as u64 * seg;
let end = if i == n - 1 {
total - 1
} else {
(i as u64 + 1) * seg - 1
};
(start, end)
})
.collect();
Ok((plan_key(&plan), plan))
} else if let Some(size) = self.opts.segment_size.filter(|s| *s > 0) {
let n = total.div_ceil(size) as usize;
let plan: Vec<(u64, u64)> = (0..n)
.map(|i| {
let start = i as u64 * size;
let end = (start + size).min(total) - 1;
(start, end)
})
.collect();
Ok((plan_key(&plan), plan))
} else {
Err(SegErr::Fallback)
}
}
fn run_chunks(
&mut self,
plan: &[(u64, u64)],
total: u64,
chunk_key: u32,
validators: &Validators,
bitmap: Vec<u8>,
meter: &Arc<ByteMeter>,
) -> std::result::Result<(), SegErr> {
let num_chunks = plan.len();
let workers = self
.opts
.parallelism
.clamp(1, num_chunks.max(1))
.min(MAX_SEGMENT_WORKERS);
let plan: Arc<Vec<(u64, u64)>> = Arc::new(plan.to_vec());
let bitmap = Arc::new(Mutex::new(bitmap));
let next = Arc::new(AtomicUsize::new(0));
let failed: Arc<Mutex<Option<Error>>> = Arc::new(Mutex::new(None));
let fallback = Arc::new(AtomicBool::new(false));
let validators = Arc::new(validators.clone());
let if_range = Arc::new(validators.if_range().unwrap_or("").to_string());
let retry = self.retry();
let mut handles = Vec::with_capacity(workers);
for _ in 0..workers {
let plan = Arc::clone(&plan);
let bitmap = Arc::clone(&bitmap);
let next = Arc::clone(&next);
let failed = Arc::clone(&failed);
let fallback = Arc::clone(&fallback);
let meter = Arc::clone(meter);
let store = Arc::clone(&self.store);
let validators = Arc::clone(&validators);
let if_range = Arc::clone(&if_range);
let base = self.base.clone();
handles.push(std::thread::spawn(move || loop {
if failed.lock().unwrap().is_some() || fallback.load(Ordering::Relaxed) {
break;
}
let i = next.fetch_add(1, Ordering::Relaxed);
if i >= num_chunks {
break;
}
if bit_get(&bitmap.lock().unwrap(), i) {
continue;
}
let (start, end) = plan[i];
match fetch_chunk_streaming(
&base,
store.as_ref(),
start,
end,
total,
&if_range,
retry,
&meter,
) {
ChunkResult::Ok => {
let mut bm = bitmap.lock().unwrap();
bit_set(&mut bm, i);
let meta = ranged_meta(chunk_key as u64, total, &validators, &bm);
store.save_state(total, Kind::HttpRanged, &meta);
drop(bm);
meter.report(true);
}
ChunkResult::Fallback => fallback.store(true, Ordering::Relaxed),
ChunkResult::Fatal(e) => *failed.lock().unwrap() = Some(e),
}
}));
}
let mut panicked = false;
for h in handles {
panicked |= h.join().is_err();
}
if panicked {
return Err(SegErr::Fatal(Error::BadResponse(
"segment worker panicked; download incomplete".into(),
)));
}
if let Some(e) = Arc::try_unwrap(failed)
.ok()
.and_then(|m| m.into_inner().ok())
.flatten()
{
return Err(SegErr::Fatal(e));
}
if fallback.load(Ordering::Relaxed) {
return Err(SegErr::Fallback);
}
Ok(())
}
fn persist_ranged(&self, chunk: u32, total: u64, validators: &Validators, bitmap: &[u8]) {
let meta = ranged_meta(chunk as u64, total, validators, bitmap);
self.store.save_state(total, Kind::HttpRanged, &meta);
}
fn validators_from(&self, reader: &crate::http::BodyReader) -> Validators {
Validators {
url: self.url_key.clone(),
etag: reader.header("etag").unwrap_or("").to_string(),
last_modified: reader.header("last-modified").unwrap_or("").to_string(),
}
}
fn verify_and_finalize(&self, real_size: u64, _total: Option<u64>) -> Result<DownloadOutcome> {
if let Some(want) = self.opts.expected_sha256 {
match self.store.hash_prefix(real_size) {
Ok(got) if got == want => {}
Ok(_) => {
self.store.discard();
return Err(Error::BadResponse(
"downloaded file failed SHA-256 verification".into(),
));
}
Err(e) => return Err(Error::Io(e)),
}
}
self.store.finalize(real_size).map_err(Error::Io)?;
Ok(DownloadOutcome {
bytes_written: real_size,
total: Some(real_size),
resumed_from: 0,
})
}
fn retry(&self) -> Retry {
Retry {
max: self.opts.max_retries,
initial: self.opts.initial_backoff,
cap: self.opts.max_backoff,
}
}
fn backoff(&self, attempt_no: u32) {
let shift = attempt_no.saturating_sub(1).min(20);
let delay = self
.opts
.initial_backoff
.saturating_mul(1u32 << shift)
.min(self.opts.max_backoff);
if !delay.is_zero() {
std::thread::sleep(delay);
}
}
}
const MIN_SEGMENT_BYTES: u64 = 1 << 20; const MAX_SEGMENT_WORKERS: usize = 16;
const MAX_SEGMENT_CHUNKS: usize = 1024;
const PROGRESS_INTERVAL: Duration = Duration::from_millis(100);
struct ByteMeter {
done: AtomicU64,
moved: AtomicU64,
total: AtomicU64,
limit_rate: Option<u64>,
low_speed: Option<(u64, u64)>,
started: Instant,
sink: Mutex<(Option<ProgressFn>, Instant)>,
}
impl ByteMeter {
fn new(cb: Option<ProgressFn>, limit_rate: Option<u64>, low_speed: Option<(u64, u64)>) -> Self {
let now = Instant::now();
ByteMeter {
done: AtomicU64::new(0),
moved: AtomicU64::new(0),
total: AtomicU64::new(0),
limit_rate: limit_rate.filter(|r| *r > 0),
low_speed: low_speed.filter(|(min, _)| *min > 0),
started: now,
sink: Mutex::new((cb, now)),
}
}
fn set_total(&self, total: u64) {
self.total.store(total, Ordering::Relaxed);
}
fn seed(&self, done: u64) {
self.done.store(done, Ordering::Relaxed);
}
fn add(&self, n: u64) {
self.done.fetch_add(n, Ordering::Relaxed);
let moved = self.moved.fetch_add(n, Ordering::Relaxed) + n;
self.report(false);
if let Some(rate) = self.limit_rate {
let target = Duration::from_secs_f64(moved as f64 / rate as f64);
let elapsed = self.started.elapsed();
if target > elapsed {
std::thread::sleep(target - elapsed);
}
}
}
fn rewind(&self, n: u64) {
self.done.fetch_sub(n, Ordering::Relaxed);
self.moved.fetch_sub(n, Ordering::Relaxed);
}
fn report(&self, force: bool) {
let mut sink = match if force {
self.sink.lock().ok()
} else {
self.sink.try_lock().ok()
} {
Some(g) => g,
None => return,
};
if sink.0.is_none() || (!force && sink.1.elapsed() < PROGRESS_INTERVAL) {
return;
}
let done = self.done.load(Ordering::Relaxed);
let total = match self.total.load(Ordering::Relaxed) {
0 => None,
t => Some(t),
};
sink.1 = Instant::now();
if let Some(cb) = sink.0.as_mut() {
cb(done, total);
}
}
fn take_cb(&self) -> Option<ProgressFn> {
self.sink.lock().ok().and_then(|mut s| s.0.take())
}
}
enum ChunkResult {
Ok,
Fallback,
Fatal(Error),
}
#[derive(Clone, Copy)]
struct Retry {
max: u32,
initial: Duration,
cap: Duration,
}
#[allow(clippy::too_many_arguments)]
fn fetch_chunk_streaming(
base: &Request,
store: &dyn PartStore,
start: u64,
end: u64,
total: u64,
if_range: &str,
retry: Retry,
meter: &ByteMeter,
) -> ChunkResult {
let want = end - start + 1;
let mut got = 0u64;
let mut budget = retry.max;
let mut attempt_no = 0u32;
loop {
let from = start + got;
let mut req = base.clone().header("Range", &format!("bytes={from}-{end}"));
if !if_range.is_empty() {
req = req.header("If-Range", if_range);
}
match stream_chunk_once(req, store, from, want - got, total, meter) {
StreamOnce::Fallback => return ChunkResult::Fallback,
StreamOnce::Fatal(e) => return ChunkResult::Fatal(e),
StreamOnce::Advanced {
wrote,
err,
too_slow,
} => {
got += wrote;
if got >= want {
return ChunkResult::Ok;
}
let err = err.unwrap_or(Error::UnexpectedEof);
if wrote > 0 && !too_slow {
budget = retry.max;
attempt_no = 0;
}
if budget == 0 {
return ChunkResult::Fatal(err);
}
budget -= 1;
attempt_no += 1;
sleep_backoff(attempt_no, retry.initial, retry.cap);
}
}
}
}
enum StreamOnce {
Advanced {
wrote: u64,
err: Option<Error>,
too_slow: bool,
},
Fallback,
Fatal(Error),
}
fn stream_chunk_once(
req: Request,
store: &dyn PartStore,
at: u64,
want: u64,
total: u64,
meter: &ByteMeter,
) -> StreamOnce {
let mut reader = match req.send_reader() {
Ok(r) => r,
Err(e) if is_transient(&e) => {
return StreamOnce::Advanced {
wrote: 0,
err: Some(e),
too_slow: false,
}
}
Err(e) => return StreamOnce::Fatal(e),
};
match reader.status() {
206 => {}
200 => return StreamOnce::Fallback,
s if (400..500).contains(&s) => return StreamOnce::Fatal(status_error(s, &reader)),
s if s >= 500 => {
return StreamOnce::Advanced {
wrote: 0,
err: Some(status_error(s, &reader)),
too_slow: false,
}
}
s => return StreamOnce::Fatal(Error::BadResponse(format!("unexpected status {s}"))),
}
match parse_content_range(reader.header("content-range")) {
Some(cr) if cr.start == at && cr.total == Some(total) => {}
_ => return StreamOnce::Fallback,
}
let p = pump_to_store(&mut reader, store, at, want, Some(meter));
match p.err {
Some(Error::Io(e)) if pump_open_failed(&e) => StreamOnce::Fatal(Error::Io(e)),
err => StreamOnce::Advanced {
wrote: p.wrote,
err,
too_slow: p.too_slow,
},
}
}
fn pump_open_failed(e: &io::Error) -> bool {
matches!(
e.kind(),
io::ErrorKind::NotFound | io::ErrorKind::PermissionDenied
)
}
struct Pumped {
wrote: u64,
err: Option<Error>,
too_slow: bool,
}
impl Pumped {
fn clean(wrote: u64) -> Self {
Pumped {
wrote,
err: None,
too_slow: false,
}
}
fn broke(wrote: u64, err: Error) -> Self {
Pumped {
wrote,
err: Some(err),
too_slow: false,
}
}
}
fn pump_to_store(
reader: &mut crate::http::BodyReader,
store: &dyn PartStore,
at: u64,
want: u64,
meter: Option<&ByteMeter>,
) -> Pumped {
let low_speed = meter.and_then(|m| m.low_speed);
let started = Instant::now();
let mut wrote = 0u64;
let mut buf = [0u8; 64 * 1024];
loop {
if wrote >= want {
return Pumped::clean(wrote);
}
if let Some((min, secs)) = low_speed {
let el = started.elapsed().as_secs();
if el >= secs && wrote / el.max(1) < min {
return Pumped {
wrote,
err: Some(low_speed_error()),
too_slow: true,
};
}
}
let cap = (want - wrote).min(buf.len() as u64) as usize;
match reader.read(&mut buf[..cap]) {
Ok(0) => return Pumped::clean(wrote),
Ok(n) => {
if let Err(e) = store.write_at(at + wrote, &buf[..n]) {
return Pumped::broke(wrote, Error::Io(e));
}
wrote += n as u64;
if let Some(m) = meter {
m.add(n as u64);
}
}
Err(e) => return Pumped::broke(wrote, Error::Io(e)),
}
}
}
fn low_speed_error() -> Error {
Error::Io(io::Error::new(
io::ErrorKind::TimedOut,
"transfer below low-speed limit",
))
}
fn sleep_backoff(attempt_no: u32, initial: Duration, max_backoff: Duration) {
let shift = attempt_no.saturating_sub(1).min(20);
let delay = initial.saturating_mul(1u32 << shift).min(max_backoff);
if !delay.is_zero() {
std::thread::sleep(delay);
}
}
fn plan_done_bytes(plan: &[(u64, u64)], bitmap: &[u8]) -> u64 {
plan.iter()
.enumerate()
.filter(|(i, _)| bit_get(bitmap, *i))
.map(|(_, (s, e))| e - s + 1)
.sum()
}
fn classify_pre_body(err: Error, have: u64) -> Attempt {
if is_transient(&err) {
Attempt::Transient {
written: have,
resumable: false,
err,
}
} else {
Attempt::Fatal(err)
}
}
fn status_error(code: u16, reader: &crate::http::BodyReader) -> Error {
Error::Status {
code,
reason: reader.head().reason.clone(),
}
}
fn is_transient(err: &Error) -> bool {
match err {
Error::Io(e) => matches!(
e.kind(),
io::ErrorKind::ConnectionReset
| io::ErrorKind::ConnectionAborted
| io::ErrorKind::BrokenPipe
| io::ErrorKind::TimedOut
| io::ErrorKind::UnexpectedEof
| io::ErrorKind::WouldBlock
| io::ErrorKind::NotConnected
),
Error::UnexpectedEof => true,
Error::BadResponse(m) => {
let m = m.to_ascii_lowercase();
m.contains("reset by server") || m.contains("goaway")
}
Error::Status { code, .. } => *code >= 500,
_ => false,
}
}
#[derive(Debug, PartialEq, Eq)]
struct ContentRange {
start: u64,
end: u64,
total: Option<u64>,
}
fn parse_content_range(v: Option<&str>) -> Option<ContentRange> {
let v = v?.trim();
let rest = v
.strip_prefix("bytes ")
.or_else(|| v.strip_prefix("bytes="))?;
let (range, total) = rest.split_once('/')?;
let (start, end) = range.split_once('-')?;
let start = start.trim().parse::<u64>().ok()?;
let end = end.trim().parse::<u64>().ok()?;
let total = match total.trim() {
"*" => None,
t => Some(t.parse::<u64>().ok()?),
};
if start > end || total.is_some_and(|t| end >= t) {
return None;
}
Some(ContentRange { start, end, total })
}
fn parse_unsatisfied_total(v: Option<&str>) -> Option<u64> {
let v = v?.trim();
let rest = v
.strip_prefix("bytes ")
.or_else(|| v.strip_prefix("bytes="))?;
rest.trim().strip_prefix("*/")?.trim().parse().ok()
}
fn plan_key(plan: &[(u64, u64)]) -> u32 {
let mut h: u32 = 0x811c_9dc5;
for (s, e) in plan {
for b in s.to_le_bytes().iter().chain(e.to_le_bytes().iter()) {
h ^= u32::from(*b);
h = h.wrapping_mul(0x0100_0193);
}
}
h
}
fn hash_file_prefix(mut f: File, len: u64) -> io::Result<[u8; 32]> {
use std::io::Seek as _;
f.seek(io::SeekFrom::Start(0))?;
let mut hasher = Sha256::new();
let mut left = len;
let mut buf = [0u8; 64 * 1024];
while left > 0 {
let want = left.min(buf.len() as u64) as usize;
let n = f.read(&mut buf[..want])?;
if n == 0 {
return Err(io::Error::from(io::ErrorKind::UnexpectedEof));
}
hasher.update(&buf[..n]);
left -= n as u64;
}
Ok(hasher.finalize())
}
fn hash_blob_prefix(blob: &TempBlob, len: u64) -> io::Result<[u8; 32]> {
let mut hasher = Sha256::new();
let mut at = 0u64;
let mut buf = [0u8; 64 * 1024];
while at < len {
let cap = (len - at).min(buf.len() as u64) as usize;
let n = blob.read_at(&mut buf[..cap], at)?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"temp blob shorter than expected while hashing",
));
}
hasher.update(&buf[..n]);
at += n as u64;
}
Ok(hasher.finalize())
}
fn hash_prefix(path: &Path, len: u64) -> io::Result<[u8; 32]> {
let mut f = std::fs::File::open(path)?;
let mut hasher = Sha256::new();
let mut remaining = len;
let mut buf = [0u8; 64 * 1024];
while remaining > 0 {
let cap = remaining.min(buf.len() as u64) as usize;
let n = f.read(&mut buf[..cap])?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"file shorter than expected while hashing",
));
}
hasher.update(&buf[..n]);
remaining -= n as u64;
}
Ok(hasher.finalize())
}
fn encode_validators(v: &Validators) -> Vec<u8> {
let mut out = Vec::new();
for s in [&v.url, &v.etag, &v.last_modified] {
let b = s.as_bytes();
out.extend_from_slice(&(b.len() as u16).to_le_bytes());
out.extend_from_slice(b);
}
out
}
fn decode_validators(p: &[u8]) -> Option<Validators> {
let mut p = p;
let mut take = || -> Option<String> {
if p.len() < 2 {
return None;
}
let n = u16::from_le_bytes([p[0], p[1]]) as usize;
let s = String::from_utf8_lossy(p.get(2..2 + n)?).into_owned();
p = &p[2 + n..];
Some(s)
};
Some(Validators {
url: take()?,
etag: take()?,
last_modified: take()?,
})
}
fn stream_meta(total: u64, done: u64, validators: &Validators) -> Vec<u8> {
let mut m = Vec::with_capacity(16);
m.extend_from_slice(&total.to_le_bytes());
m.extend_from_slice(&done.to_le_bytes());
m.extend_from_slice(&encode_validators(validators));
m
}
fn parse_stream_meta(meta: &[u8]) -> Option<(u64, Validators)> {
if meta.len() < 16 {
return None;
}
let done = u64::from_le_bytes(meta[8..16].try_into().unwrap());
let v = decode_validators(&meta[16..])?;
Some((done, v))
}
fn ranged_meta(chunk: u64, total: u64, validators: &Validators, bitmap: &[u8]) -> Vec<u8> {
let v = encode_validators(validators);
let mut m = Vec::with_capacity(16 + v.len() + bitmap.len());
m.extend_from_slice(&(chunk as u32).to_le_bytes());
m.extend_from_slice(&total.to_le_bytes());
m.extend_from_slice(&(v.len() as u32).to_le_bytes());
m.extend_from_slice(&v);
m.extend_from_slice(bitmap);
m
}
fn parse_ranged_full(meta: &[u8]) -> Option<(u32, Validators, Vec<u8>)> {
if meta.len() < 16 {
return None;
}
let chunk = u32::from_le_bytes(meta[0..4].try_into().unwrap());
let vlen = u32::from_le_bytes(meta[12..16].try_into().unwrap()) as usize;
let rest = meta.get(16..)?;
let validators = decode_validators(rest.get(..vlen)?)?;
let bitmap = rest.get(vlen..)?.to_vec();
Some((chunk, validators, bitmap))
}
fn bit_get(map: &[u8], i: usize) -> bool {
map[i / 8] & (1 << (i % 8)) != 0
}
fn bit_set(map: &mut [u8], i: usize) {
map[i / 8] |= 1 << (i % 8);
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{Read, Seek, SeekFrom, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
fn make_body(n: usize, seed: u64) -> Vec<u8> {
let mut x = seed | 1;
(0..n)
.map(|_| {
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
(x >> 33) as u8
})
.collect()
}
struct Origin {
body: Vec<u8>,
etag: String,
accept_ranges: bool,
ignore_range: bool,
kill_after: Option<u64>,
kill_range_start: Option<u64>,
kills_left: usize,
slow_left: usize,
slow_range_start: Option<u64>,
ranges: Vec<String>,
misrange_start: Option<u64>,
}
impl Origin {
fn shared(body: Vec<u8>, etag: &str) -> Arc<Mutex<Origin>> {
Arc::new(Mutex::new(Origin {
body,
etag: etag.to_string(),
accept_ranges: true,
ignore_range: false,
kill_after: None,
kill_range_start: None,
kills_left: 0,
slow_left: 0,
slow_range_start: None,
ranges: Vec::new(),
misrange_start: None,
}))
}
}
const SLOW_PIECE: usize = 1024;
const SLOW_PIECES: usize = 64;
const SLOW_GAP: Duration = Duration::from_millis(50);
fn start(origin: Arc<Mutex<Origin>>) -> u16 {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let port = listener.local_addr().unwrap().port();
thread::spawn(move || {
for conn in listener.incoming() {
let Ok(mut sock) = conn else { continue };
let origin = Arc::clone(&origin);
thread::spawn(move || handle(&mut sock, &origin));
}
});
port
}
fn read_head(sock: &mut TcpStream) -> String {
let mut buf = Vec::new();
let mut byte = [0u8; 1];
while sock.read(&mut byte).map(|n| n == 1).unwrap_or(false) {
buf.push(byte[0]);
if buf.ends_with(b"\r\n\r\n") {
break;
}
}
String::from_utf8_lossy(&buf).into_owned()
}
fn header_val<'a>(head: &'a str, name: &str) -> Option<&'a str> {
head.lines().find_map(|l| {
let (k, v) = l.split_once(':')?;
k.eq_ignore_ascii_case(name).then(|| v.trim())
})
}
fn parse_req_range(v: &str) -> Option<(u64, Option<u64>)> {
let r = v.trim().strip_prefix("bytes=")?;
let (a, b) = r.split_once('-')?;
let start = a.trim().parse::<u64>().ok()?;
let end = if b.trim().is_empty() {
None
} else {
Some(b.trim().parse::<u64>().ok()?)
};
Some((start, end))
}
fn handle(sock: &mut TcpStream, origin: &Arc<Mutex<Origin>>) {
let head = read_head(sock);
if head.is_empty() {
return;
}
let is_head = head.split_whitespace().next() == Some("HEAD");
let mut o = origin.lock().unwrap();
let range = header_val(&head, "range").map(|s| s.to_string());
let if_range = header_val(&head, "if-range").map(|s| s.to_string());
o.ranges.push(range.clone().unwrap_or_default());
let len = o.body.len() as u64;
if is_head {
let ar = if o.accept_ranges {
"Accept-Ranges: bytes\r\n"
} else {
""
};
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Length: {len}\r\n{ar}ETag: {}\r\nConnection: close\r\n\r\n",
o.etag
);
drop(o);
let _ = sock.write_all(resp.as_bytes());
let _ = sock.shutdown(Shutdown::Write);
return;
}
let if_range_ok = if_range.as_deref().map(|v| v == o.etag).unwrap_or(true);
let use_range = range.is_some() && o.accept_ranges && !o.ignore_range && if_range_ok;
let (mut start, mut end) = match (use_range, range.as_deref().and_then(parse_req_range)) {
(true, Some((s, e))) => (s, e.unwrap_or(len - 1).min(len - 1)),
_ => (0u64, len.saturating_sub(1)),
};
if use_range && o.misrange_start == Some(start) {
o.misrange_start = None;
end -= start;
start = 0;
}
if use_range && start >= len {
let resp = format!(
"HTTP/1.1 416 Range Not Satisfiable\r\nContent-Range: bytes */{len}\r\n\
Content-Length: 0\r\nConnection: close\r\n\r\n"
);
let _ = sock.write_all(resp.as_bytes());
let _ = sock.shutdown(Shutdown::Write);
return;
}
let slice = o.body[start as usize..=end as usize].to_vec();
let slice_len = slice.len() as u64;
let status_206 = use_range;
let head_bytes = if status_206 {
format!(
"HTTP/1.1 206 Partial Content\r\nContent-Range: bytes {start}-{end}/{len}\r\n\
Content-Length: {slice_len}\r\nETag: {}\r\nAccept-Ranges: bytes\r\n\
Connection: close\r\n\r\n",
o.etag
)
} else {
let ar = if o.accept_ranges {
"Accept-Ranges: bytes\r\n"
} else {
""
};
format!(
"HTTP/1.1 200 OK\r\nContent-Length: {len}\r\n{ar}ETag: {}\r\nConnection: close\r\n\r\n",
o.etag
)
};
let kill_here = o.kills_left > 0
&& o.kill_after.is_some()
&& o.kill_range_start.map(|s| s == start).unwrap_or(true);
let send_n = if kill_here {
o.kills_left -= 1;
o.kill_after.unwrap().min(slice_len) as usize
} else {
slice.len()
};
let slow_here = o.slow_left > 0 && o.slow_range_start.map(|st| st == start).unwrap_or(true);
if slow_here {
o.slow_left -= 1;
}
let payload = slice[..send_n].to_vec();
drop(o);
let _ = sock.write_all(head_bytes.as_bytes());
if slow_here {
let _ = sock.flush();
for piece in payload.chunks(SLOW_PIECE).take(SLOW_PIECES) {
if sock.write_all(piece).is_err() {
break;
}
let _ = sock.flush();
thread::sleep(SLOW_GAP);
}
} else {
let _ = sock.write_all(&payload);
}
let _ = sock.flush();
let _ = sock.shutdown(Shutdown::Write);
let mut sink = [0u8; 64];
let _ = sock.set_read_timeout(Some(Duration::from_millis(200)));
while sock.read(&mut sink).map(|n| n > 0).unwrap_or(false) {}
}
fn tmp(tag: &str) -> PathBuf {
static N: AtomicU64 = AtomicU64::new(0);
let n = N.fetch_add(1, Ordering::Relaxed);
std::env::temp_dir().join(format!("rsurl_dl_{}_{tag}_{n}", std::process::id()))
}
fn cleanup(p: &Path) {
let _ = std::fs::remove_file(p);
let _ = std::fs::remove_file(resume::part_path(p));
}
fn no_backoff() -> DownloadOptions {
DownloadOptions {
initial_backoff: Duration::ZERO,
max_backoff: Duration::ZERO,
..Default::default()
}
}
#[test]
fn resumes_after_midbody_disconnect() {
let body = make_body(120_000, 0xABCD);
let origin = Origin::shared(body.clone(), "v1");
{
let mut o = origin.lock().unwrap();
o.kill_after = Some(45_000);
o.kills_left = 1;
}
let port = start(origin.clone());
let out = tmp("resume");
let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, no_backoff())
.expect("download");
assert_eq!(outcome.bytes_written, 120_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
assert_eq!(ranges.len(), 2, "ranges: {ranges:?}");
assert_eq!(ranges[0], "");
assert_eq!(ranges[1], "bytes=45000-");
cleanup(&out);
}
#[test]
fn full_download_retry_when_no_range_support() {
let body = make_body(50_000, 0x1234);
let origin = Origin::shared(body.clone(), "");
{
let mut o = origin.lock().unwrap();
o.accept_ranges = false;
o.ignore_range = true; o.kill_after = Some(20_000);
o.kills_left = 1; }
let port = start(origin.clone());
let out = tmp("noranges");
let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, no_backoff())
.expect("download");
assert_eq!(outcome.bytes_written, 50_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
assert_eq!(ranges.len(), 2, "ranges: {ranges:?}");
cleanup(&out);
}
#[test]
fn validator_mismatch_discards_stale_partial() {
let v1 = make_body(100_000, 0x1111);
let v2 = make_body(100_000, 0x2222);
let origin = Origin::shared(v2.clone(), "v2");
let port = start(origin);
let out = tmp("validator");
let part = resume::part_path(&out);
let url = format!("http://127.0.0.1:{port}/file");
let u = Request::get(&url).unwrap();
let u = u.url();
let stale = Validators {
url: format!("{}://{}:{}{}", u.scheme, u.host, u.port, u.path),
etag: "v1".into(),
last_modified: String::new(),
};
{
let mut f = std::fs::OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(&part)
.unwrap();
f.set_len(100_000).unwrap();
f.write_all(&v1[..40_000]).unwrap();
}
resume::write_state(
&part,
100_000,
Kind::HttpStream,
&stream_meta(100_000, 40_000, &stale),
)
.unwrap();
assert_eq!(load_state_done(&part), Some(40_000));
let outcome = download(&url, &out, no_backoff()).expect("download v2");
assert_eq!(outcome.bytes_written, 100_000);
assert_eq!(std::fs::read(&out).unwrap(), v2);
assert_ne!(std::fs::read(&out).unwrap(), v1);
cleanup(&out);
}
fn load_state_done(part: &Path) -> Option<u64> {
let st = resume::read_state(part).ok()??;
parse_stream_meta(&st.meta).map(|(done, _)| done)
}
#[test]
fn segmented_retries_only_the_failing_chunk() {
let body = make_body(10_000, 0x9999);
let origin = Origin::shared(body.clone(), "seg");
{
let mut o = origin.lock().unwrap();
o.kill_after = Some(0);
o.kill_range_start = Some(3000);
o.kills_left = 1;
}
let port = start(origin.clone());
let out = tmp("segmented");
let mut opts = no_backoff();
opts.segment_size = Some(1000); let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect("segmented download");
assert_eq!(outcome.bytes_written, 10_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
assert_eq!(ranges.len(), 11, "ranges: {ranges:?}");
assert_eq!(ranges[0], "bytes=0-", "first GET doubles as the probe");
let at_3000 = ranges.iter().filter(|r| *r == "bytes=3000-3999").count();
assert_eq!(at_3000, 2, "failed chunk retried: {ranges:?}");
let at_4000 = ranges.iter().filter(|r| *r == "bytes=4000-4999").count();
assert_eq!(at_4000, 1, "completed chunks not refetched");
cleanup(&out);
}
#[test]
fn sha256_mismatch_deletes_partial() {
let body = make_body(8_000, 0x7777);
let origin = Origin::shared(body.clone(), "h");
let port = start(origin);
let out = tmp("sha");
let mut opts = no_backoff();
opts.expected_sha256 = Some([0u8; 32]); let err = download(&format!("http://127.0.0.1:{port}/file"), &out, opts).unwrap_err();
assert!(matches!(err, Error::BadResponse(_)));
assert!(!out.exists(), "final file not created");
assert!(!resume::part_path(&out).exists(), "partial deleted");
cleanup(&out);
}
#[test]
fn sha256_match_finalizes() {
let body = make_body(8_000, 0x5555);
let want = Sha256::digest(&body);
let origin = Origin::shared(body.clone(), "h");
let port = start(origin);
let out = tmp("sha_ok");
let mut opts = no_backoff();
opts.expected_sha256 = Some(want);
let outcome =
download(&format!("http://127.0.0.1:{port}/file"), &out, opts).expect("download");
assert_eq!(outcome.bytes_written, 8_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
cleanup(&out);
}
#[test]
fn is_transient_classifies_h2_reset_and_transport_faults() {
assert!(is_transient(&Error::BadResponse(
"stream 1 reset by server, error code 2".into()
)));
assert!(is_transient(&Error::BadResponse("received GOAWAY".into())));
assert!(is_transient(&Error::UnexpectedEof));
assert!(is_transient(&Error::Io(io::Error::from(
io::ErrorKind::ConnectionReset
))));
assert!(is_transient(&Error::Io(io::Error::from(
io::ErrorKind::TimedOut
))));
assert!(is_transient(&Error::Status {
code: 503,
reason: "Service Unavailable".into()
}));
assert!(!is_transient(&Error::Status {
code: 404,
reason: "Not Found".into()
}));
assert!(!is_transient(&Error::BadResponse(
"malformed header".into()
)));
assert!(!is_transient(&Error::Cancelled));
}
#[test]
fn data_uri_base64_is_written() {
let out = tmp("data_b64");
let outcome = download(
"data:text/plain;base64,SGVsbG8sIGRhdGEh",
&out,
no_backoff(),
)
.expect("data uri");
assert_eq!(outcome.bytes_written, 12);
assert_eq!(std::fs::read(&out).unwrap(), b"Hello, data!");
cleanup(&out);
}
#[test]
fn data_uri_percent_encoded_and_plain() {
let out = tmp("data_pct");
download("data:,a%20b%2Fc%00d", &out, no_backoff()).expect("data uri");
assert_eq!(std::fs::read(&out).unwrap(), b"a b/c\0d");
download("data:,plain", &out, no_backoff()).expect("data uri");
assert_eq!(std::fs::read(&out).unwrap(), b"plain");
cleanup(&out);
}
#[test]
fn data_uri_sha256_and_size_and_malformed() {
let out = tmp("data_meta");
let body = b"verify me";
let want = Sha256::digest(body);
let mut opts = no_backoff();
opts.expected_sha256 = Some(want);
let outcome = download("data:,verify me", &out, opts).expect("sha ok");
assert_eq!(outcome.bytes_written, body.len() as u64);
let mut opts = no_backoff();
opts.expected_sha256 = Some([0u8; 32]);
assert!(
download("data:,verify me", &out, opts).is_err(),
"sha mismatch"
);
let mut opts = no_backoff();
opts.max_size = Some(3);
assert!(download("data:,too long", &out, opts).is_err(), "size cap");
assert!(matches!(
download("data:text/plain", &out, no_backoff()),
Err(Error::InvalidUrl(_))
));
cleanup(&out);
}
#[test]
fn front_door_dispatches_data_and_http() {
let out = tmp("fd_data");
fetch_to_file("data:;base64,aGVsbG8=", &out, no_backoff()).expect("data");
assert_eq!(std::fs::read(&out).unwrap(), b"hello");
cleanup(&out);
let body = make_body(3_000, 0xF00D);
let origin = Origin::shared(body.clone(), "fd");
let port = start(origin);
let out = tmp("fd_http");
let outcome =
fetch_to_file(&format!("http://127.0.0.1:{port}/f"), &out, no_backoff()).expect("http");
assert_eq!(outcome.bytes_written, 3_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
cleanup(&out);
}
#[cfg(unix)]
#[test]
fn front_door_dispatches_file_scheme() {
let src = tmp("fd_src");
std::fs::write(&src, b"local file contents").unwrap();
let out = tmp("fd_file");
let url = format!("file://{}", src.display());
let outcome = fetch_to_file(&url, &out, no_backoff()).expect("file");
assert_eq!(outcome.bytes_written, 19);
assert_eq!(std::fs::read(&out).unwrap(), b"local file contents");
cleanup(&out);
cleanup(&src);
}
#[test]
fn front_door_rejects_magnet_and_bad_scheme() {
let out = tmp("fd_bad");
assert!(matches!(
fetch_to_file("magnet:?xt=urn:btih:abc", &out, no_backoff()),
Err(Error::UnsupportedScheme(_))
));
assert!(matches!(
fetch_to_file("not a url", &out, no_backoff()),
Err(Error::InvalidUrl(_))
));
cleanup(&out);
}
#[test]
fn decode_data_uri_recognition() {
assert!(decode_data_uri("http://x/").is_none());
assert!(decode_data_uri("DATA:,x").is_some()); assert_eq!(decode_data_uri("data:,hi").unwrap().unwrap(), b"hi");
assert_eq!(
decode_data_uri("data:;base64,aGk=").unwrap().unwrap(),
b"hi"
);
}
#[test]
fn parse_content_range_variants() {
assert_eq!(
parse_content_range(Some("bytes 100-199/1000")),
Some(ContentRange {
start: 100,
end: 199,
total: Some(1000)
})
);
assert_eq!(
parse_content_range(Some("bytes 0-0/*")),
Some(ContentRange {
start: 0,
end: 0,
total: None
})
);
assert_eq!(parse_content_range(Some("bytes 5-4/10")), None);
assert_eq!(parse_content_range(Some("bytes 0-10/10")), None);
assert_eq!(parse_unsatisfied_total(Some("bytes */1000")), Some(1000));
assert_eq!(parse_unsatisfied_total(Some("bytes 0-1/1000")), None);
assert_eq!(parse_content_range(Some("garbage")), None);
assert_eq!(parse_content_range(None), None);
}
#[test]
fn fresh_download_no_faults() {
let body = make_body(5_000, 0x4242);
let origin = Origin::shared(body.clone(), "e");
let port = start(origin);
let out = tmp("fresh");
let outcome =
download(&format!("http://127.0.0.1:{port}/file"), &out, no_backoff()).expect("dl");
assert_eq!(outcome.bytes_written, 5_000);
assert_eq!(outcome.resumed_from, 0);
assert_eq!(std::fs::read(&out).unwrap(), body);
cleanup(&out);
}
#[test]
fn segmented_single_chunk_streams_and_resumes_within_chunk() {
let body = make_body(50_000, 0xC0DE);
let origin = Origin::shared(body.clone(), "one");
{
let mut o = origin.lock().unwrap();
o.kill_after = Some(20_000);
o.kill_range_start = Some(0); o.kills_left = 1;
}
let port = start(origin.clone());
let out = tmp("onechunk");
let mut opts = no_backoff();
opts.segment_size = Some(1_000_000); let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect("single-chunk download");
assert_eq!(outcome.bytes_written, 50_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
assert!(
ranges.iter().any(|r| r == "bytes=20000-49999"),
"expected an in-chunk resume from 20000: {ranges:?}"
);
cleanup(&out);
}
#[test]
fn segmented_parallel_fetches_every_chunk_once() {
let body = make_body(10_000, 0xBEEF);
let origin = Origin::shared(body.clone(), "par");
let port = start(origin.clone());
let out = tmp("parallel");
let mut opts = no_backoff();
opts.segment_size = Some(1000); opts.parallelism = 4; let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect("parallel download");
assert_eq!(outcome.bytes_written, 10_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
let range_gets = ranges.iter().filter(|r| r.starts_with("bytes=")).count();
assert_eq!(range_gets, 10, "one GET per chunk: {ranges:?}");
}
#[test]
fn segmented_resume_reads_total_from_disk_without_probing() {
let body = make_body(5_000, 0x3033);
let origin = Origin::shared(body.clone(), "rv1");
let port = start(origin.clone());
let out = tmp("segresume");
let part = resume::part_path(&out);
let url = format!("http://127.0.0.1:{port}/file");
let u = Request::get(&url).unwrap();
let u = u.url();
let validators = Validators {
url: format!("{}://{}:{}{}", u.scheme, u.host, u.port, u.path),
etag: "rv1".into(),
last_modified: String::new(),
};
{
let mut f = std::fs::OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(&part)
.unwrap();
f.set_len(5_000).unwrap();
f.write_all(&body[..2_000]).unwrap();
}
let mut bitmap = vec![0u8]; bit_set(&mut bitmap, 0);
bit_set(&mut bitmap, 1);
resume::write_state(
&part,
5_000,
Kind::HttpRanged,
&ranged_meta(
plan_key(&[
(0, 999),
(1000, 1999),
(2000, 2999),
(3000, 3999),
(4000, 4999),
]) as u64,
5_000,
&validators,
&bitmap,
),
)
.unwrap();
let mut opts = no_backoff();
opts.segment_size = Some(1000);
let outcome = download(&url, &out, opts).expect("resumed segmented download");
assert_eq!(outcome.bytes_written, 5_000);
assert_eq!(outcome.resumed_from, 2_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
assert!(
!ranges.iter().any(|r| r == "bytes=0-" || r.is_empty()),
"no probe expected on resume: {ranges:?}"
);
let mut got: Vec<_> = ranges.clone();
got.sort();
assert_eq!(
got,
vec![
"bytes=2000-2999".to_string(),
"bytes=3000-3999".to_string(),
"bytes=4000-4999".to_string(),
],
"only the missing chunks fetched: {ranges:?}"
);
cleanup(&out);
}
#[test]
fn segments_mode_splits_into_n_equal_parts() {
let body = make_body(4_200_000, 0x5EED);
let origin = Origin::shared(body.clone(), "seg");
let port = start(origin.clone());
let out = tmp("segments");
let mut opts = no_backoff();
opts.segments = Some(4);
opts.parallelism = 4;
let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect("segments download");
assert_eq!(outcome.bytes_written, 4_200_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
assert_eq!(outcome.resumed_from, 0, "fresh download resumed nothing");
let ranges = origin.lock().unwrap().ranges.clone();
let range_gets = ranges.iter().filter(|r| r.starts_with("bytes=")).count();
assert_eq!(range_gets, 4, "split into 4 segments: {ranges:?}");
cleanup(&out);
}
#[test]
fn segments_plan_is_capped_independently_of_the_worker_pool() {
const CHUNKS: usize = MAX_SEGMENT_WORKERS + 4;
let body = make_body(CHUNKS << 20, 0xC0FFEE);
let origin = Origin::shared(body.clone(), "wide");
let port = start(origin.clone());
let out = tmp("wide");
let mut opts = no_backoff();
opts.segments = Some(CHUNKS);
opts.parallelism = 4;
let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect("wide segments download");
assert_eq!(outcome.bytes_written, body.len() as u64);
assert_eq!(std::fs::read(&out).unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
let range_gets = ranges.iter().filter(|r| r.starts_with("bytes=")).count();
assert_eq!(
range_gets, CHUNKS,
"{CHUNKS} chunks over 4 workers: {ranges:?}"
);
cleanup(&out);
}
#[test]
fn chunk_retry_budget_refreshes_on_forward_progress() {
let body = make_body(3 << 20, 0xFEED);
let origin = Origin::shared(body.clone(), "flaky");
{
let mut o = origin.lock().unwrap();
o.kill_after = Some(400_000);
o.kills_left = 8;
}
let port = start(origin.clone());
let out = tmp("flaky_seg");
let mut opts = no_backoff();
opts.segments = Some(2);
opts.parallelism = 2;
opts.max_retries = 1; let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect("advancing segments must not exhaust the retry budget");
assert_eq!(outcome.bytes_written, body.len() as u64);
assert_eq!(std::fs::read(&out).unwrap(), body);
assert_eq!(
origin.lock().unwrap().kills_left,
0,
"the injected breaks should all have been spent"
);
cleanup(&out);
}
#[test]
fn slow_connection_is_cut_and_retried_on_a_fresh_one() {
const TOTAL: usize = 256 * 1024;
const CHUNK: u64 = 64 * 1024;
let body = make_body(TOTAL, 0x5107);
let origin = Origin::shared(body.clone(), "slow");
{
let mut o = origin.lock().unwrap();
o.slow_range_start = Some(CHUNK);
o.slow_left = 1;
}
let port = start(origin.clone());
let out = tmp("slow_conn");
let mut opts = no_backoff();
opts.segment_size = Some(CHUNK);
opts.parallelism = 2;
opts.low_speed = Some((100_000, 1)); let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect("a slow connection should be replaced, not fatal");
assert_eq!(outcome.bytes_written, TOTAL as u64);
assert_eq!(std::fs::read(&out).unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
let gets = ranges.iter().filter(|r| r.starts_with("bytes=")).count();
assert!(
gets > TOTAL / CHUNK as usize,
"the cut chunk should have been re-requested: {ranges:?}"
);
assert!(
ranges.iter().any(|r| {
r.strip_prefix("bytes=")
.and_then(|r| r.split('-').next())
.and_then(|st| st.parse::<u64>().ok())
.is_some_and(|st| st > CHUNK && st < 2 * CHUNK)
}),
"the retry should resume inside the cut chunk, not restart it: {ranges:?}"
);
cleanup(&out);
}
#[test]
fn persistently_slow_transfer_still_gives_up() {
const TOTAL: usize = 256 * 1024;
const CHUNK: u64 = 64 * 1024;
let body = make_body(TOTAL, 0x5108);
let origin = Origin::shared(body.clone(), "slower");
origin.lock().unwrap().slow_left = usize::MAX; let port = start(origin.clone());
let out = tmp("slow_all");
let mut opts = no_backoff();
opts.segment_size = Some(CHUNK);
opts.parallelism = 1;
opts.max_retries = 1;
opts.low_speed = Some((100_000, 1));
let err = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect_err("a persistently slow link must give up");
assert!(
err.to_string().contains("low-speed"),
"expected the -Y abort the CLI maps to exit 28, got {err}"
);
cleanup(&out);
}
#[test]
fn progress_reports_bytes_in_flight_and_honours_the_rate_limit() {
const TOTAL: usize = 2 << 20;
const SEG: u64 = (TOTAL / 2) as u64;
const RATE: u64 = 4 << 20;
let body = make_body(TOTAL, 0x1234);
let origin = Origin::shared(body.clone(), "prog");
let port = start(origin.clone());
let out = tmp("prog");
type Reports = Arc<Mutex<Vec<(u64, Option<u64>)>>>;
let seen: Reports = Arc::new(Mutex::new(Vec::new()));
let sink = Arc::clone(&seen);
let mut opts = no_backoff();
opts.segments = Some(2);
opts.parallelism = 2;
opts.limit_rate = Some(RATE);
opts.progress = Some(Box::new(move |n, total| {
sink.lock().unwrap().push((n, total));
}));
let started = Instant::now();
let outcome = download(&format!("http://127.0.0.1:{port}/file"), &out, opts)
.expect("progress download");
let elapsed = started.elapsed();
assert_eq!(outcome.bytes_written, TOTAL as u64);
assert_eq!(std::fs::read(&out).unwrap(), body);
let seen = seen.lock().unwrap().clone();
assert!(!seen.is_empty(), "no progress reported");
let mut prev = 0;
for (n, total) in &seen {
assert!(*n >= prev, "progress went backwards: {seen:?}");
assert!(*n <= TOTAL as u64, "progress overshot: {seen:?}");
assert_eq!(*total, Some(TOTAL as u64));
prev = *n;
}
assert_eq!(
seen.last().map(|(n, _)| *n),
Some(TOTAL as u64),
"the final report should be the whole file: {seen:?}"
);
assert!(
seen.iter().any(|(n, _)| *n % SEG != 0),
"every report sat on a chunk boundary — in-flight bytes are not \
being counted: {seen:?}"
);
assert!(
elapsed >= Duration::from_secs_f64(TOTAL as f64 / RATE as f64 / 2.0),
"finished in {elapsed:?} — the rate limit was not applied"
);
cleanup(&out);
}
fn tmp_dir(tag: &str) -> PathBuf {
static N: AtomicU64 = AtomicU64::new(0);
let n = N.fetch_add(1, Ordering::Relaxed);
let d = std::env::temp_dir().join(format!("rsurl_tmpdl_{}_{tag}_{n}", std::process::id()));
std::fs::create_dir_all(&d).unwrap();
d
}
fn stray_entries(dir: &Path) -> Vec<String> {
std::fs::read_dir(dir)
.unwrap()
.filter_map(|e| Some(e.ok()?.file_name().to_string_lossy().into_owned()))
.filter(|name| cfg!(unix) || name.contains("rsurlpart"))
.collect()
}
#[test]
fn tmp_download_small_body_never_touches_the_filesystem() {
let body = make_body(4_000, 0x1111);
let port = start(Origin::shared(body.clone(), "mem"));
let dir = tmp_dir("mem");
let mut opts = no_backoff();
opts.tmp_dir = Some(dir.clone());
let mut blob =
download_to_tmp(&format!("http://127.0.0.1:{port}/file"), opts).expect("tmp download");
assert!(
blob.is_in_memory(),
"4 KB is under the 1 MiB spill threshold"
);
assert_eq!(blob.len(), 4_000);
assert_eq!(blob.to_vec().unwrap(), body);
assert!(stray_entries(&dir).is_empty(), "nothing should be on disk");
let mut head = [0u8; 8];
blob.read_exact(&mut head).unwrap();
assert_eq!(head, body[..8]);
assert_eq!(blob.seek(SeekFrom::End(-4)).unwrap(), 3_996);
let mut tail = Vec::new();
blob.read_to_end(&mut tail).unwrap();
assert_eq!(tail, body[3_996..]);
let mut mid = [0u8; 16];
assert_eq!(blob.read_at(&mut mid, 2_000).unwrap(), 16);
assert_eq!(mid, body[2_000..2_016]);
blob.close().unwrap();
let _ = std::fs::remove_dir(&dir);
}
#[test]
fn tmp_download_spills_to_an_anonymous_file_with_no_sidecar() {
let body = make_body(200_000, 0x2222);
let port = start(Origin::shared(body.clone(), "spill"));
let dir = tmp_dir("spill");
let mut opts = no_backoff();
opts.tmp_dir = Some(dir.clone());
opts.tmp_spill_threshold = Some(64 * 1024);
let blob =
download_to_tmp(&format!("http://127.0.0.1:{port}/file"), opts).expect("tmp download");
assert!(!blob.is_in_memory(), "200 KB is past the 64 KiB threshold");
assert_eq!(blob.len(), 200_000);
assert_eq!(blob.to_vec().unwrap(), body);
assert!(
stray_entries(&dir).is_empty(),
"the spilled file must have no name and leave no sidecar: {:?}",
stray_entries(&dir)
);
drop(blob);
assert!(stray_entries(&dir).is_empty(), "nothing survives the drop");
let _ = std::fs::remove_dir(&dir);
}
#[test]
fn tmp_download_resumes_after_a_midbody_disconnect() {
let body = make_body(120_000, 0x3333);
let origin = Origin::shared(body.clone(), "v1");
{
let mut o = origin.lock().unwrap();
o.kill_after = Some(45_000);
o.kills_left = 1;
}
let port = start(origin.clone());
let dir = tmp_dir("resume");
let mut opts = no_backoff();
opts.tmp_dir = Some(dir.clone());
opts.tmp_spill_threshold = Some(1024); let blob =
download_to_tmp(&format!("http://127.0.0.1:{port}/file"), opts).expect("tmp download");
assert_eq!(blob.to_vec().unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
assert_eq!(ranges.len(), 2, "ranges: {ranges:?}");
assert_eq!(ranges[1], "bytes=45000-", "resumed, not restarted");
assert!(stray_entries(&dir).is_empty());
let _ = std::fs::remove_dir(&dir);
}
#[test]
fn tmp_download_segmented_parallel_fetches_every_chunk_once() {
let body = make_body(40_000, 0x4444);
let origin = Origin::shared(body.clone(), "seg");
let port = start(origin.clone());
let dir = tmp_dir("seg");
let mut opts = no_backoff();
opts.tmp_dir = Some(dir.clone());
opts.tmp_spill_threshold = Some(0); opts.segment_size = Some(4_000);
opts.parallelism = 4;
let blob =
download_to_tmp(&format!("http://127.0.0.1:{port}/file"), opts).expect("tmp download");
assert_eq!(blob.len(), 40_000);
assert_eq!(blob.to_vec().unwrap(), body);
let ranges = origin.lock().unwrap().ranges.clone();
assert_eq!(ranges.len(), 10, "one GET per chunk: {ranges:?}");
assert!(stray_entries(&dir).is_empty(), "no sidecar, no named spill");
let _ = std::fs::remove_dir(&dir);
}
#[test]
fn tmp_download_checks_sha256_and_max_size() {
let body = make_body(8_000, 0x5566);
let origin = Origin::shared(body.clone(), "h");
let port = start(origin);
let url = format!("http://127.0.0.1:{port}/file");
let mut opts = no_backoff();
opts.expected_sha256 = Some(Sha256::digest(&body));
let blob = download_to_tmp(&url, opts).expect("hash matches");
assert_eq!(blob.to_vec().unwrap(), body);
let mut opts = no_backoff();
opts.expected_sha256 = Some([0u8; 32]);
assert!(matches!(
download_to_tmp(&url, opts).unwrap_err(),
Error::BadResponse(_)
));
let mut opts = no_backoff();
opts.max_size = Some(100);
assert!(matches!(
download_to_tmp(&url, opts).unwrap_err(),
Error::BadResponse(_)
));
}
#[test]
fn fetch_to_tmp_dispatches_data_file_and_http() {
let blob = fetch_to_tmp("data:text/plain;base64,aGVsbG8=", no_backoff()).unwrap();
assert_eq!(blob.to_vec().unwrap(), b"hello");
#[cfg(unix)]
{
let src = tmp("tmp_front_door");
std::fs::write(&src, b"from a local file").unwrap();
let url = format!("file://{}", src.display());
let blob = fetch_to_tmp(&url, no_backoff()).unwrap();
assert_eq!(blob.to_vec().unwrap(), b"from a local file");
cleanup(&src);
}
let body = make_body(2_000, 0x6677);
let port = start(Origin::shared(body.clone(), "front"));
let blob = fetch_to_tmp(&format!("http://127.0.0.1:{port}/file"), no_backoff()).unwrap();
assert_eq!(blob.to_vec().unwrap(), body);
assert!(matches!(
fetch_to_tmp("magnet:?xt=urn:btih:0000", no_backoff()).unwrap_err(),
Error::UnsupportedScheme(_)
));
}
#[test]
fn chunk_with_misplaced_content_range_is_not_spliced() {
let body = make_body(5_000, 0x5151);
let origin = Origin::shared(body.clone(), "mr");
origin.lock().unwrap().misrange_start = Some(2_000);
let port = start(origin.clone());
let out = tmp("misrange");
let mut opts = no_backoff();
opts.segment_size = Some(1000);
opts.parallelism = 1;
download(&format!("http://127.0.0.1:{port}/file"), &out, opts).expect("download");
assert_eq!(std::fs::read(&out).unwrap(), body);
cleanup(&out);
}
fn write_stream_partial(out: &Path, url: &str, body: &[u8], done: u64, total: u64, etag: &str) {
let part = resume::part_path(out);
let u = Request::get(url).unwrap();
let u = u.url();
let validators = Validators {
url: format!("{}://{}:{}{}", u.scheme, u.host, u.port, u.path),
etag: etag.into(),
last_modified: String::new(),
};
{
let mut f = std::fs::OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(&part)
.unwrap();
f.set_len(total).unwrap();
f.write_all(&body[..done as usize]).unwrap();
}
resume::write_state(
&part,
total,
Kind::HttpStream,
&stream_meta(total, done, &validators),
)
.unwrap();
}
#[test]
fn resume_416_for_a_shrunken_resource_restarts_instead_of_truncating() {
let old = make_body(1_000, 0x7070);
let new = make_body(400, 0x0707);
let origin = Origin::shared(new.clone(), "same");
let port = start(origin.clone());
let out = tmp("shrunk");
let url = format!("http://127.0.0.1:{port}/file");
write_stream_partial(&out, &url, &old, 600, 1_000, "same");
let outcome = download(&url, &out, no_backoff()).expect("download");
assert_eq!(outcome.bytes_written, 400);
assert_eq!(std::fs::read(&out).unwrap(), new);
cleanup(&out);
}
#[test]
fn resume_416_with_matching_total_is_complete() {
let body = make_body(1_000, 0x7171);
let origin = Origin::shared(body.clone(), "e");
let port = start(origin.clone());
let out = tmp("complete416");
let url = format!("http://127.0.0.1:{port}/file");
write_stream_partial(&out, &url, &body, 1_000, 1_000, "e");
let outcome = download(&url, &out, no_backoff()).expect("download");
assert_eq!(outcome.bytes_written, 1_000);
assert_eq!(std::fs::read(&out).unwrap(), body);
cleanup(&out);
}
#[test]
fn if_range_falls_back_to_last_modified_for_weak_or_missing_etag() {
let lm = "Wed, 21 Oct 2015 07:28:00 GMT";
let mut v = Validators {
url: String::new(),
etag: "W/\"weak\"".into(),
last_modified: lm.into(),
};
assert_eq!(v.if_range(), Some(lm));
v.etag.clear();
assert_eq!(v.if_range(), Some(lm));
v.etag = "\"strong\"".into();
assert_eq!(v.if_range(), Some("\"strong\""));
v.etag = "W/\"weak\"".into();
v.last_modified.clear();
assert_eq!(v.if_range(), None);
}
#[test]
fn plan_key_distinguishes_layouts_with_equal_chunk_size() {
let seg = [(0, 2), (3, 5), (6, 9)];
let size = [(0, 2), (3, 5), (6, 8), (9, 9)];
assert_ne!(plan_key(&seg), plan_key(&size));
assert_eq!(plan_key(&seg), plan_key(&[(0, 2), (3, 5), (6, 9)]));
}
#[test]
fn panicking_segment_worker_fails_the_download() {
let body = make_body(6_000, 0x6161);
let origin = Origin::shared(body, "p");
let port = start(origin);
let out = tmp("panic");
let mut opts = no_backoff();
opts.segment_size = Some(1000);
opts.parallelism = 2;
opts.progress = Some(Box::new(|_, _| {
if std::thread::current().name().is_none() {
panic!("progress callback panicked (expected by test)");
}
}));
let res = download(&format!("http://127.0.0.1:{port}/file"), &out, opts);
assert!(res.is_err(), "worker panic must fail the download");
assert!(!out.exists(), "no output may be finalized");
cleanup(&out);
}
fn file_url(p: &Path) -> String {
let s = p.display().to_string().replace('\\', "/");
if s.starts_with('/') {
format!("file://{s}")
} else {
format!("file:///{s}")
}
}
#[test]
fn failed_one_shot_transfer_keeps_existing_output() {
let out = tmp("keep");
std::fs::write(&out, b"precious").unwrap();
let missing = tmp("missing-src");
let url = file_url(&missing);
assert!(fetch_to_file(&url, &out, DownloadOptions::default()).is_err());
assert_eq!(std::fs::read(&out).unwrap(), b"precious");
let dir = out.parent().unwrap();
let prefix = format!(".{}.rsurltmp.", out.file_name().unwrap().to_string_lossy());
let leftovers = std::fs::read_dir(dir)
.unwrap()
.filter_map(|e| e.ok())
.filter(|e| e.file_name().to_string_lossy().starts_with(&prefix))
.count();
assert_eq!(leftovers, 0);
cleanup(&out);
}
#[test]
fn one_shot_sha256_mismatch_keeps_existing_output() {
let src = tmp("sha-src");
std::fs::write(&src, b"new contents").unwrap();
let out = tmp("sha-out");
std::fs::write(&out, b"old").unwrap();
let opts = DownloadOptions {
expected_sha256: Some([0u8; 32]),
..Default::default()
};
let url = file_url(&src);
assert!(fetch_to_file(&url, &out, opts).is_err());
assert_eq!(std::fs::read(&out).unwrap(), b"old");
fetch_to_file(&url, &out, DownloadOptions::default()).unwrap();
assert_eq!(std::fs::read(&out).unwrap(), b"new contents");
cleanup(&out);
let _ = std::fs::remove_file(&src);
}
}