onemoney-protocol 0.18.0

//! Multi-signature account utilities.
//!
//! This module provides utilities for working with multi-signature accounts:
//! - Derive multi-sig account addresses from signer configurations
//! - Aggregate multiple signatures for multi-sig transactions
//! - Build and validate multi-sig transactions

use alloy_primitives::{Address, keccak256};
use om_primitives_types::transaction::{B264, MultiSigSignatureEntry, Signature};
use thiserror::Error;

/// Domain separation tag (DST) for deriving multi-sig account addresses.
pub const DST_MULTISIG_ADDR_V1: &[u8] = b"MULTISIG_V1";

/// Signer configuration for multi-signature accounts.
///
/// Each signer has a public key (33-byte SEC1 compressed format),
/// a weight, and optional metadata.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SignerConfig {
    /// Signer's public key in SEC1 compressed format (33 bytes)
    pub public_key: Vec<u8>,
    /// Weight of this signer's vote
    pub weight: u8,
}

impl SignerConfig {
    /// Create a new signer configuration.
    ///
    /// # Arguments
    /// * `public_key` - Public key in SEC1 compressed format (must be 33 bytes)
    /// * `weight` - Weight of this signer (must be > 0)
    ///
    /// # Errors
    /// Returns error if public key length is not 33 bytes or weight is 0.
    pub fn new(public_key: Vec<u8>, weight: u8) -> Result<Self, MultiSigError> {
        if public_key.len() != 33 {
            return Err(MultiSigError::InvalidPublicKeyLength {
                expected: 33,
                actual: public_key.len(),
            });
        }
        if weight == 0 {
            return Err(MultiSigError::InvalidWeight);
        }
        Ok(Self { public_key, weight })
    }
}

/// Threshold configuration for multi-signature accounts.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ThresholdConfig {
    /// Minimum total weight required for transaction approval
    pub threshold: u16,
}

impl ThresholdConfig {
    /// Create a new threshold configuration.
    ///
    /// # Arguments
    /// * `threshold` - Minimum weight required (must be > 0)
    ///
    /// # Errors
    /// Returns error if threshold is 0.
    pub fn new(threshold: u16) -> Result<Self, MultiSigError> {
        if threshold == 0 {
            return Err(MultiSigError::InvalidThreshold);
        }
        Ok(Self { threshold })
    }
}

/// Errors that can occur during multi-signature operations.
#[derive(Debug, Error)]
pub enum MultiSigError {
    #[error("invalid public key length: expected {expected} bytes, got {actual} bytes")]
    InvalidPublicKeyLength { expected: usize, actual: usize },

    #[error("invalid weight: must be greater than 0")]
    InvalidWeight,

    #[error("invalid threshold: must be greater than 0")]
    InvalidThreshold,

    #[error("no signers provided")]
    NoSigners,

    #[error("threshold {threshold} exceeds total weight {total_weight}")]
    ThresholdExceedsTotalWeight { threshold: u16, total_weight: u16 },

    #[error("duplicate public key found")]
    DuplicatePublicKey,

    #[error("signature count mismatch: expected at least 1, got {count}")]
    InsufficientSignatures { count: usize },
}

