Multi Thread sync, ported from 5d2b85c03a

lib/src/lightwallet.rs

@@ -1,11 +1,15 @@
-use std::time::SystemTime;
+use std::time::{SystemTime, Duration};
 use std::io::{self, Read, Write};
 use std::cmp;
 use std::collections::{HashMap, HashSet};
 use std::sync::{Arc, RwLock};
 use std::io::{Error, ErrorKind};
 
+use threadpool::ThreadPool;
+use std::sync::mpsc::{channel};
+
 use rand::{Rng, rngs::OsRng};
+use subtle::{ConditionallySelectable, ConstantTimeEq, CtOption};
 
 use log::{info, warn, error};
 
@@ -21,20 +25,23 @@ use sha2::{Sha256, Digest};
 
 use zcash_client_backend::{
     encoding::{encode_payment_address, encode_extended_spending_key},
-    proto::compact_formats::CompactBlock, welding_rig::scan_block,
+    proto::compact_formats::{CompactBlock, CompactOutput},
+    wallet::{WalletShieldedOutput, WalletShieldedSpend}
 };
 
 use zcash_primitives::{
+    jubjub::fs::Fs,
     block::BlockHash,
-    merkle_tree::{CommitmentTree},
     serialize::{Vector},
     transaction::{
         builder::{Builder},
         components::{Amount, OutPoint, TxOut}, components::amount::DEFAULT_FEE,
         TxId, Transaction, 
     },
-     legacy::{Script, TransparentAddress},
-    note_encryption::{Memo, try_sapling_note_decryption, try_sapling_output_recovery},
+    sapling::Node,
+    merkle_tree::{CommitmentTree, IncrementalWitness},
+    legacy::{Script, TransparentAddress},
+    note_encryption::{Memo, try_sapling_note_decryption, try_sapling_output_recovery, try_sapling_compact_note_decryption},
     zip32::{ExtendedFullViewingKey, ExtendedSpendingKey, ChildIndex},
     JUBJUB,
     primitives::{PaymentAddress},
@@ -136,6 +143,8 @@ pub struct LightWallet {
 
     // Non-serialized fields
     config: LightClientConfig,
+
+    pub total_scan_duration: Arc<RwLock<Vec<Duration>>>,
 }
 
 impl LightWallet {
@@ -254,6 +263,7 @@ impl LightWallet {
             mempool_txs: Arc::new(RwLock::new(HashMap::new())),
             config:      config.clone(),
             birthday:    latest_block,
+            total_scan_duration: Arc::new(RwLock::new(vec![Duration::new(0, 0)])),
         };
 
         // If restoring from seed, make sure we are creating 50 addresses for users
@@ -375,6 +385,7 @@ impl LightWallet {
             mempool_txs: Arc::new(RwLock::new(HashMap::new())),
             config:      config.clone(),
             birthday,
+            total_scan_duration: Arc::new(RwLock::new(vec![Duration::new(0, 0)])),
         })
     }
 
