spark-cryptography 0.1.11

//! Spark Cryptography Module
//!
//! This module provides cryptographic functionality for the Spark wallet system, specifically
//! handling key derivation and management.
//!
//! # Overview
//!
//! The Spark wallet uses a hierarchical deterministic (HD) wallet structure based on BIP32,
//! with a custom purpose number (8797555) for Spark-specific keys. The key hierarchy is:
//!
//! ```
//! // m/8797555'/account'/key_type'[/leaf_index']
//! ```
//!
//! Where:
//! - `account`: The account index (hardened)
//! - `key_type`: The type of key (0' for identity, 1' for base signing, 2' for deposit)
//! - `leaf_index`: Optional leaf-specific index (hardened)
//!
//! # Key Types
//!
//! The wallet supports three types of keys:
//!
//! 1. **Identity Key** (0')
//!    - Used for wallet authentication and identification
//!    - Used for signature verification with Spark Operators
//!
//! 2. **Base Signing Key** (1')
//!    - Foundation for leaf-specific keys
//!    - Used for general signing operations
//!
//! 3. **Deposit Key** (2')
//!    - Used for deposit transactions. All deposit transactions are signed with this key.
//!      However, each deposit address is different. Deposit signing key is static, and it is
//!      not to be confused with one-time deposit address.
//!    - After deposit, leaves are transferred to BaseSigning key
//!
//! # Usage
//!
//! ```rust
//! use spark_cryptography::derivation_path::{derive_spark_key, SparkKeyType};
//! use bitcoin::Network;
//!
//! // Generate a seed (in practice, use a secure random number generator)
//! let seed_bytes = b"0000000000000000000000000000000000000000000000000000000000000000";
//!
//! // Derive an identity key
//! let identity_key = derive_spark_key(
//!     None,                    // No leaf ID for identity key
//!     0,                       // Account index
//!     seed_bytes,             // Seed
//!     SparkKeyType::Identity, // Key type
//!     Network::Bitcoin,       // Network
//! ).unwrap();
//!
//! // Derive a leaf-specific key
//! let leaf_key = derive_spark_key(
//!     Some("leaf-uuid".to_string()),
//!     0,
//!     seed_bytes,
//!     SparkKeyType::BaseSigning,
//!     Network::Bitcoin,
//! ).unwrap();
//! ```
//!
//! # Security Considerations
//!
//! - All derivation paths use hardened derivation (with ' suffix)
//! - Seeds must be at least 16 bytes long
//! - Leaf indices are deterministically derived from leaf IDs
//! - Different key types ensure separation of concerns
//!
//! # Error Handling
//!
//! The module provides detailed error types for various failure cases:
//! - Invalid seeds
//! - Invalid derivation paths
//! - Invalid leaf indices
//! - Invalid key types
//! - Network-specific errors

use bitcoin::{
    bip32::{ChildNumber, Xpriv},
    hashes::{sha256, Hash, HashEngine},
    key::Secp256k1,
    secp256k1::{All, PublicKey, SecretKey},
    Network,
};

use crate::error::{InvalidLeafError, SparkCryptographyError};

pub const SPARK_DERIVATION_PATH_PURPOSE: u32 = 8797555;

/// The derivation path for a Spark key.
///
/// This is a wrapper around a vector of `ChildNumber` that represents the derivation path for a Spark key.
/// It is used to derive a secret key from a seed and a derivation path.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SparkDerivationPath(pub(crate) Vec<ChildNumber>);

impl std::ops::Deref for SparkDerivationPath {
    type Target = Vec<ChildNumber>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Key types for Spark wallet.
///
/// Each key type corresponds to a specific path in the hierarchical deterministic
/// wallet structure, with the path format: m/8797555'/account'/key_type'
///
/// The key types serve different purposes in the Spark wallet:
/// - Identity key: Used for wallet authentication and identification
/// - Base signing key: Used as the foundation for leaf-specific keys
/// - Temporary signing key: Used for one-time operations like deposits
///
/// These key types are encoded as the third component in the derivation path.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SparkKeyType {
    /// Identity key with path m/8797555'/account'/0'
    ///
    /// This key is used for wallet authentication, identification,
    /// and signature verification with Spark Operators.
    Identity,

