Make handle_block explicitly return the new decrypted notes

js/pivx_shield.ts

@@ -20,6 +20,7 @@ interface Block {
 
 interface TransactionResult {
   decrypted_notes: [Note, string][];
+  decrypted_new_notes: [Note, string][];
   commitment_tree: string;
   nullifiers: string[];
 }
@@ -309,15 +310,14 @@ export class PIVXShield {
       );
     }
     for (const tx of block.txs) {
-      const { belongToWallet, decryptedNotes } = await this.decryptTransaction(
+      const { belongToWallet, decryptedNewNotes } = await this.addTransaction(
         tx.hex,
       );
-      await this.addTransaction(tx.hex);
       if (belongToWallet) {
         walletTransactions.push(tx.hex);
       }
       // Add all the decryptedNotes to the Nullifier->Note map
-      for (const note of decryptedNotes) {
+      for (const note of decryptedNewNotes) {
         const nullifier = await this.generateNullifierFromNote(note);
         const simplifiedNote = {
           value: note[0].value,
@@ -354,7 +354,7 @@ export class PIVXShield {
     );
   }
   async decryptTransactionOutputs(hex: string) {
-    const { decryptedNotes } = await this.decryptTransaction(hex);
+    const decryptedNotes = await this.decryptTransaction(hex);
     const simplifiedNotes = [];
     for (const [note, _] of decryptedNotes) {
       simplifiedNotes.push({
@@ -374,11 +374,20 @@ export class PIVXShield {
       this.unspentNotes,
     );
     this.commitmentTree = res.commitment_tree;
-    this.unspentNotes = res.decrypted_notes;
+    this.unspentNotes = res.decrypted_notes.concat(res.decrypted_new_notes);
 
     if (res.nullifiers.length > 0) {
       await this.removeSpentNotes(res.nullifiers);
     }
+    // Check if the transaction belongs to the wallet:
+    let belongToWallet = res.decrypted_new_notes.length > 0;
+    for (const nullifier of res.nullifiers) {
+      if (belongToWallet) {
+        break;
+      }
+      belongToWallet = belongToWallet || this.mapNullifierNote.has(nullifier);
+    }
+    return { belongToWallet, decryptedNewNotes: res.decrypted_new_notes };
   }
 
   async decryptTransaction(hex: string) {
@@ -390,15 +399,7 @@ export class PIVXShield {
       this.isTestnet,
       [],
     );
-    // Check if the transaction belongs to the wallet:
-    let belongToWallet = res.decrypted_notes.length > 0;
-    for (const nullifier of res.nullifiers) {
-      if (belongToWallet) {
-        break;
-      }
-      belongToWallet = belongToWallet || this.mapNullifierNote.has(nullifier);
-    }
-    return { belongToWallet, decryptedNotes: res.decrypted_notes };
+    return res.decrypted_new_notes;
   }
 
   /**
@@ -467,7 +468,7 @@ export class PIVXShield {
     if (useShieldInputs) {
       this.pendingSpentNotes.set(txid, nullifiers);
     }
-    const { decryptedNotes } = await this.decryptTransaction(txhex);
+    const decryptedNotes = await this.decryptTransaction(txhex);
     this.pendingUnspentNotes.set(
       txid,
       decryptedNotes.map((n) => n[0]),

src/transaction.rs

@@ -126,6 +126,7 @@ pub async fn load_prover() -> bool {
 #[derive(Serialize, Deserialize)]
 pub struct JSTxSaplingData {
     pub decrypted_notes: Vec<(Note, String)>,
+    pub decrypted_new_notes: Vec<(Note, String)>,
     pub nullifiers: Vec<String>,
     pub commitment_tree: String,
 }
@@ -156,17 +157,21 @@ pub fn handle_transaction(
             (note, IncrementalWitness::read(wit).unwrap())
         })
         .collect::<Vec<_>>();
-    let nullifiers = handle_transaction_internal(&mut tree, tx, &key, true, &mut comp_note)
-        .map_err(|_| "Cannot decode tx")?;
-    let mut ser_comp_note: Vec<(Note, String)> = vec![];
+    let mut new_comp_note: Vec<(Note, IncrementalWitness<Node>)> = vec![];
+    let nullifiers = handle_transaction_internal(
+        &mut tree,
+        tx,
+        &key,
+        true,
+        &mut comp_note,
+        &mut new_comp_note,
+    )
+    .map_err(|_| "Cannot decode tx")?;
+    let mut ser_comp_note: Vec<(Note, String)> =
+        serialize_comp_note(comp_note).map_err(|_| "Cannot serialize notes")?;
+    let mut ser_new_comp_note: Vec<(Note, String)> =
+        serialize_comp_note(new_comp_note).map_err(|_| "Cannot serialize notes")?;
     let mut ser_nullifiers: Vec<String> = vec![];
-    for (note, witness) in comp_note.iter() {
-        let mut buff = Vec::new();
-        witness
-            .write(&mut buff)
-            .map_err(|_| "Cannot write witness to buffer")?;
-        ser_comp_note.push((note.clone(), hex::encode(&buff)));
-    }
 
     for nullif in nullifiers.iter() {
         ser_nullifiers.push(hex::encode(nullif.0));
@@ -178,19 +183,35 @@ pub fn handle_transaction(
 
     let res: JSTxSaplingData = JSTxSaplingData {
         decrypted_notes: ser_comp_note,
+        decrypted_new_notes: ser_new_comp_note,
         nullifiers: ser_nullifiers,
         commitment_tree: hex::encode(buff),
     };
     Ok(serde_wasm_bindgen::to_value(&res).map_err(|_| "Cannot serialize tx output")?)
 }
 
+pub fn serialize_comp_note(
+    comp_note: Vec<(Note, IncrementalWitness<Node>)>,
+) -> Result<Vec<(Note, String)>, Box<dyn Error>> {
+    let mut ser_comp_note: Vec<(Note, String)> = vec![];
+    for (note, witness) in comp_note.iter() {
+        let mut buff = Vec::new();
+        witness
+            .write(&mut buff)
+            .map_err(|_| "Cannot write witness to buffer")?;
+        ser_comp_note.push((note.clone(), hex::encode(&buff)));
+    }
+    Ok(ser_comp_note)
+}
+
 //add a tx to a given commitment tree and the return a witness to each output
 pub fn handle_transaction_internal(
     tree: &mut CommitmentTree<Node>,
     tx: &str,
     key: &UnifiedFullViewingKey,
     is_testnet: bool,
     witnesses: &mut Vec<(Note, IncrementalWitness<Node>)>,
+    new_witnesses: &mut Vec<(Note, IncrementalWitness<Node>)>,
 ) -> Result<Vec<Nullifier>, Box<dyn Error>> {
     let tx = Transaction::read(
         Cursor::new(hex::decode(tx)?),
@@ -210,10 +231,9 @@ pub fn handle_transaction_internal(
         }
 
         for (i, out) in sapling.shielded_outputs().iter().enumerate() {
-            println!("note found!");
             tree.append(Node::from_cmu(out.cmu()))
                 .map_err(|_| "Failed to add cmu to tree")?;
-            for (_, witness) in witnesses.iter_mut() {
+            for (_, witness) in witnesses.iter_mut().chain(new_witnesses.iter_mut()) {
                 witness
                     .append(Node::from_cmu(out.cmu()))
                     .map_err(|_| "Failed to add cmu to witness")?;
@@ -222,7 +242,7 @@ pub fn handle_transaction_internal(
                 if note.index == i {
                     // Save witness
                     let witness = IncrementalWitness::from_tree(tree);
-                    witnesses.push((note.note.clone(), witness));
+                    new_witnesses.push((note.note.clone(), witness));
                 }
             }
         }