rhtdl/

lib.rs

#![doc = include_str!("../README.md")]

#![forbid(unsafe_code)]
// TODO: integration test with traffic control (tc(8)) to simulate poor network connection

use std::fmt;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use std::ops::Deref;
use std::net::SocketAddr;
use reqwest::{Url, Client, Response};
use headers::{ContentRange, Range, HeaderMap, HeaderMapExt};
use tracing::{trace, error, warn, info, info_span, debug};

mod integrity;
mod error;
mod output;

use crate::integrity::Integrity;
pub use crate::output::Output;
pub use crate::error::Error;

pub type Result<T> = std::result::Result<T, Error>;

// this could be made into a wrapper struct with 0xFFFF reserved as a sentinel
// value, but i don't think it's worth it just to save 3 bytes.
/// Amount of times to retry an operation.
///
/// `None` = infinite retries.
///
/// note that this is the number of retries, not the number of *attempts*.
pub type RetryCount = Option<u16>;

/// Tracks how much of a file has been downloaded.
///
/// Allows resuming downloads and showing a progress bar.
///
/// This may be written to a file or database to allow gracefully resuming if a program unexpectedly terminates.
#[derive(Default, Copy, Clone, Debug)]
pub struct Progress {
    pub bytes_read: u64,
    pub bytes_total: Option<u64>,
}

impl fmt::Display for Progress {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        if let Some(total) = self.bytes_total {
            write!(f, "{}/{}", self.bytes_read, total)?;
        } else {
            write!(f, "{}/??", self.bytes_read)?;
        }
        Ok(())
    }
}

impl Progress {
    fn complete(&self) -> bool {
        if let Some(total) = self.bytes_total {
            self.bytes_read >= total
        } else {
            false
        }
    }

    /// Recover the progress of a download based on the length of its output file.
    ///
    /// `total_len` will be set to `None`.
    pub fn from_file_len(file: &std::fs::File) -> std::io::Result<Progress> {
        Ok(Progress{
            bytes_read: file.metadata()?.len(),
            bytes_total: None,
        })
    }
}

/// parameters for how a file should be downloaded.
///
/// the default was carefully selected to be suitable for most usecases over
/// virtually any type of network.
///
/// unlike [`Options`], this may be reused for multiple downloads.
#[derive(Debug, Clone)]
// TODO: optional serde support
pub struct Config {
    /// if `Some`, abort the download after the specified duration has passed.
    pub timeout: Option<Duration>,
    /// if the specified duration passes with no progress being made,
    /// abort the download.
    pub stalled_timeout: Option<Duration>,
    /// number of times to retry the download while no progress is being made.
    ///
    /// None = infinite retries
    ///
    /// Note that this is the number of *retries*, which is one less than
    /// the number of attempts.  With a retry count of 0, exactly
    /// one attempt will be made, and it will fail if any error occurs.
    ///
    /// The larger this number, the longer a burst of spurious errors
    /// will need to be in order to promoted to a fatal error.
    pub stalled_retries: RetryCount,
    /// total number of times to retry the download
    ///
    /// Using this is not recommended unless you know you will be
    /// downloading over a stable network.  instead try configuring
    /// the stalled download timeout or overall deadline.
    pub max_retries: RetryCount,
    /// maximum number of times to restart the download from the beginning
    ///
    /// usually, if there is an error, `rhtdl` will use the http Range
    /// header to resume the download from where it left off,
    /// but if the server does not implement `Range` requests, it will
    /// restart the download from the beginning.
    ///
    /// the download will also be restarted if `rhtdl` detects that the
    /// resource has been modified (eg. due to an `ETag` change).
    pub max_restarts: RetryCount,
    /// default size of transfer chunks requested with the Range header.
    ///
    /// larger chunk sizes will usually result in faster downloads,
    /// (especially if the remote server refuses to keep the connection open
    /// after a request), but will use more memory.
    pub chunk_size: u64,
    /// base retry delay used for exponential backoff.
    ///
    /// if there is only one error, this is the `amount` of time that will be
    /// waited before retrying.
    pub backoff_base: Duration,
    // TODO: arbitrary headers
}

