nym-sdk 1.21.1

// Copyright 2022-2023 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0

use super::{connection_state::BuilderState, Config, StoragePaths};
use crate::bandwidth::{BandwidthAcquireClient, BandwidthImporter};
use crate::mixnet::socks5_client::Socks5MixnetClient;
use crate::mixnet::{CredentialStorage, MixnetClient, Recipient};
use crate::GatewayTransceiver;
use crate::NymNetworkDetails;
use crate::{Error, Result};
use log::{debug, warn};
use nym_client_core::client::base_client::storage::gateways_storage::GatewayRegistration;
use nym_client_core::client::base_client::storage::helpers::{
    get_active_gateway_identity, get_all_registered_identities, has_gateway_details,
    set_active_gateway,
};
use nym_client_core::client::base_client::storage::{
    Ephemeral, GatewaysDetailsStore, MixnetClientStorage, OnDiskPersistent,
};
use nym_client_core::client::base_client::{BaseClient, EventSender};
use nym_client_core::client::key_manager::persistence::KeyStore;
use nym_client_core::client::{
    base_client::BaseClientBuilder, replies::reply_storage::ReplyStorageBackend,
};
use nym_client_core::config::{DebugConfig, ForgetMe, RememberMe, StatsReporting};
use nym_client_core::error::ClientCoreError;
use nym_client_core::init::helpers::gateways_for_init;
use nym_client_core::init::types::{GatewaySelectionSpecification, GatewaySetup};
use nym_client_core::init::{refresh_gateway_published_data, setup_gateway};
use nym_credentials_interface::TicketType;
use nym_crypto::hkdf::DerivationMaterial;
use nym_socks5_client_core::config::Socks5;
use nym_task::ShutdownTracker;
use nym_topology::provider_trait::TopologyProvider;
use nym_topology::RoutingNode;
use nym_validator_client::{nyxd, QueryHttpRpcNyxdClient, UserAgent};
use rand::rngs::OsRng;
use std::path::Path;
use std::path::PathBuf;
#[cfg(unix)]
use std::sync::Arc;
use url::Url;
use zeroize::Zeroizing;

/// The number of reply SURBs to include in a message by default.
pub(crate) const DEFAULT_NUMBER_OF_SURBS: u32 = 10;

/// Configures and builds a [`MixnetClient`].
///
/// Use [`new_ephemeral`](Self::new_ephemeral) for in-memory keys (discarded on disconnect),
/// or `new_with_default_storage` for persistent identity that survives restarts.
///
/// # Example
///
/// ```no_run
/// use nym_sdk::mixnet::MixnetClientBuilder;
///
/// # #[tokio::main]
/// # async fn main() {
/// let client = MixnetClientBuilder::new_ephemeral()
///     .build()
///     .unwrap()
///     .connect_to_mixnet()
///     .await
///     .unwrap();
/// # }
/// ```
#[derive(Default)]
pub struct MixnetClientBuilder<S: MixnetClientStorage = Ephemeral> {
    config: Config,
    storage_paths: Option<StoragePaths>,
    socks5_config: Option<Socks5>,

    wait_for_gateway: bool,
    wait_for_initial_topology: bool,
    custom_topology_provider: Option<Box<dyn TopologyProvider + Send + Sync>>,
    custom_gateway_transceiver: Option<Box<dyn GatewayTransceiver + Send + Sync>>,
    custom_shutdown: Option<ShutdownTracker>,
    event_tx: Option<EventSender>,
    force_tls: bool,
    no_hostname: bool,
    user_agent: Option<UserAgent>,
    #[cfg(unix)]
    connection_fd_callback: Option<Arc<dyn Fn(std::os::fd::RawFd) + Send + Sync>>,

    // TODO: incorporate it properly into `MixnetClientStorage` (I will need it in wasm anyway)
    gateway_endpoint_config_path: Option<PathBuf>,

    storage: S,
    forget_me: ForgetMe,
    remember_me: RememberMe,
    derivation_material: Option<DerivationMaterial>,
    stream_idle_timeout: Option<std::time::Duration>,
}

impl MixnetClientBuilder<Ephemeral> {
    /// Creates a client builder with ephemeral storage.
    #[must_use]
    pub fn new_ephemeral() -> Self {
        MixnetClientBuilder {
            ..Default::default()
        }
    }

