spark_rust/wallet/config/

spark.rs

//! # Spark Protocol Configuration
//!
//! This module defines configuration for connecting to and interacting with
//! the Spark protocol network, including:
//!
//! - Operator connection management
//! - Threshold signing configuration
//! - Authentication session handling
//! - RPC and GraphQL client configuration
//!
//! The Spark protocol uses a threshold signing approach where multiple operators
//! collaborate with the client to create valid signatures. This module manages
//! connections to these operators and handles the related configuration.

use crate::constants::spark::connection::{DEFAULT_COORDINATOR_INDEX, DEFAULT_RETRY_COUNT};
use crate::constants::spark::{
    LIGHTSPARK_SSP_MAINNET_IDENTITY_PUBLIC_KEY, LIGHTSPARK_SSP_REGTEST_IDENTITY_PUBLIC_KEY,
};
use crate::error::CryptoError;
use crate::error::{NetworkError, SparkSdkError, ValidationError};
use crate::rpc::connections::connection::SparkConnection;
use crate::rpc::SparkRpcClient;
use crate::wallet::graphql::GraphqlClient;
use crate::SparkNetwork;
use bitcoin::secp256k1::PublicKey;
use frost_secp256k1_tr_unofficial::Identifier;
use hashbrown::HashMap;
use parking_lot::RwLock;
use spark_cryptography::signing::identifier_to_hex_string;
use spark_protos::spark::spark_service_client::SparkServiceClient;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
use tonic::transport::Channel;
use tonic::transport::Uri;
use tonic::Request;

use crate::constants::spark::connection::SPARK_DEV_OPERATORS;
use crate::constants::spark::SPARK_REGTEST_SIGNING_THRESHOLD;

use crate::rpc::traits::SparkRpcConnection;

/// Primary configuration for the Spark protocol
///
/// This struct manages the network-specific configuration for the Spark protocol,
/// including connections to operators, threshold signing parameters, and authentication.
/// It serves as the central point for all protocol-related settings and provides methods
/// for interacting with the Spark operator network.
///
/// Key responsibilities:
/// - Managing operator connections and authentication
/// - Providing access to RPC and GraphQL clients
/// - Maintaining threshold signing configuration
/// - Handling connection retries and backoff
#[derive(Clone)]
pub(crate) struct SparkConfig {
    /// Network to use for the wallet (mainnet, testnet, etc.)
    pub(crate) network: SparkNetwork,

    /// Spark operator pool
    pub(crate) operator_pool: Arc<SparkOperatorPool>,

    /// SSP identity public key
    pub(crate) ssp_identity_public_key: PublicKey,

    /// The threshold constant used for signing
    pub(crate) threshold: u32,

    /// GraphQL client for SSP API interactions
    pub(crate) ssp_graphql_client: Arc<GraphqlClient>,

    /// gRPC client used for Spark Operator RPC calls
    spark_clients: Arc<RwLock<HashMap<u32, SparkRpcClient>>>,

    /// Session
    pub(crate) sessions: Arc<parking_lot::RwLock<Vec<Session>>>,
}

/// Authentication session for Spark operator communication
///
/// Represents an authenticated session with a Spark operator, containing
/// the session token and expiration information needed for secure communication.
#[derive(Debug, Clone)]
pub struct Session {
    /// The session token
    pub(crate) session_token: String,

    /// The expiration timestamp
    pub(crate) _expiration_timestamp: i64, // TODO: add auto-refresh until Spark Authn server adds refresh token support
}

/// Pool of Spark operators for threshold signing
///
/// Manages a collection of Spark operators that participate in threshold signing
/// operations. The pool includes configuration for retry behavior and identifies
/// the coordinator operator that orchestrates multi-party signing.
#[derive(Clone)]
pub(crate) struct SparkOperatorPool {
    /// List of available Spark operators
    pub(crate) operators: Vec<Arc<SparkOperator>>,

    /// Index of the operator that acts as coordinator
    pub(crate) coordinator_index: u32,

    /// Number of times to retry failed operations
    retry_count: u32,

    /// Milliseconds to wait between retry attempts
    backoff_ms: u64,
}

impl SparkOperatorPool {}

/// Configuration for an individual Spark operator node
///
/// Represents a single Spark operator that participates in threshold signing
/// operations. Each operator has a unique identifier, network address, and
/// identity public key used for authentication and FROST threshold signatures.
#[derive(Debug, Clone)]
pub(crate) struct SparkOperator {
    /// The index of the signing operator.
    pub(crate) id: u32,

