sdk_common/src/device.rs

use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use tsify::Tsify;
use wasm_bindgen::prelude::*;

use sp_client::{
    bitcoin::{
        absolute::Height, hashes::Hash, secp256k1::PublicKey, Amount, OutPoint, Transaction,
        XOnlyPublicKey,
    },
    silentpayments::{utils::receiving::calculate_ecdh_shared_secret, SilentPaymentAddress},
    OutputSpendStatus, OwnedOutput, SpClient,
};

use crate::{pcd::Member, silentpayments::SpWallet};

#[derive(Debug, Serialize, Deserialize, Clone, Default, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub struct Device {
    sp_wallet: SpWallet,
    pairing_process_commitment: Option<OutPoint>,
    paired_member: Member,
}

impl Device {
    pub fn new(sp_client: SpClient) -> Self {
        let local_address = sp_client.get_receiving_address();
        let member = Member::new(vec![SilentPaymentAddress::try_from(local_address).unwrap()]);
        Self {
            sp_wallet: SpWallet::new(sp_client),
            pairing_process_commitment: None,
            paired_member: member,
        }
    }

    pub fn get_sp_wallet(&self) -> &SpWallet {
        &self.sp_wallet
    }

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

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

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

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

    pub fn update_outputs_with_transaction(
        &mut self,
        tx: &Transaction,
        blockheight: u32,
        partial_tweak: PublicKey,
    ) -> anyhow::Result<HashMap<OutPoint, OwnedOutput>> {
        // First check that we haven't already scanned this transaction
        let txid = tx.txid();

        for i in 0..tx.output.len() {
            if self.sp_wallet.get_outputs().contains_key(&OutPoint {
                txid,
                vout: i as u32,
            }) {
                return Err(anyhow::Error::msg("Transaction already scanned"));
            }
        }

        for input in tx.input.iter() {
            if let Some(output) = self.sp_wallet.get_outputs().get(&input.previous_output) {
                match &output.spend_status {
                    OutputSpendStatus::Spent(tx) => {
                        if *tx == txid.as_raw_hash().to_byte_array() {
                            return Err(anyhow::Error::msg("Transaction already scanned"));
                        }
                    }
                    OutputSpendStatus::Mined(_) => {
                        return Err(anyhow::Error::msg("Transaction already scanned"))
                    }
                    _ => continue,
                }
            }
        }

        let shared_secret = calculate_ecdh_shared_secret(
            &partial_tweak,
            &self.sp_wallet.get_sp_client().get_scan_key(),
        );
        let mut pubkeys_to_check: HashMap<XOnlyPublicKey, u32> = HashMap::new();
        for (vout, output) in (0u32..).zip(tx.output.iter()) {
            if output.script_pubkey.is_p2tr() {
                let xonly = XOnlyPublicKey::from_slice(&output.script_pubkey.as_bytes()[2..])?;
                pubkeys_to_check.insert(xonly, vout);
            }
        }
        let ours = self
            .sp_wallet
            .get_sp_client()
            .sp_receiver
            .scan_transaction(&shared_secret, pubkeys_to_check.keys().cloned().collect())?;
        let mut new_outputs: HashMap<OutPoint, OwnedOutput> = HashMap::new();
        for (label, map) in ours.iter() {
            for (key, scalar) in map {
                let vout = pubkeys_to_check.get(&key).unwrap().to_owned();
                let txout = tx.output.get(vout as usize).unwrap();

                let outpoint = OutPoint::new(tx.txid(), vout);
                let owned = OwnedOutput {
                    blockheight: Height::from_consensus(blockheight)?,
                    tweak: scalar.to_be_bytes(),
                    amount: txout.value,
                    script: txout.script_pubkey.to_bytes().try_into()?,
                    label: label.clone(),
                    spend_status: OutputSpendStatus::Unspent,
                };
                new_outputs.insert(outpoint, owned);
            }
        }
        let mut res = new_outputs.clone();
        self.sp_wallet.get_mut_outputs().extend(new_outputs);

        let txid = tx.txid();
        // update outputs that we own and that are spent
        for input in tx.input.iter() {
            if let Some(prevout) = self
                .sp_wallet
                .get_mut_outputs()
                .get_mut(&input.previous_output)
            {
                // This is spent by this tx
                prevout.spend_status = OutputSpendStatus::Spent(*txid.as_byte_array());
                res.insert(input.previous_output, prevout.clone());
            }
        }

        Ok(res)
    }

    pub fn get_pairing_commitment(&self) -> Option<OutPoint> {
        self.pairing_process_commitment.clone()
    }

    pub fn pair(&mut self, commitment_outpoint: OutPoint, member: Member) {
        self.pairing_process_commitment = Some(commitment_outpoint);
        self.paired_member = member;
    }

    pub fn unpair(&mut self) {
        let local_address = self.get_sp_client().get_receiving_address();
        let member = Member::new(vec![SilentPaymentAddress::try_from(local_address).unwrap()]);
        self.paired_member = member;
        self.pairing_process_commitment = None;
    }

    pub fn to_member(&self) -> Member {
        self.paired_member.clone()
    }

    pub fn get_address(&self) -> SilentPaymentAddress {
        self.get_sp_client().get_receiving_address()
    }

    pub fn get_other_addresses(&self) -> Vec<String> {
        let our_address: String = self.get_sp_client().get_receiving_address().into();
        self.to_member()
            .get_addresses()
            .into_iter()
            .filter(|a| *a != our_address)
            .collect()
    }
}