onemoney-protocol 0.18.0

//! BLS signature verification utilities for finalized transactions.

use alloy_primitives::B256;
use om_crypto_types::bls12381::{
    AggregateSignature as BlsAggregateSignature, PublicKey as BlsPublicKey, Signature as BlsSignature,
};
/// Computes the message hash that validators counter-sign for a transaction.
pub use om_crypto_types::counter_sign::signature_hash_for_counter_sign;
use om_rest_types::RestBlsAggregateSignature;

use crate::error::{Error, Result};

/// Verify a BLS aggregate signature from a finalized transaction.
///
/// This function verifies that the BLS aggregate signature is valid for the
/// given message hash. It decodes the signature and validator public keys from
/// the REST API format and performs cryptographic verification.
///
/// **Note**: Most users should use
/// `Client::get_and_verify_finalized_transaction_by_hash()` which automatically
/// fetches and verifies the transaction. This function is for advanced use
/// cases where you need fine-grained control.
///
/// # Arguments
///
/// * `message_hash` - The hash of the message that was signed (computed via
///   `signature_hash_for_counter_sign`)
/// * `aggregate_sig` - The BLS aggregate signature from the finalized
///   transaction
///
/// # Returns
///
/// `Ok(())` if the signature is valid, `Err` otherwise.
///
/// # Example
///
/// ```rust,ignore
/// use alloy_primitives::B256;
/// use onemoney_protocol::{Client, signature_hash_for_counter_sign, verify_bls_aggregate_signature};
/// use std::str::FromStr;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     let client = Client::testnet()?;
///
///     // Easier way: use get_and_verify directly
///     let finalized_tx = client
///         .get_and_verify_finalized_transaction_by_hash("0x...")
///         .await?;
///
///     // Advanced: manual verification
///     let tx_hash = B256::from_str("0x...")?;
///     let finalized_tx = client.get_finalized_transaction_by_hash("0x...").await?;
///     let message_hash = signature_hash_for_counter_sign(&tx_hash, &finalized_tx.epoch);
///     verify_bls_aggregate_signature(&message_hash, &finalized_tx.counter_signature)?;
///
///     Ok(())
/// }
/// ```
pub fn verify_bls_aggregate_signature(message_hash: &B256, aggregate_sig: &RestBlsAggregateSignature) -> Result<()> {
    // Decode the signer bitmask
    let bitmask_hex = aggregate_sig
        .signer_bitmask
        .strip_prefix("0x")
        .unwrap_or(&aggregate_sig.signer_bitmask);
    let bitmask_bytes =
        hex::decode(bitmask_hex).map_err(|e| Error::verification_error(format!("Invalid bitmask hex: {}", e)))?;

    // Decode the BLS signature
    let signature_hex = aggregate_sig
        .signature
        .strip_prefix("0x")
        .unwrap_or(&aggregate_sig.signature);
    let signature_bytes =
        hex::decode(signature_hex).map_err(|e| Error::verification_error(format!("Invalid signature hex: {}", e)))?;

    if signature_bytes.len() != 48 {
        return Err(Error::verification_error(format!(
            "BLS signature must be 48 bytes, got {}",
            signature_bytes.len()
        )));
    }

    let bls_signature = BlsSignature::from_bytes(&signature_bytes)
        .map_err(|e| Error::verification_error(format!("Invalid BLS signature bytes: {}", e)))?;

    // Reconstruct the aggregate signature
    let aggregate_signature = BlsAggregateSignature::new(bitmask_bytes, bls_signature);

    // Decode validator public keys
    let validator_pubkeys: Result<Vec<BlsPublicKey>> = aggregate_sig
        .validator_public_keys
        .iter()
        .map(|pk_hex| {
            let pk_hex = pk_hex.strip_prefix("0x").unwrap_or(pk_hex);
            let pk_bytes =
                hex::decode(pk_hex).map_err(|e| Error::verification_error(format!("Invalid public key hex: {}", e)))?;

            if pk_bytes.len() != 96 {
                return Err(Error::verification_error(format!(
                    "BLS public key must be 96 bytes, got {}",
                    pk_bytes.len()
                )));
            }

            BlsPublicKey::from_bytes(&pk_bytes)
                .map_err(|e| Error::verification_error(format!("Invalid BLS public key bytes: {}", e)))
        })
        .collect();

    let validator_pubkeys = validator_pubkeys?;

    // Verify the aggregate signature
    aggregate_signature
        .verify(message_hash.as_slice(), &validator_pubkeys)
        .map_err(|e| Error::verification_error(format!("BLS signature verification failed: {}", e)))?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use om_crypto_types::bls12381::PrivateKey as BlsPrivateKey;

    use super::*;

    fn generate_test_keypair() -> (BlsPrivateKey, BlsPublicKey) {
        let private_key = BlsPrivateKey::generate(&mut rand::thread_rng());
        let public_key = private_key.public_key();
        (private_key, public_key)
    }

    #[test]
    fn test_verify_single_validator_signature() {
        let (private_key, public_key) = generate_test_keypair();

        let tx_hash = B256::from([1u8; 32]);
        let epoch = 42u64;
        let message_hash = signature_hash_for_counter_sign(&tx_hash, &epoch);

        let signature = private_key.sign(message_hash.as_slice());
        let bitmask = vec![0b00000001u8];
        let aggregate_sig = BlsAggregateSignature::new(bitmask.clone(), signature);

        let rest_sig = RestBlsAggregateSignature {
            signer_bitmask: format!("0x{}", hex::encode(&bitmask)),
            signature: format!("0x{}", hex::encode(aggregate_sig.signature().to_bytes())),
            validator_public_keys: vec![format!("0x{}", hex::encode(public_key.to_bytes()))],
        };

        let result = verify_bls_aggregate_signature(&message_hash, &rest_sig);
        if let Err(e) = &result {
            eprintln!("Verification error: {:?}", e);
        }
        assert!(
            result.is_ok(),
            "Valid signature should verify successfully: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_verify_multiple_validators_signature() {
        let validators: Vec<_> = (0..3).map(|_| generate_test_keypair()).collect();

        let tx_hash = B256::from([2u8; 32]);
        let epoch = 100u64;
        let message_hash = signature_hash_for_counter_sign(&tx_hash, &epoch);

        let signatures: Vec<_> = validators
            .iter()
            .map(|(sk, _)| sk.sign(message_hash.as_slice()))
            .collect();
        let sig_refs: Vec<_> = signatures.iter().collect();

        let aggregated = BlsSignature::aggregate(&sig_refs).expect("Should aggregate");

        let bitmask = vec![0b00000111u8];

        let rest_sig = RestBlsAggregateSignature {
            signer_bitmask: format!("0x{}", hex::encode(&bitmask)),
            signature: format!("0x{}", hex::encode(aggregated.to_bytes())),
            validator_public_keys: validators
                .iter()
                .map(|(_, pk)| format!("0x{}", hex::encode(pk.to_bytes())))
                .collect(),
        };

        let result = verify_bls_aggregate_signature(&message_hash, &rest_sig);
        assert!(result.is_ok(), "Valid multi-validator signature should verify");
    }

    #[test]
    fn test_verify_fails_with_wrong_message() {
        let (private_key, public_key) = generate_test_keypair();

        let tx_hash = B256::from([1u8; 32]);
        let epoch = 42u64;
        let message_hash = signature_hash_for_counter_sign(&tx_hash, &epoch);

        let signature = private_key.sign(message_hash.as_slice());
        let bitmask = vec![0b00000001u8];

        let rest_sig = RestBlsAggregateSignature {
            signer_bitmask: format!("0x{}", hex::encode(&bitmask)),
            signature: format!("0x{}", hex::encode(signature.to_bytes())),
            validator_public_keys: vec![format!("0x{}", hex::encode(public_key.to_bytes()))],
        };

        let wrong_message_hash = B256::from([99u8; 32]);
        let result = verify_bls_aggregate_signature(&wrong_message_hash, &rest_sig);
        assert!(result.is_err(), "Verification should fail with wrong message");
    }

    #[test]
    fn test_verify_fails_with_invalid_signature() {
        let (_, public_key) = generate_test_keypair();

        let tx_hash = B256::from([1u8; 32]);
        let epoch = 42u64;
        let message_hash = signature_hash_for_counter_sign(&tx_hash, &epoch);

        let bitmask = vec![0b00000001u8];

        // Create invalid signature (all zeros)
        let rest_sig = RestBlsAggregateSignature {
            signer_bitmask: format!("0x{}", hex::encode(&bitmask)),
            signature:
                "0x000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
                    .to_string(),
            validator_public_keys: vec![format!("0x{}", hex::encode(public_key.to_bytes()))],
        };

        let result = verify_bls_aggregate_signature(&message_hash, &rest_sig);
        assert!(result.is_err(), "Verification should fail with invalid signature");
    }

    #[test]
    fn test_verify_fails_with_wrong_public_key() {
        let (private_key, _) = generate_test_keypair();
        let (_, wrong_public_key) = generate_test_keypair();

        let tx_hash = B256::from([1u8; 32]);
        let epoch = 42u64;
        let message_hash = signature_hash_for_counter_sign(&tx_hash, &epoch);

        let signature = private_key.sign(message_hash.as_slice());
        let bitmask = vec![0b00000001u8];

        let rest_sig = RestBlsAggregateSignature {
            signer_bitmask: format!("0x{}", hex::encode(&bitmask)),
            signature: format!("0x{}", hex::encode(signature.to_bytes())),
            validator_public_keys: vec![format!("0x{}", hex::encode(wrong_public_key.to_bytes()))],
        };

        let result = verify_bls_aggregate_signature(&message_hash, &rest_sig);
        assert!(result.is_err(), "Verification should fail with wrong public key");
    }
}