spark-cryptography 0.1.11

use bitcoin::hashes::{sha256, Hash, HashEngine};
use bitcoin::secp256k1::{schnorr::Signature, Message, PublicKey, Secp256k1, SecretKey};
use k256::elliptic_curve::{generic_array::GenericArray, group::GroupEncoding, PrimeField};
use k256::ProjectivePoint;
use rand::thread_rng;

#[derive(Debug)]
pub struct AdaptorSignature {
    pub signature: [u8; 64],
    pub adaptor_secret_key: [u8; 32],
}

#[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
pub fn generate_adaptor_from_signature(
    signature: &[u8; 64],
) -> Result<AdaptorSignature, Box<dyn std::error::Error>> {
    let mut rng = thread_rng();

    // Generate random adaptor private key
    let adaptor_secret_key = SecretKey::new(&mut rng);

    // Parse the original signature
    let sig = Signature::from_slice(signature)?;

    // Extract r and s components
    let r = sig[..32].to_vec();
    let s = sig[32..].to_vec();

    let s_ga = GenericArray::from_slice(&s);
    let adaptor_secret_key_ga = GenericArray::from_slice(adaptor_secret_key.as_ref());

    // Convert s to scalar
    let s_scalar = k256::Scalar::from_repr(*s_ga).unwrap();
    let t_scalar = k256::Scalar::from_repr(*adaptor_secret_key_ga).unwrap();

    // Calculate s - adaptor_secret_key
    let new_s = s_scalar - t_scalar;

    // Create new signature
    let mut new_signature = [0u8; 64];
    new_signature[..32].copy_from_slice(&r);
    new_signature[32..].copy_from_slice(&new_s.to_bytes());

    Ok(AdaptorSignature {
        signature: new_signature,
        adaptor_secret_key: adaptor_secret_key.secret_bytes(),
    })
}

#[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
pub fn validate_outbound_adaptor_signature(
    pubkey: &[u8],
    message: &[u8],
    signature: &[u8],
    adaptor_pubkey: &[u8],
) -> Result<bool, Box<dyn std::error::Error>> {
    schnorr_verify_with_adaptor(signature, message, pubkey, adaptor_pubkey, false)
}

#[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
pub fn apply_adaptor_to_signature(
    pubkey: &[u8],
    message: &[u8],
    signature: &[u8],
    adaptor_secret_key: &[u8],
) -> Result<[u8; 64], Box<dyn std::error::Error>> {
    let secp = Secp256k1::new();

    // Parse the signature
    let sig: [u8; 64] = signature
        .try_into()
        .map_err(|_| "Invalid signature length")?;
    let (r, s) = sig.split_at(32);

    // let r_ga = GenericArray::from_slice(r);
    let s_ga = GenericArray::from_slice(s);
    let adaptor_secret_key_ga = GenericArray::from_slice(adaptor_secret_key);

    // Convert to scalars
    let s_scalar = k256::Scalar::from_repr(*s_ga).unwrap();
    let t_scalar = k256::Scalar::from_repr(*adaptor_secret_key_ga).unwrap();

    // Try adding adaptor first
    let new_s = s_scalar + t_scalar;
    let mut new_sig = [0u8; 64];
    new_sig[..32].copy_from_slice(r);
    new_sig[32..].copy_from_slice(&new_s.to_bytes());

    // Verify the signature
    let msg = Message::from_digest_slice(message)?;
    let pk = PublicKey::from_slice(pubkey)?;

    if secp
        .verify_schnorr(
            &bitcoin::secp256k1::schnorr::Signature::from_slice(&new_sig)?,
            &msg,
            &pk.x_only_public_key().0,
        )
        .is_ok()
    {
        return Ok(new_sig);
    }

    // If adding didn't work, try subtracting
    let alt_s = s_scalar - t_scalar;
    let mut alt_sig = [0u8; 64];
    alt_sig[..32].copy_from_slice(r);
    alt_sig[32..].copy_from_slice(&alt_s.to_bytes());

    if secp
        .verify_schnorr(
            &bitcoin::secp256k1::schnorr::Signature::from_slice(&alt_sig)?,
            &msg,
            &pk.x_only_public_key().0,
        )
        .is_ok()
    {
        return Ok(alt_sig);
    }

    Err("Cannot apply adaptor to signature".into())
}

