Skip to main content

ronami_core/bot/
download.rs

1use std::path::Path;
2
3use bytes::Bytes;
4use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream};
5
6use tokio::{
7    fs::File,
8    io::{AsyncReadExt, AsyncWrite, AsyncWriteExt},
9};
10
11use crate::{
12    DownloadError,
13    bot::Bot,
14    net::{self, Download},
15};
16
17impl Download for Bot {
18    type Err<'dst> = DownloadError;
19
20    // I would like to unbox this, but my coworkers will kill me if they'll see yet
21    // another hand written `Future`. (waffle)
22    type Fut<'dst> = BoxFuture<'dst, Result<(), Self::Err<'dst>>>;
23
24    fn download_file<'dst>(
25        &self,
26        path: &'dst str,
27        destination: &'dst mut (dyn AsyncWrite + Unpin + Send),
28    ) -> Self::Fut<'dst> {
29        let is_localhost = match &self.api_url.host_str() {
30            Some(host) => ["localhost", "127.0.0.1", "::1"].contains(host),
31            None => false,
32        };
33        // If path is absolute and api_url contains localhost, it is pretty clear there
34        // is a local TBA server with --local option, we can just copy the file
35        if is_localhost && Path::new(&path).is_absolute() {
36            return copy_file(path, destination).boxed();
37        }
38
39        net::download_file(
40            &self.client,
41            reqwest::Url::clone(&*self.api_url),
42            &self.token,
43            path,
44            destination,
45        )
46        .boxed()
47    }
48
49    type StreamErr = reqwest::Error;
50
51    type Stream = BoxStream<'static, Result<Bytes, Self::StreamErr>>;
52
53    fn download_file_stream(&self, path: &str) -> Self::Stream {
54        net::download_file_stream(
55            &self.client,
56            reqwest::Url::clone(&*self.api_url),
57            &self.token,
58            path,
59        )
60        .map(|res| res.map_err(crate::errors::hide_token))
61        .boxed()
62    }
63}
64
65async fn copy_file<'o, D>(path: &'o str, dst: &'o mut D) -> Result<(), DownloadError>
66where
67    D: ?Sized + AsyncWrite + Unpin,
68{
69    let mut src_file = File::open(path).await?;
70
71    let mut buffer = [0; 128 * 1024];
72    loop {
73        let n = src_file.read(&mut buffer).await?;
74        if n == 0 {
75            break;
76        }
77
78        dst.write_all(&buffer[..n]).await?;
79    }
80
81    dst.flush().await?;
82    Ok(())
83}