sdk_common/src/silentpayments.rs

use std::collections::HashMap;

use anyhow::{Error, Result};
use serde::{Deserialize, Serialize};
use sp_client::bitcoin::absolute::Height;
use sp_client::bitcoin::secp256k1::{PublicKey, SecretKey};
use tsify::Tsify;

use rand::{thread_rng, Rng};
use sp_client::bitcoin::{Amount, OutPoint, Transaction, TxOut, XOnlyPublicKey};
use sp_client::{FeeRate, OutputSpendStatus, OwnedOutput, Recipient, SilentPaymentUnsignedTransaction, SpClient};
use sp_client::silentpayments::utils::receiving::calculate_ecdh_shared_secret;

#[derive(Debug, Default, Deserialize, Serialize, Clone)]
pub struct SpWallet {
    sp_client: SpClient,
    outputs: HashMap<OutPoint, OwnedOutput>,
    birthday: u32,
    last_scan: u32,
}

impl SpWallet {
    pub fn new(sp_client: SpClient) -> Self {
        Self {
            sp_client,
            ..Default::default()
        }
    }

    pub fn get_sp_client(&self) -> &SpClient {
        &self.sp_client
    }

    pub fn get_outputs(&self) -> &HashMap<OutPoint, OwnedOutput> {
        &self.outputs
    }

    pub fn get_mut_outputs(&mut self) -> &mut HashMap<OutPoint, OwnedOutput> {
        &mut self.outputs
    }

    pub fn get_unspent_outputs(&self) -> HashMap<OutPoint, OwnedOutput> {
        self.outputs.iter()
            .filter(|(_, output)| output.spend_status == OutputSpendStatus::Unspent)
            .map(|(outpoint, output)| (*outpoint, output.clone()))
            .collect()
    }

    pub fn get_balance(&self) -> Amount {
        self.outputs.values()
            .filter(|output| output.spend_status == OutputSpendStatus::Unspent)
            .fold(Amount::ZERO, |acc, x| acc + x.amount)
    }

    pub fn get_birthday(&self) -> u32 {
        self.birthday
    }

    pub fn set_birthday(&mut self, birthday: u32) {
        self.birthday = birthday;
    }

    pub fn get_last_scan(&self) -> u32 {
        self.last_scan
    }

    pub fn set_last_scan(&mut self, last_scan: u32) {
        self.last_scan = last_scan;
    }

    pub fn update_with_transaction(&mut self, tx: &Transaction, public_tweak: &PublicKey, height: u32) -> Result<HashMap<OutPoint, OwnedOutput>> {
        let receiver = &self.get_sp_client().sp_receiver;
        let p2tr_outs: Vec<(usize, &TxOut)> = tx
            .output
            .iter()
            .enumerate()
            .filter(|(_, o)| o.script_pubkey.is_p2tr())
            .collect();

        if p2tr_outs.is_empty() {
            return Err(Error::msg("No taproot outputs"))
        }; // That should never happen since we have a tweak_data, but anyway

        // Now we can just run sp_receiver on all the p2tr outputs
        let xonlykeys: Result<Vec<XOnlyPublicKey>> = p2tr_outs
            .iter()
            .map(|(_, o)| {
                XOnlyPublicKey::from_slice(&o.script_pubkey.as_bytes()[2..]).map_err(Error::new)
            })
            .collect();

        let tweak_data = calculate_ecdh_shared_secret(public_tweak, &self.sp_client.get_scan_key());
        let ours = receiver.scan_transaction(&tweak_data, xonlykeys?)?;

        let txid = tx.txid();
        let height = Height::from_consensus(height)?;
        let mut res = HashMap::new();
        for (label, map) in ours.iter() {
            res.extend(p2tr_outs.iter().filter_map(|(i, o)| {
                match XOnlyPublicKey::from_slice(&o.script_pubkey.as_bytes()[2..]) {
                    Ok(key) => {
                        if let Some(scalar) = map.get(&key) {
                            match SecretKey::from_slice(&scalar.to_be_bytes()) {
                                Ok(tweak) => {
                                    let outpoint = OutPoint {
                                        txid,
                                        vout: *i as u32,
                                    };
                                    return Some((
                                        outpoint,
                                        OwnedOutput {
                                            blockheight: height,
                                            tweak: tweak.secret_bytes(),
                                            amount: o.value,
                                            script: o.script_pubkey.clone(),
                                            label: label.clone(),
                                            spend_status: OutputSpendStatus::Unspent,
                                        },
                                    ));
                                }
                                Err(_) => {
                                    return None;
                                }
                            }
                        }
                        None
                    }
                    Err(_) => None,
                }
            }));
        }
        self.get_mut_outputs().extend(res.clone());
        Ok(res)
    }
}

#[derive(Debug, Serialize, Deserialize, PartialEq, Tsify)]
#[tsify(from_wasm_abi)]
pub struct TsUnsignedTransaction(SilentPaymentUnsignedTransaction); 

impl TsUnsignedTransaction {
    pub fn new(unsigned_tx: SilentPaymentUnsignedTransaction) -> Self {
        Self(unsigned_tx)
    }

    pub fn as_inner(&self) -> &SilentPaymentUnsignedTransaction {
        &self.0
    }

    pub fn to_inner(self) -> SilentPaymentUnsignedTransaction {
        self.0
    }
}

pub fn create_transaction(
    available_outpoints: Vec<(OutPoint, OwnedOutput)>,
    sp_client: &SpClient,
    mut recipients: Vec<Recipient>,
    payload: Option<Vec<u8>>,
    fee_rate: FeeRate,
) -> Result<SilentPaymentUnsignedTransaction> {
    let mut commitment = [0u8; 32];
    if let Some(ref p) = payload {
        commitment.copy_from_slice(&p);
    } else {
        thread_rng().fill(&mut commitment);
    }

    recipients.push(Recipient { 
        address: sp_client::RecipientAddress::Data(commitment.to_vec()), 
        amount: Amount::ZERO 
    });

    let new_transaction = sp_client.create_new_transaction(
        available_outpoints, 
        recipients, 
        fee_rate,
        sp_client.get_network()
    )?;

    let finalized_transaction = SpClient::finalize_transaction(new_transaction)?;

    Ok(finalized_transaction)
}

pub fn sign_transaction(sp_client: &SpClient, unsigned_transaction: SilentPaymentUnsignedTransaction) -> Result<Transaction> {
    let mut aux_rand = [0u8; 32];
    thread_rng().fill(&mut aux_rand);
    sp_client.sign_transaction(unsigned_transaction, &aux_rand)
}