    /// Identifier is the FROST identifier of the signing operator, which will be index + 1 in 32 bytes big endian hex string.
    /// Used as shamir secret share identifier in DKG key shares.
    pub(crate) frost_identifier: Identifier,

    /// Address of the signing operator
    pub(crate) address: Uri,

    /// Public key of the signing operator
    pub(crate) identity_public_key: PublicKey,
}

impl SparkOperator {
    // Parse an operator specified as a string (in the format <pubkey>@<address>) into a
    // [`SparkOperator`].
    pub(crate) fn parse(id: u32, op: &str) -> Result<Self, SparkSdkError> {
        let parts: Vec<&str> = op.split('@').collect();
        if parts.len() != 2 {
            return Err(SparkSdkError::from(ValidationError::InvalidArgument {
                argument: format!(
                    "Invalid operator string format {} (should be <pubkey>@<address>)",
                    op
                ),
            }));
        }

        let identity_public_key = PublicKey::from_str(parts[0]).map_err(|err| {
            SparkSdkError::from(ValidationError::InvalidArgument {
                argument: format!("Invalid public key for operator {}: {}", parts[0], err),
            })
        })?;

        let address = Uri::from_str(parts[1])
            .map_err(|err| SparkSdkError::from(NetworkError::InvalidUri(err)))?;

        let frost_identifier = Identifier::try_from(id as u16 + 1).map_err(|err| {
            SparkSdkError::from(ValidationError::InvalidArgument {
                argument: format!(
                    "Invalid frost identifier for operator {}: {}",
                    parts[0], err
                ),
            })
        })?;

        Ok(SparkOperator {
            id,
            frost_identifier,
            address,
            identity_public_key,
        })
    }

    pub(crate) fn frost_identifier_str(&self) -> String {
        identifier_to_hex_string(&self.frost_identifier)
    }
}

impl SparkConfig {
    /// Try and get operators from environment variables. Each operator should be in a separate
    /// environment variable in the format `SPARK_OPERATOR_<index>`.
    ///
    /// Returns `None` if operators are not specified as environmental variables. Otherwise,
    /// returns Ok with the parsed operators, or an error if parsing operators fails.
    fn operators_from_env() -> Option<Result<Vec<SparkOperator>, SparkSdkError>> {
        // If the first operator is not set, assume the rest aren't set either.
        if std::env::var("SPARK_OPERATOR_0").is_err() {
            return None;
        }

        let mut operators = vec![];
        for i in 0..10 {
            let variable = format!("SPARK_OPERATOR_{}", i);
            let operator = std::env::var(&variable);
            let operator = match operator {
                Ok(operator) => operator,
                Err(_) => break,
            };

            match SparkOperator::parse(i as u32, &operator) {
                Ok(operator) => operators.push(operator),
                Err(err) => {
                    return Some(Err(SparkSdkError::from(ValidationError::InvalidArgument {
                        argument: format!("Unable to parse operator {}: {}", variable, err),
                    })))
                }
            }
        }

        Some(Ok(operators))
    }

    /// Try to get the operators from the environment first. Otherwise, parse the defaults from
    /// [`SPARK_DEV_OPERATORS`].
    fn operators() -> Result<Vec<SparkOperator>, SparkSdkError> {
        if let Some(operators) = Self::operators_from_env() {
            return operators;
        }

        let mut spark_operators = vec![];
        for (i, operator) in SPARK_DEV_OPERATORS.iter().enumerate() {
            // TODO: Add environment
            let address = Uri::from_str(operator.0)
                .map_err(|err| SparkSdkError::from(NetworkError::InvalidUri(err)))?;

            let identity_public_key = PublicKey::from_str(operator.1).map_err(|err| {
                SparkSdkError::from(ValidationError::InvalidArgument {
                    argument: format!("Invalid public key for operator {}: {}", operator.1, err),
                })
            })?;

            let frost_identifier = Identifier::try_from(i as u16 + 1).map_err(|err| {
                SparkSdkError::from(ValidationError::InvalidArgument {
                    argument: format!(
                        "Invalid frost identifier for operator {}: {}",
                        operator.1, err
                    ),
                })
            })?;

            spark_operators.push(SparkOperator {
                id: i as u32,
                frost_identifier,
                address,
                identity_public_key,
            });
        }

        Ok(spark_operators)
    }