#[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
fn schnorr_verify_with_adaptor(
    signature: &[u8],
    message: &[u8],
    pubkey: &[u8],
    adaptor_pubkey: &[u8],
    _inbound: bool,
) -> Result<bool, Box<dyn std::error::Error>> {
    // Verify message length
    if message.len() != 32 {
        return Err(format!("wrong size for message (got {}, want 32)", message.len()).into());
    }

    // Parse public key
    let pk = PublicKey::from_slice(pubkey)?;
    let pk_bytes = pk.serialize();
    let pk_ga = GenericArray::from_slice(pk_bytes.as_slice());
    let pk_point = ProjectivePoint::from_bytes(pk_ga).unwrap();

    // Parse signature
    if signature.len() != 64 {
        return Err(format!("wrong signature length: {}", signature.len()).into());
    }

    let (r_bytes, s_bytes) = signature.split_at(32);

    // Compute tagged hash for challenge
    // Create a hasher
    let mut engine = sha256::Hash::engine();

    // Update with tagged data
    engine.input(b"BIP0340/challenge");
    engine.input(r_bytes);
    engine.input(&pk.serialize()[1..]); // Skip first byte (parity)
    engine.input(message);

    // Finalize and return the hash
    let e = sha256::Hash::from_engine(engine);

    // Convert challenge to scalar
    let e_ga = GenericArray::from_slice(e.as_byte_array());
    let e_scalar = k256::Scalar::from_repr(*e_ga).unwrap();
    let neg_e = -e_scalar;

    // Calculate R = sG - eP
    let s_ga = GenericArray::from_slice(s_bytes);
    let s_scalar = k256::Scalar::from_repr(*s_ga).unwrap();
    let base_point = ProjectivePoint::GENERATOR;
    let r_point = (base_point * s_scalar) + (pk_point * neg_e);

    // Add adaptor public key
    let adaptor_ga = GenericArray::from_slice(adaptor_pubkey);
    let adaptor_point = ProjectivePoint::from_bytes(adaptor_ga).unwrap();
    let _new_r = r_point + adaptor_point;

    // TODO: Check for point at infinity
    // if !inbound && new_r.is_identity().into() {
    //     return Err("calculated R point is the point at infinity".into());
    // }

    // TODO: Convert to affine and check y coordinate parity
    // let new_r_affine = new_r.to_affine();
    // if new_r_affine.y().is_odd().into() {
    //     return Err("calculated R y-value is odd".into());
    // }

    // TODO: Check if x coordinate matches r
    let r_ga = GenericArray::from_slice(r_bytes);
    let _r_scalar = k256::Scalar::from_repr(*r_ga).unwrap();
    // if new_r.x() != &r_scalar.into() {
    //     return Err("calculated R point was not given R".into());
    // }

    Ok(true)
}

#[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
pub fn generate_signature_from_existing_adaptor(
    signature: &[u8],
    adaptor_secret_key: &[u8],
) -> Result<[u8; 64], Box<dyn std::error::Error>> {
    // Parse signature
    if signature.len() != 64 {
        return Err(format!("wrong signature length: {}", signature.len()).into());
    }

    let (r, s) = signature.split_at(32);
    let s_ga = GenericArray::from_slice(s);
    let adaptor_secret_key_ga = GenericArray::from_slice(adaptor_secret_key);

    // Convert to scalars
    let s_scalar = k256::Scalar::from_repr(*s_ga).unwrap();
    let mut t_scalar = k256::Scalar::from_repr(*adaptor_secret_key_ga).unwrap();

    // Negate t
    t_scalar = -t_scalar;

    // Add negated t to s
    let new_s = s_scalar + t_scalar;

    // Create new signature
    let mut new_sig = [0u8; 64];
    new_sig[..32].copy_from_slice(r);
    new_sig[32..].copy_from_slice(&new_s.to_bytes());

    Ok(new_sig)
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use super::*;
    use bitcoin::secp256k1::Keypair;

    const MOCK_SECRET_KEY: [u8; 32] = [
        0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
        0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01,
        0x01, 0x01,
    ];

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

        let secret_key = SecretKey::from_slice(&MOCK_SECRET_KEY).unwrap();

        // Create test message and hash
        let message = b"test";
        let mut engine = sha256::Hash::engine();
        engine.input(message);
        let msg_hash = sha256::Hash::from_engine(engine);

        // Create original signature
        let msg = Message::from_digest_slice(msg_hash.as_byte_array()).unwrap();
        let keypair = Keypair::from_secret_key(&secp, &secret_key);

        let signature = Signature::from_str("b28a65d8651818ea04cfdf945455d64bbd7bac068532289a8cc0801a9ed283046159052e7ec6972cc113b2d6bc17883d5d1d4e12a4563f05b220e7b8533342e3")
            .unwrap();
        let verification = secp.verify_schnorr(
            &signature,
            &msg,
            &keypair.public_key().x_only_public_key().0,
        );

        assert!(verification.is_ok());
    }
}