From 6a2be13b7c0cef7b83314b3d96fcbefb38ff73ff Mon Sep 17 00:00:00 2001
From: NicolasCantu <nicolas.cantu@pm.me>
Date: Thu, 3 Apr 2025 15:44:52 +0200
Subject: [PATCH] Add outputs to Device

---
 src/device.rs | 109 ++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 106 insertions(+), 3 deletions(-)

diff --git a/src/device.rs b/src/device.rs
index ea0c52e..317aeff 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -1,11 +1,16 @@
+use std::collections::HashMap;
+
 use serde::{Deserialize, Serialize};
 use tsify::Tsify;
 use wasm_bindgen::prelude::*;
 
 use sp_client::{
-    SpClient,
-    bitcoin::OutPoint, 
-    silentpayments::SilentPaymentAddress,
+    bitcoin::{absolute::Height, secp256k1::PublicKey, OutPoint, Transaction, XOnlyPublicKey}, 
+    silentpayments::{
+        utils::receiving::calculate_ecdh_shared_secret, 
+        SilentPaymentAddress
+    }, 
+    OutputSpendStatus, OwnedOutput, SpClient
 };
 
 use crate::pcd::Member;
@@ -14,6 +19,7 @@ use crate::pcd::Member;
 #[tsify(into_wasm_abi, from_wasm_abi)]
 pub struct Device {
     sp_client: SpClient,
+    outputs: HashMap<OutPoint, OwnedOutput>,
     pairing_process_commitment: Option<OutPoint>,
     paired_member: Member,
 }
@@ -24,6 +30,7 @@ impl Device {
         let member = Member::new(vec![SilentPaymentAddress::try_from(local_address).unwrap()]);
         Self {
             sp_client,
+            outputs: HashMap::new(),
             pairing_process_commitment: None,
             paired_member: member,
         }
@@ -37,6 +44,102 @@ impl Device {
         &mut 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 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
+                .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.outputs.get(&input.previous_output) {
+                match &output.spend_status {
+                    OutputSpendStatus::Spent(tx) => {
+                        if *tx == txid.to_string() {
+                            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_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_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 {
+            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 label_str: Option<String>;
+                if let Some(ref l) = label {
+                    label_str = Some(l.as_string());
+                } else {
+                    label_str = None;
+                }
+
+                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_str,
+                    spend_status: OutputSpendStatus::Unspent,
+                };
+                new_outputs.insert(outpoint, owned);
+            }
+        }
+        let mut res = new_outputs.clone();
+        self.outputs.extend(new_outputs);
+
+        let txid = tx.txid().to_string();
+        // update outputs that we own and that are spent
+        for input in tx.input.iter() {
+            if let Some(prevout) = self.outputs.get_mut(&input.previous_output) {
+                // This is spent by this tx
+                prevout.spend_status = OutputSpendStatus::Spent(txid.clone());
+                res.insert(input.previous_output, prevout.clone());
+            }
+        }
+
+        Ok(res)
+    }
+
     pub fn get_pairing_commitment(&self) -> Option<OutPoint> {
         self.pairing_process_commitment.clone()
     }