fedimint_connectors/

lib.rs

pub mod error;
pub mod http;
pub mod iroh;
#[cfg(all(feature = "tor", not(target_family = "wasm")))]
pub mod tor;
pub mod ws;

use std::collections::{BTreeMap, HashMap};
use std::fmt::{self, Debug};
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{anyhow, bail};
use async_trait::async_trait;
use fedimint_core::envs::{FM_WS_API_CONNECT_OVERRIDES_ENV, parse_kv_list_from_env};
use fedimint_core::module::{ApiMethod, ApiRequestErased};
use fedimint_core::util::backoff_util::{FibonacciBackoff, custom_backoff};
use fedimint_core::util::{FmtCompact, FmtCompactAnyhow, SafeUrl};
use fedimint_core::{apply, async_trait_maybe_send};
use fedimint_logging::{LOG_CLIENT_NET_API, LOG_NET};
use reqwest::Method;
use serde_json::Value;
use tokio::sync::OnceCell;
use tracing::trace;

use crate::error::ServerError;
use crate::ws::WebsocketConnector;

pub type ServerResult<T> = Result<T, ServerError>;

/// Type for connector initialization functions
type ConnectorInitFn = Arc<
    dyn Fn() -> Pin<Box<dyn Future<Output = anyhow::Result<DynConnector>> + Send>> + Send + Sync,
>;

/// Builder for [`ConnectorRegistry`]
///
/// See [`ConnectorRegistry::build_from_client_env`] and similar
/// to create.
#[derive(Debug, Clone)]
#[allow(clippy::struct_excessive_bools)] // Shut up, Clippy
pub struct ConnectorRegistryBuilder {
    /// List of overrides to use when attempting to connect to given url
    ///
    /// This is useful for testing, or forcing non-default network
    /// connectivity.
    connection_overrides: BTreeMap<SafeUrl, SafeUrl>,

    /// Enable Iroh endpoints at all?
    iroh_enable: bool,
    /// Override the Iroh DNS server to use
    iroh_dns: Option<SafeUrl>,
    /// Should start the "next/unstable" Iroh stack
    iroh_next: bool,
    /// Enable Pkarr DHT discovery
    iroh_pkarr_dht: bool,

    /// Enable Websocket API handling at all?
    ws_enable: bool,
    ws_force_tor: bool,

    // Enable HTTP
    http_enable: bool,
}

impl ConnectorRegistryBuilder {
    #[allow(clippy::unused_async)] // Leave room for async in the future
    pub async fn bind(self) -> anyhow::Result<ConnectorRegistry> {
        // Create initialization functions for each connector type
        let mut connectors_lazy: BTreeMap<String, (ConnectorInitFn, OnceCell<DynConnector>)> =
            BTreeMap::new();

        // WS connector init function
        let builder_ws = self.clone();
        let ws_connector_init = Arc::new(move || {
            let builder = builder_ws.clone();
            Box::pin(async move { builder.build_ws_connector().await })
                as Pin<Box<dyn Future<Output = anyhow::Result<DynConnector>> + Send>>
        });
        connectors_lazy.insert("ws".into(), (ws_connector_init.clone(), OnceCell::new()));
        connectors_lazy.insert("wss".into(), (ws_connector_init.clone(), OnceCell::new()));

        // Iroh connector init function
        let builder_iroh = self.clone();
        connectors_lazy.insert(
            "iroh".into(),
            (
                Arc::new(move || {
                    let builder = builder_iroh.clone();
                    Box::pin(async move { builder.build_iroh_connector().await })
                        as Pin<Box<dyn Future<Output = anyhow::Result<DynConnector>> + Send>>
                }),
                OnceCell::new(),
            ),
        );

        let builder_http = self.clone();
        let http_connector_init = Arc::new(move || {
            let builder = builder_http.clone();
            Box::pin(async move { builder.build_http_connector() })
                as Pin<Box<dyn Future<Output = anyhow::Result<DynConnector>> + Send>>
        });

        connectors_lazy.insert(
            "http".into(),
            (http_connector_init.clone(), OnceCell::new()),
        );
        connectors_lazy.insert(
            "https".into(),
            (http_connector_init.clone(), OnceCell::new()),
        );

        Ok(ConnectorRegistry {
            connectors_lazy: Arc::new(connectors_lazy),
            connection_overrides: self.connection_overrides,
        })
    }