impl Default for Config {
    fn default() -> Config {
        // very generous defaults sutible for downloading massive files over
        // horrible connections.
        Self {
            timeout: None,
            stalled_timeout: Some(Duration::from_secs(3*60)),
            stalled_retries: Some(7),
            max_retries: None,
            max_restarts: Some(3),
            chunk_size: 1_024_000,
            backoff_base: Duration::from_millis(200),
        }
    }
}

impl Config {
    /// set infinite timeout and retries.
    ///
    /// useful for background downloads that should
    /// keep going unless the user cancels them,
    /// such as a podcast player.
    #[inline]
    pub fn never_give_up(&mut self) -> &mut Self {
        self.timeout = None;
        self.max_retries = None;
        self.stalled_retries = None;
        self.max_restarts = None;
        // TODO: timeouts
        self
    }
}

// TODO: bon::builder so the internal representation can be modified?
// TODO: split the immutable data (configuration options) from the mutable data like `progress` so an Options struct can be easily reused for multiple downloads.
/// structure that controls the behavior of a specific download.
///
/// this contains an `Arc<Config>`, but also other data, such as channels for
/// reporting the progress of the download.
#[derive(Default)]
pub struct Options {
    config: Arc<Config>,
    /// place to report progress to
    progress: Option<Arc<Mutex<Progress>>>,
    deadline: Option<Instant>,
    address_callback: Option<Box<dyn Fn(SocketAddr)>>,
}

impl fmt::Debug for Options {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // FIXME: do this properly
        write!(f, "{:?}", &self.config)
    }
}

impl Deref for Options {
    type Target = Config;

    fn deref(&self) -> &Self::Target {
        &self.config 
    }
}

impl From<Arc<Config>> for Options {
    #[inline]
    fn from(config: Arc<Config>) -> Self {
        Options {
            config,
            deadline: None,
            progress: None,
            address_callback: None,
        }
    }
}

impl From<Config> for Options {
    #[inline]
    fn from(config: Config) -> Self {
        Arc::new(config).into()
    }
}

impl Options {
    /// return a handle that can be used to track the progress of the download
    ///
    /// # Panics
    /// this method panics if it is called multiple times on the same `Options`.
    pub fn progress(&mut self) -> Arc<Mutex<Progress>> {
        assert!(self.progress.is_none());
        let handle = Arc::new(Mutex::new(Progress::default()));
        self.progress = Some(handle.clone());
        handle
    }

    /// registers a callback that is called as soon as an address is obtained.
    ///
    /// this mostly exists for the purpose of integration testing.
    pub fn register_address_callback(&mut self, callback: Box<dyn Fn(SocketAddr)>) {
        self.address_callback = Some(callback);
    }

    fn deadline_expired(&self) -> bool {
        if let Some(deadline) = self.deadline {
            Instant::now().saturating_duration_since(deadline) > Duration::from_secs(0)
        } else {
            false
        }
    }

    fn deadline_start(&mut self) {
        if self.deadline.is_none() && self.timeout.is_some() {
            self.deadline = Some(Instant::now() + self.timeout.unwrap());
        }
    }
}

// TODO: use Content-Digest to verify the download, if available.
// TODO: multiple download workers, each connecting to the host over a different IP, if available. two `RangeSet`s behind a mutex to coordinate parallel range downloads.  one `RangeSet` stores the blocks that are currently being downloaded by a worker, and one stores the blocks that have already been downloaded (maybe it would be better to invert this slightly, and have map of ranges that have been completed, and a map of ranges yet to be delegated)
// TODO: use If-Range to have the server restart the download if the etag changes.
/// a download task that is in progress
struct Download<O> {
    options: Options,
    url: Url,
    /// internal progress status
    progress: Progress,
    resp: Option<Response>,
    error_count: u32,
    restart_count: u16,
    /// fatal error encountered
    fatal: bool,
    client: Client,
    /// time since a data byte was received
    last_progress_made: Instant,
    /// number of errors since a data byte was received
    ///
    /// used for `stalled_retries` and exponential backoff
    // TODO: wrap all error counts in `Saturating`.
    errors_since_last_progress: u16,
    output: O,
    integrity: Integrity,
    last_error: Option<Error>,
}