    /// Base signing key with path m/8797555'/account'/1'
    ///
    /// This key serves as the foundation for leaf-specific keys,
    /// which are derived by adding a leaf-specific index.
    BaseSigning,

    /// Deposit key with path m/8797555'/account'/2'
    ///
    /// This key is used to sign deposits transactions. After the deposit is signed,
    /// the user should transfer the leaves to their BaseSigning key under the same
    /// leaf UUID. This happens automatically in `sparks-wallet`.
    Deposit,
}

/// Derives the secret key for a given account index, network, and key type.
///
/// This function constructs a complete derivation path for the key,
/// using the purpose index (8797555'), the account index, and the key type index.
/// It then derives the key using the provided seed and secp256k1 context.
///
/// # Arguments
///
/// - `seed`: The seed bytes to derive the key from. In &[`u8`] format.
/// - `account_index`: The account index in [`u32`] format.
/// - `network`: The [`Network`] to derive the key for.
/// - `secp`: The [`Secp256k1`] context to use for the derivation.
/// - `key_type`: The [`SparkKeyType`] to derive the key for. For example, to get the identity key, use [`SparkKeyType::Identity`].
///
/// # Returns
///
/// A [`SecretKey`] representing the derived key, or an error if the derivation fails.
///
/// # Errors
///
/// - [`SparkCryptographyError::InvalidSeed`]: If the seed is invalid.
/// - [`SparkCryptographyError::InvalidDerivationPath`]: If the derivation path is invalid.
pub fn get_secret_key_with_key_type(
    seed: &[u8],
    account_index: u32,
    network: Network,
    key_type: SparkKeyType,
    secp: &Secp256k1<All>,
) -> Result<SecretKey, SparkCryptographyError> {
    let seed = Xpriv::new_master(network, seed).map_err(|_| SparkCryptographyError::InvalidSeed)?;

    let identity_derivation_path = get_derivation_path_with_key_type(account_index, key_type)?;

    let identity_key = seed
        .derive_priv(secp, &*identity_derivation_path)
        .map_err(|_| {
            SparkCryptographyError::InvalidDerivationPath(format!("{:?}", identity_derivation_path))
        })?;

    Ok(identity_key.private_key)
}

/// Derives the public key for a given account index, network, and key type.
///
/// This function constructs a complete derivation path for the key,
/// using the purpose index (8797555'), the account index, and the key type index.
/// It then derives the key using the provided seed and secp256k1 context.
///
/// # Arguments
///
/// - `seed`: The seed bytes to derive the key from. In &[`u8`] format.
/// - `account_index`: The account index in [`u32`] format.
/// - `network`: The [`Network`] to derive the key for.
/// - `secp`: The [`Secp256k1`] context to use for the derivation.
/// - `key_type`: The [`SparkKeyType`] to derive the key for. For example, to get the identity key, use [`SparkKeyType::Identity`].
///
/// # Returns
///
/// A [`PublicKey`] representing the derived key, or an error if the derivation fails.
///
/// # Errors
///
/// - [`SparkCryptographyError::InvalidSeed`]: If the seed is invalid.
/// - [`SparkCryptographyError::InvalidDerivationPath`]: If the derivation path is invalid.
///
/// # Example
///
/// ```rust
/// use spark_cryptography::derivation_path::{get_public_key_with_key_type, SparkKeyType};
/// use bitcoin::Network;
///
/// let secp = bitcoin::secp256k1::Secp256k1::new();
/// let seed_bytes = "0x0000000000000000000000000000000000000000000000000000000000000000".as_bytes();
/// let public_key = get_public_key_with_key_type(seed_bytes, 0, Network::Bitcoin, SparkKeyType::Identity, &secp).unwrap();
/// ```
pub fn get_public_key_with_key_type(
    seed: &[u8],
    account_index: u32,
    network: Network,
    key_type: SparkKeyType,
    secp: &Secp256k1<All>,
) -> Result<PublicKey, SparkCryptographyError> {
    let identity_secret_key =
        get_secret_key_with_key_type(seed, account_index, network, key_type, secp)?;

    Ok(identity_secret_key.public_key(secp))
}