    pub async fn build_iroh_connector(&self) -> anyhow::Result<DynConnector> {
        if !self.iroh_enable {
            bail!("Iroh connector not enabled");
        }
        Ok(Arc::new(
            iroh::IrohConnector::new(self.iroh_dns.clone(), self.iroh_pkarr_dht, self.iroh_next)
                .await?,
        ) as DynConnector)
    }

    pub async fn build_ws_connector(&self) -> anyhow::Result<DynConnector> {
        if !self.ws_enable {
            bail!("Websocket connector not enabled");
        }

        match self.ws_force_tor {
            #[cfg(all(feature = "tor", not(target_family = "wasm")))]
            true => {
                use crate::tor::TorConnector;

                Ok(Arc::new(TorConnector::bootstrap().await?) as DynConnector)
            }

            false => Ok(Arc::new(WebsocketConnector::new()) as DynConnector),
            #[allow(unreachable_patterns)]
            _ => bail!("Tor requested, but not support not compiled in"),
        }
    }

    pub fn build_http_connector(&self) -> anyhow::Result<DynConnector> {
        if !self.http_enable {
            bail!("Http connector not enabled");
        }

        Ok(Arc::new(crate::http::HttpConnector::default()) as DynConnector)
    }

    pub fn iroh_pkarr_dht(self, enable: bool) -> Self {
        Self {
            iroh_pkarr_dht: enable,
            ..self
        }
    }

    pub fn iroh_next(self, enable: bool) -> Self {
        Self {
            iroh_next: enable,
            ..self
        }
    }

    pub fn ws_force_tor(self, enable: bool) -> Self {
        Self {
            ws_force_tor: enable,
            ..self
        }
    }

    pub fn set_iroh_dns(self, url: SafeUrl) -> Self {
        Self {
            iroh_dns: Some(url),
            ..self
        }
    }

    /// Apply overrides from env variables
    pub fn with_env_var_overrides(mut self) -> anyhow::Result<Self> {
        // TODO: read rest of the env
        for (k, v) in parse_kv_list_from_env::<_, SafeUrl>(FM_WS_API_CONNECT_OVERRIDES_ENV)? {
            self = self.with_connection_override(k, v);
        }

        Ok(Self { ..self })
    }

    pub fn with_connection_override(
        mut self,
        original_url: SafeUrl,
        replacement_url: SafeUrl,
    ) -> Self {
        self.connection_overrides
            .insert(original_url, replacement_url);
        self
    }
}

/// A set of available connectivity protocols a client can use to make
/// network API requests (typically to federation).
///
/// Maps from connection URL schema to [`Connector`] to use to connect to it.
///
/// See [`ConnectorRegistry::build_from_client_env`] and similar
/// to create.
///
/// [`ConnectorRegistry::connect_guardian`] is the main entry point for making
/// mixed-networking stack connection.
///
/// Responsibilities:
#[derive(Clone)]
pub struct ConnectorRegistry {
    /// Wrapped in Arc so clones share the same OnceCell instances
    connectors_lazy: Arc<BTreeMap<String, (ConnectorInitFn, OnceCell<DynConnector>)>>,
    /// Connection URL overrides for testing/custom routing
    connection_overrides: BTreeMap<SafeUrl, SafeUrl>,
}

impl fmt::Debug for ConnectorRegistry {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ConnectorRegistry")
            .field("connectors_lazy", &self.connectors_lazy.len())
            .field("connection_overrides", &self.connection_overrides)
            .finish()
    }
}

impl ConnectorRegistry {
    /// Create a builder with recommended defaults intended for client-side
    /// usage
    ///
    /// In particular mobile devices are considered.
    pub fn build_from_client_defaults() -> ConnectorRegistryBuilder {
        ConnectorRegistryBuilder {
            iroh_enable: true,
            iroh_dns: None,
            iroh_pkarr_dht: false,
            iroh_next: true,
            ws_enable: true,
            ws_force_tor: false,
            http_enable: true,

            connection_overrides: BTreeMap::default(),
        }
    }

