Merge pull request #197 from EspressoSystems/bugfix/key-scan-panic

Fix panic in key scan

Cargo.toml

@@ -2,7 +2,7 @@
 name = "seahorse"
 description = "A generic cap-style cryptocurrency wallet."
 authors = ["Espresso Systems <[email protected]>"]
-version = "0.2.4"
+version = "0.2.5"
 edition = "2018"
 license = "GPL-3.0-or-later"
 
@@ -44,6 +44,7 @@ strum = "0.24"
 strum_macros = "0.20.1"
 surf = "2.3.1"
 tempdir = "0.3.7"
+tracing = "0.1.35"
 zeroize = "1.3"
 
 # local dependencies

src/key_scan.rs

@@ -171,7 +171,17 @@ impl<L: Ledger> BackgroundKeyScan<L> {
             let records = self
                 .records
                 .into_values()
-                .map(|(ro, uid)| (ro, uid, mt.get_leaf(uid).expect_ok().unwrap().1.path))
+                .map(|(ro, uid)| {
+                    (
+                        ro,
+                        uid,
+                        mt.get_leaf(uid)
+                            .expect_ok()
+                            .unwrap_or_else(|_| panic!("uid {}/{}", uid, mt.num_leaves()))
+                            .1
+                            .path,
+                    )
+                })
                 .collect();
             Ok((
                 self.key,
@@ -308,7 +318,7 @@ impl<L: Ledger> BackgroundKeyScan<L> {
                         *uid,
                         &RecordCommitment::from(ro),
                     );
-                    self.records_mt.remember(*uid, proof).unwrap();
+                    self.remember_uid(*uid, proof);
                     self.records.insert(nullifier, (ro.clone(), *uid));
                 }
 
@@ -387,6 +397,18 @@ impl<L: Ledger> BackgroundKeyScan<L> {
         }
     }
 
+    fn remember_uid(&mut self, leaf: u64, proof: &MerkleLeafProof) {
+        // If we were planning to forget this leaf once a new leaf is appended, stop planning that.
+        if self.leaf_to_forget == Some(leaf) {
+            self.leaf_to_forget = None;
+            // `leaf_to_forget` is always represented in the tree, so we don't have to call
+            // `remember` in this case.
+            assert!(self.records_mt.get_leaf(leaf).expect_ok().is_ok());
+        } else if self.records_mt.remember(leaf, proof).is_err() {
+            tracing::error!("Got bad Merkle proof. Unable to save record {}", leaf);
+        }
+    }
+
     fn forget(&mut self, uid: u64) {
         if uid == self.records_mt.num_leaves() - 1 {
             // If the leaf we're trying to forget is on the frontier, we can't forget it