From b4a2439ef037f7198b4e7016b562e71dafaff144 Mon Sep 17 00:00:00 2001 From: xdustinface Date: Mon, 22 Jun 2026 21:51:52 +1000 Subject: [PATCH 1/2] feat(key-wallet): reserve receive addresses on hand-out Adds a reservation lifecycle to addresses so a receive address handed out to a caller is not re-issued before it is funded or explicitly released. This closes the hand-out race where two sequential requests returned the same address. - Model address lifecycle as an `AddressState` enum (`Available`, `Reserved { at }`, `Used { at }`) on `AddressInfo`, replacing the separate `used`/`used_at` fields. The states are mutually exclusive by construction, so the invariant "a used address is never reserved" holds structurally instead of being maintained by hand. - Add `next_unused_and_reserve`, `release_reservation`, and `sweep_expired_reservations` (TTL backstop) on `AddressPool`; reserved addresses are excluded from hand-out, count against the gap limit, and are never pruned or aged out when clockless. - Wire `next_receive_address_and_reserve`, `release_receive_reservation`, and `sweep_expired_receive_reservations` through `ManagedCoreFundsAccount`, bumping the monitor revision on change. - Add `reserved_count` to `PoolStats`. - Cover reserve/release/sweep, serde and bincode round-trips, gap-limit and prune interaction, and end-to-end promotion on funding. --- key-wallet-ffi/src/address_pool.rs | 8 +- key-wallet-manager/src/events.rs | 5 +- .../src/managed_account/address_pool.rs | 531 ++++++++++++++++-- .../managed_core_funds_account.rs | 87 +++ key-wallet/src/tests/address_pool_tests.rs | 30 +- .../src/tests/address_reservation_tests.rs | 141 +++++ key-wallet/src/tests/mod.rs | 2 + 7 files changed, 754 insertions(+), 50 deletions(-) create mode 100644 key-wallet/src/tests/address_reservation_tests.rs diff --git a/key-wallet-ffi/src/address_pool.rs b/key-wallet-ffi/src/address_pool.rs index 886dfb0ee..17a3aba11 100644 --- a/key-wallet-ffi/src/address_pool.rs +++ b/key-wallet-ffi/src/address_pool.rs @@ -264,9 +264,9 @@ fn address_info_to_ffi(info: &AddressInfo) -> FFIAddressInfo { public_key_len, index: info.index, path: path_str, - used: info.used, + used: info.is_used(), generated_at: info.generated_at, - used_at: info.used_at.unwrap_or(0), + used_at: info.used_at().unwrap_or(0), tx_count: info.tx_count, total_received: info.total_received, total_sent: info.total_sent, @@ -879,6 +879,7 @@ pub unsafe extern "C" fn address_info_array_free(infos: *mut *mut FFIAddressInfo mod tests { use super::*; use dash_network::ffi::FFINetwork; + use key_wallet::managed_account::address_pool::AddressState; #[test] fn test_address_pool_type_values() { @@ -914,9 +915,8 @@ mod tests { public_key: Some(PublicKeyType::ECDSA(vec![0x02, 0x03, 0x04])), index: 0, path: test_path, - used: false, + state: AddressState::Available, generated_at: 1234567890, - used_at: None, tx_count: 0, total_received: 0, total_sent: 0, diff --git a/key-wallet-manager/src/events.rs b/key-wallet-manager/src/events.rs index d5d597cba..a18115c6e 100644 --- a/key-wallet-manager/src/events.rs +++ b/key-wallet-manager/src/events.rs @@ -434,7 +434,7 @@ mod project_derived_addresses_tests { use super::*; use key_wallet::account::StandardAccountType; use key_wallet::bip32::{ChildNumber, DerivationPath}; - use key_wallet::managed_account::address_pool::AddressInfo; + use key_wallet::managed_account::address_pool::{AddressInfo, AddressState}; use std::collections::BTreeMap; /// Compressed encoding of the secp256k1 generator point (G). @@ -488,9 +488,8 @@ mod project_derived_addresses_tests { public_key: Some(PublicKeyType::ECDSA(pubkey_bytes)), index, path, - used: false, + state: AddressState::Available, generated_at: 0, - used_at: None, tx_count: 0, total_received: 0, total_sent: 0, diff --git a/key-wallet/src/managed_account/address_pool.rs b/key-wallet/src/managed_account/address_pool.rs index a6059667d..472b6719a 100644 --- a/key-wallet/src/managed_account/address_pool.rs +++ b/key-wallet/src/managed_account/address_pool.rs @@ -207,6 +207,31 @@ impl KeySource { } } +/// Lifecycle state of an address in the pool. +/// +/// The three states are mutually exclusive by construction, so the invariant +/// "a used address is never reserved" holds structurally rather than being +/// maintained by hand. Timestamps are caller-supplied (seconds), and `0` means +/// the caller had no clock available. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "bincode", derive(Encode, Decode))] +pub enum AddressState { + /// Free to hand out: neither reserved nor used. + #[default] + Available, + /// Handed out but not yet funded. `at` is when it was reserved. + Reserved { + /// Reservation timestamp. + at: u64, + }, + /// Funds have been seen at this address. `at` is when it was first used. + Used { + /// First-used timestamp. + at: u64, + }, +} + /// Information about a single address in the pool #[derive(Debug, Clone)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] @@ -222,12 +247,10 @@ pub struct AddressInfo { pub index: u32, /// Full derivation path pub path: DerivationPath, - /// Whether this address has been used - pub used: bool, + /// Lifecycle state: available, reserved, or used + pub state: AddressState, /// When the address was first generated (timestamp) pub generated_at: u64, - /// When the address was first used (timestamp) - pub used_at: Option, /// Transaction count for this address pub tx_count: u32, /// Total received amount @@ -257,9 +280,8 @@ impl AddressInfo { public_key: Some(public_key), index, path, - used: false, + state: AddressState::Available, generated_at: 0, // Should use actual timestamp - used_at: None, tx_count: 0, total_received: 0, total_sent: 0, @@ -291,9 +313,8 @@ impl AddressInfo { public_key: None, // Public key not available from script alone index, path, - used: false, + state: AddressState::Available, generated_at: 0, // Should use actual timestamp - used_at: None, tx_count: 0, total_received: 0, total_sent: 0, @@ -303,11 +324,53 @@ impl AddressInfo { }) } - /// Mark this address as used + /// Mark this address as used. + /// + /// Transitions to [`AddressState::Used`], which structurally drops any + /// reservation so the invariant "a used address is never reserved" holds + /// even when a reserved address is promoted directly to used. Idempotent: + /// a re-mark keeps the original used timestamp. fn mark_used(&mut self) { - if !self.used { - self.used = true; - self.used_at = Some(0); // Should use actual timestamp + if !self.is_used() { + self.state = AddressState::Used { + at: 0, + }; // Should use actual timestamp + } + } + + /// Whether this address is available to be handed out: neither used nor + /// reserved. + pub fn is_available(&self) -> bool { + matches!(self.state, AddressState::Available) + } + + /// Whether this address is currently reserved (handed out but not yet used). + pub fn is_reserved(&self) -> bool { + matches!(self.state, AddressState::Reserved { .. }) + } + + /// Whether this address has been used. + pub fn is_used(&self) -> bool { + matches!(self.state, AddressState::Used { .. }) + } + + /// Timestamp the address was first used, or `None` if not used. + pub fn used_at(&self) -> Option { + match self.state { + AddressState::Used { + at, + } => Some(at), + _ => None, + } + } + + /// Timestamp the address was reserved, or `None` if not reserved. + pub fn reserved_at(&self) -> Option { + match self.state { + AddressState::Reserved { + at, + } => Some(at), + _ => None, } } @@ -523,7 +586,7 @@ impl AddressPool { // First, try to find an already generated unused address for i in 0..=self.highest_generated.unwrap_or(0) { if let Some(info) = self.addresses.get(&i) { - if !info.used { + if info.is_available() { return Ok(info.address.clone()); } } @@ -548,7 +611,7 @@ impl AddressPool { // First, try to find an already generated unused address for i in 0..=self.highest_generated.unwrap_or(0) { if let Some(info) = self.addresses.get(&i) { - if !info.used { + if info.is_available() { return Ok(info.clone()); } } @@ -569,6 +632,109 @@ impl AddressPool { }) } + /// Reserve and return the next available address. + /// + /// Scans for the first available index (neither used nor reserved), marks it + /// reserved with `now`, and returns its address. Because the scan and the + /// reserve happen under a single exclusive `&mut self` borrow, two + /// sequential hand-outs cannot both observe the same index as available: the + /// first reserves it before the second runs. When no generated address is + /// available a fresh one is derived and always materialized into the pool, + /// since the reservation is stored on the address entry. + /// + /// `now` is a caller-supplied timestamp (seconds) that stamps the + /// reservation and is what [`AddressPool::sweep_expired_reservations`] later + /// compares against. The pool keeps no clock of its own. + /// + /// This path derives a fresh address whenever every existing address is + /// used or reserved, with no upper bound. Outstanding reservations are + /// reclaimed only by [`AddressPool::release_reservation`] and + /// [`AddressPool::sweep_expired_reservations`], so callers are assumed + /// trusted. Gate or cap this before exposing it to untrusted input. + pub fn next_unused_and_reserve(&mut self, key_source: &KeySource, now: u64) -> Result
{ + for i in 0..=self.highest_generated.unwrap_or(0) { + if let Some(info) = self.addresses.get_mut(&i) { + if info.is_available() { + info.state = AddressState::Reserved { + at: now, + }; + return Ok(info.address.clone()); + } + } + } + + if matches!(key_source, KeySource::NoKeySource) { + return Err(Error::NoKeySource); + } + let next_index = self.highest_generated.map(|h| h + 1).unwrap_or(0); + let address = self.generate_address_at_index(next_index, key_source, true)?; + // `generate_address_at_index` with `add_to_state = true` always inserts + // at `next_index`. If `highest_generated` ever drifts out of sync with + // `addresses` a silent lookup miss would return an unreserved address, + // reintroducing the hand-out race this method exists to close, so fail + // hard in release rather than fail open. + let info = self + .addresses + .get_mut(&next_index) + .expect("generate_address_at_index(add_to_state=true) must insert at next_index"); + debug_assert!(info.is_available()); + info.state = AddressState::Reserved { + at: now, + }; + Ok(address) + } + + /// Release a reservation, returning the address to the available pool. + /// + /// Idempotent: returns `false` when the index is not currently reserved + /// (never reserved, already released, already used, or unknown index), and + /// `true` when a reservation was actually cleared. + pub fn release_reservation(&mut self, index: u32) -> bool { + if let Some(info) = self.addresses.get_mut(&index) { + if info.is_reserved() { + info.state = AddressState::Available; + return true; + } + } + false + } + + /// Reclaim reservations older than `ttl`, returning their addresses to the + /// available pool. + /// + /// A reservation is expired when `now - reserved_at >= ttl`. Returns how + /// many reservations were reclaimed. `now` and `ttl` are caller-supplied + /// (seconds). The pool keeps no clock. This is the backstop against a + /// hand-out that is reserved but never paid and never explicitly released, + /// which would otherwise pin gap-limit headroom across restarts forever. + /// + /// A `now` of 0 means the caller has no clock available, so the sweep + /// declines to judge staleness and reclaims nothing. A `ttl` of 0 means no + /// expiry window, so nothing is reclaimed either. A reservation stamped with + /// `at == 0` (its caller had no clock) is likewise never aged out. + pub fn sweep_expired_reservations(&mut self, now: u64, ttl: u64) -> usize { + if now == 0 || ttl == 0 { + return 0; + } + let mut reclaimed = 0; + for info in self.addresses.values_mut() { + if let AddressState::Reserved { + at, + } = info.state + { + // A clockless stamp cannot be judged stale against a real clock. + if at == 0 { + continue; + } + if now.saturating_sub(at) >= ttl { + info.state = AddressState::Available; + reclaimed += 1; + } + } + } + reclaimed + } + /// Get multiple unused addresses at once /// /// Returns the requested number of unused addresses, generating new ones if needed. @@ -589,7 +755,7 @@ impl AddressPool { break; } if let Some(info) = self.addresses.get(&i) { - if !info.used { + if info.is_available() { addresses.push(info.address.clone()); collected += 1; } @@ -645,7 +811,7 @@ impl AddressPool { break; } if let Some(info) = self.addresses.get(&i) { - if !info.used { + if info.is_available() { result.push((info.address.clone(), info.clone())); collected += 1; } @@ -696,7 +862,7 @@ impl AddressPool { && self.highest_generated.map(|h| current_index <= h).unwrap_or(false) { if let Some(info) = self.addresses.get(¤t_index) { - if !info.used { + if info.is_available() { unused.push(info.address.clone()); } } @@ -717,7 +883,7 @@ impl AddressPool { pub fn mark_used(&mut self, address: &Address) -> bool { if let Some(&index) = self.address_index.get(address) { if let Some(info) = self.addresses.get_mut(&index) { - if !info.used { + if !info.is_used() { info.mark_used(); self.used_indices.insert(index); @@ -737,7 +903,7 @@ impl AddressPool { /// Mark an address at a specific index as used pub fn mark_index_used(&mut self, index: u32) -> bool { if let Some(info) = self.addresses.get_mut(&index) { - if !info.used { + if !info.is_used() { info.mark_used(); self.used_indices.insert(index); @@ -761,7 +927,7 @@ impl AddressPool { let mut found = Vec::new(); for (_, info) in self.addresses.iter_mut() { - if !info.used && check_fn(&info.address) { + if !info.is_used() && check_fn(&info.address) { info.mark_used(); self.used_indices.insert(info.index); found.push(info.address.clone()); @@ -789,12 +955,20 @@ impl AddressPool { /// Get only used addresses pub fn used_addresses(&self) -> Vec
{ - self.addresses.values().filter(|info| info.used).map(|info| info.address.clone()).collect() + self.addresses + .values() + .filter(|info| info.is_used()) + .map(|info| info.address.clone()) + .collect() } - /// Get only unused addresses + /// Get only available addresses (neither used nor reserved) pub fn unused_addresses(&self) -> Vec
{ - self.addresses.values().filter(|info| !info.used).map(|info| info.address.clone()).collect() + self.addresses + .values() + .filter(|info| info.is_available()) + .map(|info| info.address.clone()) + .collect() } /// Get address at specific index @@ -877,9 +1051,10 @@ impl AddressPool { /// Check if we need to generate more addresses pub fn needs_more_addresses(&self) -> bool { - let unused_count = self.addresses.values().filter(|info| !info.used).count() as u32; + let available_count = + self.addresses.values().filter(|info| info.is_available()).count() as u32; - unused_count < self.gap_limit + available_count < self.gap_limit } /// Generate addresses to maintain the gap limit. @@ -940,11 +1115,14 @@ impl AddressPool { /// Get pool statistics pub fn stats(&self) -> PoolStats { let used_count = self.used_indices.len() as u32; - let unused_count = self.addresses.len() as u32 - used_count; + let reserved_count = + self.addresses.values().filter(|info| info.is_reserved()).count() as u32; + let unused_count = self.addresses.len() as u32 - used_count - reserved_count; PoolStats { total_generated: self.addresses.len() as u32, used_count, + reserved_count, unused_count, highest_used: self.highest_used, highest_generated: self.highest_generated, @@ -956,8 +1134,7 @@ impl AddressPool { /// Reset the pool (for rescan) pub fn reset_usage(&mut self) { for info in self.addresses.values_mut() { - info.used = false; - info.used_at = None; + info.state = AddressState::Available; info.tx_count = 0; info.total_received = 0; info.total_sent = 0; @@ -977,9 +1154,11 @@ impl AddressPool { let mut pruned = 0; let indices_to_remove: Vec = self .addresses - .keys() - .filter(|&&idx| idx > keep_until && !self.used_indices.contains(&idx)) - .copied() + .iter() + .filter(|(&idx, info)| { + idx > keep_until && !self.used_indices.contains(&idx) && !info.is_reserved() + }) + .map(|(&idx, _)| idx) .collect(); for idx in indices_to_remove { @@ -1005,7 +1184,9 @@ pub struct PoolStats { pub total_generated: u32, /// Number of used addresses pub used_count: u32, - /// Number of unused addresses + /// Number of reserved addresses (handed out but not yet used) + pub reserved_count: u32, + /// Number of unused addresses (neither used nor reserved) pub unused_count: u32, /// Highest used index pub highest_used: Option, @@ -1021,7 +1202,7 @@ impl fmt::Display for PoolStats { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( f, - "{} pool: {} addresses ({} used, {} unused), gap limit: {}", + "{} pool: {} addresses ({} used, {} reserved, {} unused), gap limit: {}", if self.is_internal { "Internal" } else { @@ -1029,6 +1210,7 @@ impl fmt::Display for PoolStats { }, self.total_generated, self.used_count, + self.reserved_count, self.unused_count, self.gap_limit ) @@ -1218,6 +1400,266 @@ mod tests { assert_ne!(addr1, addr3); // Should return different address after marking used } + #[test] + fn test_reserve_returns_distinct_addresses() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + 5, + Network::Testnet, + ); + let key_source = test_key_source(); + + // Without reserving, two hand-outs return the same address (the race). + let peek1 = pool.next_unused(&key_source, true).unwrap(); + let peek2 = pool.next_unused(&key_source, true).unwrap(); + assert_eq!(peek1, peek2); + + // Reserving on hand-out makes two sequential calls return distinct + // addresses: the first is no longer available to the second. + let addr1 = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + let addr2 = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + assert_ne!(addr1, addr2); + + // A plain peek must skip the reserved addresses too. + let peek3 = pool.next_unused(&key_source, true).unwrap(); + assert_ne!(peek3, addr1); + assert_ne!(peek3, addr2); + } + + #[test] + fn test_release_returns_address_to_pool() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + 5, + Network::Testnet, + ); + let key_source = test_key_source(); + + let addr = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + let index = pool.address_index(&addr).unwrap(); + + // A reserved address is excluded from the available set. + assert!(!pool.unused_addresses().contains(&addr)); + + // After releasing, the same address becomes available again. + assert!(pool.release_reservation(index)); + assert!(pool.unused_addresses().contains(&addr)); + let reissued = pool.next_unused_and_reserve(&key_source, 200).unwrap(); + assert_eq!(addr, reissued); + } + + #[test] + fn test_release_is_idempotent() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + 5, + Network::Testnet, + ); + let key_source = test_key_source(); + + let addr = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + let index = pool.address_index(&addr).unwrap(); + + // First release clears the reservation, a second is a no-op. + assert!(pool.release_reservation(index)); + assert!(!pool.release_reservation(index)); + // Releasing a never-reserved index is a no-op too. + let other = pool.next_unused(&key_source, true).unwrap(); + let other_index = pool.address_index(&other).unwrap(); + assert!(!pool.release_reservation(other_index)); + // Unknown index is a no-op. + assert!(!pool.release_reservation(9_999)); + } + + #[test] + fn test_reserved_promoted_to_used_clears_reservation() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + 5, + Network::Testnet, + ); + let key_source = test_key_source(); + + let addr = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + let index = pool.address_index(&addr).unwrap(); + assert!(pool.addresses.get(&index).unwrap().is_reserved()); + + // Funds arriving mark the reserved address used, so the reservation is + // structurally dropped (a used address is never reserved). + assert!(pool.mark_used(&addr)); + let info = pool.addresses.get(&index).unwrap(); + assert!(info.is_used()); + assert!(!info.is_reserved()); + assert!(info.reserved_at().is_none()); + // Releasing a now-used index is a no-op. + assert!(!pool.release_reservation(index)); + } + + #[cfg(feature = "serde")] + #[test] + fn test_reservation_survives_serde_round_trip() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + 5, + Network::Testnet, + ); + let key_source = test_key_source(); + + let reserved = pool.next_unused_and_reserve(&key_source, 12_345).unwrap(); + let index = pool.address_index(&reserved).unwrap(); + + // Unlike an ephemeral reservation, the reserved state must survive a + // round trip so a handed-out address is not re-issued after a restart. + let json = serde_json::to_string(&pool).unwrap(); + let mut restored: AddressPool = serde_json::from_str(&json).unwrap(); + + assert_eq!(restored.addresses.get(&index).unwrap().reserved_at(), Some(12_345)); + assert!(restored.addresses.get(&index).unwrap().is_reserved()); + + // A fresh hand-out on the restored pool skips the still-reserved address. + let next = restored.next_unused_and_reserve(&key_source, 12_346).unwrap(); + assert_ne!(next, reserved); + } + + #[cfg(feature = "bincode")] + #[test] + fn test_reservation_survives_bincode_round_trip() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + 5, + Network::Testnet, + ); + let key_source = test_key_source(); + + let reserved = pool.next_unused_and_reserve(&key_source, 12_345).unwrap(); + let index = pool.address_index(&reserved).unwrap(); + + // The reserved state must survive a bincode round trip so a handed-out + // address is not re-issued after a restart. + let bytes = bincode::encode_to_vec(&pool, bincode::config::standard()).unwrap(); + let (mut restored, _): (AddressPool, usize) = + bincode::decode_from_slice(&bytes, bincode::config::standard()).unwrap(); + + assert_eq!(restored.addresses.get(&index).unwrap().reserved_at(), Some(12_345)); + assert!(restored.addresses.get(&index).unwrap().is_reserved()); + + // A fresh hand-out on the restored pool skips the still-reserved address. + let next = restored.next_unused_and_reserve(&key_source, 12_346).unwrap(); + assert_ne!(next, reserved); + } + + #[test] + fn test_sweep_expired_reservations() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + 5, + Network::Testnet, + ); + let key_source = test_key_source(); + + let old = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + let recent = pool.next_unused_and_reserve(&key_source, 200).unwrap(); + let clockless = pool.next_unused_and_reserve(&key_source, 0).unwrap(); + let old_index = pool.address_index(&old).unwrap(); + let recent_index = pool.address_index(&recent).unwrap(); + let clockless_index = pool.address_index(&clockless).unwrap(); + + // now == 0 means no clock: nothing is reclaimed. + assert_eq!(pool.sweep_expired_reservations(0, 100), 0); + assert!(pool.addresses.get(&old_index).unwrap().is_reserved()); + + // ttl == 0 means no expiry window: live reservations are left intact. + assert_eq!(pool.sweep_expired_reservations(250, 0), 0); + assert!(pool.addresses.get(&old_index).unwrap().is_reserved()); + assert!(pool.addresses.get(&recent_index).unwrap().is_reserved()); + + // At now=250 with ttl=100: old (age 150) expires, recent (age 50) stays. + // A clockless reservation (reserved_at == 0) is never aged out. + assert_eq!(pool.sweep_expired_reservations(250, 100), 1); + assert!(!pool.addresses.get(&old_index).unwrap().is_reserved()); + assert!(pool.addresses.get(&recent_index).unwrap().is_reserved()); + assert!(pool.addresses.get(&clockless_index).unwrap().is_reserved()); + + // The reclaimed address is available for hand-out again. + let reissued = pool.next_unused_and_reserve(&key_source, 300).unwrap(); + assert_eq!(old, reissued); + } + + #[test] + fn test_reset_usage_clears_reservations() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + 5, + Network::Testnet, + ); + let key_source = test_key_source(); + + let reserved = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + let index = pool.address_index(&reserved).unwrap(); + assert!(pool.addresses.get(&index).unwrap().is_reserved()); + + // A rescan reset must drop reservations too, otherwise the address stays + // pinned forever and is never handed out again. + pool.reset_usage(); + assert!(!pool.addresses.get(&index).unwrap().is_reserved()); + + let reissued = pool.next_unused_and_reserve(&key_source, 200).unwrap(); + assert_eq!(reserved, reissued); + } + + #[test] + fn test_prune_never_drops_reserved() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let gap_limit = 5; + let mut pool = AddressPool::new_without_generation( + base_path, + AddressPoolType::External, + gap_limit, + Network::Testnet, + ); + let key_source = test_key_source(); + pool.generate_addresses(10, &key_source, true).unwrap(); + + // Mark index 0 used so keep_until = highest_used + gap_limit = 5. + pool.mark_index_used(0); + + // Reserve indices 1..=7 (first available each time, index 0 is used). + let mut reserved = Vec::new(); + for t in 0..7 { + reserved.push(pool.next_unused_and_reserve(&key_source, 100 + t).unwrap()); + } + + // Indices 6 and 7 are reserved and beyond keep_until. Indices 8 and 9 + // are plain unused beyond keep_until and are the only prunable entries. + let pruned = pool.prune_unused(); + assert_eq!(pruned, 2); + assert!(!pool.addresses.contains_key(&8)); + assert!(!pool.addresses.contains_key(&9)); + + // Every reserved address survived the prune. + for addr in &reserved { + assert!(pool.address_index(addr).is_some()); + } + assert!(pool.addresses.get(&6).unwrap().is_reserved()); + assert!(pool.addresses.get(&7).unwrap().is_reserved()); + } + #[test] fn test_gap_limit_maintenance() { let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); @@ -1269,6 +1711,29 @@ mod tests { assert_eq!(pool.addresses.len(), gap_limit as usize + 3); } + #[test] + fn test_reserved_addresses_count_against_gap_limit() { + let base_path = DerivationPath::from(vec![ChildNumber::from_normal_idx(0).unwrap()]); + let key_source = test_key_source(); + let gap_limit = 5; + + // A pool with exactly gap_limit available addresses needs no more. + let mut pool = AddressPool::new( + base_path, + AddressPoolType::External, + gap_limit, + Network::Testnet, + &key_source, + ) + .unwrap(); + assert!(!pool.needs_more_addresses()); + + // Reserving one drops available headroom below the gap limit, so the + // pool now reports it needs more even though nothing was used. + pool.next_unused_and_reserve(&key_source, 100).unwrap(); + assert!(pool.needs_more_addresses()); + } + #[test] fn test_address_pool_builder() { let pool = AddressPoolBuilder::new() diff --git a/key-wallet/src/managed_account/managed_core_funds_account.rs b/key-wallet/src/managed_account/managed_core_funds_account.rs index cf95d78d1..8e674fa41 100644 --- a/key-wallet/src/managed_account/managed_core_funds_account.rs +++ b/key-wallet/src/managed_account/managed_core_funds_account.rs @@ -561,6 +561,93 @@ impl ManagedCoreFundsAccount { } } + /// Generate and reserve the next receive address. + /// + /// Like [`Self::next_receive_address`] but atomically reserves the returned + /// address so a concurrent hand-out cannot return the same one. The + /// reservation persists until funds arrive (promoting it to used), + /// [`Self::release_receive_reservation`] is called, or it is reclaimed by a + /// TTL sweep. `now` is a caller-supplied timestamp (seconds) stamping the + /// reservation. Only valid for Standard accounts. + /// + /// Derivation of fresh addresses on this path is unbounded, so callers are + /// assumed trusted. See [`address_pool::AddressPool::next_unused_and_reserve`]. + pub fn next_receive_address_and_reserve( + &mut self, + account_xpub: Option<&ExtendedPubKey>, + now: u64, + ) -> Result { + if let ManagedAccountType::Standard { + external_addresses, + .. + } = self.keys.managed_account_type_mut() + { + let key_source = match account_xpub { + Some(xpub) => address_pool::KeySource::Public(*xpub), + None => address_pool::KeySource::NoKeySource, + }; + + let addr = external_addresses.next_unused_and_reserve(&key_source, now).map_err( + |e| match e { + crate::error::Error::NoKeySource => { + "No unused addresses available and no key source provided" + } + _ => "Failed to generate receive address", + }, + )?; + self.keys.bump_monitor_revision(); + Ok(addr) + } else { + Err("Cannot generate receive address for non-standard account type") + } + } + + /// Release a previously reserved receive address back to the available pool. + /// + /// Idempotent: returns `false` if the address is unknown to the external + /// pool or is not currently reserved. Bumps the monitor revision only when + /// a reservation is actually cleared. + pub fn release_receive_reservation(&mut self, address: &Address) -> bool { + if let ManagedAccountType::Standard { + external_addresses, + .. + } = self.keys.managed_account_type_mut() + { + if let Some(index) = external_addresses.address_index(address) { + if external_addresses.release_reservation(index) { + self.keys.bump_monitor_revision(); + return true; + } + } + } + false + } + + /// Reclaim receive reservations older than `ttl`, returning their addresses + /// to the available pool. + /// + /// Backstop for reservations handed out but never funded and never + /// explicitly released, which would otherwise pin gap-limit headroom + /// forever. `now` and `ttl` are caller-supplied (seconds); the wallet keeps + /// no clock. Returns the number of reservations reclaimed and bumps the + /// monitor revision when that is non-zero. See + /// [`address_pool::AddressPool::sweep_expired_reservations`]. + pub fn sweep_expired_receive_reservations(&mut self, now: u64, ttl: u64) -> usize { + if let ManagedAccountType::Standard { + external_addresses, + .. + } = self.keys.managed_account_type_mut() + { + let reclaimed = external_addresses.sweep_expired_reservations(now, ttl); + if reclaimed > 0 { + self.keys.bump_monitor_revision(); + } + reclaimed + } else { + 0 + } + } + /// Generate multiple receive addresses at once using the optionally provided extended public key. /// Only valid for Standard accounts. pub fn next_receive_addresses( diff --git a/key-wallet/src/tests/address_pool_tests.rs b/key-wallet/src/tests/address_pool_tests.rs index 25ec66b1a..7da10b5cb 100644 --- a/key-wallet/src/tests/address_pool_tests.rs +++ b/key-wallet/src/tests/address_pool_tests.rs @@ -54,15 +54,20 @@ fn test_next_unused_multiple() { pool.mark_used(&addresses[0]); pool.mark_used(&addresses[2]); - // Request more addresses - should get 3 unused + 2 new + // Reserve the first remaining unused address: it must be skipped too. + let reserved = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + assert_eq!(reserved, addresses[1]); + + // Request more addresses: should skip used and reserved, returning the + // remaining generated ones plus freshly derived ones to satisfy the count. let more_addresses = pool.next_unused_multiple(5, &key_source, true); assert_eq!(more_addresses.len(), 5); - assert_eq!(more_addresses[0], addresses[1]); // First unused - assert_eq!(more_addresses[1], addresses[3]); // Second unused - assert_eq!(more_addresses[2], addresses[4]); // Third unused - // more_addresses[3] and [4] should be newly generated + assert!(!more_addresses.contains(&reserved)); + assert_eq!(more_addresses[0], addresses[3]); // First available + assert_eq!(more_addresses[1], addresses[4]); // Second available + // remaining entries are newly generated - assert_eq!(pool.highest_generated, Some(6)); // Generated 2 more + assert_eq!(pool.highest_generated, Some(7)); // Generated 3 more } #[test] @@ -84,18 +89,23 @@ fn test_next_unused_multiple_with_info() { for (i, (addr, info)) in address_infos.iter().enumerate() { assert_eq!(addr, &info.address); assert_eq!(info.index, i as u32); - assert!(!info.used); + assert!(!info.is_used()); assert!(info.public_key.is_some()); } // Mark first one as used pool.mark_used(&address_infos[0].0); - // Get more with info - should skip the used one + // Reserve the second one: it must be skipped alongside the used one. + let reserved = pool.next_unused_and_reserve(&key_source, 100).unwrap(); + assert_eq!(reserved, address_infos[1].0); + + // Get more with info: should skip both the used and the reserved address, + // still returning the requested count by deriving fresh addresses. let more_infos = pool.next_unused_multiple_with_info(3, &key_source, true); assert_eq!(more_infos.len(), 3); - assert_eq!(more_infos[0].0, address_infos[1].0); // Should skip the first (used) one - assert_eq!(more_infos[1].0, address_infos[2].0); + assert!(!more_infos.iter().any(|(addr, _)| addr == &reserved)); + assert_eq!(more_infos[0].0, address_infos[2].0); // First available after used + reserved } #[test] diff --git a/key-wallet/src/tests/address_reservation_tests.rs b/key-wallet/src/tests/address_reservation_tests.rs new file mode 100644 index 000000000..6f8e75a45 --- /dev/null +++ b/key-wallet/src/tests/address_reservation_tests.rs @@ -0,0 +1,141 @@ +//! End-to-end coverage for receive-address reservation through the managed +//! wallet: two reserved hand-outs are distinct, a reservation survives until +//! funds arrive, and the normal transaction-checking path promotes the +//! reserved address to used (clearing the reservation). + +use crate::account::{Account, AccountType, ManagedCoreFundsAccount}; +use crate::managed_account::address_pool::KeySource; +use crate::managed_account::managed_account_trait::ManagedAccountTrait; +use crate::managed_account::managed_account_type::ManagedAccountType; +use crate::mnemonic::{Language, Mnemonic}; +use crate::test_utils::TestWalletContext; +use crate::transaction_checking::TransactionContext; +use crate::{ExtendedPrivKey, Network}; +use dashcore::{Address, Transaction}; +use secp256k1::Secp256k1; +use std::str::FromStr; + +#[tokio::test] +async fn test_reserved_receive_address_promotes_on_funding() { + let mut ctx = TestWalletContext::new_random(); + let xpub = ctx.xpub; + + let revision_before = ctx + .managed_wallet + .first_bip44_managed_account_mut() + .expect("managed account") + .monitor_revision(); + + let a = ctx + .managed_wallet + .first_bip44_managed_account_mut() + .expect("managed account") + .next_receive_address_and_reserve(Some(&xpub), 1_000) + .expect("reserve a"); + let b = ctx + .managed_wallet + .first_bip44_managed_account_mut() + .expect("managed account") + .next_receive_address_and_reserve(Some(&xpub), 1_000) + .expect("reserve b"); + + // Each reservation must bump the monitor revision so observers re-read the + // address set. + let revision_after_reserve = ctx + .managed_wallet + .first_bip44_managed_account_mut() + .expect("managed account") + .monitor_revision(); + assert!(revision_after_reserve > revision_before); + + // Two sequential receive hand-outs return distinct addresses: reserving the + // first removes it from the available pool seen by the second. + assert_ne!(a, b); + + // Funds arriving at `a` promote it Reserved -> Used through the normal + // checker, which clears the reservation at the single mark_used chokepoint. + let tx = Transaction::dummy(&a, 0..1, &[150_000]); + let result = ctx.check_transaction(&tx, TransactionContext::Mempool).await; + assert!(result.is_relevant); + + let acc = ctx.managed_wallet.first_bip44_managed_account_mut().expect("managed account"); + // `a` is used now, so its reservation is already gone: releasing is a no-op. + assert!(!acc.release_receive_reservation(&a)); + // `b` is still reserved, so releasing it returns it to the available pool + // and bumps the monitor revision again. + let revision_before_release = acc.monitor_revision(); + assert!(acc.release_receive_reservation(&b)); + assert!(acc.monitor_revision() > revision_before_release); + + // Releasing an address the pool has never seen is a no-op. + let stray = Address::from_str("yTb47qEBpNmgXvYYsHEN4nh8yJwa5iC4Cs") + .unwrap() + .require_network(Network::Testnet) + .unwrap(); + assert!(!acc.release_receive_reservation(&stray)); +} + +#[test] +fn test_reserve_receive_without_key_source_errors() { + // A standard account whose external pool has no pre-generated addresses + // and no key source cannot hand out a reservation. + let mut account = ManagedCoreFundsAccount::dummy_bip44(); + let err = account + .next_receive_address_and_reserve(None, 1_000) + .expect_err("no key source and an empty pool must fail"); + assert_eq!(err, "No unused addresses available and no key source provided"); +} + +#[test] +fn test_reserve_receive_on_non_standard_account_errors() { + let network = Network::Testnet; + let mnemonic = Mnemonic::from_phrase( + "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about", + Language::English, + ) + .unwrap(); + let seed = mnemonic.to_seed(""); + let master = ExtendedPrivKey::new_master(network, &seed).unwrap(); + let secp = Secp256k1::new(); + + let account_type = AccountType::CoinJoin { + index: 0, + }; + let derivation_path = account_type.derivation_path(network).unwrap(); + let account_key = master.derive_priv(&secp, &derivation_path).unwrap(); + let account = Account::from_xpriv(Some([0u8; 32]), account_type, account_key, network).unwrap(); + + let key_source = KeySource::Public(account.account_xpub); + let managed_type = + ManagedAccountType::from_account_type(account.account_type, network, &key_source).unwrap(); + let mut managed = ManagedCoreFundsAccount::new(managed_type, network); + + let err = managed + .next_receive_address_and_reserve(Some(&account.account_xpub), 1_000) + .expect_err("non-standard account must not hand out receive reservations"); + assert_eq!(err, "Cannot generate receive address for non-standard account type"); + + // The sweep wrapper is likewise inert on a non-standard account. + assert_eq!(managed.sweep_expired_receive_reservations(10_000, 1), 0); +} + +#[test] +fn test_sweep_expired_receive_reservations_reclaims_through_account() { + let mut ctx = TestWalletContext::new_random(); + let xpub = ctx.xpub; + + let acc = ctx.managed_wallet.first_bip44_managed_account_mut().expect("managed account"); + let reserved = acc.next_receive_address_and_reserve(Some(&xpub), 1_000).expect("reserve"); + let revision_after_reserve = acc.monitor_revision(); + + // A sweep before the ttl elapses leaves the reservation intact and does not + // bump the monitor revision. + assert_eq!(acc.sweep_expired_receive_reservations(1_000, 500), 0); + assert_eq!(acc.monitor_revision(), revision_after_reserve); + + // Once the reservation is older than the ttl it is reclaimed, the revision + // bumps, and releasing the now-available address is a no-op. + assert_eq!(acc.sweep_expired_receive_reservations(2_000, 500), 1); + assert!(acc.monitor_revision() > revision_after_reserve); + assert!(!acc.release_receive_reservation(&reserved)); +} diff --git a/key-wallet/src/tests/mod.rs b/key-wallet/src/tests/mod.rs index 8db87a128..9650106f1 100644 --- a/key-wallet/src/tests/mod.rs +++ b/key-wallet/src/tests/mod.rs @@ -8,6 +8,8 @@ mod balance_tests; mod address_pool_tests; +mod address_reservation_tests; + mod advanced_transaction_tests; mod backup_restore_tests; From fc0668f015b89824aceb9f27176fdd1cac43e83b Mon Sep 17 00:00:00 2001 From: xdustinface Date: Tue, 23 Jun 2026 15:27:52 +1000 Subject: [PATCH 2/2] fix(key-wallet): return `InvalidState` instead of panicking on pool invariant miss `next_unused_and_reserve` and `maintain_gap_limit` guarded the "entry must exist after `generate_address_at_index(add_to_state=true)`" invariant with `expect()` and `panic!()`. This crate's error-handling philosophy is to never panic in library code, so both sites now propagate a new `Error::InvalidState` variant through their existing `Result` return type. The FFI maps it to `FFIErrorCode::InvalidState`. Addresses CodeRabbit review comment on PR #818 https://github.com/dashpay/rust-dashcore/pull/818#discussion_r3456738207 --- key-wallet-ffi/src/error.rs | 7 +++-- key-wallet/src/error.rs | 3 ++ .../src/managed_account/address_pool.rs | 31 ++++++++++--------- 3 files changed, 24 insertions(+), 17 deletions(-) diff --git a/key-wallet-ffi/src/error.rs b/key-wallet-ffi/src/error.rs index 683d8e900..3b0a0e189 100644 --- a/key-wallet-ffi/src/error.rs +++ b/key-wallet-ffi/src/error.rs @@ -250,9 +250,10 @@ impl From for FFIError { Error::InvalidNetwork => FFIErrorCode::InvalidNetwork, Error::InvalidAddress(_) => FFIErrorCode::InvalidAddress, Error::Serialization(_) => FFIErrorCode::SerializationError, - Error::WatchOnly | Error::CoinJoinNotEnabled | Error::NoKeySource => { - FFIErrorCode::InvalidState - } + Error::InvalidState(_) + | Error::WatchOnly + | Error::CoinJoinNotEnabled + | Error::NoKeySource => FFIErrorCode::InvalidState, Error::Bip32(_) | Error::Slip10(_) | Error::BLS(_) diff --git a/key-wallet/src/error.rs b/key-wallet/src/error.rs index 5e3105e7d..ab9831512 100644 --- a/key-wallet/src/error.rs +++ b/key-wallet/src/error.rs @@ -37,6 +37,8 @@ pub enum Error { Serialization(String), /// Invalid parameter InvalidParameter(String), + /// Internal invariant violated (inconsistent wallet state) + InvalidState(String), /// Watch-only wallet (no private keys available) WatchOnly, /// No key source available for address derivation @@ -61,6 +63,7 @@ impl fmt::Display for Error { Error::CoinJoinNotEnabled => write!(f, "CoinJoin not enabled for this account"), Error::Serialization(s) => write!(f, "Serialization error: {}", s), Error::InvalidParameter(s) => write!(f, "Invalid parameter: {}", s), + Error::InvalidState(s) => write!(f, "Invalid state: {}", s), Error::WatchOnly => write!(f, "Watch-only wallet: private keys not available"), Error::NoKeySource => write!(f, "No key source available for address derivation"), } diff --git a/key-wallet/src/managed_account/address_pool.rs b/key-wallet/src/managed_account/address_pool.rs index 472b6719a..2daa1f653 100644 --- a/key-wallet/src/managed_account/address_pool.rs +++ b/key-wallet/src/managed_account/address_pool.rs @@ -671,12 +671,15 @@ impl AddressPool { // `generate_address_at_index` with `add_to_state = true` always inserts // at `next_index`. If `highest_generated` ever drifts out of sync with // `addresses` a silent lookup miss would return an unreserved address, - // reintroducing the hand-out race this method exists to close, so fail - // hard in release rather than fail open. - let info = self - .addresses - .get_mut(&next_index) - .expect("generate_address_at_index(add_to_state=true) must insert at next_index"); + // reintroducing the hand-out race this method exists to close, so report + // the broken invariant as an error rather than fail open. + let info = self.addresses.get_mut(&next_index).ok_or_else(|| { + Error::InvalidState(format!( + "next_unused_and_reserve: generate_address_at_index({}) succeeded but \ + the entry was not stored; pool invariant broken", + next_index + )) + })?; debug_assert!(info.is_available()); info.state = AddressState::Reserved { at: now, @@ -1075,17 +1078,17 @@ impl AddressPool { let next_index = self.highest_generated.map(|h| h + 1).unwrap_or(0); self.generate_address_at_index(next_index, key_source, true)?; // `generate_address_at_index` with `add_to_state = true` always - // inserts at `next_index`. Asserting the invariant explicitly - // here turns a regression that breaks it (e.g. a refactor that - // hits the early-return branch on a re-derivation) into a loud - // panic instead of an infinite loop on the outer `while`. - let info = self.addresses.get(&next_index).cloned().unwrap_or_else(|| { - panic!( + // inserts at `next_index`. Surfacing a regression that breaks it + // (e.g. a refactor that hits the early-return branch on a + // re-derivation) as an error keeps the outer `while` from spinning + // into an infinite loop. + let info = self.addresses.get(&next_index).cloned().ok_or_else(|| { + Error::InvalidState(format!( "maintain_gap_limit: generate_address_at_index({}) succeeded but \ the entry was not stored; pool invariant broken", next_index - ) - }); + )) + })?; new_addresses.push(info); }