hdiff_update_core/

download.rs

use std::{
    path::{Path, PathBuf},
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use url::Url;

use crate::{error::io_path, fs_ops::replace_file, Error, Result};

const MAX_TEXT_RESPONSE_BYTES: u64 = 4 * 1024 * 1024;

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HttpHeader {
    pub name: String,
    pub value: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "event", content = "data", rename_all = "camelCase")]
pub enum DownloadEvent {
    Started { content_length: Option<u64> },
    Progress { chunk_length: usize },
    Finished,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct DownloadStats {
    pub path: PathBuf,
    pub bytes_written: u64,
}

pub async fn read_url_to_string(
    url_or_path: &str,
    headers: &[HttpHeader],
    timeout_secs: Option<u64>,
) -> Result<String> {
    if let Some(url) = parse_supported_url_or_path(url_or_path)? {
        match url.scheme() {
            "http" | "https" => {
                let client = client(timeout_secs)?;
                let mut request = client.get(url);
                for header in headers {
                    request = request.header(&header.name, &header.value);
                }
                let response = request.send().await?;
                ensure_success_status(&response)?;
                return response_text_limited(response).await;
            }
            "file" => {
                let path = url
                    .to_file_path()
                    .map_err(|_| Error::UnsupportedUrl(url_or_path.to_string()))?;
                return read_text_file_limited(&path).await;
            }
            _ => return Err(Error::UnsupportedUrl(url_or_path.to_string())),
        }
    }

    read_text_file_limited(Path::new(url_or_path)).await
}

async fn response_text_limited(response: reqwest::Response) -> Result<String> {
    if response
        .content_length()
        .is_some_and(|length| length > MAX_TEXT_RESPONSE_BYTES)
    {
        return Err(Error::DownloadLimitExceeded {
            limit: MAX_TEXT_RESPONSE_BYTES,
            attempted: response.content_length().unwrap_or_default(),
        });
    }
    let mut stream = response.bytes_stream();
    let mut bytes = Vec::new();
    while let Some(chunk) = stream.next().await {
        let chunk = chunk?;
        let attempted = bytes.len() as u64 + chunk.len() as u64;
        if attempted > MAX_TEXT_RESPONSE_BYTES {
            return Err(Error::DownloadLimitExceeded {
                limit: MAX_TEXT_RESPONSE_BYTES,
                attempted,
            });
        }
        bytes.extend_from_slice(&chunk);
    }
    String::from_utf8(bytes)
        .map_err(|error| Error::Message(format!("text response is not valid UTF-8: {error}")))
}

async fn read_text_file_limited(path: &Path) -> Result<String> {
    let metadata = tokio::fs::metadata(path)
        .await
        .map_err(|error| io_path(path, error))?;
    if metadata.len() > MAX_TEXT_RESPONSE_BYTES {
        return Err(Error::DownloadLimitExceeded {
            limit: MAX_TEXT_RESPONSE_BYTES,
            attempted: metadata.len(),
        });
    }
    tokio::fs::read_to_string(path)
        .await
        .map_err(|error| io_path(path, error))
}

pub async fn download_to_file<F>(
    url_or_path: &str,
    destination: impl AsRef<Path>,
    headers: &[HttpHeader],
    timeout_secs: Option<u64>,
    maximum_bytes: Option<u64>,
    on_event: F,
) -> Result<DownloadStats>
where
    F: FnMut(DownloadEvent),
{
    let destination = destination.as_ref();
    if let Some(parent) = destination.parent() {
        tokio::fs::create_dir_all(parent)
            .await
            .map_err(|error| io_path(parent, error))?;
    }

    let temporary_destination = temporary_path_for(destination);
    let result = download_to_temporary_file(
        url_or_path,
        &temporary_destination,
        headers,
        timeout_secs,
        maximum_bytes,
        on_event,
    )
    .await;

    match result {
        Ok(mut stats) => {
            let temporary = temporary_destination.clone();
            let destination_path = destination.to_path_buf();
            tokio::task::spawn_blocking(move || replace_file(&temporary, &destination_path))
                .await
                .map_err(|error| {
                    Error::Message(format!("download replace task failed: {error}"))
                })??;
            stats.path = destination.to_path_buf();
            Ok(stats)
        }
        Err(error) => {
            let _ = tokio::fs::remove_file(&temporary_destination).await;
            Err(error)
        }
    }
}

async fn download_to_temporary_file<F>(
    url_or_path: &str,
    destination: &Path,
    headers: &[HttpHeader],
    timeout_secs: Option<u64>,
    maximum_bytes: Option<u64>,
    on_event: F,
) -> Result<DownloadStats>
where
    F: FnMut(DownloadEvent),
{
    if let Some(url) = parse_supported_url_or_path(url_or_path)? {
        return match url.scheme() {
            "http" | "https" => {
                download_http(
                    url,
                    destination,
                    headers,
                    timeout_secs,
                    maximum_bytes,
                    on_event,
                )
                .await
            }
            "file" => {
                let source = url
                    .to_file_path()
                    .map_err(|_| Error::UnsupportedUrl(url_or_path.to_string()))?;
                copy_file_with_progress(&source, destination, maximum_bytes, on_event).await
            }
            _ => Err(Error::UnsupportedUrl(url_or_path.to_string())),
        };
    }

    copy_file_with_progress(url_or_path, destination, maximum_bytes, on_event).await
}

fn parse_supported_url_or_path(value: &str) -> Result<Option<Url>> {
    match Url::parse(value) {
        Ok(url) if matches!(url.scheme(), "http" | "https" | "file") => Ok(Some(url)),
        Ok(_) if value.contains("://") => Err(Error::UnsupportedUrl(value.to_string())),
        Ok(_) | Err(_) => Ok(None),
    }
}

async fn download_http<F>(
    url: Url,
    destination: &Path,
    headers: &[HttpHeader],
    timeout_secs: Option<u64>,
    maximum_bytes: Option<u64>,
    mut on_event: F,
) -> Result<DownloadStats>
where
    F: FnMut(DownloadEvent),
{
    let client = client(timeout_secs)?;
    let mut request = client.get(url);
    for header in headers {
        request = request.header(&header.name, &header.value);
    }

    let response = request.send().await?;
    ensure_success_status(&response)?;
    let content_length = response.content_length();
    if let (Some(limit), Some(content_length)) = (maximum_bytes, content_length) {
        if content_length > limit {
            return Err(Error::DownloadLimitExceeded {
                limit,
                attempted: content_length,
            });
        }
    }
    on_event(DownloadEvent::Started { content_length });

    let mut stream = response.bytes_stream();
    let mut file = tokio::fs::File::create(destination)
        .await
        .map_err(|error| io_path(destination, error))?;
    let mut written = 0_u64;

    while let Some(chunk) = stream.next().await {
        let chunk = chunk?;
        let attempted = written
            .checked_add(chunk.len() as u64)
            .ok_or_else(|| Error::Message("download size overflow".to_string()))?;
        if maximum_bytes.is_some_and(|limit| attempted > limit) {
            return Err(Error::DownloadLimitExceeded {
                limit: maximum_bytes.unwrap_or_default(),
                attempted,
            });
        }
        file.write_all(&chunk)
            .await
            .map_err(|error| io_path(destination, error))?;
        written = attempted;
        on_event(DownloadEvent::Progress {
            chunk_length: chunk.len(),
        });
    }

    file.flush()
        .await
        .map_err(|error| io_path(destination, error))?;
    file.sync_all()
        .await
        .map_err(|error| io_path(destination, error))?;
    on_event(DownloadEvent::Finished);

    Ok(DownloadStats {
        path: destination.to_path_buf(),
        bytes_written: written,
    })
}

async fn copy_file_with_progress<F>(
    source: impl AsRef<Path>,
    destination: impl AsRef<Path>,
    maximum_bytes: Option<u64>,
    mut on_event: F,
) -> Result<DownloadStats>
where
    F: FnMut(DownloadEvent),
{
    let source = source.as_ref();
    let destination = destination.as_ref();
    let mut input = tokio::fs::File::open(source)
        .await
        .map_err(|error| io_path(source, error))?;
    let metadata = input
        .metadata()
        .await
        .map_err(|error| io_path(source, error))?;
    if maximum_bytes.is_some_and(|limit| metadata.len() > limit) {
        return Err(Error::DownloadLimitExceeded {
            limit: maximum_bytes.unwrap_or_default(),
            attempted: metadata.len(),
        });
    }
    let mut output = tokio::fs::File::create(destination)
        .await
        .map_err(|error| io_path(destination, error))?;
    let mut buf = vec![0_u8; 256 * 1024];
    let mut written = 0_u64;

    on_event(DownloadEvent::Started {
        content_length: Some(metadata.len()),
    });
    loop {
        let read = input
            .read(&mut buf)
            .await
            .map_err(|error| io_path(source, error))?;
        if read == 0 {
            break;
        }
        let attempted = written
            .checked_add(read as u64)
            .ok_or_else(|| Error::Message("download size overflow".to_string()))?;
        if maximum_bytes.is_some_and(|limit| attempted > limit) {
            return Err(Error::DownloadLimitExceeded {
                limit: maximum_bytes.unwrap_or_default(),
                attempted,
            });
        }
        output
            .write_all(&buf[..read])
            .await
            .map_err(|error| io_path(destination, error))?;
        written = attempted;
        on_event(DownloadEvent::Progress { chunk_length: read });
    }

    output
        .flush()
        .await
        .map_err(|error| io_path(destination, error))?;
    output
        .sync_all()
        .await
        .map_err(|error| io_path(destination, error))?;
    on_event(DownloadEvent::Finished);

    Ok(DownloadStats {
        path: destination.to_path_buf(),
        bytes_written: written,
    })
}

fn client(timeout_secs: Option<u64>) -> Result<reqwest::Client> {
    let mut builder = reqwest::Client::builder().redirect(reqwest::redirect::Policy::none());
    if let Some(timeout_secs) = timeout_secs {
        builder = builder.timeout(Duration::from_secs(timeout_secs));
    }
    Ok(builder.build()?)
}

fn ensure_success_status(response: &reqwest::Response) -> Result<()> {
    if response.status().is_success() {
        return Ok(());
    }
    Err(Error::UnexpectedHttpStatus {
        status: response.status().as_u16(),
    })
}

fn temporary_path_for(destination: &Path) -> PathBuf {
    let file_name = destination
        .file_name()
        .map(|name| name.to_string_lossy())
        .unwrap_or_else(|| "download".into());
    let suffix = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_nanos())
        .unwrap_or_default();
    destination.with_file_name(format!(".{file_name}.part-{}-{suffix}", std::process::id()))
}

#[cfg(test)]
mod tests {
    use std::fs;

    use tempfile::tempdir;

    use super::download_to_file;
    use crate::Error;

    #[tokio::test]
    async fn local_download_stops_at_the_signed_size_limit() {
        let dir = tempdir().unwrap();
        let source = dir.path().join("source.bin");
        let destination = dir.path().join("destination.bin");
        fs::write(&source, b"0123456789").unwrap();

        let error = download_to_file(
            source.to_str().unwrap(),
            &destination,
            &[],
            None,
            Some(5),
            |_| {},
        )
        .await
        .unwrap_err();

        assert!(matches!(error, Error::DownloadLimitExceeded { .. }));
        assert!(!destination.exists());
    }
}