starbase_utils 0.14.0

use crate::fs::{self, FsError};
use async_trait::async_trait;
use reqwest::{
    Client, Response,
    header::{HeaderMap, HeaderName, HeaderValue},
};
use std::collections::HashMap;
use std::fmt::Debug;
use std::fs::File;
use std::io::Write;
use std::net::{IpAddr, Shutdown, SocketAddr, TcpStream, ToSocketAddrs};
use std::path::Path;
use std::sync::{Arc, mpsc};
use std::thread;
use std::time::Duration;
use tracing::{instrument, trace};
use url::Url;

pub use crate::net_error::NetError;

/// A contract for downloading files from a URL.
#[async_trait]
pub trait Downloader: Send {
    async fn download(&self, url: Url) -> Result<Response, NetError>;
}

pub type BoxedDownloader = Box<dyn Downloader>;

/// A default [`Downloader`] backed by `reqwest`.
#[derive(Default)]
pub struct DefaultDownloader {
    client: reqwest::Client,
    headers: Option<HeaderMap>,
}

fn headers_from_map(map: &HashMap<String, String>) -> Result<HeaderMap, NetError> {
    let mut headers = HeaderMap::new();
    for (k, v) in map {
        let name = HeaderName::try_from(k.as_str()).map_err(|e| NetError::HttpUnknown {
            url: String::new(),
            error: format!("invalid header name {:?}: {}", k, e),
        })?;
        let value = HeaderValue::try_from(v.as_str()).map_err(|e| NetError::HttpUnknown {
            url: String::new(),
            error: format!("invalid header value {:?}: {}", v, e),
        })?;
        headers.insert(name, value);
    }
    Ok(headers)
}

impl DefaultDownloader {
    /// Create a downloader that applies the given headers to each request.
    pub fn new_with_headers(headers: HashMap<String, String>) -> Result<Self, NetError> {
        Ok(Self {
            client: Client::new(),
            headers: Some(headers_from_map(&headers)?),
        })
    }

    /// Create a downloader with a custom client and headers.
    pub fn new_with_client_and_headers(
        client: Client,
        headers: HashMap<String, String>,
    ) -> Result<Self, NetError> {
        Ok(Self {
            client,
            headers: Some(headers_from_map(&headers)?),
        })
    }
}

#[async_trait]
impl Downloader for DefaultDownloader {
    async fn download(&self, url: Url) -> Result<Response, NetError> {
        let mut request = self.client.get(url.clone());

        if let Some(headers) = &self.headers {
            request = request.headers(headers.clone());
        }

        request.send().await.map_err(|error| NetError::Http {
            error: Box::new(error),
            url: url.to_string(),
        })
    }
}

/// A function that is called for each chunk in a download stream.
pub type OnChunkFn = Arc<dyn Fn(u64, u64) + Send + Sync>;

/// Options for customizing network downloads.
pub struct DownloadOptions {
    /// The downloader to use for fetching the file. If not provided, a default
    /// downloader backed by `reqwest` will be used.
    pub downloader: Option<BoxedDownloader>,

    /// Handle for the file being written to. If not provided, a new file will be
    /// created at the destination path.
    pub file: Option<File>,

    /// Whether to acquire an exclusive lock on the destination file while writing.
    /// Requires the `fs-lock` feature. Defaults to `true`.
    pub lock: bool,

    /// An optional callback that is called with the current and total size of the
    /// download for each chunk received.
    pub on_chunk: Option<OnChunkFn>,
}

impl DownloadOptions {
    pub fn new(downloader: impl Downloader + 'static) -> Self {
        Self {
            downloader: Some(Box::new(downloader)),
            file: None,
            lock: true,
            on_chunk: None,
        }
    }
}

impl Default for DownloadOptions {
    fn default() -> Self {
        Self {
            downloader: None,
            file: None,
            lock: true,
            on_chunk: None,
        }
    }
}

