curl_http_client/

collector.rs

use std::fmt::Debug;
use std::io::Read;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use std::{
    fs::{File, OpenOptions},
    io::{Seek, SeekFrom, Write},
    path::PathBuf,
};

use curl::easy::{Handler, ReadError, WriteError};
use derive_deref_rs::Deref;
use http::{HeaderMap, HeaderName, HeaderValue};
use log::trace;
use tokio::sync::mpsc::{unbounded_channel, Sender, UnboundedReceiver, UnboundedSender};

/// This is an information about the transfer(Download/Upload) speed that will be sent across tasks.
/// It is useful to get the transfer speed and displayed it according to
/// user's application.
#[derive(Clone, Debug)]
pub struct TransferSpeed(f64);

impl TransferSpeed {
    pub fn as_bytes_per_sec(&self) -> u64 {
        self.0 as u64
    }
}

impl From<u64> for TransferSpeed {
    fn from(value: u64) -> Self {
        Self(value as f64)
    }
}

impl From<usize> for TransferSpeed {
    fn from(value: usize) -> Self {
        Self(value as f64)
    }
}

impl From<i32> for TransferSpeed {
    fn from(value: i32) -> Self {
        Self(value as f64)
    }
}

impl From<i64> for TransferSpeed {
    fn from(value: i64) -> Self {
        Self(value as f64)
    }
}

impl From<f64> for TransferSpeed {
    fn from(value: f64) -> Self {
        Self(value)
    }
}

/// AbortPerform is a flag that can be safely shared across threads to be able to cancel Curl perform operation
/// via progress function of the Collector.
#[derive(Deref, Clone, Debug)]
pub struct AbortPerform {
    abort: Arc<Mutex<bool>>,
}

impl AbortPerform {
    /// Creates a new AbortPerform object with false as the default value.
    pub fn new() -> Self {
        Self {
            abort: Arc::new(Mutex::new(false)),
        }
    }
}

impl Default for AbortPerform {
    fn default() -> Self {
        Self::new()
    }
}

/// Stores the path for the downloaded file or the uploaded file.
/// Internally it will also monitor the bytes transferred and the Download/Upload speed.
#[derive(Clone, Debug)]
pub struct FileInfo {
    /// File path to download or file path of the source file to be uploaded.
    pub path: PathBuf,
    /// Sends the transfer speed information via channel to another task.
    /// This is an optional parameter depends on the user application.
    send_speed_info: Option<Sender<TransferSpeed>>,
    bytes_transferred: usize,
    transfer_started: Instant,
    transfer_speed: TransferSpeed,
    abort: Option<AbortPerform>,
}

impl FileInfo {
    /// Sets the destination file path to download or file path of the source file to be uploaded.
    pub fn path(path: PathBuf) -> Self {
        Self {
            path,
            send_speed_info: None,
            bytes_transferred: 0,
            transfer_started: Instant::now(),
            transfer_speed: TransferSpeed::from(0),
            abort: None,
        }
    }

    /// Sets the FileInfo struct with a message passing channel to send transfer speed information across user applications.
    /// It uses a tokio bounded channel to send the information across tasks.
    pub fn with_transfer_speed_sender(mut self, send_speed_info: Sender<TransferSpeed>) -> Self {
        self.send_speed_info = Some(send_speed_info);
        self
    }

    /// Set the FileInfo struct with a perform aborter.
    /// AbortPerform is a shared flag across threads to be able to switch this flag to true to abort the curl perform.
    pub fn with_perform_aborter(mut self, abort: AbortPerform) -> Self {
        self.abort = Some(abort);
        self
    }

    fn update_bytes_transferred(&mut self, transferred: usize) {
        self.bytes_transferred += transferred;

        let now = Instant::now();
        let difference = now.duration_since(self.transfer_started);

        self.transfer_speed =
            TransferSpeed::from((self.bytes_transferred) as f64 / difference.as_secs_f64());
    }

    fn bytes_transferred(&self) -> usize {
        self.bytes_transferred
    }

    fn transfer_speed(&self) -> TransferSpeed {
        self.transfer_speed.clone()
    }
}

fn send_transfer_info(info: &FileInfo) {
    if let Some(tx) = info.send_speed_info.clone() {
        let transfer_speed = info.transfer_speed();
        tokio::spawn(async move {
            tx.send(transfer_speed).await.map_err(|e| {
                trace!("{:?}", e);
            })
        });
    }
}

