1use std::{
2 path::{Path, PathBuf},
3 time::{Duration, SystemTime, UNIX_EPOCH},
4};
5
6use futures_util::StreamExt;
7use serde::{Deserialize, Serialize};
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9use url::Url;
10
11use crate::{error::io_path, Error, Result};
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14#[serde(rename_all = "camelCase")]
15pub struct HttpHeader {
16 pub name: String,
17 pub value: String,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
21#[serde(tag = "event", content = "data", rename_all = "camelCase")]
22pub enum DownloadEvent {
23 Started { content_length: Option<u64> },
24 Progress { chunk_length: usize },
25 Finished,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
29#[serde(rename_all = "camelCase")]
30pub struct DownloadStats {
31 pub path: PathBuf,
32 pub bytes_written: u64,
33}
34
35pub async fn read_url_to_string(
36 url_or_path: &str,
37 headers: &[HttpHeader],
38 timeout_secs: Option<u64>,
39) -> Result<String> {
40 if let Some(url) = parse_supported_url_or_path(url_or_path)? {
41 match url.scheme() {
42 "http" | "https" => {
43 let client = client(timeout_secs)?;
44 let mut request = client.get(url);
45 for header in headers {
46 request = request.header(&header.name, &header.value);
47 }
48 let response = request.send().await?.error_for_status()?;
49 return Ok(response.text().await?);
50 }
51 "file" => {
52 let path = url
53 .to_file_path()
54 .map_err(|_| Error::UnsupportedUrl(url_or_path.to_string()))?;
55 return tokio::fs::read_to_string(&path)
56 .await
57 .map_err(|error| io_path(path, error));
58 }
59 _ => return Err(Error::UnsupportedUrl(url_or_path.to_string())),
60 }
61 }
62
63 tokio::fs::read_to_string(url_or_path)
64 .await
65 .map_err(|error| io_path(url_or_path, error))
66}
67
68pub async fn download_to_file<F>(
69 url_or_path: &str,
70 destination: impl AsRef<Path>,
71 headers: &[HttpHeader],
72 timeout_secs: Option<u64>,
73 on_event: F,
74) -> Result<DownloadStats>
75where
76 F: FnMut(DownloadEvent),
77{
78 let destination = destination.as_ref();
79 if let Some(parent) = destination.parent() {
80 tokio::fs::create_dir_all(parent)
81 .await
82 .map_err(|error| io_path(parent, error))?;
83 }
84
85 let temporary_destination = temporary_path_for(destination);
86 let result = download_to_temporary_file(
87 url_or_path,
88 &temporary_destination,
89 headers,
90 timeout_secs,
91 on_event,
92 )
93 .await;
94
95 match result {
96 Ok(mut stats) => {
97 if tokio::fs::metadata(destination).await.is_ok() {
98 tokio::fs::remove_file(destination)
99 .await
100 .map_err(|error| io_path(destination, error))?;
101 }
102 tokio::fs::rename(&temporary_destination, destination)
103 .await
104 .map_err(|error| io_path(destination, error))?;
105 stats.path = destination.to_path_buf();
106 Ok(stats)
107 }
108 Err(error) => {
109 let _ = tokio::fs::remove_file(&temporary_destination).await;
110 Err(error)
111 }
112 }
113}
114
115async fn download_to_temporary_file<F>(
116 url_or_path: &str,
117 destination: &Path,
118 headers: &[HttpHeader],
119 timeout_secs: Option<u64>,
120 on_event: F,
121) -> Result<DownloadStats>
122where
123 F: FnMut(DownloadEvent),
124{
125 if let Some(url) = parse_supported_url_or_path(url_or_path)? {
126 return match url.scheme() {
127 "http" | "https" => {
128 download_http(url, destination, headers, timeout_secs, on_event).await
129 }
130 "file" => {
131 let source = url
132 .to_file_path()
133 .map_err(|_| Error::UnsupportedUrl(url_or_path.to_string()))?;
134 copy_file_with_progress(&source, destination, on_event).await
135 }
136 _ => Err(Error::UnsupportedUrl(url_or_path.to_string())),
137 };
138 }
139
140 copy_file_with_progress(url_or_path, destination, on_event).await
141}
142
143fn parse_supported_url_or_path(value: &str) -> Result<Option<Url>> {
144 match Url::parse(value) {
145 Ok(url) if matches!(url.scheme(), "http" | "https" | "file") => Ok(Some(url)),
146 Ok(_) if value.contains("://") => Err(Error::UnsupportedUrl(value.to_string())),
147 Ok(_) | Err(_) => Ok(None),
148 }
149}
150
151async fn download_http<F>(
152 url: Url,
153 destination: &Path,
154 headers: &[HttpHeader],
155 timeout_secs: Option<u64>,
156 mut on_event: F,
157) -> Result<DownloadStats>
158where
159 F: FnMut(DownloadEvent),
160{
161 let client = client(timeout_secs)?;
162 let mut request = client.get(url);
163 for header in headers {
164 request = request.header(&header.name, &header.value);
165 }
166
167 let response = request.send().await?.error_for_status()?;
168 let content_length = response.content_length();
169 on_event(DownloadEvent::Started { content_length });
170
171 let mut stream = response.bytes_stream();
172 let mut file = tokio::fs::File::create(destination)
173 .await
174 .map_err(|error| io_path(destination, error))?;
175 let mut written = 0_u64;
176
177 while let Some(chunk) = stream.next().await {
178 let chunk = chunk?;
179 file.write_all(&chunk)
180 .await
181 .map_err(|error| io_path(destination, error))?;
182 written += chunk.len() as u64;
183 on_event(DownloadEvent::Progress {
184 chunk_length: chunk.len(),
185 });
186 }
187
188 file.flush()
189 .await
190 .map_err(|error| io_path(destination, error))?;
191 on_event(DownloadEvent::Finished);
192
193 Ok(DownloadStats {
194 path: destination.to_path_buf(),
195 bytes_written: written,
196 })
197}
198
199async fn copy_file_with_progress<F>(
200 source: impl AsRef<Path>,
201 destination: impl AsRef<Path>,
202 mut on_event: F,
203) -> Result<DownloadStats>
204where
205 F: FnMut(DownloadEvent),
206{
207 let source = source.as_ref();
208 let destination = destination.as_ref();
209 let mut input = tokio::fs::File::open(source)
210 .await
211 .map_err(|error| io_path(source, error))?;
212 let metadata = input
213 .metadata()
214 .await
215 .map_err(|error| io_path(source, error))?;
216 let mut output = tokio::fs::File::create(destination)
217 .await
218 .map_err(|error| io_path(destination, error))?;
219 let mut buf = vec![0_u8; 256 * 1024];
220 let mut written = 0_u64;
221
222 on_event(DownloadEvent::Started {
223 content_length: Some(metadata.len()),
224 });
225 loop {
226 let read = input
227 .read(&mut buf)
228 .await
229 .map_err(|error| io_path(source, error))?;
230 if read == 0 {
231 break;
232 }
233 output
234 .write_all(&buf[..read])
235 .await
236 .map_err(|error| io_path(destination, error))?;
237 written += read as u64;
238 on_event(DownloadEvent::Progress { chunk_length: read });
239 }
240
241 output
242 .flush()
243 .await
244 .map_err(|error| io_path(destination, error))?;
245 on_event(DownloadEvent::Finished);
246
247 Ok(DownloadStats {
248 path: destination.to_path_buf(),
249 bytes_written: written,
250 })
251}
252
253fn client(timeout_secs: Option<u64>) -> Result<reqwest::Client> {
254 let mut builder = reqwest::Client::builder();
255 if let Some(timeout_secs) = timeout_secs {
256 builder = builder.timeout(Duration::from_secs(timeout_secs));
257 }
258 Ok(builder.build()?)
259}
260
261fn temporary_path_for(destination: &Path) -> PathBuf {
262 let file_name = destination
263 .file_name()
264 .map(|name| name.to_string_lossy())
265 .unwrap_or_else(|| "download".into());
266 let suffix = SystemTime::now()
267 .duration_since(UNIX_EPOCH)
268 .map(|duration| duration.as_nanos())
269 .unwrap_or_default();
270 destination.with_file_name(format!(".{file_name}.part-{}-{suffix}", std::process::id()))
271}