// TODO: this should never panic, use that binary analysis thing to assert that.
/// Used to resume a download across program restarts.
///
/// The `bytes_read` of `progress` must exactly match the amount of bytes previously written to `output`, but the `bytes_total` may be `None`.
/// The easiest way to get this is via `Progress::from_file_len`.
pub async fn resume(url: Url, options: Options, output: impl Output, progress: Progress) -> Result<()> {
    use tracing::Instrument;
    let span = info_span!("download", url = url.to_string());
    let mut dl = Download {
        options,
        output,
        url,
        progress,
        error_count: 0,
        restart_count: 0,
        last_progress_made: Instant::now(),
        errors_since_last_progress: 0,
        fatal: false,
        resp: None,
        client: Client::new(),
        integrity: Integrity::default(),
        last_error: None,
    };
    dl.resume().instrument(span).await
}

/// The central function of the crate.
///
/// Downloads `url` using `options`, saving the result to `output`.
///
/// `output` must be empty.
pub async fn download(url: Url, options: impl Into<Options>, output: impl Output) -> Result<()> {
    resume(url, options.into(), output, Progress::default()).await?;
    Ok(())
}

impl<O: Output> Download<O> {
    async fn try_step_inner(&mut self) -> Result<()> {
        if let Some(resp) = &mut self.resp {
            if resp.url() != &self.url {
                // redirect happened, save the new url,
                // since the older url might me a "latest"
                // url that gets changed to point to a new url
                // whenever a new version is published,
                // but we don't want to splice two versions together.
                self.url = resp.url().clone();
            }
            let Some(chunk) = resp.chunk().await? else {
                self.resp = None;
                return Ok(())
            };
            self.progress.bytes_read += chunk.len() as u64;
            trace!("recieved {} bytes", chunk.len());
            if let Some(prog) = &self.options.progress {
                // if the progress is currently being read, don't block,
                // just wait until next pass (maybe this should use buf2sync instead to make it lockless?)
                if let Ok(mut prog) = prog.try_lock() {
                    *prog = self.progress;
                }
            }
            self.last_progress_made = Instant::now();
            self.errors_since_last_progress = 0;
            self.output.write_all(&chunk)?;
        } else {
            let rstart = self.progress.bytes_read;
            let rend = rstart + self.options.chunk_size;
            let mut req = self.client.get(self.url.clone());
            let mut headers = HeaderMap::new();
            headers.typed_insert(Range::bytes(rstart..rend).unwrap());
            // TODO: use if_range around here? will also need to slighly tweak retry logic to avoid wasting a request restarting a download that has already restarted.
            req = req.headers(headers);
            if let Some(deadline) = self.options.deadline {
                // if deadline expires, reqwest will give an error.
                req = req.timeout(deadline.saturating_duration_since(Instant::now()));
            }
            let resp = req.send().await?;
            if let Some(ac) = &self.options.address_callback {
                if let Some(ra) = resp.remote_addr() {
                    ac(ra);
                }
            }
            let maybe_cr = resp.headers().typed_get::<ContentRange>();
            let real_rstart = maybe_cr
                .as_ref()
                .and_then(headers::ContentRange::bytes_range)
                .map_or(0, |x| x.0);
            if real_rstart != rstart {
                // if rstart != 0, and Content-Range is missing,
                // we can't continue the download.  retry from start.
                warn!("expected to recive a response starting at {rstart}, instead recived a chunk starting at {real_rstart}");
                return self.restart();
            }
            if let Err(e) = self.integrity.ingest_headers(resp.headers()) {
                error!("mid-air collision detected, restarting download ({e})");
                self.last_error = Some(e);
                return self.restart();
            }
            if self.progress.bytes_total.is_none() {
                if let Some(cr) = maybe_cr {
                    self.progress.bytes_total = cr.bytes_len();
                } else {
                    self.progress.bytes_total = resp.content_length();
                }
            } else if let Some(cr) = maybe_cr {
                let old_len = self.progress.bytes_total;
                let new_len = cr.bytes_len();
                if old_len != new_len {
                    warn!(
                        old = old_len,
                        new = new_len,
                        "total length of resource changed",
                    );
                    return self.restart();
                }
            }
            self.resp = Some(resp);
        }
        Ok(())
    }