/// Derive a multi-signature account address from signer configurations.
///
/// The address is deterministically derived by:
/// 1. Sorting signers by public key (lexicographic order)
/// 2. Concatenating: pubkey1 || weight1 || pubkey2 || weight2 || ... ||
///    threshold
/// 3. Computing keccak256 hash
/// 4. Taking last 20 bytes as address
///
/// # Arguments
/// * `signers` - List of signer configurations
/// * `threshold` - Threshold configuration
///
/// # Returns
/// The derived multi-sig account address
///
/// # Errors
/// Returns error if:
/// - No signers provided
/// - Threshold exceeds total weight
/// - Duplicate public keys found
///
/// # Example
/// ```
/// use onemoney_protocol::utils::{SignerConfig, ThresholdConfig, derive_multisig_address};
///
/// let signer1 = SignerConfig::new(vec![2; 33], 1).unwrap();
/// let signer2 = SignerConfig::new(vec![3; 33], 1).unwrap();
/// let threshold = ThresholdConfig::new(2).unwrap();
///
/// let address = derive_multisig_address(&[signer1, signer2], &threshold).unwrap();
/// ```
pub fn derive_multisig_address(
    signers: &[SignerConfig],
    threshold: &ThresholdConfig,
) -> Result<Address, MultiSigError> {
    if signers.is_empty() {
        return Err(MultiSigError::NoSigners);
    }

    if threshold.threshold == 0 {
        return Err(MultiSigError::InvalidThreshold);
    }

    for signer in signers {
        if signer.weight == 0 {
            return Err(MultiSigError::InvalidWeight);
        }
        if signer.public_key.len() != 33 {
            return Err(MultiSigError::InvalidPublicKeyLength {
                expected: 33,
                actual: signer.public_key.len(),
            });
        }
    }

    // Calculate total weight
    let total_weight: u16 = signers.iter().map(|s| s.weight as u16).sum();
    if threshold.threshold > total_weight {
        return Err(MultiSigError::ThresholdExceedsTotalWeight {
            threshold: threshold.threshold,
            total_weight,
        });
    }

    // Sort signers by public key for deterministic ordering
    let mut sorted_signers: Vec<&SignerConfig> = signers.iter().collect();
    sorted_signers.sort_by(|a, b| a.public_key.cmp(&b.public_key));

    // Check for duplicates
    for i in 1..sorted_signers.len() {
        if sorted_signers[i].public_key == sorted_signers[i - 1].public_key {
            return Err(MultiSigError::DuplicatePublicKey);
        }
    }

    // Serialize configuration: pubkey || weight || pubkey || weight || ... ||
    // threshold
    let mut data = Vec::with_capacity(DST_MULTISIG_ADDR_V1.len() + signers.len() * 34 + 2);
    data.extend_from_slice(DST_MULTISIG_ADDR_V1);
    for signer in &sorted_signers {
        data.extend_from_slice(&signer.public_key);
        data.push(signer.weight);
    }
    data.extend_from_slice(&threshold.threshold.to_be_bytes());

    // Hash and take last 20 bytes
    let hash = keccak256(&data);
    let mut address_bytes = [0u8; 20];
    address_bytes.copy_from_slice(&hash[12..32]);

    Ok(Address::from(address_bytes))
}

/// Multi-signature transaction signature collector.
///
/// Helps collect signatures from multiple signers for multi-sig transactions.
/// After collecting signatures, use `Signed::new_multi_sig` to create the
/// signed transaction.
///
/// # Example
/// ```no_run
/// use alloy_primitives::Address;
/// use onemoney_protocol::{
///     MultiSigSignatureEntry, PaymentPayload, Signature, utils::MultiSigSignatureCollector,
/// };
///
/// let mut collector = MultiSigSignatureCollector::new();
/// collector.add_signature(signer_pubkey1, signature1);
/// collector.add_signature(signer_pubkey2, signature2);
///
/// // Get collected signatures
/// let signatures = collector.signatures();
///
/// // User creates Signed transaction:
/// // let signed = Signed::new_multi_sig(payload, multisig_account, signatures);
/// ```
pub struct MultiSigSignatureCollector {
    signatures: Vec<MultiSigSignatureEntry>,
}

impl MultiSigSignatureCollector {
    /// Create a new signature collector.
    pub fn new() -> Self {
        Self { signatures: Vec::new() }
    }

    /// Add a signature from a signer.
    ///
    /// # Arguments
    /// * `signer_pubkey` - Signer's public key (33 bytes SEC1 compressed)
    /// * `signature` - Signature from this signer
    ///
    /// # Returns
    /// Mutable reference to self for method chaining
    pub fn add_signature(&mut self, signer_pubkey: B264, signature: Signature) -> &mut Self {
        self.signatures.push(MultiSigSignatureEntry {
            signer_pubkey,
            signature,
        });
        self
    }

    /// Add multiple signatures at once.
    ///
    /// # Arguments
    /// * `signatures` - Vector of signature entries
    ///
    /// # Returns
    /// Mutable reference to self for method chaining
    pub fn add_signatures(&mut self, mut signatures: Vec<MultiSigSignatureEntry>) -> &mut Self {
        self.signatures.append(&mut signatures);
        self
    }

    /// Get the number of signatures collected so far.
    pub fn signature_count(&self) -> usize {
        self.signatures.len()
    }

