zcash_client_sqlite: Add methods for fixing broken note commitment tr…

…ees.

Cargo.toml

@@ -198,3 +198,6 @@ unexpected_cfgs = { level = "warn", check-cfg = ['cfg(zcash_unstable, values("zf
 orchard = { git = "https://github.com/zcash/orchard.git", rev = "b1c22c07300db22239235d16dab096e23369948f" }
 redjubjub = { git = "https://github.com/ZcashFoundation/redjubjub", rev = "eae848c5c14d9c795d000dd9f4c4762d1aee7ee1" }
 sapling = { package = "sapling-crypto", git = "https://github.com/zcash/sapling-crypto.git", rev = "6ca338532912adcd82369220faeea31aab4720c5" }
+incrementalmerkletree = { path = "../../incrementalmerkletree/canon/incrementalmerkletree/" }
+incrementalmerkletree-testing = { path = "../../incrementalmerkletree/canon/incrementalmerkletree-testing/" }
+shardtree = { path = "../../incrementalmerkletree/canon/shardtree/" }

zcash_client_sqlite/CHANGELOG.md

@@ -9,6 +9,8 @@ and this library adheres to Rust's notion of
 
 ### Added
 - `zcash_client_sqlite::WalletDb::from_connection`
+- `zcash_client_sqlite::WalletDb::check_witnesses`
+- `zcash_client_sqlite::WalletDb::queue_rescans`
 
 ### Changed
 - MSRV is now 1.81.0.

zcash_client_sqlite/src/lib.rs

@@ -39,6 +39,7 @@ use secrecy::{ExposeSecret, SecretVec};
 use shardtree::{error::ShardTreeError, ShardTree};
 use std::{
     borrow::{Borrow, BorrowMut},
+    cmp::{max, min},
     collections::HashMap,
     convert::AsRef,
     fmt,
@@ -140,6 +141,7 @@ pub mod wallet;
 use wallet::{
     commitment_tree::{self, put_shard_roots},
     common::spendable_notes_meta,
+    scanning::replace_queue_entries,
     SubtreeProgressEstimator,
 };
 
@@ -302,6 +304,49 @@ impl<C: BorrowMut<Connection>, P: consensus::Parameters + Clone> WalletDb<C, P>
         tx.commit()?;
         Ok(result)
     }
+
+    /// Attempts to construct a witness for each note belonging to the wallet that is believed by
+    /// the wallet to currently be spendable, and returns a vector of .
+    ///
+    /// This method is intended for repairing wallets that broke due to bugs in `shardtree`.
+    ///
+    /// Returns a vector of the ranges that must be rescanned in order to correct missing witness
+    /// data.
+    pub fn check_witnesses(&mut self) -> Result<Vec<Range<BlockHeight>>, SqliteClientError> {
+        self.transactionally(|wdb| wallet::commitment_tree::check_witnesses(&wdb.conn.0))
+    }
+
+    /// Updates the scan queue by inserting scan ranges for the given range of block heights, with
+    /// the specified scanning priority.
+    pub fn queue_rescans(
+        &mut self,
+        rescan_ranges: NonEmpty<Range<BlockHeight>>,
+        priority: ScanPriority,
+    ) -> Result<(), SqliteClientError> {
+        let query_range = rescan_ranges
+            .iter()
+            .fold(None, |acc: Option<Range<BlockHeight>>, scan_range| {
+                if let Some(range) = acc {
+                    Some(min(range.start, scan_range.start)..max(range.end, scan_range.end))
+                } else {
+                    Some(scan_range.clone())
+                }
+            })
+            .expect("rescan_ranges is nonempty");
+
+        self.transactionally::<_, _, SqliteClientError>(|wdb| {
+            replace_queue_entries(
+                wdb.conn.0,
+                &query_range,
+                rescan_ranges
+                    .into_iter()
+                    .map(|r| ScanRange::from_parts(r, priority)),
+                true,
+            )
+        })?;
+
+        Ok(())
+    }
 }
 
 impl<C: Borrow<rusqlite::Connection>, P: consensus::Parameters> InputSource for WalletDb<C, P> {

zcash_client_sqlite/src/wallet.rs

@@ -80,7 +80,7 @@ use std::collections::{HashMap, HashSet};
 use std::convert::TryFrom;
 use std::io::{self, Cursor};
 use std::num::NonZeroU32;
-use std::ops::RangeInclusive;
+use std::ops::{Range, RangeInclusive};
 
 use tracing::{debug, warn};
 
@@ -3628,6 +3628,41 @@ pub(crate) fn prune_nullifier_map(
     Ok(())
 }
 
+pub(crate) fn get_block_range(
+    conn: &rusqlite::Connection,
+    protocol: ShieldedProtocol,
+    commitment_tree_address: incrementalmerkletree::Address,
+) -> Result<Option<Range<BlockHeight>>, SqliteClientError> {
+    let prefix = match protocol {
+        ShieldedProtocol::Sapling => "sapling",
+        ShieldedProtocol::Orchard => "orchard",
+    };
+    let mut stmt = conn.prepare_cached(&format!(
+        "SELECT MIN(height), MAX(height)
+         FROM blocks
+         WHERE {prefix}_commitment_tree_size BETWEEN :min_tree_size AND :max_tree_size"
+    ))?;
+
+    stmt.query_row(
+        named_params! {
+            ":min_tree_size": u64::from(commitment_tree_address.position_range_start()) + 1,
+            ":max_tree_size": u64::from(commitment_tree_address.position_range_start()) + 1,
+        },
+        |row| {
+            let min_height = row
+                .get::<_, u32>(0)
+                .map(|h| BlockHeight::from_u32(h.saturating_sub(1)))?;
+            let max_height = row
+                .get::<_, u32>(1)
+                .map(|h| BlockHeight::from_u32(h.saturating_add(1)))?;
+
+            Ok(min_height..max_height)
+        },
+    )
+    .optional()
+    .map_err(SqliteClientError::from)
+}
+
 #[cfg(any(test, feature = "test-dependencies"))]
 pub mod testing {
     use incrementalmerkletree::Position;

zcash_client_sqlite/src/wallet/commitment_tree.rs

@@ -11,17 +11,22 @@ use std::{
 
 use incrementalmerkletree::{Address, Hashable, Level, Position, Retention};
 use shardtree::{
-    error::ShardTreeError,
+    error::{QueryError, ShardTreeError},
     store::{Checkpoint, ShardStore, TreeState},
-    LocatedPrunableTree, LocatedTree, PrunableTree, RetentionFlags,
+    LocatedPrunableTree, LocatedTree, PrunableTree, RetentionFlags, ShardTree,
 };
 
 use zcash_client_backend::{
-    data_api::chain::CommitmentTreeRoot,
+    data_api::{chain::CommitmentTreeRoot, SAPLING_SHARD_HEIGHT},
     serialization::shardtree::{read_shard, write_shard},
 };
 use zcash_primitives::merkle_tree::HashSer;
-use zcash_protocol::consensus::BlockHeight;
+use zcash_protocol::{consensus::BlockHeight, ShieldedProtocol};
+
+use crate::{error::SqliteClientError, PRUNING_DEPTH, SAPLING_TABLES_PREFIX};
+
+#[cfg(feature = "orchard")]
+use {crate::ORCHARD_TABLES_PREFIX, zcash_client_backend::data_api::ORCHARD_SHARD_HEIGHT};
 
 /// Errors that can appear in SQLite-back [`ShardStore`] implementation operations.
 #[derive(Debug)]
@@ -1112,6 +1117,74 @@ pub(crate) fn put_shard_roots<
     Ok(())
 }
 
+pub(crate) fn check_witnesses(
+    conn: &rusqlite::Transaction<'_>,
+) -> Result<Vec<Range<BlockHeight>>, SqliteClientError> {
+    let chain_tip_height =
+        super::chain_tip_height(conn)?.ok_or(SqliteClientError::ChainHeightUnknown)?;
+    let wallet_birthday = super::wallet_birthday(conn)?.ok_or(SqliteClientError::AccountUnknown)?;
+    let unspent_sapling_note_meta =
+        super::sapling::select_unspent_note_meta(conn, chain_tip_height, wallet_birthday)?;
+
+    let mut scan_ranges = vec![];
+    let mut sapling_incomplete = vec![];
+    let sapling_tree =
+        ShardTree::<_, { sapling::NOTE_COMMITMENT_TREE_DEPTH }, SAPLING_SHARD_HEIGHT>::new(
+            SqliteShardStore::<_, sapling::Node, SAPLING_SHARD_HEIGHT>::from_connection(
+                conn,
+                SAPLING_TABLES_PREFIX,
+            )
+            .map_err(|e| ShardTreeError::Storage(Error::Query(e)))?,
+            PRUNING_DEPTH.try_into().unwrap(),
+        );
+    for m in unspent_sapling_note_meta.iter() {
+        match sapling_tree.witness_at_checkpoint_depth(m.commitment_tree_position(), 0) {
+            Ok(_) => {}
+            Err(ShardTreeError::Query(QueryError::TreeIncomplete(mut addrs))) => {
+                sapling_incomplete.append(&mut addrs);
+            }
+            Err(other) => {
+                return Err(SqliteClientError::CommitmentTree(other));
+            }
+        }
+    }
+
+    for addr in sapling_incomplete {
+        let range = super::get_block_range(conn, ShieldedProtocol::Sapling, addr)?;
+        scan_ranges.extend(range.into_iter());
+    }
+
+    #[cfg(feature = "orchard")]
+    {
+        let unspent_orchard_note_meta =
+            super::orchard::select_unspent_note_meta(conn, chain_tip_height, wallet_birthday)?;
+        let mut orchard_incomplete = vec![];
+        let orchard_tree = ShardTree::<_, {orchard::NOTE_COMMITMENT_TREE_DEPTH as u8}, ORCHARD_SHARD_HEIGHT>::new(
+        SqliteShardStore::<_, orchard::tree::MerkleHashOrchard, ORCHARD_SHARD_HEIGHT>::from_connection(conn, ORCHARD_TABLES_PREFIX)
+            .map_err(|e| ShardTreeError::Storage(Error::Query(e)))?,
+        PRUNING_DEPTH.try_into().unwrap(),
+    );
+        for m in unspent_orchard_note_meta.iter() {
+            match orchard_tree.witness_at_checkpoint_depth(m.commitment_tree_position(), 0) {
+                Ok(_) => {}
+                Err(ShardTreeError::Query(QueryError::TreeIncomplete(mut addrs))) => {
+                    orchard_incomplete.append(&mut addrs);
+                }
+                Err(other) => {
+                    return Err(SqliteClientError::CommitmentTree(other));
+                }
+            }
+        }
+
+        for addr in orchard_incomplete {
+            let range = super::get_block_range(conn, ShieldedProtocol::Orchard, addr)?;
+            scan_ranges.extend(range.into_iter());
+        }
+    }
+
+    Ok(scan_ranges)
+}
+
 #[cfg(test)]
 mod tests {
     use tempfile::NamedTempFile;

zcash_client_sqlite/src/wallet/common.rs

@@ -1,5 +1,6 @@
 //! Functions common to Sapling and Orchard support in the wallet.
 
+use incrementalmerkletree::Position;
 use rusqlite::{named_params, types::Value, Connection, Row};
 use std::{num::NonZeroU64, rc::Rc};
 
@@ -236,6 +237,98 @@ where
         .collect::<Result<_, _>>()
 }
 
+#[allow(dead_code)]
+pub(crate) struct UnspentNoteMeta {
+    note_id: ReceivedNoteId,
+    txid: TxId,
+    output_index: u32,
+    commitment_tree_position: Position,
+    value: Zatoshis,
+}
+
+#[allow(dead_code)]
+impl UnspentNoteMeta {
+    pub(crate) fn note_id(&self) -> ReceivedNoteId {
+        self.note_id
+    }
+
+    pub(crate) fn txid(&self) -> TxId {
+        self.txid
+    }
+
+    pub(crate) fn output_index(&self) -> u32 {
+        self.output_index
+    }
+
+    pub(crate) fn commitment_tree_position(&self) -> Position {
+        self.commitment_tree_position
+    }
+
+    pub(crate) fn value(&self) -> Zatoshis {
+        self.value
+    }
+}
+
+pub(crate) fn select_unspent_note_meta(
+    conn: &rusqlite::Connection,
+    protocol: ShieldedProtocol,
+    chain_tip_height: BlockHeight,
+    wallet_birthday: BlockHeight,
+) -> Result<Vec<UnspentNoteMeta>, SqliteClientError> {
+    let (table_prefix, index_col, _) = per_protocol_names(protocol);
+    let mut stmt = conn.prepare_cached(&format!("
+        SELECT {table_prefix}_received_notes.id AS id, txid, {index_col}, 
+               commitment_tree_position, value
+        FROM {table_prefix}_received_notes
+        INNER JOIN transactions
+           ON transactions.id_tx = {table_prefix}_received_notes.tx
+        WHERE value > 5000 -- FIXME #1316, allow selection of dust inputs
+        AND recipient_key_scope IS NOT NULL
+        AND nf IS NOT NULL
+        AND commitment_tree_position IS NOT NULL
+        AND {table_prefix}_received_notes.id NOT IN (
+          SELECT {table_prefix}_received_note_id
+          FROM {table_prefix}_received_note_spends
+          JOIN transactions stx ON stx.id_tx = transaction_id
+          WHERE stx.block IS NOT NULL -- the spending tx is mined
+          OR stx.expiry_height IS NULL -- the spending tx will not expire
+          OR stx.expiry_height > :anchor_height -- the spending tx is unexpired
+        )
+        AND NOT EXISTS (
+           SELECT 1 FROM v_{table_prefix}_shard_unscanned_ranges unscanned
+           -- select all the unscanned ranges involving the shard containing this note
+           WHERE {table_prefix}_received_notes.commitment_tree_position >= unscanned.start_position
+           AND {table_prefix}_received_notes.commitment_tree_position < unscanned.end_position_exclusive
+           -- exclude unscanned ranges that start above the anchor height (they don't affect spendability)
+           AND unscanned.block_range_start <= :anchor_height
+           -- exclude unscanned ranges that end below the wallet birthday
+           AND unscanned.block_range_end > :wallet_birthday
+        )
+    "))?;
+
+    let res = stmt
+        .query_and_then::<_, SqliteClientError, _, _>(
+            named_params![
+                ":anchor_height": u32::from(chain_tip_height),
+                ":wallet_birthday": u32::from(wallet_birthday),
+            ],
+            |row| {
+                Ok(UnspentNoteMeta {
+                    note_id: row.get("id").map(|id| ReceivedNoteId(protocol, id))?,
+                    txid: row.get("txid").map(TxId::from_bytes)?,
+                    output_index: row.get(index_col)?,
+                    commitment_tree_position: row
+                        .get::<_, u64>("commitment_tree_position")
+                        .map(Position::from)?,
+                    value: Zatoshis::from_nonnegative_i64(row.get("value")?)?,
+                })
+            },
+        )?
+        .collect::<Result<Vec<_>, _>>()?;
+
+    Ok(res)
+}
+
 pub(crate) fn spendable_notes_meta(
     conn: &rusqlite::Connection,
     protocol: ShieldedProtocol,

zcash_client_sqlite/src/wallet/orchard.rs

@@ -24,7 +24,7 @@ use zip32::Scope;
 
 use crate::{error::SqliteClientError, AccountUuid, ReceivedNoteId, TxRef};
 
-use super::{get_account_ref, memo_repr, parse_scope, scope_code};
+use super::{common::UnspentNoteMeta, get_account_ref, memo_repr, parse_scope, scope_code};
 
 /// This trait provides a generalization over shielded output representations.
 pub(crate) trait ReceivedOrchardOutput {
@@ -226,6 +226,19 @@ pub(crate) fn select_spendable_orchard_notes<P: consensus::Parameters>(
     )
 }
 
+pub(crate) fn select_unspent_note_meta(
+    conn: &Connection,
+    chain_tip_height: BlockHeight,
+    wallet_birthday: BlockHeight,
+) -> Result<Vec<UnspentNoteMeta>, SqliteClientError> {
+    super::common::select_unspent_note_meta(
+        conn,
+        ShieldedProtocol::Orchard,
+        chain_tip_height,
+        wallet_birthday,
+    )
+}
+
 /// Records the specified shielded output as having been received.
 ///
 /// This implementation relies on the facts that: