grammers_client/client/

files.rs

// Copyright 2020 - developers of the `grammers` project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::mem;
use std::sync::Arc;
#[cfg(feature = "fs")]
use std::{io::SeekFrom, path::Path};

use futures_util::stream::{FuturesUnordered, StreamExt as _};
use grammers_mtsender::InvocationError;
use grammers_tl_types as tl;
use tokio::io::{self, AsyncRead, AsyncReadExt};
use tokio::sync::Mutex as AsyncMutex;
#[cfg(feature = "fs")]
use tokio::{
    fs,
    io::{AsyncSeekExt, AsyncWriteExt},
    sync::mpsc::unbounded_channel,
};

use super::Client;
use crate::media::{Downloadable, Uploaded};
use crate::utils::generate_random_id;

pub const MIN_CHUNK_SIZE: i32 = 4 * 1024;
pub const MAX_CHUNK_SIZE: i32 = 512 * 1024;
const FILE_MIGRATE_ERROR: i32 = 303;
const BIG_FILE_SIZE: usize = 10 * 1024 * 1024;
const WORKER_COUNT: usize = 4;

/// Iterator returned by [`Client::iter_download`].
pub struct DownloadIter {
    client: Client,
    done: bool,
    size: Option<usize>,
    variant: DownloadIterVariant,
}

enum DownloadIterVariant {
    Request(tl::functions::upload::GetFile),
    PreDownloaded(Vec<u8>),
    PreFailed(io::Error),
    Empty,
}

impl DownloadIter {
    /// Changes the chunk size, in bytes, used to make requests. Useful if you only need to get a
    /// small part of a file. By default, `MAX_CHUNK_SIZE` is used.
    ///
    /// # Panics
    ///
    /// Panics if `size` is not divisible by `MIN_CHUNK_SIZE`, or if `size` is not in contained in
    /// the range `MIN_CHUNK_SIZE..=MAX_CHUNK_SIZE`.
    pub fn chunk_size(mut self, size: i32) -> Self {
        assert!((MIN_CHUNK_SIZE..=MAX_CHUNK_SIZE).contains(&size) && size % MIN_CHUNK_SIZE == 0);
        match &mut self.variant {
            DownloadIterVariant::Request(request) => request.limit = size,
            _ => {}
        }
        self
    }

    /// Skips `n` chunks to start downloading a different offset from the file. If you want to
    /// skip less data, modify the `chunk_size` before calling this method, and then reset it to
    /// any value you want.
    pub fn skip_chunks(mut self, n: i32) -> Self {
        match &mut self.variant {
            DownloadIterVariant::Request(request) => {
                request.offset += request.limit as i64 * (n as i64)
            }
            _ => {}
        }
        self
    }

    /// Fetch and return the next chunk.
    pub async fn next(&mut self) -> Result<Option<Vec<u8>>, InvocationError> {
        if self.done {
            return Ok(None);
        }

        let variant = mem::replace(&mut self.variant, DownloadIterVariant::Empty);

        let mut request = match variant {
            DownloadIterVariant::Request(r) => r,
            DownloadIterVariant::PreDownloaded(data) => {
                self.done = true;
                return Ok(Some(data.clone()));
            }
            DownloadIterVariant::PreFailed(error) => {
                return Err(InvocationError::Io(error));
            }
            DownloadIterVariant::Empty => return Ok(None),
        };

        use tl::enums::upload::File;

        // TODO handle maybe FILEREF_UPGRADE_NEEDED
        let mut dc = self.client.0.session.home_dc_id()?;
        loop {
            break match self.client.invoke_in_dc(dc, &request).await {
                Ok(File::File(f)) => {
                    let reached_known_size = self
                        .size
                        .is_some_and(|size| request.offset as usize + f.bytes.len() >= size);

                    if reached_known_size || f.bytes.len() < request.limit as usize {
                        self.done = true;
                        if f.bytes.is_empty() {
                            break Ok(None);
                        }
                    }

                    request.offset += request.limit as i64;
                    self.variant = DownloadIterVariant::Request(request);

                    Ok(Some(f.bytes))
                }
                Ok(File::CdnRedirect(_)) => {
                    panic!("API returned File::CdnRedirect even though cdn_supported = false");
                }
                Err(InvocationError::Rpc(err)) if &err.name == "AUTH_KEY_UNREGISTERED" => {
                    match self.client.copy_auth_to_dc(dc).await {
                        Ok(_) => continue,
                        Err(e) => Err(e),
                    }
                }
                Err(InvocationError::Rpc(err)) if err.code == FILE_MIGRATE_ERROR => {
                    dc = err.value.unwrap() as _;
                    continue;
                }
                Err(e) => Err(e),
            };
        }
    }
}

