Merge pull request #81 from PIVX-Labs/nullifier_map

Shield Activity Support

js/pivx_shield.ts

@@ -88,6 +88,14 @@ export class PIVXShield {
    */
   private pendingUnspentNotes: Map<string, Note[]> = new Map();
 
+  /**
+   *
+   * @private
+   * Map nullifier -> Note
+   * It contains all notes in the history of the wallet, both spent and unspent
+   */
+  private mapNullifierNote: Map<string, SimplifiedNote> = new Map();
+
   private promises: Map<
     string,
     { res: (...args: any) => void; rej: (...args: any) => void }
@@ -248,6 +256,7 @@ export class PIVXShield {
       diversifierIndex: this.diversifierIndex,
       unspentNotes: this.unspentNotes,
       isTestnet: this.isTestnet,
+      mapNullifierNote: Object.fromEntries(this.mapNullifierNote),
     });
   }
   /**
@@ -273,46 +282,123 @@ export class PIVXShield {
       shieldData.lastProcessedBlock,
       shieldData.commitmentTree,
     );
+    pivxShield.mapNullifierNote = new Map(
+      Object.entries(shieldData.mapNullifierNote ?? {}),
+    );
     pivxShield.diversifierIndex = shieldData.diversifierIndex;
     pivxShield.unspentNotes = shieldData.unspentNotes;
-    return pivxShield;
+
+    // Shield activity update: mapNullifierNote must be present in the shieldData
+    let success = true;
+    if (!shieldData.mapNullifierNote) {
+      success = false;
+    }
+    return { pivxShield, success };
   }
 
   /**
    * Loop through the txs of a block and update useful shield data
    * @param block - block outputted from any PIVX node
+   * @returns list of transactions belonging to the wallet
    */
   async handleBlock(block: Block) {
+    let walletTransactions: string[] = [];
     if (this.lastProcessedBlock > block.height) {
       throw new Error(
         "Blocks must be processed in a monotonically increasing order!",
       );
     }
     for (const tx of block.txs) {
+      const { belongToWallet, decryptedNotes } = await this.decryptTransaction(
+        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) {
+        const nullifier = await this.generateNullifierFromNote(note);
+        const simplifiedNote = {
+          value: note[0].value,
+          recipient: await this.getShieldAddressFromNote(note[0]),
+        };
+        this.mapNullifierNote.set(nullifier, simplifiedNote);
+      }
+      // Delete the corresponding pending transaction
       this.pendingUnspentNotes.delete(tx.txid);
     }
     this.lastProcessedBlock = block.height;
+    return walletTransactions;
   }
 