    async fn try_step(&mut self) -> Result<()> {
        if let Some(deadline) = self.stall_deadline() {
            tokio::time::timeout_at(deadline.into(), self.try_step_inner()).await?
        } else {
            self.try_step_inner().await
        }
    }
    
    // only returns the last error
    async fn resume(&mut self) -> Result<()> {
        self.options.deadline_start();
        debug!("downloading from offset {}", self.progress.bytes_read);
        while !self.progress.complete() {
            match self.try_step().await {
                Ok(()) => {},
                Err(err) => {
                    if err.is_fatal() {
                        error!("fatal error: {err}");
                        return Err(err);
                    }
                    if err.is_timeout() {
                        // since there is no way to construct
                        // reqwest::Error, we need to get reqwest to
                        // make one for us, by using its timeout
                        // TODO: now that we have our own error type, perhaps reconsider?
                        if self.options.deadline_expired() {
                            error!("deadline expired");
                            return Err(err);
                        }
                    }
                    // exponential backoff.
                    // is computed before incrementing error_count to
                    // avoid a (mostly insignificant) off-by-one error.
                    // TODO: "recovery time", if there are no errors for a long period, reduce the retry delay
                    let delay = self.options.backoff_base * 2_u16.saturating_pow(self.error_count).into();
                    self.error_count += 1;
                    self.errors_since_last_progress += 1;
                    warn!(error_count = self.error_count, "non-fatal error: {err:?}");
                    if let Some(max) = self.options.max_retries {
                        if self.error_count >= max.into() {
                            error!("max_retries exceeded");
                            return Err(err);
                        }
                    }
                    if let Some(max) = self.options.stalled_retries {
                        if self.error_count >= max.into() {
                            error!(stalled_retries = self.options.stalled_retries, "stalled_retries exceeded");
                            return Err(err);
                        }
                    }
                    self.last_error = Some(err);
                    info!("retrying in {} ms", delay.as_millis());
                    tokio::time::sleep(delay).await;
                },
            }
        }
        if self.fatal {
            if let Some(e) = self.last_error.take() {
                error!("previous error upgraded to fatal: {e}");
                return Err(e);
            }
        }
        Ok(())
    }

    /// rewind to the start and retry the download from there.
    fn restart(&mut self) -> Result<()> {
        warn!("restarting download");
        if let Some(max) = self.options.max_restarts {
            error!("maximum restart count exceeded");
            if self.restart_count >= max {
                self.fatal = true;
                return Ok(());
            }
        }
        self.restart_count = self.restart_count.saturating_add(1);
        self.progress.bytes_read = 0;
        self.resp = None;
        self.output.restart()?;
        self.integrity = Default::default();
        Ok(())
    }

    fn stall_deadline(&self) -> Option<Instant> {
        self.options.stalled_timeout.map(|x| self.last_progress_made + x)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /*#[tokio::test]
    async fn failure() {
        // for some reason, my ISP will redirect *any* unknown domain to
        // the router login page, making this test impossible.
        return;
        // TODO: check the actual error
        assert_eq!(
            resume("http://bad-domain.example/".parse().unwrap(),
                   Options::default(), &std::io::sink(),
                   Progress::default()).await.unwrap_err().to_string(),
            "error sending request for url (http://bad-domain.example/)");
    }*/

    #[tokio::test]
    async fn success() {
        let mut out = Vec::<u8>::new();
        resume("https://raw.githubusercontent.com/unicode-org/cldr-json/44.0.1/cldr-json/cldr-annotations-modern/annotations/en-001/annotations.json".parse().unwrap(), Options::default(), &mut out, Progress::default()).await.unwrap();
    }

    #[tokio::test]
    async fn short_timeout() {
        // TODO: capture the "tracing" output to make sure it mentions a timeout/deadline
        let mut cfg = Config::default();
        cfg.timeout = Some(Duration::from_nanos(10));
        let result = download(
            "https://static.rust-lang.org/dist/2024-07-26/cargo-nightly-powerpc64-unknown-linux-gnu.tar.xz".parse().unwrap(),
            cfg, &std::io::sink())
            .await;
        assert!(result.is_err());
        assert!(result.unwrap_err().is_timeout());
    }
}