mithril_client/file_downloader/
http.rs

1use std::{
2    io::{BufReader, Read, Write},
3    path::Path,
4};
5
6use anyhow::{Context, anyhow};
7use async_trait::async_trait;
8use flate2::read::GzDecoder;
9use flume::{Receiver, Sender};
10use futures::StreamExt;
11use reqwest::{Response, StatusCode, Url};
12use slog::{Logger, debug};
13use tar::Archive;
14use tokio::fs::File;
15use tokio::io::AsyncReadExt;
16
17use mithril_common::{StdResult, logging::LoggerExtensions};
18
19use crate::common::CompressionAlgorithm;
20use crate::feedback::FeedbackSender;
21use crate::utils::StreamReader;
22
23use super::{FileDownloader, FileDownloaderUri, interface::DownloadEvent};
24
25/// A file downloader that only handles download through HTTP.
26pub struct HttpFileDownloader {
27    http_client: reqwest::Client,
28    feedback_sender: FeedbackSender,
29    logger: Logger,
30}
31
32impl HttpFileDownloader {
33    /// Constructs a new `HttpFileDownloader`.
34    pub fn new(feedback_sender: FeedbackSender, logger: Logger) -> StdResult<Self> {
35        let http_client = reqwest::ClientBuilder::new()
36            .build()
37            .with_context(|| "Building http client for HttpFileDownloader failed")?;
38
39        Ok(Self {
40            http_client,
41            feedback_sender,
42            logger: logger.new_with_component_name::<Self>(),
43        })
44    }
45
46    async fn get(&self, location: &str) -> StdResult<Response> {
47        debug!(self.logger, "GET Snapshot location='{location}'.");
48        let request_builder = self.http_client.get(location);
49        let response = request_builder.send().await.with_context(|| {
50            format!("Cannot perform a GET for the snapshot (location='{location}')")
51        })?;
52
53        match response.status() {
54            StatusCode::OK => Ok(response),
55            StatusCode::NOT_FOUND => Err(anyhow!("Location='{location} not found")),
56            status_code => Err(anyhow!("Unhandled error {status_code}")),
57        }
58    }
59
60    fn file_scheme_to_local_path(file_url: &str) -> Option<String> {
61        Url::parse(file_url)
62            .ok()
63            .filter(|url| url.scheme() == "file")
64            .and_then(|url| url.to_file_path().ok())
65            .map(|path| path.to_string_lossy().into_owned())
66    }
67
68    /// Stream the `location` directly from the local filesystem
69    async fn download_local_file(
70        &self,
71        local_path: &str,
72        sender: &Sender<Vec<u8>>,
73        download_event_type: DownloadEvent,
74        file_size: u64,
75    ) -> StdResult<()> {
76        let mut downloaded_bytes: u64 = 0;
77        let mut file = File::open(local_path).await?;
78        let size = match file.metadata().await {
79            Ok(metadata) => metadata.len(),
80            Err(_) => file_size,
81        };
82
83        self.feedback_sender
84            .send_event(download_event_type.build_download_started_event(size))
85            .await;
86
87        loop {
88            // We can either allocate here each time, or clone a shared buffer into sender.
89            // A larger read buffer is faster, less context switches:
90            let mut buffer = vec![0; 16 * 1024 * 1024];
91            let bytes_read = file.read(&mut buffer).await?;
92            if bytes_read == 0 {
93                break;
94            }
95            buffer.truncate(bytes_read);
96            sender.send_async(buffer).await.with_context(|| {
97                format!("Local file read: could not write {bytes_read} bytes to stream.")
98            })?;
99            downloaded_bytes += bytes_read as u64;
100            let event = download_event_type.build_download_progress_event(downloaded_bytes, size);
101            self.feedback_sender.send_event(event).await;
102        }
103
104        self.feedback_sender
105            .send_event(download_event_type.build_download_completed_event())
106            .await;
107
108        Ok(())
109    }
110
111    /// Stream the `location` remotely
112    async fn download_remote_file(
113        &self,
114        location: &str,
115        sender: &Sender<Vec<u8>>,
116        download_event_type: DownloadEvent,
117        file_size: u64,
118    ) -> StdResult<()> {
119        let mut downloaded_bytes: u64 = 0;
120        let response = self.get(location).await?;
121        let size = response.content_length().unwrap_or(file_size);
122        let mut remote_stream = response.bytes_stream();
123
124        self.feedback_sender
125            .send_event(download_event_type.build_download_started_event(size))
126            .await;
127
128        while let Some(item) = remote_stream.next().await {
129            let chunk = item.with_context(|| "Download: Could not read from byte stream")?;
130            sender.send_async(chunk.to_vec()).await.with_context(|| {
131                format!("Download: could not write {} bytes to stream.", chunk.len())
132            })?;
133            downloaded_bytes += chunk.len() as u64;
134            let event = download_event_type.build_download_progress_event(downloaded_bytes, size);
135            self.feedback_sender.send_event(event).await;
136        }
137
138        self.feedback_sender
139            .send_event(download_event_type.build_download_completed_event())
140            .await;
141
142        Ok(())
143    }
144
145    fn unpack_file(
146        stream: Receiver<Vec<u8>>,
147        compression_algorithm: Option<CompressionAlgorithm>,
148        unpack_dir: &Path,
149        download_id: String,
150    ) -> StdResult<()> {
151        let input = StreamReader::new(stream);
152        match compression_algorithm {
153            Some(CompressionAlgorithm::Gzip) => {
154                let gzip_decoder = GzDecoder::new(input);
155                let mut file_archive = Archive::new(gzip_decoder);
156                file_archive.unpack(unpack_dir).with_context(|| {
157                    format!(
158                        "Could not unpack with 'Gzip' from streamed data to directory '{}'",
159                        unpack_dir.display()
160                    )
161                })?;
162            }
163            Some(CompressionAlgorithm::Zstandard) => {
164                let zstandard_decoder = zstd::Decoder::new(input)
165                    .with_context(|| "Unpack failed: Create Zstandard decoder error")?;
166                let mut file_archive = Archive::new(zstandard_decoder);
167                file_archive.unpack(unpack_dir).with_context(|| {
168                    format!(
169                        "Could not unpack with 'Zstd' from streamed data to directory '{}'",
170                        unpack_dir.display()
171                    )
172                })?;
173            }
174            None => {
175                let file_path = unpack_dir.join(download_id);
176                if file_path.exists() {
177                    std::fs::remove_file(file_path.clone())?;
178                }
179                let mut file = std::fs::File::create(file_path)?;
180                let input_buffered = BufReader::new(input);
181                for byte in input_buffered.bytes() {
182                    file.write_all(&[byte?])?;
183                }
184                file.flush()?;
185            }
186        };
187
188        Ok(())
189    }
190}
191
192#[async_trait]
193impl FileDownloader for HttpFileDownloader {
194    async fn download_unpack(
195        &self,
196        location: &FileDownloaderUri,
197        file_size: u64,
198        target_dir: &Path,
199        compression_algorithm: Option<CompressionAlgorithm>,
200        download_event_type: DownloadEvent,
201    ) -> StdResult<()> {
202        if !target_dir.is_dir() {
203            Err(
204                anyhow!("target path is not a directory or does not exist: `{target_dir:?}`")
205                    .context("Download-Unpack: prerequisite error"),
206            )?;
207        }
208
209        let (sender, receiver) = flume::bounded(32);
210        let dest_dir = target_dir.to_path_buf();
211        let download_id = download_event_type.download_id().to_owned();
212        let unpack_thread = tokio::task::spawn_blocking(move || -> StdResult<()> {
213            Self::unpack_file(receiver, compression_algorithm, &dest_dir, download_id)
214        });
215        if let Some(local_path) = Self::file_scheme_to_local_path(location.as_str()) {
216            self.download_local_file(&local_path, &sender, download_event_type, file_size)
217                .await?;
218        } else {
219            self.download_remote_file(location.as_str(), &sender, download_event_type, file_size)
220                .await?;
221        }
222        drop(sender);
223        unpack_thread
224            .await
225            .with_context(|| {
226                format!(
227                    "Unpack: panic while unpacking to dir '{}'",
228                    target_dir.display()
229                )
230            })?
231            .with_context(|| {
232                format!("Unpack: could not unpack to dir '{}'", target_dir.display())
233            })?;
234
235        Ok(())
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use std::sync::Arc;
242
243    use httpmock::MockServer;
244
245    use mithril_common::{entities::FileUri, test_utils::TempDir};
246
247    use crate::{
248        feedback::{MithrilEvent, MithrilEventCardanoDatabase, StackFeedbackReceiver},
249        test_utils::TestLogger,
250    };
251
252    use super::*;
253
254    #[cfg(not(target_family = "windows"))]
255    fn local_file_uri(path: &Path) -> FileDownloaderUri {
256        FileDownloaderUri::FileUri(FileUri(format!(
257            "file://{}",
258            path.canonicalize().unwrap().to_string_lossy()
259        )))
260    }
261
262    #[cfg(target_family = "windows")]
263    fn local_file_uri(path: &Path) -> FileDownloaderUri {
264        // We need to transform `\\?\C:\data\Temp\mithril_test\snapshot.txt` to `file://C:/data/Temp/mithril_test/snapshot.txt`
265        FileDownloaderUri::FileUri(FileUri(format!(
266            "file:/{}",
267            path.canonicalize()
268                .unwrap()
269                .to_string_lossy()
270                .replace("\\", "/")
271                .replace("?/", ""),
272        )))
273    }
274
275    #[tokio::test]
276    async fn test_download_http_file_send_feedback() {
277        let target_dir = TempDir::create(
278            "client-http-downloader",
279            "test_download_http_file_send_feedback",
280        );
281        let content = "Hello, world!";
282        let size = content.len() as u64;
283        let server = MockServer::start();
284        server.mock(|when, then| {
285            when.method(httpmock::Method::GET).path("/snapshot.tar");
286            then.status(200)
287                .body(content)
288                .header(reqwest::header::CONTENT_LENGTH.as_str(), size.to_string());
289        });
290        let feedback_receiver = Arc::new(StackFeedbackReceiver::new());
291        let http_file_downloader = HttpFileDownloader::new(
292            FeedbackSender::new(&[feedback_receiver.clone()]),
293            TestLogger::stdout(),
294        )
295        .unwrap();
296        let download_id = "id".to_string();
297
298        http_file_downloader
299            .download_unpack(
300                &FileDownloaderUri::FileUri(FileUri(server.url("/snapshot.tar"))),
301                0,
302                &target_dir,
303                None,
304                DownloadEvent::Digest {
305                    download_id: download_id.clone(),
306                },
307            )
308            .await
309            .unwrap();
310
311        let expected_events = vec![
312            MithrilEvent::CardanoDatabase(MithrilEventCardanoDatabase::DigestDownloadStarted {
313                download_id: download_id.clone(),
314                size,
315            }),
316            MithrilEvent::CardanoDatabase(MithrilEventCardanoDatabase::DigestDownloadProgress {
317                download_id: download_id.clone(),
318                downloaded_bytes: size,
319                size,
320            }),
321            MithrilEvent::CardanoDatabase(MithrilEventCardanoDatabase::DigestDownloadCompleted {
322                download_id: download_id.clone(),
323            }),
324        ];
325        assert_eq!(expected_events, feedback_receiver.stacked_events());
326    }
327
328    #[tokio::test]
329    async fn test_download_local_file_send_feedback() {
330        let target_dir = TempDir::create(
331            "client-http-downloader",
332            "test_download_local_file_send_feedback",
333        );
334        let content = "Hello, world!";
335        let size = content.len() as u64;
336
337        let source_file_path = target_dir.join("snapshot.txt");
338        let mut file = std::fs::File::create(&source_file_path).unwrap();
339        file.write_all(content.as_bytes()).unwrap();
340
341        let feedback_receiver = Arc::new(StackFeedbackReceiver::new());
342        let http_file_downloader = HttpFileDownloader::new(
343            FeedbackSender::new(&[feedback_receiver.clone()]),
344            TestLogger::stdout(),
345        )
346        .unwrap();
347        let download_id = "id".to_string();
348
349        http_file_downloader
350            .download_unpack(
351                &local_file_uri(&source_file_path),
352                0,
353                &target_dir,
354                None,
355                DownloadEvent::Digest {
356                    download_id: download_id.clone(),
357                },
358            )
359            .await
360            .unwrap();
361
362        let expected_events = vec![
363            MithrilEvent::CardanoDatabase(MithrilEventCardanoDatabase::DigestDownloadStarted {
364                download_id: download_id.clone(),
365                size,
366            }),
367            MithrilEvent::CardanoDatabase(MithrilEventCardanoDatabase::DigestDownloadProgress {
368                download_id: download_id.clone(),
369                downloaded_bytes: size,
370                size,
371            }),
372            MithrilEvent::CardanoDatabase(MithrilEventCardanoDatabase::DigestDownloadCompleted {
373                download_id: download_id.clone(),
374            }),
375        ];
376        assert_eq!(expected_events, feedback_receiver.stacked_events());
377    }
378}
mithril_client/file_downloader/http.rs

mithril_client/file_downloader/
http.rs