    /// Create a client builder with default values.
    #[must_use]
    pub fn new() -> Self {
        Self::new_ephemeral()
    }
}

impl MixnetClientBuilder<OnDiskPersistent> {
    pub async fn new_with_default_storage(storage_paths: StoragePaths) -> Result<Self> {
        Ok(MixnetClientBuilder {
            config: Default::default(),
            storage_paths: None,
            socks5_config: None,
            wait_for_gateway: false,
            wait_for_initial_topology: false,
            custom_topology_provider: None,
            storage: storage_paths
                .initialise_default_persistent_storage()
                .await?,
            gateway_endpoint_config_path: None,
            custom_shutdown: None,
            event_tx: None,
            custom_gateway_transceiver: None,
            force_tls: false,
            no_hostname: false,
            user_agent: None,
            #[cfg(unix)]
            connection_fd_callback: None,
            forget_me: Default::default(),
            remember_me: Default::default(),
            derivation_material: None,
            stream_idle_timeout: None,
        })
    }
}

impl<S> MixnetClientBuilder<S>
where
    S: MixnetClientStorage + Clone + 'static,
    S::ReplyStore: Send + Sync,
    S::GatewaysDetailsStore: Sync,
    <S::ReplyStore as ReplyStorageBackend>::StorageError: Sync + Send,
    <S::CredentialStore as CredentialStorage>::StorageError: Send + Sync,
    <S::KeyStore as KeyStore>::StorageError: Send + Sync,
    <S::GatewaysDetailsStore as GatewaysDetailsStore>::StorageError: Send + Sync,
{
    /// Creates a client builder with the provided client storage implementation.
    #[must_use]
    pub fn new_with_storage(storage: S) -> MixnetClientBuilder<S> {
        MixnetClientBuilder {
            config: Default::default(),
            storage_paths: None,
            socks5_config: None,
            wait_for_gateway: false,
            wait_for_initial_topology: false,
            custom_topology_provider: None,
            custom_gateway_transceiver: None,
            custom_shutdown: None,
            event_tx: None,
            force_tls: false,
            no_hostname: false,
            user_agent: None,
            #[cfg(unix)]
            connection_fd_callback: None,
            gateway_endpoint_config_path: None,
            storage,
            forget_me: Default::default(),
            remember_me: Default::default(),
            derivation_material: None,
            stream_idle_timeout: None,
        }
    }

    /// Change the underlying storage implementation.
    #[must_use]
    pub fn set_storage<T: MixnetClientStorage>(self, storage: T) -> MixnetClientBuilder<T> {
        MixnetClientBuilder {
            config: self.config,
            storage_paths: self.storage_paths,
            socks5_config: self.socks5_config,
            wait_for_gateway: self.wait_for_gateway,
            wait_for_initial_topology: self.wait_for_initial_topology,
            custom_topology_provider: self.custom_topology_provider,
            custom_gateway_transceiver: self.custom_gateway_transceiver,
            custom_shutdown: self.custom_shutdown,
            event_tx: self.event_tx,
            force_tls: self.force_tls,
            no_hostname: self.no_hostname,
            user_agent: self.user_agent,
            #[cfg(unix)]
            connection_fd_callback: self.connection_fd_callback,
            gateway_endpoint_config_path: self.gateway_endpoint_config_path,
            storage,
            forget_me: self.forget_me,
            remember_me: self.remember_me,
            derivation_material: self.derivation_material,
            stream_idle_timeout: self.stream_idle_timeout,
        }
    }

    #[must_use]
    pub fn with_derivation_material(mut self, derivation_material: DerivationMaterial) -> Self {
        self.derivation_material = Some(derivation_material);
        self
    }

    /// Change the underlying storage of this builder to use default implementation of on-disk disk_persistence.
    #[must_use]
    pub fn set_default_storage(
        self,
        storage: OnDiskPersistent,
    ) -> MixnetClientBuilder<OnDiskPersistent> {
        self.set_storage(storage)
    }

    #[must_use]
    pub fn with_forget_me(mut self, forget_me: ForgetMe) -> Self {
        self.forget_me = forget_me;
        self
    }

    #[must_use]
    pub fn with_remember_me(mut self, remember_me: RememberMe) -> Self {
        self.remember_me = remember_me;
        self
    }

    /// Set the idle timeout for streams. Streams with no activity for this
    /// duration are automatically cleaned up by the router.
    /// Defaults to 30 minutes if not set.
    #[must_use]
    pub fn with_stream_idle_timeout(mut self, timeout: std::time::Duration) -> Self {
        self.stream_idle_timeout = Some(timeout);
        self
    }

    /// Request a specific gateway instead of a random one.
    #[must_use]
    pub fn request_gateway(mut self, user_chosen_gateway: String) -> Self {
        self.config.user_chosen_gateway = Some(user_chosen_gateway);
        self
    }

    #[must_use]
    pub fn with_extended_topology(mut self, use_extended_topology: bool) -> Self {
        self.config.debug_config.topology.use_extended_topology = use_extended_topology;
        self
    }

    #[must_use]
    pub fn with_ignore_epoch_roles(mut self, ignore_epoch_roles: bool) -> Self {
        self.config.debug_config.topology.ignore_egress_epoch_role = ignore_epoch_roles;
        self
    }

    /// Use a specific network instead of the default (mainnet) one.
    #[must_use]
    pub fn network_details(mut self, network_details: NymNetworkDetails) -> Self {
        self.config.network_details = network_details;
        self
    }

    /// Attempt to only choose a gateway that supports wss protocol.
    #[must_use]
    pub fn force_tls(mut self, must_use_tls: bool) -> Self {
        self.force_tls = must_use_tls;
        self
    }

    /// Attempt to only choose a gateway with its IP address only, ignored if force_tls is set
    #[must_use]
    pub fn no_hostname(mut self, no_hostname: bool) -> Self {
        self.no_hostname = no_hostname;
        self
    }

    /// Enable paid coconut bandwidth credentials mode.
    #[must_use]
    pub fn enable_credentials_mode(mut self) -> Self {
        self.config.enabled_credentials_mode = true;
        self
    }

    /// Enable paid coconut bandwidth credentials mode.
    #[must_use]
    pub fn credentials_mode(mut self, credentials_mode: bool) -> Self {
        self.config.enabled_credentials_mode = credentials_mode;
        self
    }

    /// Use a custom debugging configuration.
    #[must_use]
    pub fn debug_config(mut self, debug_config: DebugConfig) -> Self {
        self.config.debug_config = debug_config;
        self
    }

    /// Configure the SOCKS5 mode.
    #[must_use]
    pub fn socks5_config(mut self, socks5_config: Socks5) -> Self {
        self.socks5_config = Some(socks5_config);
        self
    }

    /// Use a custom topology provider.
    #[must_use]
    pub fn custom_topology_provider(
        mut self,
        topology_provider: Box<dyn TopologyProvider + Send + Sync>,
    ) -> Self {
        self.custom_topology_provider = Some(topology_provider);
        self
    }

    /// Use an externally managed shutdown mechanism.
    #[must_use]
    pub fn custom_shutdown(mut self, shutdown: ShutdownTracker) -> Self {
        self.custom_shutdown = Some(shutdown);
        self
    }

    /// Use an externally managed shutdown mechanism.
    #[must_use]
    pub fn event_tx(mut self, event_tx: EventSender) -> Self {
        self.event_tx = Some(event_tx);
        self
    }

    /// Attempt to wait for the selected gateway (if applicable) to come online if it's currently not bonded.
    #[must_use]
    pub fn with_wait_for_gateway(mut self, wait_for_gateway: bool) -> Self {
        self.wait_for_gateway = wait_for_gateway;
        self
    }

    /// Attempt to wait for initial network topology to become online before finalizing client setup
    /// this is useful during network bootstrapping phases
    #[must_use]
    pub fn with_wait_for_initial_topology(mut self, wait_for_initial_topology: bool) -> Self {
        self.wait_for_initial_topology = wait_for_initial_topology;
        self
    }

    #[must_use]
    pub fn with_user_agent(mut self, user_agent: UserAgent) -> Self {
        self.user_agent = Some(user_agent);
        self
    }

    #[must_use]
    pub fn with_statistics_reporting(mut self, config: StatsReporting) -> Self {
        self.config.debug_config.stats_reporting = config;
        self
    }

    #[cfg(unix)]
    #[must_use]
    pub fn with_connection_fd_callback(
        mut self,
        connection_fd_callback: Arc<dyn Fn(std::os::fd::RawFd) + Send + Sync>,
    ) -> Self {
        self.connection_fd_callback = Some(connection_fd_callback);
        self
    }

    /// Use custom mixnet sender that might not be the default websocket gateway connection.
    /// only for advanced use
    #[must_use]
    pub fn custom_gateway_transceiver(
        mut self,
        gateway_transceiver: Box<dyn GatewayTransceiver + Send + Sync>,
    ) -> Self {
        self.custom_gateway_transceiver = Some(gateway_transceiver);
        self
    }

    /// Use specified file for storing gateway configuration.
    pub fn gateway_endpoint_config_path<P: AsRef<Path>>(mut self, path: P) -> Self {
        self.gateway_endpoint_config_path = Some(path.as_ref().to_owned());
        self
    }

    /// Construct a [`DisconnectedMixnetClient`] from the setup specified.
    #[allow(clippy::result_large_err)]
    pub fn build(self) -> Result<DisconnectedMixnetClient<S>> {
        let mut client = DisconnectedMixnetClient::new(
            self.config,
            self.socks5_config,
            self.storage,
            self.event_tx,
        )?;

        client.custom_gateway_transceiver = self.custom_gateway_transceiver;
        client.custom_topology_provider = self.custom_topology_provider;
        client.custom_shutdown = self.custom_shutdown;
        client.wait_for_gateway = self.wait_for_gateway;
        client.wait_for_initial_topology = self.wait_for_initial_topology;
        client.force_tls = self.force_tls;
        client.no_hostname = self.no_hostname;
        client.user_agent = self.user_agent;
        #[cfg(unix)]
        if self.connection_fd_callback.is_some() {
            client.connection_fd_callback = self.connection_fd_callback;
        }
        client.forget_me = self.forget_me;
        client.remember_me = self.remember_me;
        client.derivation_material = self.derivation_material;
        client.stream_idle_timeout = self.stream_idle_timeout;
        Ok(client)
    }
}

