bin/drk: revert to specific height functionality added

bin/drk/src/cli_util.rs

@@ -362,7 +362,7 @@ pub fn generate_completions(shell: &str) -> Result<()> {
     // Scan
     let reset = Arg::with_name("reset")
         .long("reset")
-        .help("Reset Merkle tree and start scanning from first block");
+        .help("Reset wallet state to provided block height and start scanning");
 
     let scan = SubCommand::with_name("scan")
         .about("Scan the blockchain and parse relevant transactions")

bin/drk/src/lib.rs

@@ -107,6 +107,23 @@ impl Drk {
         Ok(())
     }
 
+    /// Auxiliary function to completely reset wallet state.
+    pub async fn reset(&self) -> WalletDbResult<()> {
+        println!("Resetting full wallet state");
+        self.reset_scanned_blocks()?;
+        self.reset_money_tree().await?;
+        self.reset_money_smt()?;
+        self.reset_money_coins()?;
+        self.reset_mint_authorities()?;
+        self.reset_dao_trees().await?;
+        self.reset_daos().await?;
+        self.reset_dao_proposals().await?;
+        self.reset_dao_votes()?;
+        self.reset_tx_history()?;
+        println!("Successfully reset full wallet state");
+        Ok(())
+    }
+
     /// Auxiliary function to reset `walletdb` inverse cache state.
     /// Additionally, set current trees state inverse queries.
     /// We keep the entire trees state as two distinct inverse queries,

bin/drk/src/main.rs

@@ -216,8 +216,8 @@ enum Subcmd {
     /// Scan the blockchain and parse relevant transactions
     Scan {
         #[structopt(long)]
-        /// Reset Merkle tree and start scanning from first block
-        reset: bool,
+        /// Reset wallet state to provided block height and start scanning
+        reset: Option<u32>,
     },
 
     /// Explorer related subcommands
@@ -1827,18 +1827,14 @@ async fn realmain(args: Args, ex: Arc<smol::Executor<'static>>) -> Result<()> {
             )
             .await;
 
-            if reset {
-                println!("Reset requested.");
-                if let Err(e) = drk.scan_blocks(true).await {
-                    eprintln!("Failed during scanning: {e:?}");
+            if let Some(height) = reset {
+                if let Err(e) = drk.reset_to_height(height).await {
+                    eprintln!("Failed during wallet reset: {e:?}");
                     exit(2);
                 }
-                println!("Finished scanning blockchain");
-
-                return drk.stop_rpc_client().await
             }
 
-            if let Err(e) = drk.scan_blocks(false).await {
+            if let Err(e) = drk.scan_blocks().await {
                 eprintln!("Failed during scanning: {e:?}");
                 exit(2);
             }

bin/drk/src/rpc.rs

@@ -62,7 +62,7 @@ impl Drk {
 
         // Handle genesis(0) block
         if last_known == 0 {
-            if let Err(e) = self.scan_blocks(true).await {
+            if let Err(e) = self.scan_blocks().await {
                 return Err(Error::DatabaseError(format!(
                     "[subscribe_blocks] Scanning from genesis block failed: {e:?}"
                 )))
@@ -252,31 +252,18 @@ impl Drk {
         Ok(())
     }
 
-    /// Scans the blockchain starting from the last scanned block, for relevant
-    /// money transfer transactions. If reset flag is provided, Merkle tree state
-    /// and coins are reset, and start scanning from beginning. Alternatively,
-    /// it looks for a checkpoint in the wallet to reset and start scanning from.
-    pub async fn scan_blocks(&self, reset: bool) -> WalletDbResult<()> {
+    /// Scans the blockchain for wallet relevant transactions,
+    /// starting from the last scanned block.
+    pub async fn scan_blocks(&self) -> WalletDbResult<()> {
         // Grab last scanned block height
         let (mut height, _) = self.get_last_scanned_block()?;
-        // If last scanned block is genesis (0) or reset flag
-        // has been provided we reset, otherwise continue with
-        // the next block height
-        if height == 0 || reset {
-            self.reset_scanned_blocks()?;
-            self.reset_money_tree().await?;
-            self.reset_money_smt()?;
-            self.reset_money_coins()?;
-            self.reset_mint_authorities()?;
-            self.reset_dao_trees().await?;
-            self.reset_daos().await?;
-            self.reset_dao_proposals().await?;
-            self.reset_dao_votes()?;
-            self.reset_tx_history()?;
-            height = 0;
+        // If last scanned block is genesis(0) we reset,
+        // otherwise continue with the next block height.
+        if height == 0 {
+            self.reset().await?;
         } else {
             height += 1;
-        };
+        }
 
         loop {
             let rep = match self

bin/drk/src/scanned_blocks.rs

@@ -18,8 +18,6 @@
 
 use rusqlite::types::Value;
 
-use darkfi::{Error, Result};
-
 use crate::{
     convert_named_params,
     error::{WalletDbError, WalletDbResult},
@@ -52,35 +50,26 @@ impl Drk {
     }
 
     /// Get a scanned block information record.
-    pub fn get_scanned_block_record(&self, height: u32) -> Result<(u32, String, String)> {
-        let row = match self.wallet.query_single(
+    pub fn get_scanned_block_record(&self, height: u32) -> WalletDbResult<(u32, String, String)> {
+        let row = self.wallet.query_single(
             WALLET_SCANNED_BLOCKS_TABLE,
             &[],
             convert_named_params! {(WALLET_SCANNED_BLOCKS_COL_HEIGH, height)},
-        ) {
-            Ok(r) => r,
-            Err(e) => {
-                return Err(Error::DatabaseError(format!(
-                    "[get_scanned_block_record] Scanned block information record retrieval failed: {e:?}"
-                )))
-            }
-        };
+        )?;
 
         let Value::Integer(height) = row[0] else {
-            return Err(Error::ParseFailed("[get_scanned_block_record] Block height parsing failed"))
+            return Err(WalletDbError::ParseColumnValueError);
         };
         let Ok(height) = u32::try_from(height) else {
-            return Err(Error::ParseFailed("[get_scanned_block_record] Block height parsing failed"))
+            return Err(WalletDbError::ParseColumnValueError);
         };
 
         let Value::Text(ref hash) = row[1] else {
-            return Err(Error::ParseFailed("[get_scanned_block_record] Hash parsing failed"))
+            return Err(WalletDbError::ParseColumnValueError);
         };
 
         let Value::Text(ref rollback_query) = row[2] else {
-            return Err(Error::ParseFailed(
-                "[get_scanned_block_record] Rollback query parsing failed",
-            ))
+            return Err(WalletDbError::ParseColumnValueError);
         };
 
         Ok((height, hash.clone(), rollback_query.clone()))
@@ -125,4 +114,40 @@ impl Drk {
 
         Ok(())
     }
+
+    /// Reset state to provided block height.
+    /// If genesis block height(0) was provided, perform a full reset.
+    pub async fn reset_to_height(&self, height: u32) -> WalletDbResult<()> {
+        println!("Resetting wallet state to block: {height}");
+
+        // If genesis block height(0) was provided,
+        // perform a full reset.
+        if height == 0 {
+            return self.reset().await
+        }
+
+        // Grab last scanned block height
+        let (last, _) = self.get_last_scanned_block()?;
+
+        // Check if requested height is after it
+        if last <= height {
+            println!("Requested block height is greater or equal to last scanned block");
+            return Ok(())
+        }
+
+        // Iterate the range (height, last] in reverse to grab the corresponding blocks
+        for height in (height + 1..=last).rev() {
+            let (height, hash, query) = self.get_scanned_block_record(height)?;
+            println!("Reverting block: {height} - {hash}");
+            self.wallet.exec_batch_sql(&query)?;
+            let query = format!(
+                "DELETE FROM {} WHERE {} = {};",
+                WALLET_SCANNED_BLOCKS_TABLE, WALLET_SCANNED_BLOCKS_COL_HEIGH, height
+            );
+            self.wallet.exec_batch_sql(&query)?;
+        }
+
+        println!("Successfully reset wallet state");
+        Ok(())
+    }
 }