Browse Source

Fix wallet restore gaps (#1452)

* test: add restore with counter gap integration test

Add test to verify wallet can continue operating after restoring
from a seed where counter values were non-sequential (e.g., due to
failed operations or multi-device usage)

* fix(wallet): set counter to max position after restore

After restore, set the counter to the highest counter position that had
a signature plus one, rather than the count of proofs found. This prevents
"blinded message already signed" errors when restoring wallets with
non-sequential counter values.
tsk 1 month ago
parent
commit
fea75c6767

+ 143 - 0
crates/cdk-integration-tests/tests/happy_path_mint_wallet.rs

@@ -355,6 +355,149 @@ async fn test_restore() {
     }
 }
 
+/// Tests that wallet restore correctly handles non-sequential counter values
+///
+/// This test verifies that after restoring a wallet where there were gaps in the
+/// counter sequence (e.g., due to failed operations or multi-device usage), the
+/// wallet can continue to operate without errors.
+///
+/// Test scenario:
+/// 1. Wallet mints proofs using counters 0-N
+/// 2. Counter is incremented to simulate failed operations that consumed counter values
+/// 3. Wallet mints more proofs using counters at higher values
+/// 4. New wallet restores from seed and finds proofs at non-sequential counter positions
+/// 5. Wallet should be able to continue normal operations (swaps) after restore
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_restore_with_counter_gap() {
+    let seed = Mnemonic::generate(12).unwrap().to_seed_normalized("");
+    let wallet = Wallet::new(
+        &get_mint_url_from_env(),
+        CurrencyUnit::Sat,
+        Arc::new(memory::empty().await.unwrap()),
+        seed,
+        None,
+    )
+    .expect("failed to create new wallet");
+
+    // Mint first batch of proofs (uses counters starting at 0)
+    let mint_quote = wallet.mint_quote(100.into(), None).await.unwrap();
+    let invoice = Bolt11Invoice::from_str(&mint_quote.request).unwrap();
+    pay_if_regtest(&get_test_temp_dir(), &invoice)
+        .await
+        .unwrap();
+
+    let _proofs1 = wallet
+        .wait_and_mint_quote(
+            mint_quote.clone(),
+            SplitTarget::default(),
+            None,
+            tokio::time::Duration::from_secs(60),
+        )
+        .await
+        .expect("first mint failed");
+
+    assert_eq!(wallet.total_balance().await.unwrap(), 100.into());
+
+    // Get the active keyset ID to increment counter
+    let active_keyset = wallet.fetch_active_keyset().await.unwrap();
+    let keyset_id = active_keyset.id;
+
+    // Create a gap in the counter sequence
+    // This simulates failed operations or multi-device usage where counter values
+    // were consumed but no signatures were obtained
+    let gap_size = 50u32;
+    {
+        let mut tx = wallet.localstore.begin_db_transaction().await.unwrap();
+        tx.increment_keyset_counter(&keyset_id, gap_size)
+            .await
+            .unwrap();
+        tx.commit().await.unwrap();
+    }
+
+    // Mint second batch of proofs (uses counters after the gap)
+    let mint_quote2 = wallet.mint_quote(100.into(), None).await.unwrap();
+    let invoice2 = Bolt11Invoice::from_str(&mint_quote2.request).unwrap();
+    pay_if_regtest(&get_test_temp_dir(), &invoice2)
+        .await
+        .unwrap();
+
+    let _proofs2 = wallet
+        .wait_and_mint_quote(
+            mint_quote2.clone(),
+            SplitTarget::default(),
+            None,
+            tokio::time::Duration::from_secs(60),
+        )
+        .await
+        .expect("second mint failed");
+
+    assert_eq!(wallet.total_balance().await.unwrap(), 200.into());
+
+    // Create a new wallet with the same seed (simulating wallet restore scenario)
+    let wallet_restored = Wallet::new(
+        &get_mint_url_from_env(),
+        CurrencyUnit::Sat,
+        Arc::new(memory::empty().await.unwrap()),
+        seed,
+        None,
+    )
+    .expect("failed to create restored wallet");
+
+    assert_eq!(wallet_restored.total_balance().await.unwrap(), 0.into());
+
+    // Restore the wallet - this should find proofs at non-sequential counter positions
+    let restored = wallet_restored.restore().await.unwrap();
+    assert_eq!(restored, 200.into());
+
+    let proofs = wallet_restored.get_unspent_proofs().await.unwrap();
+    assert!(!proofs.is_empty());
+
+    // Swap the restored proofs to verify they are valid
+    let expected_fee = wallet_restored.get_proofs_fee(&proofs).await.unwrap().total;
+    wallet_restored
+        .swap(None, SplitTarget::default(), proofs, None, false)
+        .await
+        .expect("first swap after restore failed");
+
+    let balance_after_first_swap = Amount::from(200) - expected_fee;
+    assert_eq!(
+        wallet_restored.total_balance().await.unwrap(),
+        balance_after_first_swap
+    );
+
+    // Perform multiple swaps to verify the wallet can continue operating
+    // after restore with non-sequential counter values
+    for i in 0..gap_size {
+        let proofs = wallet_restored.get_unspent_proofs().await.unwrap();
+        if proofs.is_empty() {
+            break;
+        }
+
+        let swap_result = wallet_restored
+            .swap(None, SplitTarget::default(), proofs.clone(), None, false)
+            .await;
+
+        match swap_result {
+            Ok(_) => {
+                // Swap succeeded, continue
+            }
+            Err(e) => {
+                let error_str = format!("{:?}", e);
+                if error_str.contains("BlindedMessageAlreadySigned") {
+                    panic!(
+                        "Got 'blinded message already signed' error on swap {} after restore. \
+                         Counter was not correctly set after restoring with non-sequential values.",
+                        i + 1
+                    );
+                } else {
+                    // Some other error - might be expected (e.g., insufficient funds due to fees)
+                    break;
+                }
+            }
+        }
+    }
+}
+
 /// Tests that the melt quote status can be checked after a melt has completed
 ///
 /// This test verifies:

+ 28 - 7
crates/cdk/src/wallet/mod.rs

@@ -458,6 +458,8 @@ impl Wallet {
             let keys = self.load_keyset_keys(keyset.id).await?;
             let mut empty_batch = 0;
             let mut start_counter = 0;
+            // Track the highest counter value that had a signature
+            let mut highest_counter: Option<u32> = None;
 
             while empty_batch.lt(&3) {
                 let premint_secrets = PreMintSecrets::restore_batch(
@@ -487,12 +489,23 @@ impl Wallet {
                     continue;
                 }
 
-                let premint_secrets: Vec<_> = premint_secrets
+                // Enumerate secrets to track their original index (which corresponds to counter value)
+                let matched_secrets: Vec<_> = premint_secrets
                     .secrets
                     .iter()
-                    .filter(|p| response.outputs.contains(&p.blinded_message))
+                    .enumerate()
+                    .filter(|(_, p)| response.outputs.contains(&p.blinded_message))
                     .collect();
 
+                // Update highest counter based on matched indices
+                if let Some(&(max_idx, _)) = matched_secrets.last() {
+                    let counter_value = start_counter + max_idx as u32;
+                    highest_counter =
+                        Some(highest_counter.map_or(counter_value, |c| c.max(counter_value)));
+                }
+
+                let premint_secrets: Vec<_> = matched_secrets.into_iter().map(|(_, p)| p).collect();
+
                 // the response outputs and premint secrets should be the same after filtering
                 // blinded messages the mint did not have signatures for
                 assert_eq!(response.outputs.len(), premint_secrets.len());
@@ -506,11 +519,6 @@ impl Wallet {
 
                 tracing::debug!("Restored {} proofs", proofs.len());
 
-                let mut tx = self.localstore.begin_db_transaction().await?;
-                tx.increment_keyset_counter(&keyset.id, proofs.len() as u32)
-                    .await?;
-                tx.commit().await?;
-
                 let states = self.check_proofs_spent(proofs.clone()).await?;
 
                 let unspent_proofs: Vec<Proof> = proofs
@@ -542,6 +550,19 @@ impl Wallet {
                 empty_batch = 0;
                 start_counter += 100;
             }
+
+            // Set counter to highest found + 1 to avoid reusing any counter values
+            // that already have signatures at the mint
+            if let Some(highest) = highest_counter {
+                let mut tx = self.localstore.begin_db_transaction().await?;
+                tx.increment_keyset_counter(&keyset.id, highest + 1).await?;
+                tx.commit().await?;
+                tracing::debug!(
+                    "Set keyset {} counter to {} after restore",
+                    keyset.id,
+                    highest + 1
+                );
+            }
         }
         Ok(restored_value)
     }