/// Calculates the derivation path component for a leaf key based on a leaf ID.
///
/// This function hashes the leaf ID using SHA-256 and processes the hash to derive
/// a hardened index for use in leaf key derivation paths.
///
/// The index is calculated by:
/// 1. Computing SHA-256 hash of the leaf ID
/// 2. Interpreting the first 4 bytes of the hash as u32 value
/// 3. Modulo the u32 value with 0x80000000 to ensure it fits within valid index range
/// 4. Converting to a hardened ChildNumber
///
/// # Arguments
///
/// - `leaf_id`: The leaf ID in &[`str`] format to derive the index for.
///
/// # Returns
///
/// A [`ChildNumber`] representing the derived index, or an error if the derivation fails.
pub fn get_leaf_index(leaf_id: &str) -> Result<ChildNumber, SparkCryptographyError> {
    // Compute SHA-256 hash of the leaf ID
    let mut engine = sha256::Hash::engine();
    engine.input(leaf_id.as_bytes());
    let hash = sha256::Hash::from_engine(engine);

    let chunk = &hash[0..4];
    let amount = u32::from_be_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]) % 0x80000000;

    // Return the hardened path component
    get_child_number(amount, true)
}

/// Creates a ChildNumber from an index value and hardened flag.
///
/// This function handles validation of the provided index and creation of the
/// appropriate ChildNumber type (normal or hardened).
///
/// # Parameters
///
/// - `index`: The index value (must be less than 0x80000000)
/// - `hardened`: Whether the index should be hardened
///
/// # Returns
///
/// A ChildNumber representing the index, or an error if the index is invalid.
pub fn get_child_number(index: u32, hardened: bool) -> Result<ChildNumber, SparkCryptographyError> {
    if index > 0x7FFFFFFF {
        return Err(SparkCryptographyError::InvalidLeaf(
            InvalidLeafError::CannotDeriveChildNumberFromIndex(index),
        ));
    }

    let child_number = if hardened {
        ChildNumber::from_hardened_idx(index).unwrap()
    } else {
        ChildNumber::from_normal_idx(index).unwrap()
    };

    Ok(child_number)
}

/// Derives a Spark key from a seed and derivation path.
///
/// This function constructs a complete derivation path for a Spark key, including
/// the purpose index, account index, key type index, and optional leaf index.
/// It then uses the seed to derive the key using the secp256k1 library.
///
/// # Arguments
///
/// - `leaf_id`: The leaf ID to derive the key for.
/// - `account`: The account index to derive the key for.
/// - `seed_bytes`: The seed bytes to derive the key from.
/// - `key_type`: The key type to derive the key for.
/// - `network`: The network to derive the key for.
///
/// # Returns
///
/// A SecretKey representing the derived key, or an error if the derivation fails.
///
/// # Errors
///
/// - `InvalidSeed`: If the seed is invalid.
/// - `InvalidDerivationPath`: If the derivation path is invalid.
/// - `InvalidLeaf`: If the leaf index is invalid.
/// - `InvalidKeyType`: If the key type is invalid.
/// - `InvalidNetwork`: If the network is invalid.
///
/// # Example
///
/// ```rust
/// use spark_cryptography::derivation_path::{derive_spark_key, SparkKeyType};
/// use bitcoin::Network;
///
/// let seed_bytes = "0x0000000000000000000000000000000000000000000000000000000000000000".as_bytes();
/// let key_type = SparkKeyType::Identity;
/// let network = Network::Bitcoin;
///
/// let key = derive_spark_key(None, 0, seed_bytes, key_type, network).unwrap();
/// let key_bytes = key.secret_bytes();
/// ```
pub fn derive_spark_key(
    leaf_id: Option<String>,
    account: u32,
    seed_bytes: &[u8],
    key_type: SparkKeyType,
    network: Network,
) -> Result<SecretKey, SparkCryptographyError> {
    let path = prepare_derivation_path(account, key_type, leaf_id)?;

    // Create the master key from seed
    let seed = match Xpriv::new_master(network, seed_bytes) {
        Ok(seed) => seed,
        Err(_) => return Err(SparkCryptographyError::InvalidSeed),
    };

    // Derive the key for the given path
    let secp = bitcoin::secp256k1::Secp256k1::new();
    let extended_key = seed
        .derive_priv(&secp, &path)
        .map_err(|_| SparkCryptographyError::InvalidDerivationPath(format!("{:?}", path)))?;

    Ok(extended_key.private_key)
}