/// Download a file from the provided source URL, to the destination file path,
/// using custom options.
#[instrument(skip(options))]
pub async fn download_from_url_with_options<S: AsRef<str> + Debug, D: AsRef<Path> + Debug>(
    source_url: S,
    dest_file: D,
    options: DownloadOptions,
) -> Result<(), NetError> {
    let source_url = source_url.as_ref();
    let dest_file = dest_file.as_ref();
    let DownloadOptions {
        downloader,
        file: file_handle,
        lock: lock_file,
        on_chunk,
    } = options;

    let downloader = downloader.unwrap_or_else(|| Box::new(DefaultDownloader::default()));

    let handle_fs_error = |error: std::io::Error| FsError::Write {
        path: dest_file.to_path_buf(),
        error: Box::new(error),
    };
    let handle_net_error = |error: reqwest::Error| NetError::Http {
        error: Box::new(error),
        url: source_url.to_owned(),
    };

    trace!(
        source_url,
        dest_file = ?dest_file,
        "Downloading file from remote URL to local file",
    );

    // Fetch the file from the HTTP source
    let mut response = downloader
        .download(
            Url::parse(source_url).map_err(|error| NetError::UrlParseFailed {
                url: source_url.to_owned(),
                error: Box::new(error),
            })?,
        )
        .await?;
    let status = response.status();

    if status.as_u16() == 404 {
        return Err(NetError::UrlNotFound {
            url: source_url.to_owned(),
        });
    }

    if !status.is_success() {
        return Err(NetError::DownloadFailed {
            url: source_url.to_owned(),
            status: status.to_string(),
        });
    }

    let mut file = match file_handle {
        Some(file) => file,
        // Don't truncate the file in case there's a parallel process
        // interacting with this file! This happens more than you think!
        None => fs::create_file_if_missing(dest_file)?,
    };

    #[cfg(feature = "fs-lock")]
    if lock_file {
        fs::acquire_exclusive_lock(dest_file, &file)?;
    }

    let write_result = async {
        // Now truncate the file after we have a lock
        fs::truncate_file_handle(dest_file, &mut file)?;

        let total_size = response.content_length().unwrap_or(0);
        let mut current_size: u64 = 0;

        if let Some(on_chunk) = &on_chunk {
            on_chunk(0, total_size);
        }

        while let Some(chunk) = response.chunk().await.map_err(handle_net_error)? {
            file.write_all(&chunk).map_err(handle_fs_error)?;

            if let Some(on_chunk) = &on_chunk {
                current_size = current_size.saturating_add(chunk.len() as u64);

                on_chunk(
                    if total_size > 0 {
                        current_size.min(total_size)
                    } else {
                        current_size
                    },
                    total_size,
                );
            }
        }

        Ok::<(), NetError>(())
    }
    .await;

    match write_result {
        Ok(_) => {
            // The data is already committed; if unlock fails here the kernel
            // will release it when `file` drops anyway, so don't fail the call.
            #[cfg(feature = "fs-lock")]
            if lock_file {
                let _ = fs::release_lock(dest_file, &file);
            }

            Ok(())
        }
        Err(error) => {
            // Release the lock before removing the file so a concurrent caller
            // can't slip in between, create a new inode at the same path, and
            // acquire its own lock while ours is still held on the orphan inode.
            #[cfg(feature = "fs-lock")]
            if lock_file {
                let _ = fs::release_lock(dest_file, &file);
            }

            // Cleanup on failure, otherwise the file may only be partially written to
            let _ = fs::remove_file(dest_file);

            Err(error)
        }
    }
}

/// Download a file from the provided source URL, to the destination file path,
/// using a custom `reqwest` [`Client`].
pub async fn download_from_url_with_client<S: AsRef<str> + Debug, D: AsRef<Path> + Debug>(
    source_url: S,
    dest_file: D,
    client: &Client,
) -> Result<(), NetError> {
    download_from_url_with_options(
        source_url,
        dest_file,
        DownloadOptions::new(DefaultDownloader {
            client: client.to_owned(),
            headers: None,
        }),
    )
    .await
}

/// Download a file from the provided source URL, to the destination file path.
pub async fn download_from_url<S: AsRef<str> + Debug, D: AsRef<Path> + Debug>(
    source_url: S,
    dest_file: D,
) -> Result<(), NetError> {
    download_from_url_with_options(source_url, dest_file, DownloadOptions::default()).await
}

mod offline {
    use super::*;

    pub fn check_connection(address: SocketAddr, timeout: u64) -> bool {
        trace!("Resolving {address}");

        if let Ok(stream) = TcpStream::connect_timeout(&address, Duration::from_millis(timeout)) {
            let _ = stream.shutdown(Shutdown::Both);

            return true;
        }

        false
    }

    pub fn check_connection_from_host(host: String, timeout: u64) -> bool {
        // Wrap in a thread because resolving a host to an IP address
        // may take an unknown amount of time. If longer than our timeout,
        // exit early.
        let (sender, receiver) = mpsc::channel();

        thread::spawn(move || {
            let _ = sender.send(host.to_socket_addrs().ok());
        });

        if let Ok(Some(addresses)) = receiver.recv_timeout(Duration::from_millis(timeout)) {
            for address in addresses {
                if check_connection(address, timeout) {
                    return true;
                }
            }
        }

        false
    }
}