/// Method implementations related to uploading or downloading files.
impl Client {
    /// Returns a new iterator over the contents of a media document that will be downloaded.
    ///
    /// # Examples
    ///
    /// ```
    /// # async fn f(downloadable: grammers_client::media::Media, client: grammers_client::Client) -> Result<(), Box<dyn std::error::Error>> {
    /// let mut file_bytes = Vec::new();
    /// let mut download = client.iter_download(&downloadable);
    ///
    /// while let Some(chunk) = download.next().await? {
    ///     file_bytes.extend(chunk);
    /// }
    ///
    /// // The file is now downloaded in-memory, inside `file_bytes`!
    /// # Ok(())
    /// # }
    /// ```
    pub fn iter_download<D: Downloadable>(&self, downloadable: &D) -> DownloadIter {
        if let Some(data) = downloadable.to_data() {
            DownloadIter {
                client: self.clone(),
                done: false,
                size: Some(data.len()),
                variant: DownloadIterVariant::PreDownloaded(data),
            }
        } else if let Some(location) = downloadable.to_raw_input_location() {
            DownloadIter {
                client: self.clone(),
                done: false,
                size: downloadable.size(),
                variant: DownloadIterVariant::Request(tl::functions::upload::GetFile {
                    precise: false,
                    cdn_supported: false,
                    location,
                    offset: 0,
                    limit: MAX_CHUNK_SIZE,
                }),
            }
        } else {
            DownloadIter {
                client: self.clone(),
                done: false,
                size: None,
                variant: DownloadIterVariant::PreFailed(io::Error::new(
                    io::ErrorKind::Other,
                    "media not downloadable",
                )),
            }
        }
    }

    /// Downloads a media file into the specified path.
    ///
    /// If the file already exists, it will be overwritten.
    ///
    /// This is a small wrapper around [`Client::iter_download`] for the common case of
    /// wanting to save the file locally.
    ///
    /// # Examples
    ///
    /// ```
    /// # async fn f(downloadable: grammers_client::media::Media, client: grammers_client::Client) -> Result<(), Box<dyn std::error::Error>> {
    /// client.download_media(&downloadable, "/home/username/photos/holidays.jpg").await?;
    /// # Ok(())
    /// # }
    /// ```
    #[cfg(feature = "fs")]
    pub async fn download_media<D: Downloadable, P: AsRef<Path>>(
        &self,
        downloadable: &D,
        path: P,
    ) -> Result<(), InvocationError> {
        // Concurrent downloader
        if let Some((location, size)) = downloadable
            .to_raw_input_location()
            .zip(downloadable.size())
        {
            if size > BIG_FILE_SIZE {
                return self
                    .download_media_concurrent(location, size, path, WORKER_COUNT)
                    .await;
            }
        }

        let mut download = self.iter_download(downloadable);
        Ok(Client::load(path, &mut download).await?)
    }

    #[cfg(feature = "fs")]
    async fn load<P: AsRef<Path>>(path: P, download: &mut DownloadIter) -> Result<(), io::Error> {
        let mut file = fs::File::create(path).await?;
        while let Some(chunk) = download.next().await.map_err(io::Error::other)? {
            file.write_all(&chunk).await?;
        }

        Ok(())
    }