/// Converts a SparkKeyType to its corresponding child index number.
///
/// Returns a hardened ChildNumber representing the key type index.
fn get_key_type_index(key_type: SparkKeyType) -> Result<ChildNumber, SparkCryptographyError> {
    let index = match key_type {
        SparkKeyType::Identity => 0,
        SparkKeyType::BaseSigning => 1,
        SparkKeyType::Deposit => 2,
    };

    // Key type index should be hardened
    get_child_number(index, true)
}

fn prepare_derivation_path(
    account_index: u32,
    key_type: SparkKeyType,
    leaf_index: Option<String>,
) -> Result<Vec<ChildNumber>, SparkCryptographyError> {
    let purpose_index = get_child_number(SPARK_DERIVATION_PATH_PURPOSE, true)?;

    // Returns the hardened account index
    let account_index = get_child_number(account_index, true)?;

    // Returns the hardened key type index
    let key_type_index = get_key_type_index(key_type)?;

    // If leaf_id is provided and non-empty, calculate the leaf index
    let leaf_index = if let Some(leaf_id) = leaf_index {
        Some(get_leaf_index(leaf_id.as_str())?)
    } else {
        None
    };

    Ok(prepare_path(
        purpose_index,
        account_index,
        key_type_index,
        leaf_index,
    ))
}

/// Prepares a complete derivation path for a Spark key.
///
/// Constructs a derivation path with the following components:
/// - Purpose index (8797555')
/// - Account index
/// - Key type index (0' for identity, 1' for base signing, 2' for temporary signing)
/// - Optional leaf index (for leaf keys)
///
/// The purpose index is the first element of the path, followed by account index,
/// key type, and optionally a leaf index.
fn prepare_path(
    purpose_index: ChildNumber,
    account_index: ChildNumber,
    key_type_index: ChildNumber,
    // If leaf index is NOT provided, it means this is for the identity key.
    leaf_index: Option<ChildNumber>,
) -> Vec<ChildNumber> {
    let mut path = vec![purpose_index, account_index, key_type_index];

    if let Some(leaf_index) = leaf_index {
        path.push(leaf_index);
    }

    path
}

/// Derives the identity derivation path for a given account index.
///
/// This function constructs a complete derivation path for the identity key,
/// using the purpose index (8797555'), the account index, and the key type index (0').
///
/// # Arguments
///
/// - `account_index`: The account index in [`u32`] format.
///
/// # Returns
///
/// A [`SparkDerivationPath`] representing the derivation path for the identity key, or an error if the derivation fails.
///
/// # Errors
///
/// - [`SparkCryptographyError::InvalidLeaf`]: If the leaf index is invalid.
///
/// # Example
///
/// ```rust
/// use spark_cryptography::derivation_path::{get_derivation_path_with_key_type, SparkKeyType};
/// use bitcoin::Network;
///
/// let path = get_derivation_path_with_key_type(0, SparkKeyType::Identity).unwrap();
/// ```
pub fn get_derivation_path_with_key_type(
    account_index: u32,
    key_type: SparkKeyType,
) -> Result<SparkDerivationPath, SparkCryptographyError> {
    let path = prepare_derivation_path(account_index, key_type, None)?;
    Ok(SparkDerivationPath(path))
}

#[cfg(test)]
mod derivation_path_tests {
    use super::*;
    use bitcoin::Network;

