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());
}
}