    /// Downloads a `Document` to specified path using multiple connections
    #[cfg(feature = "fs")]
    async fn download_media_concurrent<P: AsRef<Path>>(
        &self,
        location: tl::enums::InputFileLocation,
        size: usize,
        path: P,
        workers: usize,
    ) -> Result<(), InvocationError> {
        // Allocate
        let mut file = fs::File::create(path).await.map_err(InvocationError::Io)?;
        file.set_len(size as u64)
            .await
            .map_err(InvocationError::Io)?;
        file.seek(SeekFrom::Start(0))
            .await
            .map_err(InvocationError::Io)?;

        // Start workers
        let (tx, mut rx) = unbounded_channel();
        let part_index = Arc::new(tokio::sync::Mutex::<i64>::new(0));
        let mut tasks = vec![];
        let home_dc_id = self.0.session.home_dc_id()?;
        for _ in 0..workers {
            let location = location.clone();
            let tx = tx.clone();
            let part_index = part_index.clone();
            let client = self.clone();
            let task = tokio::task::spawn(async move {
                let mut retry_offset = None;
                let mut dc = home_dc_id;
                loop {
                    // Calculate file offset
                    let offset: i64 = {
                        if let Some(offset) = retry_offset {
                            retry_offset = None;
                            offset
                        } else {
                            let mut i = part_index.lock().await;
                            *i += 1;
                            MAX_CHUNK_SIZE as i64 * (*i - 1)
                        }
                    };
                    if offset >= size as i64 {
                        break;
                    }
                    // Fetch from telegram
                    let request = &tl::functions::upload::GetFile {
                        precise: true,
                        cdn_supported: false,
                        location: location.clone(),
                        offset,
                        limit: MAX_CHUNK_SIZE,
                    };
                    match client.invoke_in_dc(dc, request).await {
                        Ok(tl::enums::upload::File::File(file)) => {
                            tx.send((offset as u64, file.bytes)).unwrap();
                        }
                        Ok(tl::enums::upload::File::CdnRedirect(_)) => {
                            panic!(
                                "API returned File::CdnRedirect even though cdn_supported = false"
                            );
                        }
                        Err(InvocationError::Rpc(err)) if &err.name == "AUTH_KEY_UNREGISTERED" => {
                            match client.copy_auth_to_dc(dc).await {
                                Ok(_) => {
                                    retry_offset = Some(offset);
                                    continue;
                                }
                                Err(e) => return Err(e),
                            }
                        }
                        Err(InvocationError::Rpc(err)) => {
                            if err.code == FILE_MIGRATE_ERROR {
                                dc = err.value.unwrap() as _;
                                retry_offset = Some(offset);
                                continue;
                            }
                            return Err(InvocationError::Rpc(err));
                        }
                        Err(e) => return Err(e),
                    }
                }
                Ok::<(), InvocationError>(())
            });
            tasks.push(task);
        }
        drop(tx);

        // File write loop
        let mut pos = 0;
        while let Some((offset, data)) = rx.recv().await {
            if offset != pos {
                file.seek(SeekFrom::Start(offset)).await?;
            }
            file.write_all(&data).await?;
            pos = offset + data.len() as u64;
        }

        // Check if all tasks finished succesfully
        for task in tasks {
            let res = task.await.map_err(io::Error::other)?;
            res?;
        }
        Ok(())
    }

    /// Uploads an async stream to Telegram servers.
    ///
    /// The file is not sent to any peer, but can be used as media when sending messages for a
    /// certain period of time (less than a day). You can use this uploaded file multiple times.
    ///
    /// Refer to [`InputMessage`] to learn more uses for `uploaded_file`.
    ///
    /// The stream size must be known beforehand. If this is not possible, you might need to
    /// process the entire async stream to determine its size, and then use the size and the
    /// downloaded buffer.
    ///
    /// The stream size may be less or equal to the actual length of the stream, but not more.
    /// If it's less, you may continue to read from the stream after the method returns.
    /// If it's more, the method will fail because it does not have enough data to read.
    ///
    /// Note that Telegram uses the file name in certain methods, for example, to make sure the
    /// file is an image when trying to use send the file as photo media, so it is important that
    /// the file name at least uses the right extension, even if the name is a dummy value.
    /// If the input file name is empty, the non-empty dummy value "a" will be used instead.
    /// Because it has no extension, you may not be able to use the file in certain methods.
    ///
    /// # Examples
    ///
    /// ```
    /// # async fn f(peer: grammers_session::types::PeerRef, client: grammers_client::Client, some_vec: &[u8]) -> Result<(), Box<dyn std::error::Error>> {
    /// use grammers_client::message::InputMessage;
    ///
    /// // In-memory `Vec<u8>` buffers can be used as async streams
    /// let size = some_vec.len();
    /// let mut stream = std::io::Cursor::new(some_vec);
    /// let uploaded_file = client.upload_stream(&mut stream, size, "sleep.jpg".to_string()).await?;
    ///
    /// client.send_message(peer, InputMessage::new().text("Zzz...").photo(uploaded_file)).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// [`InputMessage`]: crate::message::InputMessage
    pub async fn upload_stream<S: AsyncRead + Unpin>(
        &self,
        stream: &mut S,
        size: usize,
        name: String,
    ) -> Result<Uploaded, io::Error> {
        let file_id = generate_random_id();
        let name = if name.is_empty() {
            "a".to_string()
        } else {
            name
        };

        let big_file = size > BIG_FILE_SIZE;
        let parts = PartStream::new(stream, size);
        let total_parts = parts.total_parts();

        if big_file {
            let parts = Arc::new(parts);
            let mut tasks = FuturesUnordered::new();
            for _ in 0..WORKER_COUNT {
                let handle = self.clone();
                let parts = Arc::clone(&parts);
                let task = async move {
                    while let Some((part, bytes)) = parts.next_part().await? {
                        let ok = handle
                            .invoke(&tl::functions::upload::SaveBigFilePart {
                                file_id,
                                file_part: part,
                                file_total_parts: total_parts,
                                bytes,
                            })
                            .await
                            .map_err(io::Error::other)?;

                        if !ok {
                            return Err(io::Error::new(
                                io::ErrorKind::Other,
                                "server failed to store uploaded data",
                            ));
                        }
                    }
                    Ok(())
                };
                tasks.push(task);
            }

            while let Some(res) = tasks.next().await {
                res?;
            }

            Ok(Uploaded::from_raw(
                tl::types::InputFileBig {
                    id: file_id,
                    parts: total_parts,
                    name,
                }
                .into(),
            ))
        } else {
            let mut md5 = md5::Context::new();
            while let Some((part, bytes)) = parts.next_part().await? {
                md5.consume(&bytes);
                let ok = self
                    .invoke(&tl::functions::upload::SaveFilePart {
                        file_id,
                        file_part: part,
                        bytes,
                    })
                    .await
                    .map_err(io::Error::other)?;

                if !ok {
                    return Err(io::Error::new(
                        io::ErrorKind::Other,
                        "server failed to store uploaded data",
                    ));
                }
            }
            Ok(Uploaded::from_raw(
                tl::types::InputFile {
                    id: file_id,
                    parts: total_parts,
                    name,
                    md5_checksum: format!("{:x}", md5.finalize()),
                }
                .into(),
            ))
        }
    }