/// `StreamHandler` is a lightweight helper for managing streamed responses.
///
/// It provides two main components:
/// - `chunk_sender`: An asynchronous channel sender used to forward each
///   received chunk of data to a consumer in real time.
/// - `abort`: An optional abort flag that can be used to stop an ongoing
///   curl perform operation from another thread.
///
/// This struct is typically used in combination with the `Collector::Streaming`
/// variant to enable progressive consumption of large or continuous responses
/// without buffering everything in memory.
#[derive(Clone, Debug)]
pub struct StreamHandler {
    /// Asynchronous sender that delivers streamed response chunks (`Vec<u8>`)
    /// to a receiving end. Each received chunk is sent immediately, enabling
    /// consumers to process data progressively.
    chunk_sender: UnboundedSender<Vec<u8>>,

    /// Optional abort handle (`AbortPerform`) that can be set to signal
    /// cancellation of the ongoing curl perform operation.
    ///
    /// When triggered, this flag instructs the underlying transfer to stop early.
    abort: Option<AbortPerform>,
}

impl StreamHandler {
    /// Creates a new `StreamHandler` along with its corresponding receiver.
    ///
    /// This convenience method sets up a channel for streaming data,
    /// returning both the `StreamHandler` (containing the sender) and
    /// the `UnboundedReceiver` for consuming chunks.
    ///
    /// # Returns
    /// A tuple of:
    /// - `StreamHandler`: The handler containing the sender and optional abort flag.
    /// - `UnboundedReceiver<Vec<u8>>`: The receiving end for streamed chunks.
    pub fn new() -> (Self, UnboundedReceiver<Vec<u8>>) {
        let (tx, rx) = unbounded_channel();
        (
            Self {
                chunk_sender: tx,
                abort: None,
            },
            rx,
        )
    }

    /// Associates an abort handle with this `StreamHandler`.
    ///
    /// `AbortPerform` is a shared flag across threads that can be toggled
    /// to `true` in order to abort the curl perform operation prematurely.
    ///
    /// # Example
    /// ```
    /// use curl_http_client::AbortPerform;
    /// use curl_http_client::StreamHandler;
    ///
    /// let (handler, rx) = StreamHandler::new();
    /// let aborter = AbortPerform::new();
    /// let handler = handler.with_perform_aborter(aborter);
    /// ```
    pub fn with_perform_aborter(mut self, abort: AbortPerform) -> Self {
        self.abort = Some(abort);
        self
    }
}

fn send_stream_data(stream: &StreamHandler, data: Vec<u8>) {
    let tx = stream.chunk_sender.clone();

    let _ = tx.send(data).map_err(|e| {
        trace!("{:?}", e);
    });
}

/// This is an extended trait for the curl::easy::Handler trait.
pub trait ExtendedHandler: Handler {
    // Return the response body if the Collector is available.
    fn get_response_body(&self) -> Option<Vec<u8>> {
        None
    }
    // Return the response body if the Collector is available with complete headers.
    fn get_response_body_and_headers(&self) -> (Option<Vec<u8>>, Option<HeaderMap>) {
        (None, None)
    }
}

/// Collector::File(FileInfo) is used to be able to download and upload files.
/// Collector::Ram(`Vec<u8>`) is used to store response body into Memory.
/// Collector::RamWithHeaders(`Vec<u8>`, `Vec<u8>`) is used to store response body into Memory and with complete headers.
/// Collector::FileAndHeaders(`FileInfo`, `Vec<u8>`) is used to be able to download and upload files and with complete headers.
/// Collector::Streaming(`StreamHandler`, `Vec<u8>`) is used to process the response body as a **stream of chunks** instead of buffering the entire
/// response in memory or writing it directly to a file.
#[derive(Clone, Debug)]
pub enum Collector {
    /// Collector::File(`FileInfo`) is used to be able to download and upload files.
    File(FileInfo),
    /// Collector::Ram(`Vec<u8>`) is used to store response body into Memory.
    Ram(Vec<u8>),
    /// Collector::RamWithHeaders(`Vec<u8>`, `Vec<u8>`) is used to store response body into Memory and with complete headers.
    RamAndHeaders(Vec<u8>, Vec<u8>),
    /// Collector::FileAndHeaders(`FileInfo`, `Vec<u8>`) is used to be able to download and upload files and with complete headers.
    FileAndHeaders(FileInfo, Vec<u8>),
    /// Collector::Streaming(`StreamHandler`, `Vec<u8>`) is used to process
    /// the response body as a **stream of chunks** instead of buffering the entire
    /// response in memory or writing it directly to a file with complete headers.
    /// This is useful for large responses, continuous data feeds, or situations
    /// where you don’t want to block until the full body is received.
    Streaming(StreamHandler, Vec<u8>),
}