/// Represents a client that is not yet connected to the mixnet.
///
/// Represents a client that is not yet connected to the mixnet. You typically create one when you
/// want to have a separate configuration and connection phase. Once the mixnet client builder is
/// configured, call [`connect_to_mixnet()`](Self::connect_to_mixnet) or
/// [`connect_to_mixnet_via_socks5()`](Self::connect_to_mixnet_via_socks5) to transition to a connected
/// client.
pub struct DisconnectedMixnetClient<S>
where
    S: MixnetClientStorage + Clone,
{
    /// Client configuration
    config: Config,

    /// Socks5 configuration
    socks5_config: Option<Socks5>,

    /// The client can be in one of multiple states, depending on how it is created and if it's
    /// connected to the mixnet.
    state: BuilderState,

    /// Underlying storage of this client.
    storage: S,

    /// In the case of enabled credentials, a client instance responsible for querying the state of the
    /// dkg and coconut contracts
    dkg_query_client: Option<QueryHttpRpcNyxdClient>,

    /// Alternative provider of network topology used for constructing sphinx packets.
    custom_topology_provider: Option<Box<dyn TopologyProvider + Send + Sync>>,

    /// advanced usage of custom gateways
    custom_gateway_transceiver: Option<Box<dyn GatewayTransceiver + Send + Sync>>,

    /// Attempt to wait for the selected gateway (if applicable) to come online if it's currently not bonded.
    wait_for_gateway: bool,

    /// Attempt to wait for initial network topology to become online before finalizing client setup
    /// this is useful during network bootstrapping phases
    wait_for_initial_topology: bool,

    /// Force the client to connect using wss protocol with the gateway.
    force_tls: bool,

    /// Force the client to pick gateway IP and not hostname, ignored if force_tls is set
    no_hostname: bool,

    /// Allows passing an externally controlled shutdown handle.
    custom_shutdown: Option<ShutdownTracker>,

    /// Sender of mixnet client events to the SDK caller
    event_tx: Option<EventSender>,

    user_agent: Option<UserAgent>,

    /// Callback on the websocket fd as soon as the connection has been established
    #[cfg(unix)]
    connection_fd_callback: Option<Arc<dyn Fn(std::os::fd::RawFd) + Send + Sync>>,

    forget_me: ForgetMe,

    remember_me: RememberMe,
    /// The derivation material to use for the client keys, its up to the caller to save this for rederivation later
    derivation_material: Option<DerivationMaterial>,

    stream_idle_timeout: Option<std::time::Duration>,
}