-  async addTransaction(hex: string, decryptOnly = false) {
+  /**
+   *
+   * @param note - Note and its corresponding witness
+   * Generate the nullifier for a given pair note, witness
+   */
+  private async generateNullifierFromNote(note: [Note, String]) {
+    return await this.callWorker<string>(
+      "get_nullifier_from_note",
+      note,
+      this.extfvk,
+      this.isTestnet,
+    );
+  }
+
+  private async getShieldAddressFromNote(note: Note) {
+    return await this.callWorker<string>(
+      "encode_payment_address",
+      this.isTestnet,
+      note.recipient,
+    );
+  }
+  async decryptTransactionOutputs(hex: string) {
+    const { decryptedNotes } = await this.decryptTransaction(hex);
+    const simplifiedNotes = [];
+    for (const [note, _] of decryptedNotes) {
+      simplifiedNotes.push({
+        value: note.value,
+        recipient: await this.getShieldAddressFromNote(note),
+      });
+    }
+    return simplifiedNotes;
+  }
+  async addTransaction(hex: string) {
     const res = await this.callWorker<TransactionResult>(
       "handle_transaction",
       this.commitmentTree,
       hex,
       this.extfvk,
       this.isTestnet,
-      decryptOnly ? [] : this.unspentNotes,
+      this.unspentNotes,
     );
-    if (!decryptOnly) {
-      this.commitmentTree = res.commitment_tree;
-      this.unspentNotes = res.decrypted_notes;
+    this.commitmentTree = res.commitment_tree;
+    this.unspentNotes = res.decrypted_notes;
 
-      if (res.nullifiers.length > 0) {
-        await this.removeSpentNotes(res.nullifiers);
+    if (res.nullifiers.length > 0) {
+      await this.removeSpentNotes(res.nullifiers);
+    }
+  }
+
+  async decryptTransaction(hex: string) {
+    const res = await this.callWorker<TransactionResult>(
+      "handle_transaction",
+      this.commitmentTree,
+      hex,
+      this.extfvk,
+      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 res.decrypted_notes;
+    return { belongToWallet, decryptedNotes: res.decrypted_notes };
   }
 
   /**
@@ -381,16 +467,10 @@ export class PIVXShield {
     if (useShieldInputs) {
       this.pendingSpentNotes.set(txid, nullifiers);
     }
-    const decryptedNewNotes = (await this.addTransaction(txhex, true)).filter(
-      (note) =>
-        !this.unspentNotes.some(
-          (note2) => JSON.stringify(note2[0]) === JSON.stringify(note[0]),
-        ),
-    );
-
+    const { decryptedNotes } = await this.decryptTransaction(txhex);
     this.pendingUnspentNotes.set(
       txid,
-      decryptedNewNotes.map((n) => n[0]),
+      decryptedNotes.map((n) => n[0]),
     );
     return {
       hex: txhex,
@@ -465,6 +545,13 @@ export class PIVXShield {
     return this.lastProcessedBlock;
   }
 
+  /**
+   * @param nullifier - A sapling nullifier
+   * @returns the Note corresponding to a given nullifier
+   */
+  getNoteFromNullifier(nullifier: string) {
+    return this.mapNullifierNote.get(nullifier);
+  }
   /**
    * @returns sapling root
    */
@@ -487,6 +574,7 @@ export class PIVXShield {
     this.unspentNotes = [];
     this.pendingSpentNotes = new Map();
     this.pendingUnspentNotes = new Map();
+    this.mapNullifierNote = new Map();
   }
 }
 
@@ -504,6 +592,11 @@ export interface Note {
   rseed: number[];
 }
 
+export interface SimplifiedNote {
+  recipient: string;
+  value: number;
+}
+
 export interface ShieldData {
   extfvk: string;
   lastProcessedBlock: number;

src/keys.rs

@@ -71,7 +71,7 @@ pub fn encode_extsk(extsk: &ExtendedSpendingKey, is_testnet: bool) -> String {
     encoding::encode_extended_spending_key(enc_str, extsk)
 }
 
-pub fn encode_payment_address(addr: &PaymentAddress, is_testnet: bool) -> String {
+pub fn encode_payment_address_internal(addr: &PaymentAddress, is_testnet: bool) -> String {
     let enc_str: &str = if is_testnet {
         TEST_NETWORK.hrp_sapling_payment_address()
     } else {
@@ -80,6 +80,23 @@ pub fn encode_payment_address(addr: &PaymentAddress, is_testnet: bool) -> String
     encoding::encode_payment_address(enc_str, addr)
 }
 
+#[wasm_bindgen]
+pub fn encode_payment_address(
+    is_testnet: bool,
+    ser_payment_address: &[u8],
+) -> Result<JsValue, JsValue> {
+    let enc_payment_address = encode_payment_address_internal(
+        &PaymentAddress::from_bytes(
+            &ser_payment_address
+                .try_into()
+                .map_err(|_| "Bad ser_payment_address")?,
+        )
+        .ok_or("Failed to deserialize payment address")?,
+        is_testnet,
+    );
+    Ok(serde_wasm_bindgen::to_value(&enc_payment_address)?)
+}
+
 pub fn decode_extended_full_viewing_key(
     enc_extfvk: &str,
     is_testnet: bool,
@@ -136,7 +153,7 @@ pub fn generate_default_payment_address(
         decode_extended_full_viewing_key(&enc_extfvk, is_testnet).map_err(|e| e.to_string())?;
     let (def_index, def_address) = extfvk.to_diversifiable_full_viewing_key().default_address();
     Ok(serde_wasm_bindgen::to_value(&NewAddress {
-        address: encode_payment_address(&def_address, is_testnet),
+        address: encode_payment_address_internal(&def_address, is_testnet),
         diversifier_index: def_index.0.to_vec(),
     })?)
 }
@@ -163,7 +180,7 @@ pub fn generate_next_shielding_payment_address(
         .ok_or("No valid indeces left")?; // There are so many valid addresses this should never happen
 
     Ok(serde_wasm_bindgen::to_value(&NewAddress {
-        address: encode_payment_address(&address, is_testnet),
+        address: encode_payment_address_internal(&address, is_testnet),
         diversifier_index: new_index.0.to_vec(),
     })?)
 }

src/transaction.rs

@@ -25,6 +25,7 @@ use pivx_primitives::transaction::components::{OutPoint, TxOut};
 pub use pivx_primitives::transaction::fees::fixed::FeeRule;
 pub use pivx_primitives::transaction::Transaction;
 pub use pivx_primitives::zip32::AccountId;
+use pivx_primitives::zip32::ExtendedFullViewingKey;
 pub use pivx_primitives::zip32::ExtendedSpendingKey;
 pub use pivx_primitives::zip32::Scope;
 pub use pivx_proofs::prover::LocalTxProver;
@@ -264,6 +265,36 @@ pub fn remove_spent_notes(
     Ok(serde_wasm_bindgen::to_value(&unspent_notes)?)
 }
 
+#[wasm_bindgen]
+pub fn get_nullifier_from_note(
+    note_data: JsValue,
+    enc_extfvk: String,
+    is_testnet: bool,
+) -> Result<JsValue, JsValue> {
+    let extfvk =
+        decode_extended_full_viewing_key(&enc_extfvk, is_testnet).map_err(|e| e.to_string())?;
+    let (note, hex_witness): (Note, String) = serde_wasm_bindgen::from_value(note_data)?;
+    let ser_nullifiers =
+        get_nullifier_from_note_internal(extfvk, note, hex_witness).map_err(|e| e.to_string())?;
+    Ok(serde_wasm_bindgen::to_value(&ser_nullifiers)?)
+}
+
+pub fn get_nullifier_from_note_internal(
+    extfvk: ExtendedFullViewingKey,
+    note: Note,
+    hex_witness: String,
+) -> Result<String, Box<dyn Error>> {
+    let nullif_key = extfvk
+        .to_diversifiable_full_viewing_key()
+        .to_nk(Scope::External);
+    let witness = Cursor::new(hex::decode(hex_witness).map_err(|e| e.to_string())?);
+    let path = IncrementalWitness::<Node>::read(witness)
+        .map_err(|_| "Cannot read witness from buffer")?
+        .path()
+        .ok_or("Cannot find witness path")?;
+    Ok(hex::encode(note.nf(&nullif_key, path.position).0))
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct JSTransaction {
     pub txid: String,

src/transaction/test.rs

@@ -1,6 +1,6 @@
 #![cfg(test)]
 
-use crate::transaction::create_transaction_internal;
+use crate::transaction::{create_transaction_internal, get_nullifier_from_note_internal};
 
 use super::handle_transaction_internal;
 use either::Either;
@@ -15,6 +15,7 @@ use pivx_primitives::sapling::value::NoteValue;
 use pivx_primitives::sapling::Node;
 use pivx_primitives::sapling::Note;
 use pivx_primitives::sapling::Rseed::BeforeZip212;
+use pivx_primitives::zip32::Scope;
 use std::error::Error;
 use std::io::Cursor;
 
@@ -92,7 +93,7 @@ pub async fn test_create_transaction() -> Result<(), Box<dyn Error>> {
     path.write(&mut path_vec)?;
     let path = hex::encode(path_vec);
     let tx = create_transaction_internal(
-        Either::Left(vec![(note.clone(), path)]),
+        Either::Left(vec![(note.clone(), path.clone())]),
         &extended_spending_key,
         output,
         address,
@@ -108,7 +109,15 @@ pub async fn test_create_transaction() -> Result<(), Box<dyn Error>> {
         nullifier,
         "5269442d8022af933774f9f22775566d92089a151ba733f6d751f5bb65a7f56d"
     );
-    // When we implement mempool, test that new notes work correctly
+    // Verify that get_nullifier_from_note_internal yields the same nullifier
+    assert_eq!(
+        get_nullifier_from_note_internal(
+            extended_spending_key.to_extended_full_viewing_key(),
+            note.clone(),
+            path
+        )?,
+        "5269442d8022af933774f9f22775566d92089a151ba733f6d751f5bb65a7f56d"
+    );
 
     Ok(())
 }