    #[test]
    fn test_get_leaf_index() {
        let leaf_id_1 = "019534f0-f4e2-7845-87fe-c6ea2fa69f80";
        let leaf_id_2 = "019534f0-f4e2-7868-b3fa-d06dc10b79e7";
        let leaf_id_3 = "dbb5c090-dca4-47ec-9f20-41edd4594dcf";

        let child_number_1 = get_leaf_index(leaf_id_1).unwrap();
        let child_number_2 = get_leaf_index(leaf_id_2).unwrap();
        let child_number_3 = get_leaf_index(leaf_id_3).unwrap();

        assert_eq!(
            child_number_1,
            ChildNumber::from_hardened_idx(1137822116).unwrap()
        );
        assert_eq!(
            child_number_2,
            ChildNumber::from_hardened_idx(1199130649).unwrap()
        );
        assert_eq!(
            child_number_3,
            ChildNumber::from_hardened_idx(1743780874).unwrap()
        );
    }

    #[test]
    fn test_get_identity_derivation_path() {
        let path = get_derivation_path_with_key_type(0, SparkKeyType::Identity).unwrap();
        assert_eq!(path.0.len(), 3);

        // Test with different account indices
        let path_account_1 = get_derivation_path_with_key_type(1, SparkKeyType::Identity).unwrap();
        let path_account_2 = get_derivation_path_with_key_type(2, SparkKeyType::Identity).unwrap();

        assert_eq!(path_account_1.0.len(), 3);
        assert_eq!(path_account_2.0.len(), 3);

        // Verify the components
        assert_eq!(
            path.0[0],
            ChildNumber::from_hardened_idx(SPARK_DERIVATION_PATH_PURPOSE).unwrap()
        );
        assert_eq!(path.0[1], ChildNumber::from_hardened_idx(0).unwrap());
        assert_eq!(path.0[2], ChildNumber::from_hardened_idx(0).unwrap()); // Identity key type
    }

    #[test]
    fn test_derive_spark_key() {
        let seed_bytes = b"0000000000000000000000000000000000000000000000000000000000000000";
        let network = Network::Bitcoin;
        let secp = Secp256k1::new();

        // Test identity key derivation
        let identity_key =
            derive_spark_key(None, 0, seed_bytes, SparkKeyType::Identity, network).unwrap();
        let identity_pubkey = identity_key.public_key(&secp);

        // Test base signing key derivation
        let base_signing_key =
            derive_spark_key(None, 0, seed_bytes, SparkKeyType::BaseSigning, network).unwrap();
        let base_signing_pubkey = base_signing_key.public_key(&secp);

        // Test deposit key derivation
        let deposit_key =
            derive_spark_key(None, 0, seed_bytes, SparkKeyType::Deposit, network).unwrap();
        let deposit_pubkey = deposit_key.public_key(&secp);

        // Verify all keys are different
        assert_ne!(identity_pubkey, base_signing_pubkey);
        assert_ne!(identity_pubkey, deposit_pubkey);
        assert_ne!(base_signing_pubkey, deposit_pubkey);
    }

    #[test]
    fn test_derive_spark_key_with_leaf() {
        let seed_bytes = b"0000000000000000000000000000000000000000000000000000000000000000";
        let network = Network::Bitcoin;
        let leaf_id = "test-leaf-id";

        // Test base signing key with leaf
        let key_with_leaf = derive_spark_key(
            Some(leaf_id.to_string()),
            0,
            seed_bytes,
            SparkKeyType::BaseSigning,
            network,
        )
        .unwrap();

        // Test base signing key without leaf
        let key_without_leaf =
            derive_spark_key(None, 0, seed_bytes, SparkKeyType::BaseSigning, network).unwrap();

        // Verify keys are different
        assert_ne!(key_with_leaf, key_without_leaf);
    }

    #[test]
    fn test_get_child_number() {
        assert_eq!(2147483648, 0x7FFFFFFF_u32 + 1);

        // Test valid hardened indices
        assert!(get_child_number(0, true).is_ok());
        assert!(get_child_number(0x7FFFFFFF, true).is_ok());

        // Test valid normal indices
        assert!(get_child_number(0, false).is_ok());
        assert!(get_child_number(0x7FFFFFFF, false).is_ok());

        assert_eq!(
            get_child_number(0x80000000, false).unwrap_err(),
            SparkCryptographyError::InvalidLeaf(
                InvalidLeafError::CannotDeriveChildNumberFromIndex(0x80000000)
            )
        );
    }