    /// Get all collected signatures.
    ///
    /// # Returns
    /// Vector of all signature entries
    pub fn signatures(self) -> Vec<MultiSigSignatureEntry> {
        self.signatures
    }

    /// Check if any signatures have been collected.
    pub fn is_empty(&self) -> bool {
        self.signatures.is_empty()
    }
}

impl Default for MultiSigSignatureCollector {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_signer_config_new_valid() {
        let pubkey = vec![2; 33];
        let signer = SignerConfig::new(pubkey.clone(), 1).unwrap();
        assert_eq!(signer.public_key, pubkey);
        assert_eq!(signer.weight, 1);
    }

    #[test]
    fn test_signer_config_invalid_pubkey_length() {
        let pubkey = vec![2; 32]; // Wrong length
        let result = SignerConfig::new(pubkey, 1);
        assert!(matches!(result, Err(MultiSigError::InvalidPublicKeyLength { .. })));
    }

    #[test]
    fn test_signer_config_zero_weight() {
        let pubkey = vec![2; 33];
        let result = SignerConfig::new(pubkey, 0);
        assert!(matches!(result, Err(MultiSigError::InvalidWeight)));
    }

    #[test]
    fn test_threshold_config_new_valid() {
        let threshold = ThresholdConfig::new(2).unwrap();
        assert_eq!(threshold.threshold, 2);
    }

    #[test]
    fn test_threshold_config_zero() {
        let result = ThresholdConfig::new(0);
        assert!(matches!(result, Err(MultiSigError::InvalidThreshold)));
    }

    #[test]
    fn test_derive_multisig_address_2_of_3() {
        let signer1 = SignerConfig::new(vec![2; 33], 1).unwrap();
        let signer2 = SignerConfig::new(vec![3; 33], 1).unwrap();
        let signer3 = SignerConfig::new(vec![4; 33], 1).unwrap();
        let threshold = ThresholdConfig::new(2).unwrap();

        let address =
            derive_multisig_address(&[signer1.clone(), signer2.clone(), signer3.clone()], &threshold).unwrap();

        // Address should be deterministic
        let address2 = derive_multisig_address(&[signer3, signer1, signer2], &threshold).unwrap();
        assert_eq!(
            address, address2,
            "Address should be deterministic regardless of input order"
        );
    }

    #[test]
    fn test_derive_multisig_address_no_signers() {
        let threshold = ThresholdConfig::new(1).unwrap();
        let result = derive_multisig_address(&[], &threshold);
        assert!(matches!(result, Err(MultiSigError::NoSigners)));
    }

    #[test]
    fn test_derive_multisig_address_threshold_exceeds_weight() {
        let signer1 = SignerConfig::new(vec![2; 33], 1).unwrap();
        let signer2 = SignerConfig::new(vec![3; 33], 1).unwrap();
        let threshold = ThresholdConfig::new(3).unwrap(); // Total weight is 2

        let result = derive_multisig_address(&[signer1, signer2], &threshold);
        assert!(matches!(result, Err(MultiSigError::ThresholdExceedsTotalWeight { .. })));
    }

    #[test]
    fn test_derive_multisig_address_duplicate_pubkey() {
        let signer1 = SignerConfig::new(vec![2; 33], 1).unwrap();
        let signer2 = SignerConfig::new(vec![2; 33], 2).unwrap(); // Same pubkey
        let threshold = ThresholdConfig::new(2).unwrap();

        let result = derive_multisig_address(&[signer1, signer2], &threshold);
        assert!(matches!(result, Err(MultiSigError::DuplicatePublicKey)));
    }

    #[test]
    fn test_multisig_collector() {
        let mut collector = MultiSigSignatureCollector::new();

        assert_eq!(collector.signature_count(), 0);
        assert!(collector.is_empty());

        let sig1 = Signature::test_signature();
        let sig2 = Signature::test_signature();

        collector.add_signature(B264::repeat_byte(2), sig1);
        collector.add_signature(B264::repeat_byte(3), sig2);

        assert_eq!(collector.signature_count(), 2);
        assert!(!collector.is_empty());

        let signatures = collector.signatures();
        assert_eq!(signatures.len(), 2);
    }

    #[test]
    fn test_multisig_collector_default() {
        let collector = MultiSigSignatureCollector::default();
        assert!(collector.is_empty());
        assert_eq!(collector.signature_count(), 0);
    }
}