    /// Creates a new wallet configuration with the specified network.
    ///
    /// This function:
    /// 1. Creates a vector of 5 Spark operators with their configurations
    /// 2. Sets up gRPC connections to each operator
    /// 3. Initializes the wallet configuration with default values
    ///
    /// # Arguments
    ///
    /// * `network` - The Spark network to use (mainnet, testnet, regtest)
    ///
    /// # Returns
    ///
    /// Returns a Result containing either:
    /// * The initialized SparkConfig struct on success
    /// * A [`SparkSdkError`] on failure (e.g. connection errors, invalid URLs)
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// * Failed to establish gRPC connections to operators
    /// * Failed to parse operator URLs
    /// * Failed to decode operator public keys
    pub async fn new(network: SparkNetwork) -> Result<Self, SparkSdkError> {
        // set the threshold
        let threshold = SPARK_REGTEST_SIGNING_THRESHOLD;

        // set the default coordinator index
        let coordinator_index = DEFAULT_COORDINATOR_INDEX;

        // set the operators
        let spark_operators = Self::operators()?;

        // set the connection pool for operators
        let mut spark_clients = HashMap::new();
        for operator in &spark_operators {
            // establish the secure connection
            // since this uses rustls, self-signed certificates will fail
            let spark_rpc_client =
                SparkConnection::establish_connection(operator.address.clone()).await?;
            spark_clients.insert(operator.id, spark_rpc_client);
        }

        let ssp_identity_public_key = match network {
            SparkNetwork::Regtest => PublicKey::from_str(
                std::env::var("LIGHTSPARK_SSP_REGTEST_IDENTITY_PUBLIC_KEY")
                    .as_deref()
                    .unwrap_or(LIGHTSPARK_SSP_REGTEST_IDENTITY_PUBLIC_KEY),
            )
            .map_err(|err| SparkSdkError::from(CryptoError::Secp256k1(err)))?,
            SparkNetwork::Mainnet => PublicKey::from_str(
                std::env::var("LIGHTSPARK_SSP_MAINNET_IDENTITY_PUBLIC_KEY")
                    .as_deref()
                    .unwrap_or(LIGHTSPARK_SSP_MAINNET_IDENTITY_PUBLIC_KEY),
            )
            .map_err(|err| SparkSdkError::from(CryptoError::Secp256k1(err)))?,
        };

        let ssp_graphql_client = GraphqlClient::new().map(Arc::new)?;

        #[cfg(feature = "telemetry")]
        tracing::info!(
            ssp_identity_public_key = ssp_identity_public_key.to_string(),
            "ssp public key"
        );

        let operator_pool = SparkOperatorPool {
            operators: spark_operators.into_iter().map(Arc::new).collect(),
            coordinator_index,
            retry_count: DEFAULT_RETRY_COUNT,
            backoff_ms: 200,
        };

        let wallet_config = Self {
            network,
            operator_pool: Arc::new(operator_pool),
            threshold,
            spark_clients: Arc::new(RwLock::new(spark_clients)),
            ssp_graphql_client,
            ssp_identity_public_key,
            sessions: Arc::new(parking_lot::RwLock::new(vec![])),
        };

        Ok(wallet_config)
    }

