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, Error, Result};
#[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?.error_for_status()?;
return Ok(response.text().await?);
}
"file" => {
let path = url
.to_file_path()
.map_err(|_| Error::UnsupportedUrl(url_or_path.to_string()))?;
return tokio::fs::read_to_string(&path)
.await
.map_err(|error| io_path(path, error));
}
_ => return Err(Error::UnsupportedUrl(url_or_path.to_string())),
}
}
tokio::fs::read_to_string(url_or_path)
.await
.map_err(|error| io_path(url_or_path, error))
}
pub async fn download_to_file<F>(
url_or_path: &str,
destination: impl AsRef<Path>,
headers: &[HttpHeader],
timeout_secs: 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,
on_event,
)
.await;
match result {
Ok(mut stats) => {
if tokio::fs::metadata(destination).await.is_ok() {
tokio::fs::remove_file(destination)
.await
.map_err(|error| io_path(destination, error))?;
}
tokio::fs::rename(&temporary_destination, destination)
.await
.map_err(|error| io_path(destination, 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>,
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, 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, on_event).await
}
_ => Err(Error::UnsupportedUrl(url_or_path.to_string())),
};
}
copy_file_with_progress(url_or_path, destination, 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>,
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?.error_for_status()?;
let content_length = response.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?;
file.write_all(&chunk)
.await
.map_err(|error| io_path(destination, error))?;
written += chunk.len() as u64;
on_event(DownloadEvent::Progress {
chunk_length: chunk.len(),
});
}
file.flush()
.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>,
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))?;
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;
}
output
.write_all(&buf[..read])
.await
.map_err(|error| io_path(destination, error))?;
written += read as u64;
on_event(DownloadEvent::Progress { chunk_length: read });
}
output
.flush()
.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();
if let Some(timeout_secs) = timeout_secs {
builder = builder.timeout(Duration::from_secs(timeout_secs));
}
Ok(builder.build()?)
}
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()))
}