fix: update_last_revealed

Previously the call to `update_last_revealed` was conditional on whether the descriptor was present in the changeset, but it may be the case that the derivation index increases while the descriptor field of the wallet ChangeSet is None. We fix this by separating the calls to `insert_descriptor` and `update_last_revealed` and also changing the signature of `update_last_revealed` to accept a descriptor id rather than a keychain.
bitcoindevkit · Sep 25, 2024 · 86ada17 · 86ada17
1 parent 212aa65
commit 86ada17
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 54 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -8,7 +8,8 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use bdk_chain::{
-    local_chain, miniscript, tx_graph, Anchor, ConfirmationBlockTime, DescriptorExt, Merge,
+    local_chain, miniscript, tx_graph, Anchor, ConfirmationBlockTime, DescriptorExt, DescriptorId,
+    Merge,
 };
 use bdk_wallet::bitcoin::{
     self,
@@ -216,30 +217,23 @@ impl Store {
 
         if let Some(ref descriptor) = changeset.descriptor {
             insert_descriptor(&mut tx, wallet_name, descriptor, External).await?;
-            if let Some(last_revealed) = changeset
-                .indexer
-                .last_revealed
-                .get(&descriptor.descriptor_id())
-            {
-                update_last_revealed(&mut tx, wallet_name, *last_revealed, External).await?;
-            }
         }
 
-        if let Some(ref change_descriptor) = changeset.clone().change_descriptor {
+        if let Some(ref change_descriptor) = changeset.change_descriptor {
             insert_descriptor(&mut tx, wallet_name, change_descriptor, Internal).await?;
-            if let Some(last_revealed) = changeset
-                .indexer
-                .last_revealed
-                .get(&change_descriptor.descriptor_id())
-            {
-                update_last_revealed(&mut tx, wallet_name, *last_revealed, Internal).await?;
-            }
         }
 
         if let Some(network) = changeset.network {
             insert_network(&mut tx, wallet_name, network).await?;
         }
 
+        let last_revealed_indices = &changeset.indexer.last_revealed;
+        if !last_revealed_indices.is_empty() {
+            for (desc_id, index) in last_revealed_indices {
+                update_last_revealed(&mut tx, wallet_name, *desc_id, *index).await?;
+            }
+        }
+
         local_chain_changeset_persist_to_postgres(&mut tx, wallet_name, &changeset.local_chain)
             .await?;
         tx_graph_changeset_persist_to_postgres(&mut tx, wallet_name, &changeset.tx_graph).await?;
@@ -302,20 +296,17 @@ async fn insert_network(
 async fn update_last_revealed(
     tx: &mut Transaction<'_, Postgres>,
     wallet_name: &str,
+    descriptor_id: DescriptorId,
     last_revealed: u32,
-    keychain: KeychainKind,
 ) -> Result<(), BdkSqlxError> {
     info!("update last revealed");
-    let keychain = match keychain {
-        External => "External",
-        Internal => "Internal",
-    };
+
     sqlx::query(
-        "UPDATE keychain SET last_revealed = $1 WHERE wallet_name = $2 AND keychainkind = $3",
+        "UPDATE keychain SET last_revealed = $1 WHERE wallet_name = $2 AND descriptor_id = $3",
     )
     .bind(last_revealed as i32)
     .bind(wallet_name)
-    .bind(keychain)
+    .bind(descriptor_id.to_byte_array())
     .execute(&mut **tx)
     .await?;
 

diff --git a/src/test.rs b/src/test.rs
@@ -1,8 +1,11 @@
 use crate::{drop_all, Store};
 use bdk_chain::bitcoin::secp256k1::Secp256k1;
 use bdk_chain::bitcoin::Network;
-use bdk_chain::bitcoin::Network::Signet;
-use bdk_wallet::{wallet_name_from_descriptor, KeychainKind, Wallet};
+use bdk_wallet::{
+    wallet_name_from_descriptor,
+    KeychainKind::{self, *},
+    Wallet,
+};
 use better_panic::Settings;
 use sqlx::PgPool;
 use std::env;
@@ -15,7 +18,7 @@ pub fn get_test_tr_single_sig_xprv_with_change_desc() -> (&'static str, &'static
         "tr(tprv8ZgxMBicQKsPdDArR4xSAECuVxeX1jwwSXR4ApKbkYgZiziDc4LdBy2WvJeGDfUSE4UT4hHhbgEwbdq8ajjUHiKDegkwrNU6V55CxcxonVN/1/*)")
 }
 
-const NETWORK: Network = Signet;
+const NETWORK: Network = Network::Signet;
 
 #[tracing::instrument]
 #[tokio::test]
@@ -34,16 +37,9 @@ async fn wallet_is_persisted() -> anyhow::Result<()> {
 
     // Set up the database URL (you might want to use a test-specific database)
     let url = env::var("DATABASE_TEST_URL").expect("DATABASE_TEST_URL must be set for tests");
-
     let pg = PgPool::connect(&url.clone()).await?;
-    match drop_all(pg).await {
-        Ok(_) => {
-            dbg!("tables dropped")
-        }
-        Err(_) => {
-            dbg!("Error dropping tables")
-        }
-    };
+    drop_all(pg).await?;
+    println!("tables dropped");
 
     // Define descriptors (you may need to adjust these based on your exact requirements)
     let (external_desc, internal_desc) = get_test_tr_single_sig_xprv_with_change_desc();
@@ -56,26 +52,24 @@ async fn wallet_is_persisted() -> anyhow::Result<()> {
     )?;
 
     // Create a new wallet
-    let wallet_spk_index = {
-        let mut store = Store::new_with_url(url.clone(), Some(wallet_name.clone())).await?;
-        let mut wallet = Wallet::create(external_desc, internal_desc)
-            .network(NETWORK)
-            .create_wallet_async(&mut store)
-            .await?;
-
-        let deposit_address = wallet.reveal_next_address(KeychainKind::External);
-        let change_address = wallet.reveal_next_address(KeychainKind::Internal);
-        dbg!(deposit_address.address);
-        dbg!(change_address.address);
+    let mut store = Store::new_with_url(url.clone(), Some(wallet_name.clone())).await?;
+    let mut wallet = Wallet::create(external_desc, internal_desc)
+        .network(NETWORK)
+        .create_wallet_async(&mut store)
+        .await?;
+
+    let external_addr0 = wallet.reveal_next_address(KeychainKind::External);
+    for keychain in [External, Internal] {
+        let _ = wallet.reveal_addresses_to(keychain, 2);
+    }
 
-        assert!(wallet.persist_async(&mut store).await?);
-        wallet.spk_index().clone()
-    };
+    assert!(wallet.persist_async(&mut store).await?);
+    let wallet_spk_index = wallet.spk_index();
 
     {
         // Recover the wallet
-        let mut store = Store::new_with_url(url.clone(), Some(wallet_name.clone())).await?;
-        let mut wallet = Wallet::load()
+        let mut store = Store::new_with_url(url.clone(), Some(wallet_name)).await?;
+        let wallet = Wallet::load()
             .descriptor(KeychainKind::External, Some(external_desc))
             .descriptor(KeychainKind::Internal, Some(internal_desc))
             .load_wallet_async(&mut store)
@@ -92,8 +86,8 @@ async fn wallet_is_persisted() -> anyhow::Result<()> {
             wallet_spk_index.last_revealed_indices()
         );
 
-        let recovered_address = wallet.reveal_next_address(KeychainKind::External);
-        println!("Recovered next address: {}", recovered_address.address);
+        let recovered_addr = wallet.peek_address(KeychainKind::External, 0);
+        assert_eq!(recovered_addr, external_addr0, "failed to recover address");
 
         assert_eq!(
             wallet.public_descriptor(KeychainKind::External).to_string(),
@@ -104,7 +98,7 @@ async fn wallet_is_persisted() -> anyhow::Result<()> {
     // Clean up (optional, depending on your test database strategy)
     // You might want to delete the test wallet from the database here
     let db = PgPool::connect(&url).await?;
-    drop_all(db).await.expect("hope its not mainet");
+    drop_all(db).await.expect("hope its not mainnet");
 
     Ok(())
 }