    /// Create a builder with recommended defaults intended for the server-side
    /// usage
    pub fn build_from_server_defaults() -> ConnectorRegistryBuilder {
        ConnectorRegistryBuilder {
            iroh_enable: true,
            iroh_dns: None,
            iroh_pkarr_dht: true,
            iroh_next: true,
            ws_enable: true,
            ws_force_tor: false,
            http_enable: false,

            connection_overrides: BTreeMap::default(),
        }
    }

    /// Create a builder with recommended defaults intended for testing
    /// usage
    pub fn build_from_testing_defaults() -> ConnectorRegistryBuilder {
        ConnectorRegistryBuilder {
            iroh_enable: true,
            iroh_dns: None,
            iroh_pkarr_dht: false,
            iroh_next: false,
            ws_enable: true,
            ws_force_tor: false,
            http_enable: true,

            connection_overrides: BTreeMap::default(),
        }
    }

    /// Like [`Self::build_from_client_defaults`] build will apply
    /// environment-provided overrides.
    pub fn build_from_client_env() -> anyhow::Result<ConnectorRegistryBuilder> {
        let builder = Self::build_from_client_defaults().with_env_var_overrides()?;
        Ok(builder)
    }

    /// Like [`Self::build_from_server_defaults`] build will apply
    /// environment-provided overrides.
    pub fn build_from_server_env() -> anyhow::Result<ConnectorRegistryBuilder> {
        let builder = Self::build_from_server_defaults().with_env_var_overrides()?;
        Ok(builder)
    }

    /// Like [`Self::build_from_testing_defaults`] build will apply
    /// environment-provided overrides.
    pub fn build_from_testing_env() -> anyhow::Result<ConnectorRegistryBuilder> {
        let builder = Self::build_from_testing_defaults().with_env_var_overrides()?;
        Ok(builder)
    }

    /// Connect to a given `url` using matching [`Connector`]
    ///
    /// This is the main function consumed by the downstream use for making
    /// connection.
    pub async fn connect_guardian(
        &self,
        url: &SafeUrl,
        api_secret: Option<&str>,
    ) -> ServerResult<DynGuaridianConnection> {
        trace!(
            target: LOG_NET,
            %url,
            "Connection requested"
        );
        let url = match self.connection_overrides.get(url) {
            Some(replacement) => {
                trace!(
                    target: LOG_NET,
                    original_url = %url,
                    replacement_url = %replacement,
                    "Using a connectivity override for connection"
                );

                replacement
            }
            None => url,
        };

        let connector_key = url.scheme();

        let Some(connector_lazy) = self.connectors_lazy.get(connector_key) else {
            return Err(ServerError::InvalidEndpoint(anyhow!(
                "Unsupported scheme: {}; missing endpoint handler",
                url.scheme()
            )));
        };

        // Clone the init function to use in the async block
        let init_fn = connector_lazy.0.clone();

        let conn = connector_lazy
            .1
            .get_or_try_init(|| async move { init_fn().await })
            .await
            .map_err(|e| {
                ServerError::Transport(anyhow!(
                    "Connector failed to initialize: {}",
                    e.fmt_compact_anyhow()
                ))
            })?
            .connect_guardian(url, api_secret)
            .await
            .inspect_err(|err| {
                trace!(
                    target: LOG_NET,
                    %url,
                    err = %err.fmt_compact(),
                    "Connection failed"
                );
            })?;

        trace!(
            target: LOG_NET,
            %url,
            "Connection returned"
        );
        Ok(conn)
    }

