spark_cryptography/

key_arithmetic.rs

use bitcoin::secp256k1::{Error as Secp256k1Error, PublicKey, Secp256k1, SecretKey};

/// Add two public keys
/// Input keys must be 33-byte compressed secp256k1 public keys
pub fn add_public_keys(key1: &[u8], key2: &[u8]) -> Result<Vec<u8>, Secp256k1Error> {
    if key1.len() != 33 || key2.len() != 33 {
        return Err(Secp256k1Error::InvalidPublicKey);
    }

    let pub1 = PublicKey::from_slice(key1)?;
    let pub2 = PublicKey::from_slice(key2)?;

    // Combine the public keys
    let combined = pub1.combine(&pub2)?;

    // Return compressed format (33 bytes)
    Ok(combined.serialize().to_vec())
}

pub fn apply_public_key_tweak(pubkey: &[u8], tweak: &[u8]) -> Result<Vec<u8>, Secp256k1Error> {
    if pubkey.len() != 33 || tweak.len() != 32 {
        return Err(Secp256k1Error::InvalidPublicKey);
    }

    let secp = Secp256k1::new();
    let pub_key = PublicKey::from_slice(pubkey)?;

    // Apply the tweak
    let tweak_key = SecretKey::from_slice(tweak)?;
    let tweaked_key = pub_key.add_exp_tweak(&secp, &tweak_key.into())?;

    // Return compressed format
    Ok(tweaked_key.serialize().to_vec())
}

/// Subtract public keys
/// Input keys must be 33-byte compressed secp256k1 public keys
pub fn subtract_public_keys(key1: &[u8], key2: &[u8]) -> Result<Vec<u8>, Secp256k1Error> {
    if key1.len() != 33 || key2.len() != 33 {
        return Err(Secp256k1Error::InvalidPublicKey);
    }

    let secp = Secp256k1::new();
    let pub1 = PublicKey::from_slice(key1)?;
    let pub2 = PublicKey::from_slice(key2)?;

    // Negate the second key and add
    let negated = pub2.negate(&secp);
    let result = pub1.combine(&negated)?;

    Ok(result.serialize().to_vec())
}

/// Add private keys
/// Input keys must be 32 bytes
pub fn add_private_keys(key1: &[u8], key2: &[u8]) -> Result<Vec<u8>, Secp256k1Error> {
    if key1.len() != 32 || key2.len() != 32 {
        return Err(Secp256k1Error::InvalidSecretKey);
    }

    let sec1 = SecretKey::from_slice(key1)?;
    let sec2 = SecretKey::from_slice(key2)?;

    // Add the keys modulo curve order
    let combined = sec1.add_tweak(&sec2.into())?;

    Ok(combined.secret_bytes().to_vec())
}

/// Subtract private keys
/// Input keys must be 32 bytes
pub fn subtract_secret_keys<T: AsRef<[u8]>, U: AsRef<[u8]>>(
    key1: T,
    key2: U,
) -> Result<Vec<u8>, Secp256k1Error> {
    if key1.as_ref().len() != 32 || key2.as_ref().len() != 32 {
        return Err(Secp256k1Error::InvalidSecretKey);
    }

    let sec1 = SecretKey::from_slice(key1.as_ref())?;
    let sec2 = SecretKey::from_slice(key2.as_ref())?;

    // Negate the second key and add
    let negated = sec2.negate();
    let result = sec1.add_tweak(&negated.into())?;

    Ok(result.secret_bytes().to_vec())
}

/// Sum of private keys
/// Returns the sum of the given private keys modulo the curve order
pub fn sum_private_keys(keys: &[&[u8]]) -> Result<Vec<u8>, Secp256k1Error> {
    if keys.is_empty() {
        return Err(Secp256k1Error::InvalidSecretKey);
    }

    let mut result = SecretKey::from_slice(keys[0])?;

    // Add all subsequent keys
    for key in keys.iter().skip(1) {
        if key.len() != 32 {
            return Err(Secp256k1Error::InvalidSecretKey);
        }
        let next_key = SecretKey::from_slice(key)?;
        result = result.add_tweak(&next_key.into())?;
    }

    Ok(result.secret_bytes().to_vec())
}

/// Last key with target
/// Tweaks the given keys so their sum equals the target
pub fn last_key_with_target(keys: &[&[u8]], target: &[u8]) -> Result<Vec<u8>, Secp256k1Error> {
    if target.len() != 32 {
        return Err(Secp256k1Error::InvalidSecretKey);
    }

    let target_key = SecretKey::from_slice(target)?;
    let current_sum = sum_private_keys(keys)?;
    let current_sum_key = SecretKey::from_slice(&current_sum)?;

    // Calculate target - sum(keys)
    let result = subtract_secret_keys(&target_key.secret_bytes(), &current_sum_key.secret_bytes())?;

    Ok(result)
}

#[cfg(test)]
mod tests {
    use super::*;
    use bitcoin::secp256k1::{rand::rngs::OsRng, Secp256k1};