impl<S> DisconnectedMixnetClient<S>
where
    S: MixnetClientStorage + Clone + 'static,
    S::ReplyStore: Send + Sync,
    S::GatewaysDetailsStore: Sync,
    <S::ReplyStore as ReplyStorageBackend>::StorageError: Sync + Send,
    <S::CredentialStore as CredentialStorage>::StorageError: Send + Sync,
    <S::KeyStore as KeyStore>::StorageError: Send + Sync,
    <S::GatewaysDetailsStore as GatewaysDetailsStore>::StorageError: Send + Sync,
{
    /// Create a new mixnet client in a disconnected state. The default configuration,
    /// creates a new mainnet client with ephemeral keys stored in RAM, which will be discarded at
    /// application close.
    ///
    /// Callers have the option of supplying further parameters to:
    /// - store persistent identities at a location on-disk, if desired;
    /// - use SOCKS5 mode
    #[allow(clippy::result_large_err)]
    fn new(
        config: Config,
        socks5_config: Option<Socks5>,
        storage: S,
        event_tx: Option<EventSender>,
    ) -> Result<DisconnectedMixnetClient<S>> {
        // don't create dkg client for the bandwidth controller if credentials are disabled
        let dkg_query_client = if config.enabled_credentials_mode {
            let client_config =
                nyxd::Config::try_from_nym_network_details(&config.network_details)?;
            let client = QueryHttpRpcNyxdClient::connect(
                client_config,
                config.network_details.endpoints[0].nyxd_url.as_str(),
            )?;
            Some(client)
        } else {
            None
        };

        let forget_me = config.debug_config.forget_me;
        let remember_me = config.debug_config.remember_me;

        Ok(DisconnectedMixnetClient {
            config,
            socks5_config,
            state: BuilderState::New,
            dkg_query_client,
            storage,
            custom_topology_provider: None,
            custom_gateway_transceiver: None,
            wait_for_gateway: false,
            wait_for_initial_topology: false,
            force_tls: false,
            no_hostname: false,
            custom_shutdown: None,
            event_tx,
            user_agent: None,
            #[cfg(unix)]
            connection_fd_callback: None,
            forget_me,
            remember_me,
            derivation_material: None,
            stream_idle_timeout: None,
        })
    }

    fn get_api_endpoints(&self) -> Vec<Url> {
        self.config
            .network_details
            .endpoints
            .iter()
            .filter_map(|details| details.api_url.as_ref())
            .filter_map(|s| Url::parse(s).ok())
            .collect()
    }

    fn get_nyxd_endpoints(&self) -> Vec<Url> {
        self.config
            .network_details
            .endpoints
            .iter()
            .map(|details| details.nyxd_url.as_ref())
            .filter_map(|s| Url::parse(s).ok())
            .collect()
    }

    pub async fn setup_client_keys(&self) -> Result<()> {
        let mut rng = OsRng;
        let key_store = self.storage.key_store();

        if key_store.load_keys().await.is_err() {
            debug!("Generating new client keys");
            nym_client_core::init::generate_new_client_keys(&mut rng, key_store).await?;
        }

        Ok(())
    }

    async fn print_all_registered_gateway_identities(&self) {
        match get_all_registered_identities(self.storage.gateway_details_store()).await {
            Err(err) => {
                warn!("failed to query for all registered gateways: {err}")
            }
            Ok(all_ids) => {
                if !all_ids.is_empty() {
                    debug!("this client is already registered with the following gateways:");
                    for id in all_ids {
                        debug!("{id}")
                    }
                }
            }
        }
    }

    async fn print_selected_gateway(&self) {
        match self.storage.gateway_details_store().active_gateway().await {
            Err(err) => {
                warn!("failed to query for the current active gateway: {err}")
            }
            Ok(active) => {
                if let Some(active) = active.registration {
                    let id = active.details.gateway_id();
                    debug!("currently selected gateway: {id}");
                }
            }
        }
    }

    async fn set_active_gateway_if_previously_registered(
        &self,
        user_chosen_gateway: &str,
    ) -> Result<bool> {
        let storage = self.storage.gateway_details_store();
        // Strictly speaking, `set_active_gateway` does this check internally as well, but since the
        // error is boxed away and we're using a generic storage, it's not so easy to match on it.
        // This function is at least less likely to fail on something unrelated to the existence of
        // the gateway in the set of registered gateways
        if has_gateway_details(storage, user_chosen_gateway).await? {
            set_active_gateway(storage, user_chosen_gateway).await?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Attempt to retrieve list of all gateways available for registration
    async fn available_gateways(&mut self) -> Result<Vec<RoutingNode>, ClientCoreError> {
        if let Some(ref mut custom_provider) = self.custom_topology_provider {
            if let Some(topology) = custom_provider.get_new_topology().await {
                // Use entry_capable_nodes() instead of entry_gateways() to include
                // all entry-capable nodes, not just actively assigned ones
                return Ok(topology.entry_capable_nodes().cloned().collect());
            }
        }

        let nym_api_endpoints = self.get_api_endpoints();
        let topology_cfg = &self.config.debug_config.topology;
        let user_agent = self.user_agent.clone();

        gateways_for_init(
            &nym_api_endpoints,
            user_agent,
            topology_cfg.minimum_gateway_performance,
            topology_cfg.ignore_ingress_epoch_role,
            None,
        )
        .await
    }

    async fn new_gateway_setup(&mut self) -> Result<GatewaySetup, ClientCoreError> {
        let selection_spec = GatewaySelectionSpecification::new(
            self.config.user_chosen_gateway.clone(),
            None,
            self.force_tls,
            self.no_hostname,
        );

        let available_gateways = self.available_gateways().await?;
        for node in available_gateways.iter() {
            debug!(
                "node_id={}, identity_key={}",
                node.node_id,
                node.identity_key.to_base58_string()
            );
        }

        Ok(GatewaySetup::New {
            specification: selection_spec,
            available_gateways,
            #[cfg(unix)]
            connection_fd_callback: self.connection_fd_callback.clone(),
        })
    }

    async fn refresh_gateway_published_data(
        &mut self,
        gateway_registration: GatewayRegistration,
    ) -> Result<(), ClientCoreError> {
        let available_gateways = self.available_gateways().await?;
        refresh_gateway_published_data(
            self.storage.gateway_details_store(),
            gateway_registration,
            available_gateways,
            self.force_tls,
            self.no_hostname,
        )
        .await
    }

    /// Register with a gateway. If a gateway is provided in the config then that will try to be
    /// used. If none is specified, a gateway at random will be picked. The used gateway is saved
    /// as the active gateway.
    ///
    /// # Errors
    ///
    /// This function will return an error if you try to re-register when in an already registered
    /// state.
    pub async fn setup_gateway(&mut self) -> Result<()> {
        if !matches!(self.state, BuilderState::New) {
            return Err(Error::ReregisteringGatewayNotSupported);
        }

        self.print_all_registered_gateway_identities().await;
        self.print_selected_gateway().await;

        // Try to set active gateway to the same as the user chosen one, if it's in the set of
        // gateways that is already registered.
        if let Some(ref user_chosen_gateway) = self.config.user_chosen_gateway {
            if self
                .set_active_gateway_if_previously_registered(user_chosen_gateway)
                .await?
            {
                debug!("user chosen gateway is already registered, set as active");
            }
        }

        let active_gateway =
            get_active_gateway_identity(self.storage.gateway_details_store()).await?;

        // Determine the gateway setup based on the currently active gateway and the user-chosen
        // gateway.
        let gateway_setup = match (self.config.user_chosen_gateway.as_ref(), active_gateway) {
            // When a user-chosen gateway exists and matches the active one.
            (Some(user_chosen_gateway), Some(active_gateway))
                if &active_gateway.to_base58_string() == user_chosen_gateway =>
            {
                GatewaySetup::MustLoad { gateway_id: None }
            }
            // When a user-chosen gateway exists but there's no active gateway, or it doesn't match the active one.
            (Some(_), _) => self.new_gateway_setup().await?,
            // When no user-chosen gateway exists but there's an active gateway.
            (None, Some(_)) => GatewaySetup::MustLoad { gateway_id: None },
            // When there's no user-chosen gateway and no active gateway.
            (None, None) => self.new_gateway_setup().await?,
        };

        // this will perform necessary key and details load and optional store
        let init_results = setup_gateway(
            gateway_setup,
            self.storage.key_store(),
            self.storage.gateway_details_store(),
        )
        .await?;

        // update gateway setup if needed
        if init_results.exipred_details() {
            self.refresh_gateway_published_data(init_results.gateway_registration.clone())
                .await?;
        }

        set_active_gateway(
            self.storage.gateway_details_store(),
            &init_results.gateway_id().to_base58_string(),
        )
        .await?;

        self.state = BuilderState::Registered {};
        Ok(())
    }

    /// Creates an associated [`BandwidthAcquireClient`] that can be used to acquire bandwidth
    /// credentials of particular type for this client to consume.
    pub async fn create_bandwidth_client(
        &self,
        mnemonic: String,
        ticketbook_type: TicketType,
    ) -> Result<BandwidthAcquireClient<S::CredentialStore>> {
        if !self.config.enabled_credentials_mode {
            return Err(Error::DisabledCredentialsMode);
        }
        let client_id_array = Zeroizing::new(
            self.storage
                .key_store()
                .load_keys()
                .await
                .map_err(|e| Error::KeyStorageError {
                    source: Box::new(e),
                })?
                .identity_keypair()
                .private_key()
                .to_bytes(),
        );
        let client_id = client_id_array.to_vec();

        BandwidthAcquireClient::new(
            self.config.network_details.clone(),
            mnemonic,
            self.storage.credential_store().clone(),
            client_id,
            ticketbook_type,
        )
    }

    pub fn begin_bandwidth_import(&self) -> BandwidthImporter<'_, S::CredentialStore> {
        BandwidthImporter::new(self.storage.credential_store())
    }

    async fn connect_to_mixnet_common(mut self) -> Result<(BaseClient, Recipient)> {
        self.setup_client_keys().await?;
        self.setup_gateway().await?;

        let nyxd_endpoints = self.get_nyxd_endpoints();
        let nym_api_endpoints = self.get_api_endpoints();

        // a temporary workaround
        let base_config = self
            .config
            .as_base_client_config(nyxd_endpoints, nym_api_endpoints.clone());

        tracing::debug!(
            "SDK: Passing nym_api_urls to BaseClientBuilder (has {} nym_api_urls)",
            self.config
                .network_details
                .nym_api_urls
                .as_ref()
                .map(|urls| urls.len())
                .unwrap_or(0)
        );

        let mut base_builder: BaseClientBuilder<_, _> =
            BaseClientBuilder::new(base_config, self.storage, self.dkg_query_client)
                .with_wait_for_gateway(self.wait_for_gateway)
                .with_wait_for_initial_topology(self.wait_for_initial_topology)
                .with_forget_me(&self.forget_me)
                .with_remember_me(&self.remember_me)
                .with_derivation_material(self.derivation_material);

        // Add nym_api_urls if available in network_details
        if let Some(nym_api_urls) = self.config.network_details.nym_api_urls.clone() {
            base_builder = base_builder.with_nym_api_urls(nym_api_urls);
        }

        if let Some(user_agent) = self.user_agent {
            base_builder = base_builder.with_user_agent(user_agent);
        }

        if let Some(topology_provider) = self.custom_topology_provider {
            base_builder = base_builder.with_topology_provider(topology_provider);
        }

        // Use custom shutdown if provided, otherwise the sdk one will be used later down the line
        if let Some(shutdown_tracker) = self.custom_shutdown {
            base_builder = base_builder.with_shutdown(shutdown_tracker);
        }

        if let Some(event_tx) = self.event_tx {
            base_builder = base_builder.with_event_tx(event_tx);
        }

        if let Some(gateway_transceiver) = self.custom_gateway_transceiver {
            base_builder = base_builder.with_gateway_transceiver(gateway_transceiver);
        }

        #[cfg(unix)]
        if let Some(connection_fd_callback) = self.connection_fd_callback {
            base_builder = base_builder.with_connection_fd_callback(connection_fd_callback);
        }

        let started_client = base_builder.start_base().await?;
        self.state = BuilderState::Registered {};
        let nym_address = started_client.address;

        Ok((started_client, nym_address))
    }

    /// Connect the client to the mixnet via SOCKS5. A SOCKS5 configuration must be specified
    /// before attempting to connect.
    ///
    /// - If the client is already registered with a gateway, use that gateway.
    /// - If no gateway is registered, but there is an existing configuration and key, use that.
    /// - If no gateway is registered, and there is no pre-existing configuration or key, try to
    ///   register a new gateway.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use nym_sdk::mixnet;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let receiving_client = mixnet::MixnetClient::connect_new().await.unwrap();
    ///     let socks5_config = mixnet::Socks5::new(receiving_client.nym_address().to_string());
    ///     let client = mixnet::MixnetClientBuilder::new_ephemeral()
    ///         .socks5_config(socks5_config)
    ///         .build()
    ///         .unwrap();
    ///     let client = client.connect_to_mixnet_via_socks5().await.unwrap();
    /// }
    /// ```
    pub async fn connect_to_mixnet_via_socks5(self) -> Result<Socks5MixnetClient> {
        let socks5_config = self
            .socks5_config
            .clone()
            .ok_or(Error::Socks5Config { set: false })?;
        let debug_config = self.config.debug_config;
        let packet_type = self.config.debug_config.traffic.packet_type;
        let (mut started_client, nym_address) = self.connect_to_mixnet_common().await?;

        let client_input = started_client.client_input.register_producer();
        let client_output = started_client.client_output.register_consumer();
        let client_state = started_client.client_state;

        nym_socks5_client_core::NymClient::<S>::start_socks5_listener(
            &socks5_config,
            debug_config,
            client_input,
            client_output,
            client_state.clone(),
            nym_address,
            started_client.shutdown_handle.clone(),
            packet_type,
        );

        Ok(Socks5MixnetClient {
            nym_address,
            client_state,
            task_handle: started_client.shutdown_handle,
            socks5_config,
        })
    }

    /// Connect the client to the mixnet.
    ///
    /// - If the client is already registered with a gateway, use that gateway.
    /// - If no gateway is registered, but there is an existing configuration and key, use that.
    /// - If no gateway is registered, and there is no pre-existing configuration or key, try to
    ///   register a new gateway.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use nym_sdk::mixnet;
    ///
    /// #[tokio::main]
    /// async fn main() {
    ///     let client = mixnet::MixnetClientBuilder::new_ephemeral()
    ///         .build()
    ///         .unwrap();
    ///     let client = client.connect_to_mixnet().await.unwrap();
    /// }
    /// ```
    pub async fn connect_to_mixnet(self) -> Result<MixnetClient> {
        if self.socks5_config.is_some() {
            return Err(Error::Socks5Config { set: true });
        }
        let stream_idle_timeout = self.stream_idle_timeout;
        let (mut started_client, nym_address) = self.connect_to_mixnet_common().await?;
        let client_input = started_client.client_input.register_producer();
        let mut client_output = started_client.client_output.register_consumer();
        let client_state: nym_client_core::client::base_client::ClientState =
            started_client.client_state;
        let stats_events_reporter = started_client.stats_reporter;

        let identity_keys = started_client.identity_keys.clone();
        let reconstructed_receiver = client_output.register_receiver()?;

        let mut client = MixnetClient::new(
            nym_address,
            identity_keys,
            client_input,
            client_output,
            client_state,
            reconstructed_receiver,
            stats_events_reporter,
            started_client.shutdown_handle,
            None,
            started_client.forget_me,
            started_client.remember_me,
        );
        if let Some(timeout) = stream_idle_timeout {
            client.stream_idle_timeout = timeout;
        }
        Ok(client)
    }
}

#[derive(Debug, Clone)]
pub enum IncludedSurbs {
    Amount(u32),
    ExposeSelfAddress,
}

impl Default for IncludedSurbs {
    fn default() -> Self {
        Self::Amount(DEFAULT_NUMBER_OF_SURBS)
    }
}

impl IncludedSurbs {
    pub fn new(reply_surbs: u32) -> Self {
        Self::Amount(reply_surbs)
    }

    pub fn none() -> Self {
        Self::Amount(0)
    }

    pub fn expose_self_address() -> Self {
        Self::ExposeSelfAddress
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mixnet_builder_default_no_custom_client() {
        let builder = MixnetClientBuilder::new_ephemeral();
        assert!(
            builder.build().is_ok(),
            "Builder should succeed without custom client"
        );
    }
}