sdk_client/crates/sp_client/src/silentpayments.rs

// This file should move to common

use std::collections::{HashMap, HashSet};
use std::io::Write;
use std::iter::Once;
use std::sync::{Mutex, MutexGuard, OnceLock};

use anyhow::{Error, Result};

use log::debug;
use rand::Rng;
use sdk_common::crypto::{Aes256Encryption, Aes256Gcm, AeadCore, Purpose};
use serde::{Deserialize, Serialize};
use sp_client::bitcoin::policy::DUST_RELAY_TX_FEE;
use sp_client::bitcoin::secp256k1::ecdh::shared_secret_point;
use sp_client::bitcoin::{block, Amount, OutPoint, Txid};
use sp_client::silentpayments::sending::SilentPaymentAddress;
use sp_client::silentpayments::utils::receiving::calculate_shared_secret;
use sp_client::spclient::{OutputList, OwnedOutput, Recipient, SpClient, SpWallet};
use sp_client::silentpayments::bitcoin_hashes::{sha256, Hash, HashEngine};
use sp_client::{
    bitcoin::{
        secp256k1::{PublicKey, Scalar, XOnlyPublicKey},
        Transaction,
    },
    silentpayments::receiving::Label,
};
use tsify::Tsify;

use crate::user::{lock_connected_user, CONNECTED_USER};
use crate::MutexExt;

#[derive(Debug, Serialize, Deserialize, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub struct ScannedTransaction(HashMap<Txid, Vec<SharedSecret>>);

impl ScannedTransaction {
    pub fn new() -> Self {
        Self(HashMap::new())
    }

    pub fn get_mut(&mut self) -> &mut HashMap<Txid, Vec<SharedSecret>> {
        &mut self.0
    }
}

pub static TRANSACTIONCACHE: OnceLock<Mutex<ScannedTransaction>> = OnceLock::new();

pub fn lock_scanned_transactions() -> Result<MutexGuard<'static, ScannedTransaction>> {
    TRANSACTIONCACHE
        .get_or_init(|| Mutex::new(ScannedTransaction::new()))
        .lock_anyhow()
}

type FoundOutputs = HashMap<Option<Label>, HashMap<XOnlyPublicKey, Scalar>>;

#[derive(Debug)]
pub struct SharedPoint([u8; 64]);

impl SharedPoint {
    pub fn as_inner(&self) -> &[u8; 64] {
        &self.0
    }
}

#[derive(Debug, Serialize, Deserialize, Tsify)]
#[tsify(from_wasm_abi, into_wasm_abi)]
pub struct SharedSecret {
    secret: [u8; 32],
    shared_with: Option<String>, // SilentPaymentAddress
}

impl SharedSecret {
    pub fn new(secret: [u8;32], shared_with: Option<String>) -> Result<Self> {
        if let Some(ref shared) = shared_with {
            if let Ok(_) = SilentPaymentAddress::try_from(shared.as_str()) {
                Ok(Self {
                    secret,
                    shared_with
                })
            } else {
                Err(Error::msg("Invalid silent payment address"))
            }
        } else {
            Ok(Self {
                secret, 
                shared_with: None
            })
        }
    }
}

pub fn create_transaction_for_address_with_shared_secret(
    sp_address: SilentPaymentAddress,
    message: String,
) -> Result<(Transaction, SharedSecret)> {
    let connected_user = lock_connected_user()?;

    let sp_wallet = if sp_address.is_testnet() {
        connected_user.try_get_recover()?
    } else {
        if let Ok(main) = connected_user.try_get_main() {
            main
        } else {
            return Err(Error::msg("Can't spend on mainnet"));
        }
    };

    let available_outpoints = sp_wallet.get_outputs().to_spendable_list();

    // Here we need to add more heuristics about which outpoint we spend
    // For now let's keep it simple

    let mut inputs: HashMap<OutPoint, OwnedOutput> = HashMap::new();

    let total_available =
        available_outpoints
            .into_iter()
            .try_fold(Amount::from_sat(0), |acc, (outpoint, output)| {
                let new_total = acc + output.amount;
                inputs.insert(outpoint, output);
                if new_total > Amount::from_sat(1000) {
                    Err(new_total)
                } else {
                    Ok(new_total)
                }
            });

    match total_available {
        Err(total) => log::debug!("Spending {} outputs totaling {} sats", inputs.len(), total),
        Ok(_) => return Err(Error::msg("Not enought fund available")),
    }

    let recipient = Recipient {
        address: sp_address.into(),
        amount: Amount::from_sat(1000),
        nb_outputs: 1,
    };

    let mut new_psbt = sp_wallet
        .get_client()
        .create_new_psbt(inputs, vec![recipient], Some(message.as_bytes()))?;
    log::debug!("Created psbt: {}", new_psbt);
    SpClient::set_fees(&mut new_psbt, Amount::from_sat(1000), sp_address.into())?;

    let partial_secret = sp_wallet
        .get_client()
        .get_partial_secret_from_psbt(&new_psbt)?;

    // This wouldn't work with many recipients in the same transaction
    // each address (or more precisely each scan public key) would have its own point
    let shared_point = shared_secret_point(&sp_address.get_scan_key(), &partial_secret);

    // The shared_point *must* be hashed to produce a proper ecdh secret
    let mut eng = sha256::HashEngine::default();
    eng.write_all(&shared_point);
    let shared_secret = sha256::Hash::from_engine(eng);

    // encrypt the message with the new shared_secret
    let message_encryption = Aes256Encryption::import_key(
        Purpose::Arbitrary,
        message.into_bytes(),
        shared_secret.to_byte_array(),
        Aes256Gcm::generate_nonce(&mut rand::thread_rng()).into(),
    )?;

    let cipher = message_encryption.encrypt_with_aes_key()?;

    sp_client::spclient::SpClient::replace_op_return_with(&mut new_psbt, &cipher)?;

    sp_wallet
        .get_client()
        .fill_sp_outputs(&mut new_psbt, partial_secret)?;
    log::debug!("Definitive psbt: {}", new_psbt);
    let mut aux_rand = [0u8; 32];
    rand::thread_rng().fill(&mut aux_rand);
    let mut signed = sp_wallet.get_client().sign_psbt(new_psbt, &aux_rand)?;
    log::debug!("signed psbt: {}", signed);
    SpClient::finalize_psbt(&mut signed)?;

    let final_tx = signed.extract_tx()?;

    Ok((final_tx, SharedSecret {
        secret: shared_secret.to_byte_array(),
        shared_with: Some(sp_address.into())
    }))
}

// This need to go
pub fn check_transaction(
    tx: &Transaction,
    blockheight: u32,
    tweak_data: PublicKey,
) -> Result<String> {
    // first check that we haven't scanned this transaction yet
    if let Some((txid, _)) = lock_scanned_transactions()?.0.get_key_value(&tx.txid()) {
        let err_msg = format!("Already scanned tx {}", txid);
        return Err(Error::msg(err_msg));
    }

    let mut connected_user = lock_connected_user()?;

    let txid = tx.txid().to_string();
    if connected_user.try_get_mut_recover()?.update_wallet_with_transaction(tx, blockheight, tweak_data)? > 0 {
        return Ok(txid);
    }

    if connected_user.try_get_mut_main()?.update_wallet_with_transaction(tx, blockheight, tweak_data)? > 0 {
        return Ok(txid);
    }

    if connected_user.try_get_mut_revoke()?.update_wallet_with_transaction(tx, blockheight, tweak_data)? > 0 {
        return Ok(txid);
    }

    return Err(Error::msg("No new outputs found"));
}