    /// Uploads a local file to Telegram servers.
    ///
    /// The file is not sent to any peer, but can be used as media when sending messages for a
    /// certain period of time (less than a day). You can use this uploaded file multiple times.
    ///
    /// Refer to [`InputMessage`] to learn more uses for `uploaded_file`.
    ///
    /// If you need more control over the uploaded data, such as performing only a partial upload
    /// or with a different name, use [`Client::upload_stream`] instead.
    ///
    /// # Examples
    ///
    /// ```
    /// # async fn f(peer: grammers_session::types::PeerRef, client: grammers_client::Client) -> Result<(), Box<dyn std::error::Error>> {
    /// use grammers_client::message::InputMessage;
    ///
    /// let uploaded_file = client.upload_file("/home/username/photos/holidays.jpg").await?;
    ///
    /// client.send_message(peer, InputMessage::new().text("Check this out!").photo(uploaded_file)).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// [`InputMessage`]: crate::message::InputMessage
    #[cfg(feature = "fs")]
    pub async fn upload_file<P: AsRef<Path>>(&self, path: P) -> Result<Uploaded, io::Error> {
        let path = path.as_ref();

        let mut file = fs::File::open(path).await?;
        let size = file.seek(SeekFrom::End(0)).await? as usize;
        file.seek(SeekFrom::Start(0)).await?;

        // File name will only be `None` for `..` path, and directories cannot be uploaded as
        // files, so it's fine to unwrap.
        let name = path.file_name().unwrap().to_string_lossy().to_string();

        self.upload_stream(&mut file, size, name).await
    }
}

struct PartStreamInner<'a, S: AsyncRead + Unpin> {
    stream: &'a mut S,
    current_part: i32,
}

struct PartStream<'a, S: AsyncRead + Unpin> {
    inner: AsyncMutex<PartStreamInner<'a, S>>,
    total_parts: i32,
}

impl<'a, S: AsyncRead + Unpin> PartStream<'a, S> {
    fn new(stream: &'a mut S, size: usize) -> Self {
        let total_parts = ((size + MAX_CHUNK_SIZE as usize - 1) / MAX_CHUNK_SIZE as usize) as i32;
        Self {
            inner: AsyncMutex::new(PartStreamInner {
                stream,
                current_part: 0,
            }),
            total_parts,
        }
    }

    fn total_parts(&self) -> i32 {
        self.total_parts
    }

    async fn next_part(&self) -> Result<Option<(i32, Vec<u8>)>, io::Error> {
        let mut lock = self.inner.lock().await;
        if lock.current_part >= self.total_parts {
            return Ok(None);
        }
        let mut read = 0;
        let mut buffer = vec![0; MAX_CHUNK_SIZE as usize];

        while read != buffer.len() {
            let n = lock.stream.read(&mut buffer[read..]).await?;
            if n == 0 {
                if lock.current_part == self.total_parts - 1 {
                    break;
                } else {
                    return Err(io::Error::new(
                        io::ErrorKind::UnexpectedEof,
                        "reached EOF before reaching the last file part",
                    ));
                }
            }
            read += n;
        }

        let bytes = if read == buffer.len() {
            buffer
        } else {
            buffer[..read].to_vec()
        };

        let res = Ok(Some((lock.current_part, bytes)));
        lock.current_part += 1;
        res
    }
}