    /// Helper function to generate n random private keys
    fn generate_test_keys(n: usize) -> Vec<Vec<u8>> {
        let secp = Secp256k1::new();
        let mut keys = Vec::with_capacity(n);
        for _ in 0..n {
            let (secret_key, _) = secp.generate_keypair(&mut OsRng);
            keys.push(secret_key.secret_bytes().to_vec());
        }
        keys
    }

    #[test]
    fn test_key_additions_multiple() {
        let secp = Secp256k1::new();

        // Test with multiple pairs of keys
        for _ in 0..10 {
            let (priv_a, pub_a) = secp.generate_keypair(&mut OsRng);
            let (priv_b, pub_b) = secp.generate_keypair(&mut OsRng);

            let priv_sum =
                add_private_keys(&priv_a.secret_bytes(), &priv_b.secret_bytes()).unwrap();

            let pub_sum = add_public_keys(&pub_a.serialize(), &pub_b.serialize()).unwrap();

            let sum_key = SecretKey::from_slice(&priv_sum).unwrap();
            let sum_pub = PublicKey::from_secret_key(&secp, &sum_key);

            assert_eq!(sum_pub.serialize().to_vec(), pub_sum);
        }
    }

    #[test]
    fn test_subtract_keys() {
        let secp = Secp256k1::new();

        // Test that (A + B) - B = A for both public and private keys
        let (priv_a, pub_a) = secp.generate_keypair(&mut OsRng);
        let (priv_b, pub_b) = secp.generate_keypair(&mut OsRng);

        // Test private key subtraction
        let sum = add_private_keys(&priv_a.secret_bytes(), &priv_b.secret_bytes()).unwrap();

        let diff = subtract_secret_keys(&sum, &priv_b.secret_bytes()).unwrap();

        assert_eq!(diff, priv_a.secret_bytes());

        // Test public key subtraction
        let pub_sum = add_public_keys(&pub_a.serialize(), &pub_b.serialize()).unwrap();

        let pub_diff = subtract_public_keys(&pub_sum, &pub_b.serialize()).unwrap();

        assert_eq!(pub_diff, pub_a.serialize());
    }

    #[test]
    #[should_panic(expected = "InvalidSecretKey")]
    fn test_invalid_private_key() {
        // Test with invalid private key (all zeros)
        let zeros = vec![0u8; 32];
        let (valid_key, _) = Secp256k1::new().generate_keypair(&mut OsRng);

        add_private_keys(&zeros, &valid_key.secret_bytes()).unwrap();
    }

    #[test]
    #[should_panic(expected = "InvalidPublicKey")]
    fn test_invalid_public_key() {
        // Test with invalid public key (all zeros)
        let zeros = vec![0u8; 33];
        let (_, valid_pub) = Secp256k1::new().generate_keypair(&mut OsRng);

        add_public_keys(&zeros, &valid_pub.serialize()).unwrap();
    }

    #[test]
    fn test_sum_of_private_keys_large_set() {
        // Test with a larger set of keys (100 keys)
        let keys = generate_test_keys(100);
        let key_slices: Vec<&[u8]> = keys.iter().map(|k| k.as_slice()).collect();

        let sum = sum_private_keys(&key_slices).unwrap();

        // Verify by adding one by one
        let mut manual_sum = keys[0].clone();
        for key in keys.iter().skip(1) {
            manual_sum = add_private_keys(&manual_sum, key).unwrap();
        }

        assert_eq!(sum, manual_sum);
    }

    #[test]
    fn test_last_key_with_target_properties() {
        let secp = Secp256k1::new();

        // Generate random target and keys
        let (target_key, _) = secp.generate_keypair(&mut OsRng);
        let keys = generate_test_keys(5);
        let key_slices: Vec<&[u8]> = keys.iter().map(|k| k.as_slice()).collect();

        // Calculate tweak
        let tweak = last_key_with_target(&key_slices, &target_key.secret_bytes()).unwrap();

        // Properties to verify:
        // 1. Tweak should be a valid private key
        assert!(SecretKey::from_slice(&tweak).is_ok());

        // 2. Sum of keys + tweak should equal target
        let mut keys_with_tweak = keys.clone();
        keys_with_tweak.push(tweak);
        let key_slices_with_tweak: Vec<&[u8]> =
            keys_with_tweak.iter().map(|k| k.as_slice()).collect();

        let sum = sum_private_keys(&key_slices_with_tweak).unwrap();
        assert_eq!(sum, target_key.secret_bytes());
    }

