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 Ok(url) = Url::parse(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 Ok(url) = Url::parse(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
143async fn download_http<F>(
144 url: Url,
145 destination: &Path,
146 headers: &[HttpHeader],
147 timeout_secs: Option<u64>,
148 mut on_event: F,
149) -> Result<DownloadStats>
150where
151 F: FnMut(DownloadEvent),
152{
153 let client = client(timeout_secs)?;
154 let mut request = client.get(url);
155 for header in headers {
156 request = request.header(&header.name, &header.value);
157 }
158
159 let response = request.send().await?.error_for_status()?;
160 let content_length = response.content_length();
161 on_event(DownloadEvent::Started { content_length });
162
163 let mut stream = response.bytes_stream();
164 let mut file = tokio::fs::File::create(destination)
165 .await
166 .map_err(|error| io_path(destination, error))?;
167 let mut written = 0_u64;
168
169 while let Some(chunk) = stream.next().await {
170 let chunk = chunk?;
171 file.write_all(&chunk)
172 .await
173 .map_err(|error| io_path(destination, error))?;
174 written += chunk.len() as u64;
175 on_event(DownloadEvent::Progress {
176 chunk_length: chunk.len(),
177 });
178 }
179
180 file.flush()
181 .await
182 .map_err(|error| io_path(destination, error))?;
183 on_event(DownloadEvent::Finished);
184
185 Ok(DownloadStats {
186 path: destination.to_path_buf(),
187 bytes_written: written,
188 })
189}
190
191async fn copy_file_with_progress<F>(
192 source: impl AsRef<Path>,
193 destination: impl AsRef<Path>,
194 mut on_event: F,
195) -> Result<DownloadStats>
196where
197 F: FnMut(DownloadEvent),
198{
199 let source = source.as_ref();
200 let destination = destination.as_ref();
201 let mut input = tokio::fs::File::open(source)
202 .await
203 .map_err(|error| io_path(source, error))?;
204 let metadata = input
205 .metadata()
206 .await
207 .map_err(|error| io_path(source, error))?;
208 let mut output = tokio::fs::File::create(destination)
209 .await
210 .map_err(|error| io_path(destination, error))?;
211 let mut buf = vec![0_u8; 256 * 1024];
212 let mut written = 0_u64;
213
214 on_event(DownloadEvent::Started {
215 content_length: Some(metadata.len()),
216 });
217 loop {
218 let read = input
219 .read(&mut buf)
220 .await
221 .map_err(|error| io_path(source, error))?;
222 if read == 0 {
223 break;
224 }
225 output
226 .write_all(&buf[..read])
227 .await
228 .map_err(|error| io_path(destination, error))?;
229 written += read as u64;
230 on_event(DownloadEvent::Progress { chunk_length: read });
231 }
232
233 output
234 .flush()
235 .await
236 .map_err(|error| io_path(destination, error))?;
237 on_event(DownloadEvent::Finished);
238
239 Ok(DownloadStats {
240 path: destination.to_path_buf(),
241 bytes_written: written,
242 })
243}
244
245fn client(timeout_secs: Option<u64>) -> Result<reqwest::Client> {
246 let mut builder = reqwest::Client::builder();
247 if let Some(timeout_secs) = timeout_secs {
248 builder = builder.timeout(Duration::from_secs(timeout_secs));
249 }
250 Ok(builder.build()?)
251}
252
253fn temporary_path_for(destination: &Path) -> PathBuf {
254 let file_name = destination
255 .file_name()
256 .map(|name| name.to_string_lossy())
257 .unwrap_or_else(|| "download".into());
258 let suffix = SystemTime::now()
259 .duration_since(UNIX_EPOCH)
260 .map(|duration| duration.as_nanos())
261 .unwrap_or_default();
262 destination.with_file_name(format!(".{file_name}.part-{}-{suffix}", std::process::id()))
263}