/// Options for detecting an online/offline connection.
#[derive(Debug)]
pub struct OfflineOptions {
    pub check_default_hosts: bool,
    pub check_default_ips: bool,
    pub custom_hosts: Vec<String>,
    pub custom_ips: Vec<IpAddr>,
    pub ip_v4: bool,
    pub ip_v6: bool,
    pub timeout: u64,
}

impl Default for OfflineOptions {
    fn default() -> Self {
        Self {
            check_default_hosts: true,
            check_default_ips: true,
            custom_hosts: vec![],
            custom_ips: vec![],
            ip_v4: true,
            ip_v6: true,
            timeout: 1000,
        }
    }
}

/// Detect if there is an internet connection, or the user is offline.
#[instrument]
pub fn is_offline(timeout: u64) -> bool {
    is_offline_with_options(OfflineOptions {
        timeout,
        ..Default::default()
    })
}

/// Detect if there is an internet connection, or the user is offline,
/// using a custom list of hosts.
#[instrument]
pub fn is_offline_with_hosts(timeout: u64, custom_hosts: Vec<String>) -> bool {
    is_offline_with_options(OfflineOptions {
        timeout,
        custom_hosts,
        ..Default::default()
    })
}

/// Detect if there is an internet connection, or the user is offline.
/// This will first ping Cloudflare and Google DNS IP addresses, which
/// is the fastest approach as they do not need to parse host names.
/// If all of these fail, then we will ping Google, Mozilla, and custom
/// hosts, which is slower, so we wrap them in a timeout.
#[instrument]
pub fn is_offline_with_options(options: OfflineOptions) -> bool {
    trace!(
        timeout = options.timeout,
        "Checking for an internet connection"
    );

    // Check these first as they do not need to resolve IP addresses!
    // These typically happen in milliseconds.
    let mut ips = vec![];

    if options.check_default_ips {
        if options.ip_v4 {
            ips.extend([
                // Cloudflare DNS: https://1.1.1.1/dns/
                SocketAddr::from(([1, 1, 1, 1], 53)),
                SocketAddr::from(([1, 0, 0, 1], 53)),
                // Google DNS: https://developers.google.com/speed/public-dns
                SocketAddr::from(([8, 8, 8, 8], 53)),
                SocketAddr::from(([8, 8, 4, 4], 53)),
            ]);
        }

        if options.ip_v6 {
            ips.extend([
                // Cloudflare DNS: https://1.1.1.1/dns/
                SocketAddr::from(([0x2606, 0x4700, 0x4700, 0, 0, 0, 0, 0x1111], 53)),
                SocketAddr::from(([0x2606, 0x4700, 0x4700, 0, 0, 0, 0, 0x1001], 53)),
                // Google DNS: https://developers.google.com/speed/public-dns
                SocketAddr::from(([0x2001, 0x4860, 0x4860, 0, 0, 0, 0, 0x8888], 53)),
                SocketAddr::from(([0x2001, 0x4860, 0x4860, 0, 0, 0, 0, 0x8844], 53)),
            ]);
        }
    }

    for custom_ip in options.custom_ips {
        ips.push(SocketAddr::new(custom_ip, 53));
    }

    let (sender, receiver) = mpsc::channel();
    let checks = ips.len();

    for address in ips {
        let sender = sender.clone();
        let timeout = options.timeout;

        thread::spawn(move || {
            let _ = sender.send(offline::check_connection(address, timeout));
        });
    }

    drop(sender);

    let online = (0..checks).any(|_| receiver.recv().is_ok_and(|online| online));

    if online {
        trace!("Online!");

        return false;
    }

    // Check these second as they need to resolve IP addresses,
    // which adds unnecessary time and overhead that can't be
    // controlled with a native timeout.
    let mut hosts = vec![];

    if options.check_default_hosts {
        hosts.extend([
            "clients3.google.com:80".to_owned(),
            "detectportal.firefox.com:80".to_owned(),
            "google.com:80".to_owned(),
        ]);
    }

    if !options.custom_hosts.is_empty() {
        hosts.extend(options.custom_hosts);
    }

    let (sender, receiver) = mpsc::channel();
    let checks = hosts.len();

    for host in hosts {
        let sender = sender.clone();
        let timeout = options.timeout;

        thread::spawn(move || {
            let _ = sender.send(offline::check_connection_from_host(host, timeout));
        });
    }

    drop(sender);

    let online = (0..checks).any(|_| receiver.recv().is_ok_and(|online| online));

    if online {
        trace!("Online!");

        return false;
    }

    trace!("Offline!!!");

    true
}