    #[test]
    fn test_apply_tweak_associative_property() {
        let secp = Secp256k1::new();

        // Test that (P + t1) + t2 = P + (t1 + t2)
        let (_, base_pub) = secp.generate_keypair(&mut OsRng);
        let (tweak1, _) = secp.generate_keypair(&mut OsRng);
        let (tweak2, _) = secp.generate_keypair(&mut OsRng);

        // Method 1: (P + t1) + t2
        let intermediate =
            apply_public_key_tweak(&base_pub.serialize(), &tweak1.secret_bytes()).unwrap();

        let result1 = apply_public_key_tweak(&intermediate, &tweak2.secret_bytes()).unwrap();

        // Method 2: P + (t1 + t2)
        let combined_tweak =
            add_private_keys(&tweak1.secret_bytes(), &tweak2.secret_bytes()).unwrap();

        let result2 = apply_public_key_tweak(&base_pub.serialize(), &combined_tweak).unwrap();

        assert_eq!(result1, result2);
    }

    #[test]
    fn test_key_additions() {
        let secp = Secp256k1::new();

        // Generate first key pair
        let (priv_a, pub_a) = secp.generate_keypair(&mut OsRng);
        let priv_a_bytes = priv_a.secret_bytes();
        let pub_a_bytes = pub_a.serialize();

        // Generate second key pair
        let (priv_b, pub_b) = secp.generate_keypair(&mut OsRng);
        let priv_b_bytes = priv_b.secret_bytes();
        let pub_b_bytes = pub_b.serialize();

        // Test that public key of private key addition equals the public key addition
        let priv_sum =
            add_private_keys(&priv_a_bytes, &priv_b_bytes).expect("Failed to add private keys");
        let pub_sum =
            add_public_keys(&pub_a_bytes, &pub_b_bytes).expect("Failed to add public keys");

        let target_key =
            SecretKey::from_slice(&priv_sum).expect("Failed to create secret key from sum");
        let target_pub = PublicKey::from_secret_key(&secp, &target_key);

        assert_eq!(
            target_pub.serialize().to_vec(),
            pub_sum,
            "Public key of private key addition does not equal the public key addition"
        );
    }

    #[test]
    fn test_sum_of_private_keys() {
        let secp = Secp256k1::new();

        // Generate 10 random keys
        let mut keys: Vec<Vec<u8>> = Vec::new();
        for _ in 0..10 {
            let (secret_key, _) = secp.generate_keypair(&mut OsRng);
            keys.push(secret_key.secret_bytes().to_vec());
        }

        // Calculate sum using sum_private_keys
        let key_slices: Vec<&[u8]> = keys.iter().map(|k| k.as_slice()).collect();
        let sum = sum_private_keys(&key_slices).expect("Failed to sum private keys");

        // Calculate sum manually using add_private_keys
        let mut sum2 = keys[0].clone();
        for i in 1..keys.len() {
            sum2 = add_private_keys(&sum2, &keys[i]).expect("Failed to add private keys");
        }

        assert_eq!(sum, sum2, "Sum of private keys does not match");
    }

    #[test]
    fn test_private_key_tweak_with_target() {
        let secp = Secp256k1::new();

        // Generate target key
        let (target_key, _) = secp.generate_keypair(&mut OsRng);
        let target_bytes = target_key.secret_bytes();

        // Generate 10 random keys
        let mut keys: Vec<Vec<u8>> = Vec::new();
        for _ in 0..10 {
            let (secret_key, _) = secp.generate_keypair(&mut OsRng);
            keys.push(secret_key.secret_bytes().to_vec());
        }

        // Get key slices for the function call
        let key_slices: Vec<&[u8]> = keys.iter().map(|k| k.as_slice()).collect();

        // Calculate tweak
        let tweak =
            last_key_with_target(&key_slices, &target_bytes).expect("Failed to calculate tweak");

        // Add tweak to keys
        keys.push(tweak);

        // Calculate sum with tweak
        let key_slices_with_tweak: Vec<&[u8]> = keys.iter().map(|k| k.as_slice()).collect();
        let sum = sum_private_keys(&key_slices_with_tweak)
            .expect("Failed to sum private keys with tweak");

        assert_eq!(
            sum,
            target_bytes.to_vec(),
            "Private key tweak with target does not match"
        );
    }

    #[test]
    fn test_apply_additive_tweak_to_public_key() {
        let secp = Secp256k1::new();

        // Generate initial keypair
        let (priv_key, pub_key) = secp.generate_keypair(&mut OsRng);
        let priv_key_bytes = priv_key.secret_bytes();
        let pub_key_bytes = pub_key.serialize();

        // Generate tweak
        let (tweak_key, _) = secp.generate_keypair(&mut OsRng);
        let tweak_bytes = tweak_key.secret_bytes();

        // Calculate new private key by adding
        let new_priv =
            add_private_keys(&priv_key_bytes, &tweak_bytes).expect("Failed to add private keys");
        let target_key =
            SecretKey::from_slice(&new_priv).expect("Failed to create secret key from sum");
        let target_pub = PublicKey::from_secret_key(&secp, &target_key);

        // Apply tweak to public key
        let new_pub_key = apply_public_key_tweak(&pub_key_bytes, &tweak_bytes)
            .expect("Failed to apply tweak to public key");

        assert_eq!(
            new_pub_key,
            target_pub.serialize(),
            "Apply additive tweak to public key does not match"
        );
    }
}