impl Handler for Collector {
    /// This will store the response from the server
    /// to the data vector or into a file depends on the
    /// Collector being used.
    fn write(&mut self, data: &[u8]) -> Result<usize, WriteError> {
        match self {
            Collector::File(info) => {
                let mut file = OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(info.path.clone())
                    .map_err(|e| {
                        trace!("{}", e);
                        WriteError::Pause
                    })?;

                file.write_all(data).map_err(|e| {
                    trace!("{}", e);
                    WriteError::Pause
                })?;

                info.update_bytes_transferred(data.len());

                send_transfer_info(info);
                Ok(data.len())
            }
            Collector::Ram(container) => {
                container.extend_from_slice(data);
                Ok(data.len())
            }
            Collector::RamAndHeaders(container, _) => {
                container.extend_from_slice(data);
                Ok(data.len())
            }
            Collector::FileAndHeaders(info, _) => {
                let mut file = OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(info.path.clone())
                    .map_err(|e| {
                        trace!("{}", e);
                        WriteError::Pause
                    })?;

                file.write_all(data).map_err(|e| {
                    trace!("{}", e);
                    WriteError::Pause
                })?;

                info.update_bytes_transferred(data.len());

                send_transfer_info(info);
                Ok(data.len())
            }
            Collector::Streaming(stream, _) => {
                send_stream_data(stream, data.to_vec());
                Ok(data.len())
            }
        }
    }
    /// This will read the chunks of data from a file that will be uploaded
    /// to the server. This will be use if the Collector is Collector::File(FileInfo).
    fn read(&mut self, data: &mut [u8]) -> Result<usize, ReadError> {
        match self {
            Collector::File(info) => {
                let mut file = File::open(info.path.clone()).map_err(|e| {
                    trace!("{}", e);
                    ReadError::Abort
                })?;

                file.seek(SeekFrom::Start(info.bytes_transferred() as u64))
                    .map_err(|e| {
                        trace!("{}", e);
                        ReadError::Abort
                    })?;

                let read_size = file.read(data).map_err(|e| {
                    trace!("{}", e);
                    ReadError::Abort
                })?;

                info.update_bytes_transferred(read_size);

                send_transfer_info(info);
                Ok(read_size)
            }
            Collector::Ram(_) => Ok(0),
            Collector::RamAndHeaders(_, _) => Ok(0),
            Collector::FileAndHeaders(info, _) => {
                let mut file = File::open(info.path.clone()).map_err(|e| {
                    trace!("{}", e);
                    ReadError::Abort
                })?;

                file.seek(SeekFrom::Start(info.bytes_transferred() as u64))
                    .map_err(|e| {
                        trace!("{}", e);
                        ReadError::Abort
                    })?;

                let read_size = file.read(data).map_err(|e| {
                    trace!("{}", e);
                    ReadError::Abort
                })?;

                info.update_bytes_transferred(read_size);

                send_transfer_info(info);
                Ok(read_size)
            }
            Collector::Streaming(_, _) => Ok(0),
        }
    }

    fn header(&mut self, data: &[u8]) -> bool {
        match self {
            Collector::File(_) => {}
            Collector::Ram(_) => {}
            Collector::RamAndHeaders(_, headers) => {
                headers.extend_from_slice(data);
            }
            Collector::FileAndHeaders(_, headers) => {
                headers.extend_from_slice(data);
            }
            Collector::Streaming(_, headers) => {
                headers.extend_from_slice(data);
            }
        }
        true
    }

