From cc1a43c495639f0dbe6da52501cb48f17b8cbbd2 Mon Sep 17 00:00:00 2001
From: Daniela Brozzoni <danielabrozzoni@protonmail.com>
Date: Thu, 24 Aug 2023 16:53:50 +0200
Subject: [PATCH] fix: SpkIterator::new_with_range takes wildcards.. ..into
 account

When you pass a non-wildcard descriptor in `new_with_range`, we make
sure that the range length is at most 1; if that's not the case, we
shorten it.
We would previously use `new_with_range` without this check and with a
non-wildcard descriptor in `spks_of_all_keychains`, this meant creating
a spkiterator that would go on producing the same spks over and over
again, causing some issues with syncing on electrum/esplora.

To reproduce the bug, run in `example-crates/example_electrum`:
```
cargo run "sh(wsh(or_d(c:pk_k(cPGudvRLDSgeV4hH9NUofLvYxYBSRjju3cpiXmBg9K8G9k1ikCMp),c:pk_k(cSBSBHRrzqSXFmrBhLkZMzQB9q4P9MnAq92v8d9a5UveBc9sLX32))))#zp9pcfs9" scan
```
---
 crates/chain/src/spk_iter.rs                  | 47 ++++++++++++++-----
 .../chain/tests/test_keychain_txout_index.rs  |  8 ++++
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/crates/chain/src/spk_iter.rs b/crates/chain/src/spk_iter.rs
index 1e09df36..bce436ed 100644
--- a/crates/chain/src/spk_iter.rs
+++ b/crates/chain/src/spk_iter.rs
@@ -45,34 +45,41 @@ where
 {
     /// Creates a new script pubkey iterator starting at 0 from a descriptor.
     pub fn new(descriptor: D) -> Self {
-        let end = if descriptor.borrow().has_wildcard() {
-            BIP32_MAX_INDEX
-        } else {
-            0
-        };
-
-        SpkIterator::new_with_range(descriptor, 0..=end)
+        SpkIterator::new_with_range(descriptor, 0..=BIP32_MAX_INDEX)
     }
 
     // Creates a new script pubkey iterator from a descriptor with a given range.
+    // If the descriptor doesn't have a wildcard, we shorten whichever range you pass in
+    // to have length <= 1. This means that if you pass in 0..0 or 0..1 the range will
+    // remain the same, but if you pass in 0..10, we'll shorten it to 0..1
     pub(crate) fn new_with_range<R>(descriptor: D, range: R) -> Self
     where
         R: RangeBounds<u32>,
     {
+        let start = match range.start_bound() {
+            Bound::Included(start) => *start,
+            Bound::Excluded(start) => *start + 1,
+            Bound::Unbounded => u32::MIN,
+        };
+
         let mut end = match range.end_bound() {
             Bound::Included(end) => *end + 1,
             Bound::Excluded(end) => *end,
             Bound::Unbounded => u32::MAX,
         };
+
         // Because `end` is exclusive, we want the maximum value to be BIP32_MAX_INDEX + 1.
         end = end.min(BIP32_MAX_INDEX + 1);
 
+        if !descriptor.borrow().has_wildcard() {
+            // The length of the range should be at most 1
+            if end != start {
+                end = start + 1;
+            }
+        }
+
         Self {
-            next_index: match range.start_bound() {
-                Bound::Included(start) => *start,
-                Bound::Excluded(start) => *start + 1,
-                Bound::Unbounded => u32::MIN,
-            },
+            next_index: start,
             end,
             descriptor,
             secp: Secp256k1::verification_only(),
@@ -195,6 +202,22 @@ mod test {
 
         let mut external_spk = SpkIterator::new(&no_wildcard_descriptor);
 
+        assert_eq!(external_spk.nth(0).unwrap(), (0, external_spk_0.clone()));
+        assert_eq!(external_spk.nth(0), None);
+
+        let mut external_spk = SpkIterator::new_with_range(&no_wildcard_descriptor, 0..0);
+
+        assert_eq!(external_spk.next(), None);
+
+        let mut external_spk = SpkIterator::new_with_range(&no_wildcard_descriptor, 0..1);
+
+        assert_eq!(external_spk.nth(0).unwrap(), (0, external_spk_0.clone()));
+        assert_eq!(external_spk.next(), None);
+
+        // We test that using new_with_range with range_len > 1 gives back an iterator with
+        // range_len = 1
+        let mut external_spk = SpkIterator::new_with_range(&no_wildcard_descriptor, 0..10);
+
         assert_eq!(external_spk.nth(0).unwrap(), (0, external_spk_0));
         assert_eq!(external_spk.nth(0), None);
     }
diff --git a/crates/chain/tests/test_keychain_txout_index.rs b/crates/chain/tests/test_keychain_txout_index.rs
index 96a1afd1..817b0349 100644
--- a/crates/chain/tests/test_keychain_txout_index.rs
+++ b/crates/chain/tests/test_keychain_txout_index.rs
@@ -384,4 +384,12 @@ fn test_non_wildcard_derivations() {
         txout_index.reveal_to_target(&TestKeychain::External, 200);
     assert_eq!(revealed_spks.count(), 0);
     assert!(revealed_changeset.is_empty());
+
+    // we check that spks_of_keychain returns a SpkIterator with just one element
+    assert_eq!(
+        txout_index
+            .spks_of_keychain(&TestKeychain::External)
+            .count(),
+        1,
+    );
 }