    /// Connect to a given `url` using matching [`Connector`] to a gateway
    ///
    /// This is the main function consumed by the downstream use for making
    /// connection.
    pub async fn connect_gateway(&self, url: &SafeUrl) -> anyhow::Result<DynGatewayConnection> {
        let url = match self.connection_overrides.get(url) {
            Some(replacement) => {
                trace!(
                    target: LOG_NET,
                    original_url = %url,
                    replacement_url = %replacement,
                    "Using a connectivity override for connection"
                );

                replacement
            }
            None => url,
        };

        let connector_key = url.scheme();

        let Some(connector_lazy) = self.connectors_lazy.get(connector_key) else {
            return Err(anyhow!(
                "Unsupported scheme: {}; missing endpoint handler",
                url.scheme()
            ));
        };

        // Clone the init function to use in the async block
        let init_fn = connector_lazy.0.clone();

        connector_lazy
            .1
            .get_or_try_init(|| async move { init_fn().await })
            .await
            .map_err(|e| {
                ServerError::Transport(anyhow!(
                    "Connector failed to initialize: {}",
                    e.fmt_compact_anyhow()
                ))
            })?
            .connect_gateway(url)
            .await
    }
}
pub type DynConnector = Arc<dyn Connector>;

#[async_trait]
pub trait Connector: Send + Sync + 'static + Debug {
    async fn connect_guardian(
        &self,
        url: &SafeUrl,
        api_secret: Option<&str>,
    ) -> ServerResult<DynGuaridianConnection>;

    async fn connect_gateway(&self, url: &SafeUrl) -> anyhow::Result<DynGatewayConnection>;
}

/// Generic connection trait shared between [`IGuardianConnection`] and
/// [`IGatewayConnection`]
#[apply(async_trait_maybe_send!)]
pub trait IConnection: Debug + Send + Sync + 'static {
    fn is_connected(&self) -> bool;

    async fn await_disconnection(&self);
}

/// A connection from api client to a federation guardian (type erased)
pub type DynGuaridianConnection = Arc<dyn IGuardianConnection>;

/// A connection from api client to a federation guardian
#[async_trait]
pub trait IGuardianConnection: IConnection + Debug + Send + Sync + 'static {
    async fn request(&self, method: ApiMethod, request: ApiRequestErased) -> ServerResult<Value>;

    fn into_dyn(self) -> DynGuaridianConnection
    where
        Self: Sized,
    {
        Arc::new(self)
    }
}

/// A connection from api client to a gateway (type erased)
pub type DynGatewayConnection = Arc<dyn IGatewayConnection>;

/// A connection from a client to a gateway
#[apply(async_trait_maybe_send!)]
pub trait IGatewayConnection: IConnection + Debug + Send + Sync + 'static {
    async fn request(
        &self,
        password: Option<String>,
        method: Method,
        route: &str,
        payload: Option<Value>,
    ) -> ServerResult<Value>;

    fn into_dyn(self) -> DynGatewayConnection
    where
        Self: Sized,
    {
        Arc::new(self)
    }
}

#[derive(Debug)]
pub struct ConnectionPool<T: IConnection + ?Sized> {
    /// Available connectors which we can make connections
    connectors: ConnectorRegistry,

    /// Connection pool
    ///
    /// Every entry in this map will be created on demand and correspond to a
    /// single outgoing connection to a certain URL that is in the process
    /// of being established, or we already established.
    #[allow(clippy::type_complexity)]
    connections: Arc<tokio::sync::Mutex<HashMap<SafeUrl, Arc<ConnectionState<T>>>>>,
}

impl<T: IConnection + ?Sized> Clone for ConnectionPool<T> {
    fn clone(&self) -> Self {
        Self {
            connectors: self.connectors.clone(),
            connections: self.connections.clone(),
        }
    }
}

impl<T: IConnection + ?Sized> ConnectionPool<T> {
    pub fn new(connectors: ConnectorRegistry) -> Self {
        Self {
            connectors,
            connections: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
        }
    }

