From b1ca9721c842fa5e636b24da2575452790e8551f Mon Sep 17 00:00:00 2001
From: aggstam <aggelosstam13@gmail.com>
Date: Sat, 10 Dec 2022 18:50:54 +0200
Subject: [PATCH] drk: mark spent coins on block scanning

---
 bin/drk/src/rpc_blockchain.rs | 11 +++++++++++
 bin/drk/src/rpc_wallet.rs     | 16 ++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/bin/drk/src/rpc_blockchain.rs b/bin/drk/src/rpc_blockchain.rs
index 71b57cb97..6733c6d60 100644
--- a/bin/drk/src/rpc_blockchain.rs
+++ b/bin/drk/src/rpc_blockchain.rs
@@ -132,6 +132,7 @@ impl Drk {
     async fn scan_block(&self, block: &BlockInfo) -> Result<()> {
         eprintln!("Iterating over {} transactions", block.txs.len());
 
+        let mut nullifiers: Vec<Nullifier> = vec![];
         let mut outputs: Vec<Output> = vec![];
 
         // TODO: FIXME: This shouldn't be hardcoded here obviously.
@@ -143,6 +144,9 @@ impl Drk {
                 {
                     eprintln!("Found Money::Transfer in call {} in tx {}", j, i);
                     let params: MoneyTransferParams = deserialize(&call.data[1..])?;
+                    for input in params.inputs {
+                        nullifiers.push(input.nullifier);
+                    }
                     for output in params.outputs {
                         outputs.push(output);
                     }
@@ -152,6 +156,9 @@ impl Drk {
                 if call.contract_id == contract_id && call.data[0] == MoneyFunction::OtcSwap as u8 {
                     eprintln!("Found Money::OtcSwap in call {} in tx {}", j, i);
                     let params: MoneyTransferParams = deserialize(&call.data[1..])?;
+                    for input in params.inputs {
+                        nullifiers.push(input.nullifier);
+                    }
                     for output in params.outputs {
                         outputs.push(output);
                     }
@@ -205,6 +212,10 @@ impl Drk {
         self.put_tree(&tree).await?;
         eprintln!("Merkle tree written successfully");
 
+        eprintln!("Marking spent coins");
+        self.mark_spent_coins(nullifiers).await?;
+        eprintln!("Spent coins marked successfully");
+
         // This is the SQL query we'll be executing to insert coins into the wallet
         let query = format!(
             "INSERT INTO {} ({}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12);",
diff --git a/bin/drk/src/rpc_wallet.rs b/bin/drk/src/rpc_wallet.rs
index b8f10078c..70573455c 100644
--- a/bin/drk/src/rpc_wallet.rs
+++ b/bin/drk/src/rpc_wallet.rs
@@ -446,6 +446,22 @@ impl Drk {
         Ok(())
     }
 
+    /// Marks all coins in the wallet as spent, if their nullifier is
+    /// in the provided set
+    pub async fn mark_spent_coins(&self, nullifiers: Vec<Nullifier>) -> Result<()> {
+        if nullifiers.is_empty() {
+            return Ok(())
+        }
+
+        for (coin, _) in self.wallet_coins(false).await? {
+            if nullifiers.contains(&coin.nullifier) {
+                self.mark_spent_coin(&coin.coin).await?;
+            }
+        }
+
+        Ok(())
+    }
+
     /// Mark a given coin in the wallet as unspent
     pub async fn unspend_coin(&self, coin: &Coin) -> Result<()> {
         let query = format!(