    fn progress(&mut self, dltotal: f64, dlnow: f64, ultotal: f64, ulnow: f64) -> bool {
        trace!("dltotal: {dltotal} dlnow: {dlnow} ultotal: {ultotal} ulnow: {ulnow}");
        match self {
            Collector::File(file_info) | Collector::FileAndHeaders(file_info, _) => {
                if let Some(abort) = &file_info.abort {
                    abort.lock().map(|a| !*a).unwrap_or_else(|err| {
                        log::error!("{:?}", err);
                        false
                    })
                } else {
                    true
                }
            }
            Collector::Ram(_) | Collector::RamAndHeaders(_, _) => true,
            Collector::Streaming(stream, _) => {
                if let Some(abort) = &stream.abort {
                    abort.lock().map(|a| !*a).unwrap_or_else(|err| {
                        log::error!("{:?}", err);
                        false
                    })
                } else {
                    true
                }
            }
        }
    }
}

impl ExtendedHandler for Collector {
    /// If Collector::File(FileInfo) is set, there will be no response body since the response
    /// will be stored into a file.
    ///
    /// If Collector::Ram(`Vec<u8>`) is set, the response body can be obtain here.
    fn get_response_body(&self) -> Option<Vec<u8>> {
        match self {
            Collector::File(_) => None,
            Collector::Ram(container) => {
                if container.is_empty() {
                    None
                } else {
                    Some(container.clone())
                }
            }
            Collector::RamAndHeaders(container, _) => {
                if container.is_empty() {
                    None
                } else {
                    Some(container.clone())
                }
            }
            Collector::FileAndHeaders(_, _) => None,
            Collector::Streaming(_, _) => None,
        }
    }

    /// If Collector::File(`FileInfo`) is set, there will be no response body since the response will be stored into a file.
    /// If Collector::Ram(`Vec<u8>`) is set, the response body can be obtain here.
    /// If Collector::RamAndHeaders(`Vec<u8>`, `Vec<u8>`) is set, the response body and the complete headers are generated.
    /// If Collector::FileAndHeaders(`FileInfo`, `Vec<u8>`) is set, there will be no response body since the response will be stored into a file but a complete headers are generated.
    fn get_response_body_and_headers(&self) -> (Option<Vec<u8>>, Option<HeaderMap>) {
        match self {
            Collector::File(_) => (None, None),
            Collector::Ram(container) => {
                if container.is_empty() {
                    (None, None)
                } else {
                    (Some(container.clone()), None)
                }
            }
            Collector::RamAndHeaders(container, headers) => {
                let header_str = std::str::from_utf8(headers).unwrap();
                let mut header_map = HeaderMap::new();

                for line in header_str.lines() {
                    // Split each line into key-value pairs
                    if let Some((key, value)) = line.split_once(": ").to_owned() {
                        if let Ok(header_name) = HeaderName::from_bytes(key.as_bytes()) {
                            if let Ok(header_value) = HeaderValue::from_str(value) {
                                // Insert the key-value pair into the HeaderMap
                                header_map.insert(header_name, header_value);
                            }
                        }
                    }
                }
                if container.is_empty() {
                    (None, Some(header_map))
                } else {
                    (Some(container.clone()), Some(header_map))
                }
            }
            Collector::FileAndHeaders(_, headers) => {
                let header_str = std::str::from_utf8(headers).unwrap();
                let mut header_map = HeaderMap::new();

                for line in header_str.lines() {
                    // Split each line into key-value pairs
                    if let Some((key, value)) = line.split_once(": ").to_owned() {
                        if let Ok(header_name) = HeaderName::from_bytes(key.as_bytes()) {
                            if let Ok(header_value) = HeaderValue::from_str(value) {
                                // Insert the key-value pair into the HeaderMap
                                header_map.insert(header_name, header_value);
                            }
                        }
                    }
                }
                (None, Some(header_map))
            }
            Collector::Streaming(_, headers) => {
                let header_str = std::str::from_utf8(headers).unwrap();
                let mut header_map = HeaderMap::new();

                for line in header_str.lines() {
                    // Split each line into key-value pairs
                    if let Some((key, value)) = line.split_once(": ").to_owned() {
                        if let Ok(header_name) = HeaderName::from_bytes(key.as_bytes()) {
                            if let Ok(header_value) = HeaderValue::from_str(value) {
                                // Insert the key-value pair into the HeaderMap
                                header_map.insert(header_name, header_value);
                            }
                        }
                    }
                }
                (None, Some(header_map))
            }
        }
    }
}