    /// Gets a connection to a Spark operator service.
    ///
    /// This function manages connections to Spark operators, creating new connections if needed
    /// and reusing existing ones. It handles both connecting to a specific operator or defaulting
    /// to the coordinator.
    ///
    /// # Arguments
    ///
    /// * `operator_id` - Optional ID of the specific operator to connect to. If None, connects to
    ///                   the default coordinator.
    ///
    /// # Returns
    ///
    /// * [`Result<SparkServiceClient<Channel>, SparkSdkError>`] - A client for the Spark service on success,
    ///    or an error if the connection fails
    ///
    /// # Errors
    ///
    /// Returns [`SparkSdkError::InvalidArgument`] if:
    /// * The operator_id is out of bounds for the available operators
    ///
    /// May also return errors from:
    /// * Channel creation/connection
    /// * URI parsing
    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) async fn get_spark_connection(
        &self,
        operator_id: Option<u32>,
    ) -> Result<SparkServiceClient<Channel>, SparkSdkError> {
        let operator_id = operator_id.unwrap_or(self.operator_pool.coordinator_index);

        // if doesn't exist, create the connection for the operator
        if !self.spark_clients.read().contains_key(&operator_id) {
            let spark_operators = self.operator_pool.operators.clone();
            if operator_id >= spark_operators.len() as u32 {
                drop(spark_operators);
                return Err(SparkSdkError::from(ValidationError::InvalidArgument {
                    argument: format!("Operator index {} is out of bounds", operator_id),
                }));
            }

            let uri = spark_operators[operator_id as usize].address.clone();
            let spark_rpc_instance = SparkConnection::establish_connection(uri).await?;
            self.spark_clients
                .write()
                .insert(operator_id, spark_rpc_instance);
        }

        // get the connection
        let client = self.spark_clients.read().get(&operator_id).unwrap().clone();
        let spark_client = client.get_new_spark_service_connection()?;

        Ok(spark_client)
    }

    /// Helper method to execute a gRPC call with automatic retries
    ///
    /// This provides a simple wrapper around the more complex retry logic
    ///
    /// # Example
    /// ```
    /// let response = config.execute_with_retry(
    ///     || Request::new(GenerateDepositAddressRequest { ... }),
    ///     |client, req| client.generate_deposit_address(req)
    /// ).await?;
    /// ```
    pub(crate) async fn execute_with_retry<T, R>(
        &self,
        create_request: impl Fn() -> Request<T>,
        rpc_call: impl Fn(
            SparkServiceClient<Channel>,
            Request<T>,
        ) -> std::pin::Pin<
            Box<dyn std::future::Future<Output = Result<tonic::Response<R>, tonic::Status>> + Send>,
        >,
        operator_id: Option<u32>,
    ) -> Result<R, SparkSdkError> {
        let operator_id = operator_id.unwrap_or(self.operator_pool.coordinator_index);
        let max_retries = self.operator_pool.retry_count;
        let backoff = Duration::from_millis(self.operator_pool.backoff_ms);
        let mut attempt = 0;

        loop {
            // Get a fresh client
            let client = self.get_spark_connection(Some(operator_id)).await?;

            // Create a fresh request for this attempt
            let mut request = create_request();

            // Add auth header
            request.metadata_mut().insert(
                "authorization",
                self.sessions.read().clone()[operator_id as usize]
                    .session_token
                    .clone()
                    .to_string()
                    .parse()
                    .unwrap(),
            );

            match rpc_call(client, request).await {
                Ok(response) => return Ok(response.into_inner()),
                Err(err) => {
                    if attempt >= max_retries - 1 {
                        return Err(SparkSdkError::from(NetworkError::Status(err)));
                    }

                    sleep(backoff).await;
                    attempt += 1;

                    #[cfg(feature = "telemetry")]
                    tracing::debug!(
                        "Retrying RPC call to operator {} (attempt {}/{})",
                        operator_id,
                        attempt + 1,
                        max_retries
                    );
                }
            }
        }
    }

    /// A super-simple way to use retries with cloneable request types.
    ///
    /// This further simplifies retry usage when the request inner type (T) implements Clone.
    ///
    /// # Example
    /// ```
    /// let request = GenerateDepositAddressRequest { ... };
    /// let response = config.call_with_retry(
    ///     request,
    ///     |client, req| Box::pin(async move { client.generate_deposit_address(req).await })
    /// ).await?;
    /// ```
    pub(crate) async fn call_with_retry<T: Clone, R>(
        &self,
        request_data: T,
        rpc_call: impl Fn(
            SparkServiceClient<Channel>,
            Request<T>,
        ) -> std::pin::Pin<
            Box<dyn std::future::Future<Output = Result<tonic::Response<R>, tonic::Status>> + Send>,
        >,
        operator_id: Option<u32>,
    ) -> Result<R, SparkSdkError> {
        let request_factory = move || Request::new(request_data.clone());
        self.execute_with_retry(request_factory, rpc_call, operator_id)
            .await
    }
}

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

    #[test]
    fn test_parse_spark_operator() {
        let operator = SparkOperator::parse(
            0,
            "0322ca18fc489ae25418a0e768273c2c61cabb823edfb14feb891e9bec62016510@http://localhost:8535",
        )
        .unwrap();

        assert_eq!(operator.id, 0);
        assert_eq!(
            operator.frost_identifier_str(),
            "0000000000000000000000000000000000000000000000000000000000000001"
        );
        assert_eq!(operator.address.to_string(), "http://localhost:8535/");
        assert_eq!(
            operator.identity_public_key.to_string(),
            "0322ca18fc489ae25418a0e768273c2c61cabb823edfb14feb891e9bec62016510"
        );
    }
}