    #[test]
    fn test_get_key_type_index() {
        // Test all key types
        let identity_idx = get_key_type_index(SparkKeyType::Identity).unwrap();
        let base_signing_idx = get_key_type_index(SparkKeyType::BaseSigning).unwrap();
        let deposit_idx = get_key_type_index(SparkKeyType::Deposit).unwrap();

        assert_eq!(identity_idx, ChildNumber::from_hardened_idx(0).unwrap());
        assert_eq!(base_signing_idx, ChildNumber::from_hardened_idx(1).unwrap());
        assert_eq!(deposit_idx, ChildNumber::from_hardened_idx(2).unwrap());
    }

    #[test]
    fn test_network_specific() {
        let seed_bytes = b"0000000000000000000000000000000000000000000000000000000000000000";
        let secp = Secp256k1::new();

        // Test with different networks
        let bitcoin_key = derive_spark_key(
            None,
            0,
            seed_bytes,
            SparkKeyType::Identity,
            Network::Bitcoin,
        )
        .unwrap();
        let testnet_key = derive_spark_key(
            None,
            0,
            seed_bytes,
            SparkKeyType::Identity,
            Network::Testnet,
        )
        .unwrap();
        let regtest_key = derive_spark_key(
            None,
            0,
            seed_bytes,
            SparkKeyType::Identity,
            Network::Regtest,
        )
        .unwrap();
        let signet_key =
            derive_spark_key(None, 0, seed_bytes, SparkKeyType::Identity, Network::Signet).unwrap();

        // Verify keys are the same for different networks. This will change in the future as the derivation path becomes different for different networks.
        assert_eq!(bitcoin_key.public_key(&secp), testnet_key.public_key(&secp));
        assert_eq!(bitcoin_key.public_key(&secp), regtest_key.public_key(&secp));
        assert_eq!(bitcoin_key.public_key(&secp), signet_key.public_key(&secp));
    }

    #[test]
    fn test_get_public_and_secret_keys_with_key_type() {
        // Setup test environment with a zero seed and Bitcoin network
        let seed_bytes = b"0000000000000000000000000000000000000000000000000000000000000000";
        let network = Network::Bitcoin;
        let secp = Secp256k1::new();

        // Test identity key derivation using derive_spark_key
        let identity_key =
            derive_spark_key(None, 0, seed_bytes, SparkKeyType::Identity, network).unwrap();
        let identity_pubkey = identity_key.public_key(&secp);

        // Test base signing key derivation using derive_spark_key
        let base_signing_key =
            derive_spark_key(None, 0, seed_bytes, SparkKeyType::BaseSigning, network).unwrap();
        let base_signing_pubkey = base_signing_key.public_key(&secp);

        // Verify that get_secret_key_with_key_type produces the same identity key
        let identity_secret_key =
            get_secret_key_with_key_type(seed_bytes, 0, network, SparkKeyType::Identity, &secp)
                .unwrap();
        assert_eq!(identity_secret_key, identity_key);

        // Verify that get_public_key_with_key_type produces the same identity public key
        let identity_public_key =
            get_public_key_with_key_type(seed_bytes, 0, network, SparkKeyType::Identity, &secp)
                .unwrap();
        assert_eq!(identity_public_key, identity_pubkey);

        // Verify that get_secret_key_with_key_type produces the same base signing key
        let base_signing_secret_key =
            get_secret_key_with_key_type(seed_bytes, 0, network, SparkKeyType::BaseSigning, &secp)
                .unwrap();
        assert_eq!(base_signing_secret_key, base_signing_key);

        // Verify that get_public_key_with_key_type produces the same base signing public key
        let base_signing_public_key =
            get_public_key_with_key_type(seed_bytes, 0, network, SparkKeyType::BaseSigning, &secp)
                .unwrap();
        assert_eq!(base_signing_public_key, base_signing_pubkey);
    }
}