feat(key-wallet): owner-tagged reservations to close the broadcast-release TOCTOU (platform#4185)

key-wallet/src/lib.rs

@@ -60,6 +60,7 @@ pub use managed_account::address_pool::{AddressInfo, AddressPool, KeySource, Poo
 pub use managed_account::managed_account_type::ManagedAccountType;
 pub use managed_account::managed_platform_account::ManagedPlatformAccount;
 pub use managed_account::platform_address::PlatformP2PKHAddress;
+pub use managed_account::ReservationToken;
 pub use mnemonic::{Language, Mnemonic};
 pub use seed::Seed;
 pub use signer::{ExtendedPubKeySigner, Signer, SignerMethod, TransactionCategory};

key-wallet/src/managed_account/managed_core_funds_account.rs

@@ -19,7 +19,7 @@ use crate::managed_account::address_pool;
 use crate::managed_account::managed_account_trait::ManagedAccountTrait;
 use crate::managed_account::managed_account_type::ManagedAccountType;
 use crate::managed_account::managed_core_keys_account::ManagedCoreKeysAccount;
-use crate::managed_account::reservation::ReservationSet;
+use crate::managed_account::reservation::{ReservationSet, ReservationToken};
 use crate::managed_account::transaction_record::{
     InputDetail, OutputDetail, OutputRole, TransactionDirection,
 };
@@ -130,6 +130,37 @@ impl ManagedCoreFundsAccount {
         self.reservations.release(tx.input.iter().map(|input| &input.previous_output));
     }
 
+    /// Owner-guarded release of `tx`'s input reservations: releases an input
+    /// only if it is *still owned by* `token`, the [`ReservationToken`] returned
+    /// when this build reserved its inputs (from
+    /// [`build_unsigned_reserved`]/[`build_signed_reserved`]).
+    ///
+    /// This is the release a caller must use when it abandons a transaction
+    /// *after having `.await`ed something* between reserving and releasing —
+    /// above all the platform broadcast path, which reserves an
+    /// asset-lock/deferred send's inputs, awaits the broadcast, and on a
+    /// `Rejected` result releases them. During that await key-wallet's TTL sweep
+    /// can invisibly reclaim the reservation and a different concurrent build can
+    /// re-reserve the same outpoint (under a new token, same wallet generation).
+    /// The unconditional [`Self::release_reservation`] would then free the other
+    /// build's inputs, letting coin selection hand them to a second transaction —
+    /// a double-spend window. Passing the original token makes the release a
+    /// no-op for any input that has since changed owners, closing that window.
+    /// The check is atomic under the reservation set's mutex, so no sweep or
+    /// re-reserve can interleave between "is it still mine?" and the removal.
+    ///
+    /// The platform layer cannot make this safe on its own because the sweep
+    /// happens inside key-wallet where it has no visibility; the owner check must
+    /// live here. See `dashpay/platform#4185`.
+    ///
+    /// [`build_unsigned_reserved`]: crate::wallet::managed_wallet_info::transaction_builder::TransactionBuilder::build_unsigned_reserved
+    /// [`build_signed_reserved`]: crate::wallet::managed_wallet_info::transaction_builder::TransactionBuilder::build_signed_reserved
+    pub fn release_reservation_if_owner(&self, tx: &Transaction, token: ReservationToken) {
+        let outpoints: Vec<OutPoint> =
+            tx.input.iter().map(|input| input.previous_output).collect();
+        self.reservations.release_if_owner(&outpoints, token);
+    }
+
     /// Get a reference to the inner keys-account state.
     pub fn keys(&self) -> &ManagedCoreKeysAccount {
         &self.keys

key-wallet/src/managed_account/mod.rs

@@ -27,3 +27,4 @@ pub mod transaction_record;
 pub use managed_account_ref::{ManagedAccountRef, ManagedAccountRefMut, OwnedManagedCoreAccount};
 pub use managed_core_funds_account::ManagedCoreFundsAccount;
 pub use managed_core_keys_account::ManagedCoreKeysAccount;
+pub use reservation::ReservationToken;

key-wallet/src/managed_account/reservation.rs

@@ -12,6 +12,26 @@
 //! held survives the lock being released and is observed by the next build. It
 //! is never persisted: after a restart it is empty, where chain and mempool
 //! sync are the source of truth for which coins are already spent.
+//!
+//! # Why reservations carry an owner token
+//!
+//! Releasing a reservation by outpoint alone is unsafe once *releasing* and
+//! *re-reserving* can interleave. Every reservation is therefore stamped with a
+//! [`ReservationToken`] identifying the build that made it, and the release path
+//! that a rejected/abandoned broadcast uses ([`ReservationSet::release_if_owner`])
+//! removes an outpoint only if it is *still owned by the releasing build*.
+//!
+//! The concrete hazard (see `dashpay/platform#4185`): the platform broadcast
+//! path reserves an asset-lock/deferred send's inputs, `.await`s the broadcast,
+//! and on a `Rejected` result releases those inputs so they become spendable
+//! again. During that await the TTL sweep below can reclaim the reservation, and
+//! a *different* concurrent build can re-reserve the very same outpoint (same
+//! wallet generation, so any `Arc::ptr_eq` generation guard still matches). An
+//! unconditional release-by-outpoint would then free the *other* build's inputs,
+//! letting coin selection hand them to a second transaction — a double-spend
+//! window. The platform layer cannot detect this because the sweep happens
+//! inside key-wallet invisibly; the "release only if still mine" check must be
+//! atomic under this set's own mutex, which is exactly what an owner token buys.
 
 use std::collections::{HashMap, HashSet};
 use std::sync::{Arc, Mutex, MutexGuard};
@@ -31,59 +51,161 @@ use dashcore::blockdata::transaction::OutPoint;
 /// the very double-spend this guards against.
 const RESERVATION_TTL_BLOCKS: u32 = 24;
 
+/// Opaque, per-`reserve`-call identity stamped onto every outpoint that call
+/// reserves.
+///
+/// A token is the *proof of ownership* a build presents to the reservation
+/// set's `release_if_owner` so a release only removes the build's own
+/// reservation. Each `reserve` call mints a fresh token from a monotonic
+/// counter, so two reservations never share one — not even two reservations
+/// taken at the same block height, which is why the height (which collides
+/// freely) cannot serve as the identity.
+///
+/// The inner counter is private and there is no public constructor: a token can
+/// only originate from a real `reserve` call. That is deliberate — it prevents a
+/// caller from forging a token that happens to match another build's ownership
+/// and releasing inputs out from under it.
+///
+/// Copy semantics let a build hold its token cheaply across an `.await` (e.g.
+/// the platform broadcast path in `dashpay/platform#4185`) and present it again
+/// when releasing.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct ReservationToken(u64);
+
+/// A single reserved outpoint: when it was reserved (for the TTL backstop) and
+/// which build owns it (for owner-guarded release).
+#[derive(Debug, Clone, Copy)]
+struct Reservation {
+    /// Block height at which the reservation was taken; drives the TTL sweep.
+    reserved_at_height: u32,
+    /// The build that reserved this outpoint. Only this owner may release it via
+    /// [`ReservationSet::release_if_owner`].
+    owner: ReservationToken,
+}
+
+/// Mutex-guarded interior: the reservations keyed by outpoint plus the counter
+/// that mints the next [`ReservationToken`].
+#[derive(Debug, Default)]
+struct Reserved {
+    entries: HashMap<OutPoint, Reservation>,
+    /// Monotonically increasing source of unique tokens. Never persisted and
+    /// only ever incremented, so within a process every issued token is unique.
+    /// (Wraparound would need ~2^64 reserves in one process lifetime, which is
+    /// unreachable in practice.)
+    next_token: u64,
+}
+
 /// Ephemeral, in-memory set of reserved outpoints. Cloning shares the
 /// underlying state, which is what lets a build's reservation outlive the
 /// wallet lock and be seen by the next build.
 #[derive(Debug, Clone, Default)]
 pub(crate) struct ReservationSet {
-    inner: Arc<Mutex<HashMap<OutPoint, u32>>>,
+    inner: Arc<Mutex<Reserved>>,
 }
 
 impl ReservationSet {
     /// Recovers from a poisoned mutex rather than panicking: the guarded data is
-    /// a plain `HashMap` with no invariant a partial write could break, and
-    /// panicking here would strand all later coin selection in a long-running
-    /// node.
-    fn lock(&self) -> MutexGuard<'_, HashMap<OutPoint, u32>> {
+    /// a plain map plus a counter with no invariant a partial write could break,
+    /// and panicking here would strand all later coin selection in a
+    /// long-running node.
+    fn lock(&self) -> MutexGuard<'_, Reserved> {
         self.inner.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
     }
 
-    fn sweep(reserved: &mut HashMap<OutPoint, u32>, current_height: u32) {
+    fn sweep(reserved: &mut Reserved, current_height: u32) {
         // Height 0 means the wallet has no processed height yet, so the elapsed
         // span is unknown and no entry can be reliably judged stale. An entry
         // stamped at height 0 therefore relies on a later non-zero height being
         // presented before it can ever be reclaimed by the TTL backstop.
         if current_height == 0 {
             return;
         }
-        reserved.retain(|_, reserved_at| {
-            current_height.saturating_sub(*reserved_at) < RESERVATION_TTL_BLOCKS
+        reserved.entries.retain(|_, reservation| {
+            current_height.saturating_sub(reservation.reserved_at_height) < RESERVATION_TTL_BLOCKS
         });
     }
 
     /// Reserve `outpoints` as of `current_height`, dropping expired entries
-    /// first. Re-reserving an outpoint refreshes its height.
-    pub(crate) fn reserve(&self, outpoints: &[OutPoint], current_height: u32) {
+    /// first, and return the [`ReservationToken`] stamped onto all of them.
+    ///
+    /// Every outpoint in a single call shares the one returned token: the caller
+    /// keeps it and later presents it to [`Self::release_if_owner`] to release
+    /// only what this call reserved. Re-reserving an outpoint (a later `reserve`
+    /// naming it again) refreshes its height *and* transfers ownership to the new
+    /// token — the previous owner's [`Self::release_if_owner`] then becomes a
+    /// no-op for it, which is precisely the behavior that closes the
+    /// release/re-reserve race described in the module docs (`platform#4185`).
+    pub(crate) fn reserve(&self, outpoints: &[OutPoint], current_height: u32) -> ReservationToken {
         let mut reserved = self.lock();
         Self::sweep(&mut reserved, current_height);
+
+        let owner = ReservationToken(reserved.next_token);
+        // Increment even on an empty `outpoints` slice so a token is never
+        // reissued; wrapping is documented-unreachable but avoids a debug panic.
+        reserved.next_token = reserved.next_token.wrapping_add(1);
+
         for outpoint in outpoints {
-            reserved.insert(*outpoint, current_height);
+            reserved.entries.insert(
+                *outpoint,
+                Reservation {
+                    reserved_at_height: current_height,
+                    owner,
+                },
+            );
         }
+        owner
     }
 
     /// Return the currently reserved outpoints, dropping expired entries first.
     pub(crate) fn reserved(&self, current_height: u32) -> HashSet<OutPoint> {
         let mut reserved = self.lock();
         Self::sweep(&mut reserved, current_height);
-        reserved.keys().copied().collect()
+        reserved.entries.keys().copied().collect()
     }
 
-    /// Release the given outpoints. Idempotent: releasing an outpoint that is
-    /// not reserved does nothing.
+    /// Unconditionally release the given outpoints, regardless of owner.
+    /// Idempotent: releasing an outpoint that is not reserved does nothing.
+    ///
+    /// This is correct only where the coins are *known spent* and so must leave
+    /// the set no matter which build reserved them — e.g. when a processed spend
+    /// hands its inputs from this ephemeral set to the durable spent set. A
+    /// caller that is merely abandoning an in-flight build (rejected broadcast,
+    /// cancelled send) must use [`Self::release_if_owner`] instead, so it cannot
+    /// free a reservation another build has since taken over.
     pub(crate) fn release<'a>(&self, outpoints: impl IntoIterator<Item = &'a OutPoint>) {
         let mut reserved = self.lock();
         for outpoint in outpoints {
-            reserved.remove(outpoint);
+            reserved.entries.remove(outpoint);
+        }
+    }
+
+    /// Release only the `outpoints` still owned by `token`, atomically under this
+    /// set's mutex. An outpoint reserved by a different owner — because the TTL
+    /// sweep reclaimed this build's reservation and another build re-reserved it
+    /// meanwhile — is left untouched.
+    ///
+    /// This is the owner-guarded counterpart to [`Self::release`] and the whole
+    /// reason reservations carry a [`ReservationToken`]. It is the release a
+    /// build must use when it abandons an in-flight transaction after having
+    /// `.await`ed something (a broadcast, an external signer): during that await
+    /// its reservation may have been swept and re-taken by a concurrent build,
+    /// and an unconditional release-by-outpoint would free that other build's
+    /// inputs, opening the double-spend window described in the module docs
+    /// (`dashpay/platform#4185`). Because the check and the removal happen
+    /// together while the mutex is held, no sweep or re-reserve can interleave
+    /// between them.
+    ///
+    /// Idempotent and a no-op for any outpoint that is unreserved or owned by a
+    /// different token — including one this build's own reservation already lost
+    /// to a sweep, so a late release after reclamation is harmless.
+    pub(crate) fn release_if_owner(&self, outpoints: &[OutPoint], token: ReservationToken) {
+        let mut reserved = self.lock();
+        for outpoint in outpoints {
+            let owned_by_token =
+                reserved.entries.get(outpoint).is_some_and(|entry| entry.owner == token);
+            if owned_by_token {
+                reserved.entries.remove(outpoint);
+            }
         }
     }
 }
@@ -175,4 +297,86 @@ mod tests {
         set.reserve(&[a], 1);
         assert!(clone.reserved(1).contains(&a));
     }
+
+    #[test]
+    fn each_reserve_call_mints_a_distinct_token() {
+        let set = ReservationSet::default();
+        // Two reserves at the SAME height must still get different tokens — the
+        // whole point of not keying ownership on height, which collides.
+        let token_a = set.reserve(&[outpoint(0x10, 0)], 100);
+        let token_b = set.reserve(&[outpoint(0x11, 0)], 100);
+        assert_ne!(token_a, token_b);
+    }
+
+    #[test]
+    fn release_if_owner_releases_only_the_owning_builds_outpoints() {
+        let set = ReservationSet::default();
+        let mine = outpoint(0x20, 0);
+        let theirs = outpoint(0x21, 0);
+
+        let my_token = set.reserve(&[mine], 100);
+        let _their_token = set.reserve(&[theirs], 100);
+
+        // Releasing with my token frees only my outpoint; theirs is untouched.
+        set.release_if_owner(&[mine, theirs], my_token);
+        assert!(!set.reserved(100).contains(&mine));
+        assert!(set.reserved(100).contains(&theirs));
+    }
+
+    #[test]
+    fn release_if_owner_is_a_noop_for_unreserved_or_wrong_token() {
+        let set = ReservationSet::default();
+        let a = outpoint(0x22, 0);
+        let stale_token = set.reserve(&[a], 100);
+
+        // Simulate the outpoint being released and re-reserved by someone else.
+        set.release([&a]);
+        let _new_token = set.reserve(&[a], 100);
+
+        // The stale token no longer owns `a`, so its release must not remove it.
+        set.release_if_owner(&[a], stale_token);
+        assert!(set.reserved(100).contains(&a));
+
+        // A token for an outpoint that was never reserved is simply ignored.
+        let never = outpoint(0x23, 0);
+        set.release_if_owner(&[never], stale_token);
+        assert!(!set.reserved(100).contains(&never));
+    }
+
+    /// The core TOCTOU regression: reserve X under token A, then simulate the
+    /// TTL sweep reclaiming X and a *different* concurrent build re-reserving X
+    /// under token B. Token A's late `release_if_owner` (the rejected-broadcast
+    /// cleanup) must NOT remove X, because X now belongs to build B — releasing
+    /// it would hand B's input to coin selection and open a double-spend window.
+    /// See `dashpay/platform#4185`.
+    #[test]
+    fn release_if_owner_does_not_free_a_reservation_taken_over_after_a_sweep() {
+        let set = ReservationSet::default();
+        let x = outpoint(0x30, 0);
+
+        // Build A reserves X.
+        let token_a = set.reserve(&[x], 100);
+        assert!(set.reserved(100).contains(&x));
+
+        // TTL sweep reclaims A's reservation mid-await (modeled by advancing the
+        // height past the TTL so the next reserve's sweep drops A's entry)...
+        let swept_height = 100 + RESERVATION_TTL_BLOCKS;
+        // ...and build B re-reserves the very same outpoint under a new token.
+        let token_b = set.reserve(&[x], swept_height);
+        assert_ne!(token_a, token_b);
+        assert!(set.reserved(swept_height).contains(&x));
+
+        // Build A's rejected-broadcast cleanup releases by (outpoint, token A).
+        set.release_if_owner(&[x], token_a);
+
+        // X must still be reserved — it is now owned by build B.
+        assert!(
+            set.reserved(swept_height).contains(&x),
+            "release_if_owner with the stale owner token must not free B's reservation"
+        );
+
+        // And B can still release its own reservation normally.
+        set.release_if_owner(&[x], token_b);
+        assert!(!set.reserved(swept_height).contains(&x));
+    }
 }