@@ -424,12 +435,19 @@ impl LightWallet {
 
         Vector::write(&mut writer, &self.blocks.read().unwrap(), |w, b| b.write(w))?;
                 
-        // The hashmap, write as a set of tuples
-        Vector::write(&mut writer, &self.txs.read().unwrap().iter().collect::<Vec<(&TxId, &WalletTx)>>(),
-                        |w, (k, v)| {
-                            w.write_all(&k.0)?;
-                            v.write(w)
-                        })?;
+        // The hashmap, write as a set of tuples. Store them sorted so that wallets are
+        // deterministically saved
+        {
+            let txlist = self.txs.read().unwrap();
+            let mut txns = txlist.iter().collect::<Vec<(&TxId, &WalletTx)>>();
+            txns.sort_by(|a, b| a.0.partial_cmp(b.0).unwrap());
+
+            Vector::write(&mut writer, &txns,
+                            |w, (k, v)| {
+                                w.write_all(&k.0)?;
+                                v.write(w)
+                            })?;
+        }
         utils::write_string(&mut writer, &self.config.chain_name)?;
 
         // While writing the birthday, get it from the fn so we recalculate it properly
@@ -1272,7 +1290,7 @@ impl LightWallet {
             // Trim all witnesses for the invalidated blocks
             for tx in txs.values_mut() {
                 for nd in tx.notes.iter_mut() {
-                    nd.witnesses.split_off(nd.witnesses.len().saturating_sub(num_invalidated));
+                    let _discard = nd.witnesses.split_off(nd.witnesses.len().saturating_sub(num_invalidated));
                 }
             }
         }
@@ -1280,8 +1298,234 @@ impl LightWallet {
         num_invalidated as u64
     }
 
-    // Scan a block. Will return an error with the block height that failed to scan
+       /// Scans a [`CompactOutput`] with a set of [`ExtendedFullViewingKey`]s.
+    ///
+    /// Returns a [`WalletShieldedOutput`] and corresponding [`IncrementalWitness`] if this
+    /// output belongs to any of the given [`ExtendedFullViewingKey`]s.
+    ///
+    /// The given [`CommitmentTree`] and existing [`IncrementalWitness`]es are incremented
+    /// with this output's commitment.
+    fn scan_output_internal(
+        &self,
+        (index, output): (usize, CompactOutput),
+        ivks: &[Fs],
+        tree: &mut CommitmentTree<Node>,
+        existing_witnesses: &mut [&mut IncrementalWitness<Node>],
+        block_witnesses: &mut [&mut IncrementalWitness<Node>],
+        new_witnesses: &mut [&mut IncrementalWitness<Node>],
+        pool: &ThreadPool
+    ) -> Option<WalletShieldedOutput> {
+        let cmu = output.cmu().ok()?;
+        let epk = output.epk().ok()?;
+        let ct = output.ciphertext;
+
+        let (tx, rx) = channel();
+        ivks.iter().enumerate().for_each(|(account, ivk)| {
+            // Clone all values for passing to the closure
+            let ivk = ivk.clone();
+            let epk = epk.clone();
+            let ct = ct.clone();
+            let tx = tx.clone();
+
+            pool.execute(move || {
+                let m = try_sapling_compact_note_decryption(&ivk, &epk, &cmu, &ct);
+                let r = match m {
+                    Some((note, to)) => {
+                        tx.send(Some(Some((note, to, account))))
+                    },
+                    None => {
+                        tx.send(Some(None))
+                    }
+                };
+
+                match r {
+                    Ok(_) => {},
+                    Err(e) => println!("Send error {:?}", e)
+                }
+            });
+        });
+
+        // Increment tree and witnesses
+        let node = Node::new(cmu.into());
+        for witness in existing_witnesses {
+            witness.append(node).unwrap();
+        }
+        for witness in block_witnesses {
+            witness.append(node).unwrap();
+        }
+        for witness in new_witnesses {
+            witness.append(node).unwrap();
+        }
+        tree.append(node).unwrap();
+
+        // Collect all the RXs and fine if there was a valid result somewhere
+        let mut wsos =  vec![];
+        for _i in 0..ivks.len() {
+            let n = rx.recv().unwrap();
+            let epk = epk.clone();
+
+            let wso = match n {
+                None => panic!("Got a none!"),
+                Some(None) => None,
+                Some(Some((note, to, account))) => {
+                    // A note is marked as "change" if the account that received it
+                    // also spent notes in the same transaction. This will catch,
+                    // for instance:
+                    // - Change created by spending fractions of notes.
+                    // - Notes created by consolidation transactions.
+                    // - Notes sent from one account to itself.
+                    //let is_change = spent_from_accounts.contains(&account);
+
+                    Some(WalletShieldedOutput {
+                        index, cmu, epk, account, note, to, is_change: false,
+                        witness: IncrementalWitness::from_tree(tree),
+                    })
+                }
+            };
+            wsos.push(wso);
+        }
+
+        match wsos.into_iter().find(|wso| wso.is_some()) {
+            Some(Some(wso)) => Some(wso),
+            _ => None
+        }
+    }
+
+    /// Scans a [`CompactBlock`] with a set of [`ExtendedFullViewingKey`]s.
+    ///
+    /// Returns a vector of [`WalletTx`]s belonging to any of the given
+    /// [`ExtendedFullViewingKey`]s, and the corresponding new [`IncrementalWitness`]es.
+    ///
+    /// The given [`CommitmentTree`] and existing [`IncrementalWitness`]es are
+    /// incremented appropriately.
+    pub fn scan_block_internal(
+        &self,
+        block: CompactBlock,
+        extfvks: &[ExtendedFullViewingKey],
+        nullifiers: Vec<(Vec<u8>, usize)>,
+        tree: &mut CommitmentTree<Node>,
+        existing_witnesses: &mut [&mut IncrementalWitness<Node>],
+        pool: &ThreadPool
+    ) -> Vec<zcash_client_backend::wallet::WalletTx> {
+        let mut wtxs: Vec<zcash_client_backend::wallet::WalletTx> = vec![];
+        let ivks = extfvks.iter().map(|extfvk| extfvk.fvk.vk.ivk()).collect::<Vec<_>>();
+
+        for tx in block.vtx.into_iter() {
+            let num_spends = tx.spends.len();
+            let num_outputs = tx.outputs.len();
+
+            let (ctx, crx) = channel();
+            {
+                let nullifiers = nullifiers.clone();
+                let tx = tx.clone();
+                pool.execute(move || {
+                    // Check for spent notes
+                    // The only step that is not constant-time is the filter() at the end.
+                    let shielded_spends: Vec<_> = tx
+                        .spends
+                        .into_iter()
+                        .enumerate()
+                        .map(|(index, spend)| {
+                            // Find the first tracked nullifier that matches this spend, and produce
+                            // a WalletShieldedSpend if there is a match, in constant time.
+                            nullifiers
+                                .iter()
+                                .map(|(nf, account)| CtOption::new(*account as u64, nf.ct_eq(&spend.nf[..])))
+                                .fold(CtOption::new(0, 0.into()), |first, next| {
+                                    CtOption::conditional_select(&next, &first, first.is_some())
+                                })
+                                .map(|account| WalletShieldedSpend {
+                                    index,
+                                    nf: spend.nf,
+                                    account: account as usize,
+                                })
+                        })
+                        .filter(|spend| spend.is_some().into())
+                        .map(|spend| spend.unwrap())
+                        .collect();
+
+                    // Collect the set of accounts that were spent from in this transaction
+                    let spent_from_accounts: HashSet<_> =
+                        shielded_spends.iter().map(|spend| spend.account).collect();
+
+                    ctx.send((shielded_spends, spent_from_accounts)).unwrap();
+
+                    drop(ctx);
+                });
+            }
+
+
+            // Check for incoming notes while incrementing tree and witnesses
+            let mut shielded_outputs: Vec<WalletShieldedOutput> = vec![];
+            {
+                // Grab mutable references to new witnesses from previous transactions
+                // in this block so that we can update them. Scoped so we don't hold
+                // mutable references to wtxs for too long.
+                let mut block_witnesses: Vec<_> = wtxs
+                    .iter_mut()
+                    .map(|tx| {
+                        tx.shielded_outputs
+                            .iter_mut()
+                            .map(|output| &mut output.witness)
+                    })
+                    .flatten()
+                    .collect();
+
+                for to_scan in tx.outputs.into_iter().enumerate() {
+                    // Grab mutable references to new witnesses from previous outputs
+                    // in this transaction so that we can update them. Scoped so we
+                    // don't hold mutable references to shielded_outputs for too long.
+                    let mut new_witnesses: Vec<_> = shielded_outputs
+                        .iter_mut()
+                        .map(|output| &mut output.witness)
+                        .collect();
+
+                    if let Some(output) = self.scan_output_internal(
+                        to_scan,
+                        &ivks,
+                        tree,
+                        existing_witnesses,
+                        &mut block_witnesses,
+                        &mut new_witnesses,
+                        pool
+                    ) {
+                        shielded_outputs.push(output);
+                    }
+                }
+            }
+
+            let (shielded_spends, spent_from_accounts) = crx.recv().unwrap();
+
+            // Identify change outputs
+            shielded_outputs.iter_mut().for_each(|output| {
+                if spent_from_accounts.contains(&output.account) {
+                    output.is_change = true;
+                }
+            });
+
+            // Update wallet tx
+            if !(shielded_spends.is_empty() && shielded_outputs.is_empty()) {
+                let mut txid = TxId([0u8; 32]);
+                txid.0.copy_from_slice(&tx.hash);
+                wtxs.push(zcash_client_backend::wallet::WalletTx {
+                    txid,
+                    index: tx.index as usize,
+                    num_spends,
+                    num_outputs,
+                    shielded_spends,
+                    shielded_outputs,
+                });
+            }
+        }
+
+        wtxs
+    }
     pub fn scan_block(&self, block_bytes: &[u8]) -> Result<Vec<TxId>, i32> {
+        self.scan_block_with_pool(&block_bytes, &ThreadPool::new(1))
+    }
+
+    // Scan a block. Will return an error with the block height that failed to scan
+    pub fn scan_block_with_pool(&self, block_bytes: &[u8], pool: &ThreadPool) -> Result<Vec<TxId>, i32> {
         let block: CompactBlock = match parse_from_bytes(block_bytes) {
             Ok(block) => block,
             Err(e) => {
@@ -1372,7 +1616,7 @@ impl LightWallet {
             }
 
             new_txs = {
-                let nf_refs: Vec<_> = nfs.iter().map(|(nf, acc, _)| (&nf[..], *acc)).collect();
+                let nf_refs = nfs.iter().map(|(nf, account, _)| (nf.to_vec(), *account)).collect::<Vec<_>>();
 
                 // Create a single mutable slice of all the newly-added witnesses.
                 let mut witness_refs: Vec<_> = txs
@@ -1381,12 +1625,13 @@ impl LightWallet {
                     .flatten()
                     .collect();
 
-                scan_block(
+                self.scan_block_internal(
                     block.clone(),
                     &self.extfvks.read().unwrap(),
-                    &nf_refs[..],
+                    nf_refs,
                     &mut block_data.tree,
                     &mut witness_refs[..],
+                    pool,
                 )
             };
         }