    pub async fn get_or_create_connection<F, Fut>(
        &self,
        url: &SafeUrl,
        api_secret: Option<&str>,
        create_connection: F,
    ) -> ServerResult<Arc<T>>
    where
        F: Fn(SafeUrl, Option<String>, ConnectorRegistry) -> Fut + Clone + Send + Sync + 'static,
        Fut: Future<Output = ServerResult<Arc<T>>> + Send + 'static,
    {
        let mut pool_locked = self.connections.lock().await;

        let pool_entry_arc = pool_locked
            .entry(url.to_owned())
                        .and_modify(|entry_arc| {
                // Check if existing connection is disconnected and reset the whole entry.
                //
                // This resets the state (like connectivity backoff), which is what we want.
                // Since the (`OnceCell`) was already initialized, it means connection was successfully
                // before, and disconnected afterwards.
                if let Some(existing_conn) = entry_arc.connection.get()
                    && !existing_conn.is_connected(){
                        trace!(target: LOG_CLIENT_NET_API, %url, "Existing connection is disconnected, removing from pool");
                        *entry_arc = Arc::new(ConnectionState::new_reconnecting());
                    }
            })
            .or_insert_with(|| Arc::new(ConnectionState::new_initial()))
            .clone();

        // Drop the pool lock so other connections can work in parallel
        drop(pool_locked);

        let conn = pool_entry_arc
            .connection
            // This serializes all the connection attempts. If one attempt to connect (including
            // waiting for the reconnect backoff) succeeds, all waiting ones will use it. If it
            // fails, any already pending/next will attempt it right afterwards.
            // Nit: if multiple calls are trying to connect to the same host that is offline, it
            // will take some of them multiples of maximum retry delay to actually return with
            // an error. This should be fine in practice and hard to avoid without a lot of
            // complexity.
            .get_or_try_init(|| async {
                let retry_delay = pool_entry_arc.pre_reconnect_delay();
                fedimint_core::runtime::sleep(retry_delay).await;

                trace!(target: LOG_CLIENT_NET_API, %url, "Attempting to create a new connection");
                create_connection(
                    url.clone(),
                    api_secret.map(std::string::ToString::to_string),
                    self.connectors.clone(),
                )
                .await
            })
            .await?;

        trace!(target: LOG_CLIENT_NET_API, %url, "Connection ready");
        Ok(conn.clone())
    }
}

/// Inner part of [`ConnectionState`] preserving state between attempts to
/// initialize [`ConnectionState::connection`]
#[derive(Debug)]
struct ConnectionStateInner {
    fresh: bool,
    backoff: FibonacciBackoff,
}

#[derive(Debug)]
pub struct ConnectionState<T: ?Sized> {
    /// Connection we are trying to or already established
    pub connection: tokio::sync::OnceCell<Arc<T>>,
    /// State that technically is protected every time by
    /// the serialization of `OnceCell::get_or_try_init`, but
    /// for Rust purposes needs to be locked.
    inner: std::sync::Mutex<ConnectionStateInner>,
}

impl<T: ?Sized> ConnectionState<T> {
    /// Create a new connection state for a first time connection
    pub fn new_initial() -> Self {
        Self {
            connection: OnceCell::new(),
            inner: std::sync::Mutex::new(ConnectionStateInner {
                fresh: true,
                backoff: custom_backoff(
                    // First time connections start quick
                    Duration::from_millis(5),
                    Duration::from_secs(30),
                    None,
                ),
            }),
        }
    }

    /// Create a new connection state for a connection that already failed, and
    /// is being reset
    pub fn new_reconnecting() -> Self {
        Self {
            connection: OnceCell::new(),
            inner: std::sync::Mutex::new(ConnectionStateInner {
                // set the attempts to 1, indicating that
                fresh: false,
                backoff: custom_backoff(
                    // Connections after a disconnect start with some minimum delay
                    Duration::from_millis(500),
                    Duration::from_secs(30),
                    None,
                ),
            }),
        }
    }

    /// Record the fact that an attempt to connect is being made, and return
    /// time the caller should wait.
    pub fn pre_reconnect_delay(&self) -> Duration {
        let mut backoff_locked = self.inner.lock().expect("Locking failed");
        let fresh = backoff_locked.fresh;

        backoff_locked.fresh = false;

        if fresh {
            Duration::default()
        } else {
            backoff_locked.backoff.next().expect("Keeps retrying")
        }
    }
}