Skip to main content

fedimint_wallet_client/
lib.rs

1#![deny(clippy::pedantic)]
2#![allow(clippy::cast_possible_truncation)]
3#![allow(clippy::missing_errors_doc)]
4#![allow(clippy::missing_panics_doc)]
5#![allow(clippy::module_name_repetitions)]
6#![allow(clippy::must_use_candidate)]
7
8pub mod api;
9#[cfg(feature = "cli")]
10mod cli;
11
12mod backup;
13
14pub mod client_db;
15/// Legacy, state-machine based peg-ins, replaced by `pegin_monitor`
16/// but retained for time being to ensure existing peg-ins complete.
17mod deposit;
18pub mod events;
19use events::SendPaymentEvent;
20/// Peg-in monitor: a task monitoring deposit addresses for peg-ins.
21mod pegin_monitor;
22mod withdraw;
23
24use std::collections::{BTreeMap, BTreeSet};
25use std::future;
26use std::sync::Arc;
27use std::time::{Duration, SystemTime};
28
29use anyhow::{Context as AnyhowContext, anyhow, bail, ensure};
30use async_stream::{stream, try_stream};
31use backup::WalletModuleBackup;
32use bitcoin::address::NetworkUnchecked;
33use bitcoin::secp256k1::{All, SECP256K1, Secp256k1};
34use bitcoin::{Address, Network, ScriptBuf};
35use client_db::{DbKeyPrefix, PegInTweakIndexKey, SupportsSafeDepositKey, TweakIdx};
36use fedimint_api_client::api::{DynModuleApi, FederationResult};
37use fedimint_bitcoind::{BitcoindTracked, DynBitcoindRpc, IBitcoindRpc, create_esplora_rpc};
38use fedimint_client_module::module::init::{
39    ClientModuleInit, ClientModuleInitArgs, ClientModuleRecoverArgs,
40};
41use fedimint_client_module::module::recovery::RecoveryProgress;
42use fedimint_client_module::module::{ClientContext, ClientModule, IClientModule, OutPointRange};
43use fedimint_client_module::oplog::UpdateStreamOrOutcome;
44use fedimint_client_module::sm::{Context, DynState, ModuleNotifier, State, StateTransition};
45use fedimint_client_module::transaction::{
46    ClientOutput, ClientOutputBundle, ClientOutputSM, TransactionBuilder,
47};
48use fedimint_client_module::{DynGlobalClientContext, sm_enum_variant_translation};
49use fedimint_core::core::{Decoder, IntoDynInstance, ModuleInstanceId, ModuleKind, OperationId};
50use fedimint_core::db::{
51    AutocommitError, Database, DatabaseTransaction, IDatabaseTransactionOpsCoreTyped,
52};
53use fedimint_core::encoding::{Decodable, Encodable};
54use fedimint_core::envs::{BitcoinRpcConfig, is_running_in_test_env};
55use fedimint_core::module::{
56    Amounts, ApiAuth, ApiVersion, CommonModuleInit, ModuleCommon, ModuleConsensusVersion,
57    ModuleInit, MultiApiVersion,
58};
59use fedimint_core::task::{MaybeSend, MaybeSync, TaskGroup, sleep};
60use fedimint_core::util::backoff_util::background_backoff;
61use fedimint_core::util::{BoxStream, backoff_util, retry};
62use fedimint_core::{
63    BitcoinHash, OutPoint, TransactionId, apply, async_trait_maybe_send, push_db_pair_items,
64    runtime, secp256k1,
65};
66use fedimint_derive_secret::{ChildId, DerivableSecret};
67use fedimint_logging::LOG_CLIENT_MODULE_WALLET;
68pub use fedimint_wallet_common as common;
69use fedimint_wallet_common::config::{FeeConsensus, WalletClientConfig};
70use fedimint_wallet_common::tweakable::Tweakable;
71pub use fedimint_wallet_common::*;
72use futures::{Stream, StreamExt};
73use rand::{Rng, thread_rng};
74use secp256k1::Keypair;
75use serde::{Deserialize, Serialize};
76use strum::IntoEnumIterator;
77use tokio::sync::watch;
78use tracing::{debug, instrument};
79
80use crate::api::WalletFederationApi;
81use crate::backup::{FEDERATION_RECOVER_MAX_GAP, RecoveryStateV2, WalletRecovery};
82use crate::client_db::{
83    ClaimedPegInData, ClaimedPegInKey, ClaimedPegInPrefix, NextPegInTweakIndexKey,
84    PegInPoolCursorKey, PegInTweakIndexData, PegInTweakIndexPrefix, RecoveryFinalizedKey,
85    RecoveryStateKey, SupportsSafeDepositPrefix,
86};
87use crate::deposit::DepositStateMachine;
88use crate::withdraw::{CreatedWithdrawState, WithdrawStateMachine, WithdrawStates};
89
90const WALLET_TWEAK_CHILD_ID: ChildId = ChildId(0);
91
92#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
93pub struct BitcoinTransactionData {
94    /// The bitcoin transaction is saved as soon as we see it so the transaction
95    /// can be re-transmitted if it's evicted from the mempool.
96    pub btc_transaction: bitcoin::Transaction,
97    /// Index of the deposit output
98    pub out_idx: u32,
99}
100
101#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
102pub enum DepositStateV1 {
103    WaitingForTransaction,
104    WaitingForConfirmation(BitcoinTransactionData),
105    Confirmed(BitcoinTransactionData),
106    Claimed(BitcoinTransactionData),
107    Failed(String),
108}
109
110#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
111pub enum DepositStateV2 {
112    WaitingForTransaction,
113    WaitingForConfirmation {
114        #[serde(with = "bitcoin::amount::serde::as_sat")]
115        btc_deposited: bitcoin::Amount,
116        btc_out_point: bitcoin::OutPoint,
117    },
118    Confirmed {
119        #[serde(with = "bitcoin::amount::serde::as_sat")]
120        btc_deposited: bitcoin::Amount,
121        btc_out_point: bitcoin::OutPoint,
122    },
123    Claimed {
124        #[serde(with = "bitcoin::amount::serde::as_sat")]
125        btc_deposited: bitcoin::Amount,
126        btc_out_point: bitcoin::OutPoint,
127    },
128    Failed(String),
129}
130
131/// Deposit address allocated by this client.
132#[derive(Debug, Clone, PartialEq, Eq)]
133pub struct DepositAddressInfo {
134    pub operation_id: OperationId,
135    pub address: Address,
136    pub tweak_idx: TweakIdx,
137}
138
139/// Result of [`WalletClientModule::allocate_deposit_address_pooled_stateless`].
140///
141/// Callers that need a custom address-picking strategy can use
142/// [`MaybeNewAddress::TooManyUnusedAddresses`] to see all currently reusable
143/// addresses instead of relying on this crate's round-robin policy.
144#[allow(clippy::enum_variant_names)]
145#[derive(Debug, Clone, PartialEq, Eq)]
146pub enum MaybeNewAddress {
147    /// A new tweak/operation was created on this call.
148    NewAddress(DepositAddressInfo),
149    /// The unused-address gap is full. No new address was allocated.
150    ///
151    /// Reusable unused addresses are ordered by `creation_time` ascending.
152    TooManyUnusedAddresses(Vec<DepositAddressInfo>),
153}
154
155/// Outcome of [`WalletClientModule::allocate_deposit_address_pooled`], so
156/// callers can decide whether to perform per-operation initialization (notes,
157/// fee bookkeeping, metadata writes) — for `Reused` returns, that work was
158/// already done at the time of the original `Fresh` allocation and must not be
159/// repeated.
160#[derive(Debug, Clone, Copy, PartialEq, Eq)]
161pub enum AllocateDepositOutcome {
162    /// A new tweak/operation was created on this call.
163    Fresh,
164    /// An existing unused address was returned. The carried `TweakIdx` is the
165    /// same as the one returned in the tuple and refers to the original
166    /// allocation; it's exposed here as well for explicit diagnostic use.
167    Reused { original_tweak_idx: TweakIdx },
168}
169
170#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
171pub enum WithdrawState {
172    Created,
173    Succeeded(bitcoin::Txid),
174    Failed(String),
175    // TODO: track refund
176    // Refunded,
177    // RefundFailed(String),
178}
179
180async fn next_withdraw_state<S>(stream: &mut S) -> Option<WithdrawStates>
181where
182    S: Stream<Item = WalletClientStates> + Unpin,
183{
184    loop {
185        if let WalletClientStates::Withdraw(ds) = stream.next().await? {
186            return Some(ds.state);
187        }
188        tokio::task::yield_now().await;
189    }
190}
191
192#[derive(Debug, Clone, Default)]
193// TODO: should probably move to DB
194pub struct WalletClientInit(pub Option<DynBitcoindRpc>);
195
196const SLICE_SIZE: u64 = 1000;
197
198impl WalletClientInit {
199    pub fn new(rpc: DynBitcoindRpc) -> Self {
200        Self(Some(rpc))
201    }
202
203    async fn recover_from_slices(
204        &self,
205        args: &ClientModuleRecoverArgs<Self>,
206    ) -> anyhow::Result<()> {
207        let data = WalletClientModuleData {
208            cfg: args.cfg().clone(),
209            module_root_secret: args.module_root_secret().clone(),
210        };
211
212        let total_items = args.module_api().fetch_recovery_count().await?;
213
214        let mut state = RecoveryStateV2::new();
215
216        state.refill_pending_pool_up_to(&data, TweakIdx(FEDERATION_RECOVER_MAX_GAP));
217
218        for start in (0..total_items).step_by(SLICE_SIZE as usize) {
219            let end = std::cmp::min(start + SLICE_SIZE, total_items);
220
221            let items = args.module_api().fetch_recovery_slice(start, end).await?;
222
223            for item in &items {
224                match item {
225                    RecoveryItem::Input { outpoint, script } => {
226                        state.handle_item(*outpoint, script, &data);
227                    }
228                }
229            }
230
231            args.update_recovery_progress(RecoveryProgress {
232                complete: end.try_into().unwrap_or(u32::MAX),
233                total: total_items.try_into().unwrap_or(u32::MAX),
234            });
235        }
236
237        let mut dbtx = args.db().begin_transaction().await;
238
239        for tweak_idx in 0..state.new_start_idx().0 {
240            let operation_id = data.derive_peg_in_script(TweakIdx(tweak_idx)).3;
241
242            let claimed = state
243                .claimed_outpoints
244                .get(&TweakIdx(tweak_idx))
245                .cloned()
246                .unwrap_or_default();
247
248            dbtx.insert_new_entry(
249                &PegInTweakIndexKey(TweakIdx(tweak_idx)),
250                &PegInTweakIndexData {
251                    operation_id,
252                    creation_time: fedimint_core::time::now(),
253                    last_check_time: None,
254                    next_check_time: Some(fedimint_core::time::now()),
255                    claimed,
256                },
257            )
258            .await;
259        }
260
261        dbtx.insert_new_entry(&NextPegInTweakIndexKey, &state.new_start_idx())
262            .await;
263
264        dbtx.commit_tx().await;
265
266        Ok(())
267    }
268}
269
270impl ModuleInit for WalletClientInit {
271    type Common = WalletCommonInit;
272
273    async fn dump_database(
274        &self,
275        dbtx: &mut DatabaseTransaction<'_>,
276        prefix_names: Vec<String>,
277    ) -> Box<dyn Iterator<Item = (String, Box<dyn erased_serde::Serialize + Send>)> + '_> {
278        let mut wallet_client_items: BTreeMap<String, Box<dyn erased_serde::Serialize + Send>> =
279            BTreeMap::new();
280        let filtered_prefixes = DbKeyPrefix::iter().filter(|f| {
281            prefix_names.is_empty() || prefix_names.contains(&f.to_string().to_lowercase())
282        });
283
284        for table in filtered_prefixes {
285            match table {
286                DbKeyPrefix::NextPegInTweakIndex => {
287                    if let Some(index) = dbtx.get_value(&NextPegInTweakIndexKey).await {
288                        wallet_client_items
289                            .insert("NextPegInTweakIndex".to_string(), Box::new(index));
290                    }
291                }
292                DbKeyPrefix::PegInTweakIndex => {
293                    push_db_pair_items!(
294                        dbtx,
295                        PegInTweakIndexPrefix,
296                        PegInTweakIndexKey,
297                        PegInTweakIndexData,
298                        wallet_client_items,
299                        "Peg-In Tweak Index"
300                    );
301                }
302                DbKeyPrefix::ClaimedPegIn => {
303                    push_db_pair_items!(
304                        dbtx,
305                        ClaimedPegInPrefix,
306                        ClaimedPegInKey,
307                        ClaimedPegInData,
308                        wallet_client_items,
309                        "Claimed Peg-In"
310                    );
311                }
312                DbKeyPrefix::RecoveryFinalized => {
313                    if let Some(val) = dbtx.get_value(&RecoveryFinalizedKey).await {
314                        wallet_client_items.insert("RecoveryFinalized".to_string(), Box::new(val));
315                    }
316                }
317                DbKeyPrefix::SupportsSafeDeposit => {
318                    push_db_pair_items!(
319                        dbtx,
320                        SupportsSafeDepositPrefix,
321                        SupportsSafeDepositKey,
322                        (),
323                        wallet_client_items,
324                        "Supports Safe Deposit"
325                    );
326                }
327                DbKeyPrefix::PegInPoolCursor => {
328                    if let Some(cursor) = dbtx.get_value(&PegInPoolCursorKey).await {
329                        wallet_client_items.insert("PegInPoolCursor".to_string(), Box::new(cursor));
330                    }
331                }
332                DbKeyPrefix::RecoveryState
333                | DbKeyPrefix::ExternalReservedStart
334                | DbKeyPrefix::CoreInternalReservedStart
335                | DbKeyPrefix::CoreInternalReservedEnd => {}
336            }
337        }
338
339        Box::new(wallet_client_items.into_iter())
340    }
341}
342
343#[apply(async_trait_maybe_send!)]
344impl ClientModuleInit for WalletClientInit {
345    type Module = WalletClientModule;
346
347    fn supported_api_versions(&self) -> MultiApiVersion {
348        MultiApiVersion::try_from_iter([ApiVersion { major: 0, minor: 0 }])
349            .expect("no version conflicts")
350    }
351
352    async fn init(&self, args: &ClientModuleInitArgs<Self>) -> anyhow::Result<Self::Module> {
353        let data = WalletClientModuleData {
354            cfg: args.cfg().clone(),
355            module_root_secret: args.module_root_secret().clone(),
356        };
357
358        let db = args.db().clone();
359
360        let rpc_config = WalletClientModule::get_rpc_config(args.cfg());
361
362        // Priority:
363        // 1. user-provided bitcoind RPC from ClientBuilder::with_bitcoind_rpc
364        // 2. user-provided no-chain-id factory from
365        //    ClientBuilder::with_bitcoind_rpc_no_chain_id
366        // 3. WalletClientInit constructor
367        // 4. create from config (esplora)
368        let btc_rpc = if let Some(user_rpc) = args.user_bitcoind_rpc() {
369            user_rpc.clone()
370        } else if let Some(factory) = args.user_bitcoind_rpc_no_chain_id() {
371            if let Some(rpc) = factory(rpc_config.url.clone()).await {
372                rpc
373            } else {
374                self.0
375                    .clone()
376                    .unwrap_or(create_esplora_rpc(&rpc_config.url)?)
377            }
378        } else {
379            self.0
380                .clone()
381                .unwrap_or(create_esplora_rpc(&rpc_config.url)?)
382        };
383        let btc_rpc = BitcoindTracked::new(btc_rpc, "wallet-client").into_dyn();
384
385        let module_api = args.module_api().clone();
386
387        let (pegin_claimed_sender, pegin_claimed_receiver) = watch::channel(());
388        let (pegin_monitor_wakeup_sender, pegin_monitor_wakeup_receiver) = watch::channel(());
389
390        Ok(WalletClientModule {
391            db,
392            data,
393            module_api,
394            notifier: args.notifier().clone(),
395            rpc: btc_rpc,
396            client_ctx: args.context(),
397            pegin_monitor_wakeup_sender,
398            pegin_monitor_wakeup_receiver,
399            pegin_claimed_receiver,
400            pegin_claimed_sender,
401            task_group: args.task_group().clone(),
402            admin_auth: args.admin_auth().cloned(),
403        })
404    }
405
406    /// Wallet recovery
407    ///
408    /// Uses slice-based recovery if supported by the federation, otherwise
409    /// falls back to session-based history recovery.
410    async fn recover(
411        &self,
412        args: &ClientModuleRecoverArgs<Self>,
413        snapshot: Option<&<Self::Module as ClientModule>::Backup>,
414    ) -> anyhow::Result<()> {
415        // Check if V1 (session-based) recovery state exists (resuming interrupted
416        // recovery)
417        if args
418            .db()
419            .begin_transaction_nc()
420            .await
421            .get_value(&RecoveryStateKey)
422            .await
423            .is_some()
424        {
425            return args
426                .recover_from_history::<WalletRecovery>(self, snapshot)
427                .await;
428        }
429
430        // Determine which method to use based on endpoint availability
431        if args.module_api().fetch_recovery_count().await.is_ok() {
432            self.recover_from_slices(args).await
433        } else {
434            args.recover_from_history::<WalletRecovery>(self, snapshot)
435                .await
436        }
437    }
438
439    fn used_db_prefixes(&self) -> Option<BTreeSet<u8>> {
440        Some(
441            DbKeyPrefix::iter()
442                .map(|p| p as u8)
443                .chain(
444                    DbKeyPrefix::ExternalReservedStart as u8
445                        ..=DbKeyPrefix::CoreInternalReservedEnd as u8,
446                )
447                .collect(),
448        )
449    }
450}
451
452#[derive(Debug, Clone, Serialize, Deserialize)]
453pub struct WalletOperationMeta {
454    pub variant: WalletOperationMetaVariant,
455    pub extra_meta: serde_json::Value,
456}
457
458#[derive(Debug, Clone, Serialize, Deserialize)]
459#[serde(rename_all = "snake_case")]
460pub enum WalletOperationMetaVariant {
461    Deposit {
462        address: Address<NetworkUnchecked>,
463        /// Added in 0.4.2, can be `None` for old deposits or `Some` for ones
464        /// using the pegin monitor. The value is the child index of the key
465        /// used to generate the address, so we can re-generate the secret key
466        /// from our root secret.
467        #[serde(default)]
468        tweak_idx: Option<TweakIdx>,
469        #[serde(default, skip_serializing_if = "Option::is_none")]
470        expires_at: Option<SystemTime>,
471    },
472    Withdraw {
473        address: Address<NetworkUnchecked>,
474        #[serde(with = "bitcoin::amount::serde::as_sat")]
475        amount: bitcoin::Amount,
476        fee: PegOutFees,
477        change: Vec<OutPoint>,
478    },
479
480    RbfWithdraw {
481        rbf: Rbf,
482        change: Vec<OutPoint>,
483    },
484}
485
486/// The non-resource, just plain-data parts of [`WalletClientModule`]
487#[derive(Debug, Clone)]
488pub struct WalletClientModuleData {
489    cfg: WalletClientConfig,
490    module_root_secret: DerivableSecret,
491}
492
493impl WalletClientModuleData {
494    fn derive_deposit_address(
495        &self,
496        idx: TweakIdx,
497    ) -> (Keypair, secp256k1::PublicKey, Address, OperationId) {
498        let idx = ChildId(idx.0);
499
500        let secret_tweak_key = self
501            .module_root_secret
502            .child_key(WALLET_TWEAK_CHILD_ID)
503            .child_key(idx)
504            .to_secp_key(fedimint_core::secp256k1::SECP256K1);
505
506        let public_tweak_key = secret_tweak_key.public_key();
507
508        let address = self
509            .cfg
510            .peg_in_descriptor
511            .tweak(&public_tweak_key, bitcoin::secp256k1::SECP256K1)
512            .address(self.cfg.network.0)
513            .unwrap();
514
515        // TODO: make hash?
516        let operation_id = OperationId(public_tweak_key.x_only_public_key().0.serialize());
517
518        (secret_tweak_key, public_tweak_key, address, operation_id)
519    }
520
521    fn derive_peg_in_script(
522        &self,
523        idx: TweakIdx,
524    ) -> (ScriptBuf, bitcoin::Address, Keypair, OperationId) {
525        let (secret_tweak_key, _, address, operation_id) = self.derive_deposit_address(idx);
526
527        (
528            self.cfg
529                .peg_in_descriptor
530                .tweak(&secret_tweak_key.public_key(), SECP256K1)
531                .script_pubkey(),
532            address,
533            secret_tweak_key,
534            operation_id,
535        )
536    }
537}
538
539#[derive(Debug)]
540pub struct WalletClientModule {
541    data: WalletClientModuleData,
542    db: Database,
543    module_api: DynModuleApi,
544    notifier: ModuleNotifier<WalletClientStates>,
545    rpc: DynBitcoindRpc,
546    client_ctx: ClientContext<Self>,
547    /// Updated to wake up pegin monitor
548    pegin_monitor_wakeup_sender: watch::Sender<()>,
549    pegin_monitor_wakeup_receiver: watch::Receiver<()>,
550    /// Called every time a peg-in was claimed
551    pegin_claimed_sender: watch::Sender<()>,
552    pegin_claimed_receiver: watch::Receiver<()>,
553    task_group: TaskGroup,
554    admin_auth: Option<ApiAuth>,
555}
556
557#[apply(async_trait_maybe_send!)]
558impl ClientModule for WalletClientModule {
559    type Init = WalletClientInit;
560    type Common = WalletModuleTypes;
561    type Backup = WalletModuleBackup;
562    type ModuleStateMachineContext = WalletClientContext;
563    type States = WalletClientStates;
564
565    fn context(&self) -> Self::ModuleStateMachineContext {
566        WalletClientContext {
567            rpc: self.rpc.clone(),
568            wallet_descriptor: self.cfg().peg_in_descriptor.clone(),
569            wallet_decoder: self.decoder(),
570            secp: Secp256k1::default(),
571            client_ctx: self.client_ctx.clone(),
572        }
573    }
574
575    async fn start(&self) {
576        self.task_group.spawn_cancellable("peg-in monitor", {
577            let client_ctx = self.client_ctx.clone();
578            let db = self.db.clone();
579            let btc_rpc = self.rpc.clone();
580            let module_api = self.module_api.clone();
581            let data = self.data.clone();
582            let pegin_claimed_sender = self.pegin_claimed_sender.clone();
583            let pegin_monitor_wakeup_receiver = self.pegin_monitor_wakeup_receiver.clone();
584            pegin_monitor::run_peg_in_monitor(
585                client_ctx,
586                db,
587                btc_rpc,
588                module_api,
589                data,
590                pegin_claimed_sender,
591                pegin_monitor_wakeup_receiver,
592            )
593        });
594
595        self.task_group
596            .spawn_cancellable("supports-safe-deposit-version", {
597                let db = self.db.clone();
598                let module_api = self.module_api.clone();
599
600                poll_supports_safe_deposit_version(db, module_api)
601            });
602    }
603
604    fn supports_backup(&self) -> bool {
605        true
606    }
607
608    async fn backup(&self) -> anyhow::Result<backup::WalletModuleBackup> {
609        // fetch consensus height first
610        let session_count = self.client_ctx.global_api().session_count().await?;
611
612        let mut dbtx = self.db.begin_transaction_nc().await;
613        let next_pegin_tweak_idx = dbtx
614            .get_value(&NextPegInTweakIndexKey)
615            .await
616            .unwrap_or_default();
617        let claimed = dbtx
618            .find_by_prefix(&PegInTweakIndexPrefix)
619            .await
620            .filter_map(|(k, v)| async move {
621                if v.claimed.is_empty() {
622                    None
623                } else {
624                    Some(k.0)
625                }
626            })
627            .collect()
628            .await;
629        Ok(backup::WalletModuleBackup::new_v1(
630            session_count,
631            next_pegin_tweak_idx,
632            claimed,
633        ))
634    }
635
636    fn input_fee(
637        &self,
638        _amount: &Amounts,
639        _input: &<Self::Common as ModuleCommon>::Input,
640    ) -> Option<Amounts> {
641        Some(Amounts::new_bitcoin(self.cfg().fee_consensus.peg_in_abs))
642    }
643
644    fn output_fee(
645        &self,
646        _amount: &Amounts,
647        _output: &<Self::Common as ModuleCommon>::Output,
648    ) -> Option<Amounts> {
649        Some(Amounts::new_bitcoin(self.cfg().fee_consensus.peg_out_abs))
650    }
651
652    async fn handle_rpc(
653        &self,
654        method: String,
655        request: serde_json::Value,
656    ) -> BoxStream<'_, anyhow::Result<serde_json::Value>> {
657        Box::pin(try_stream! {
658            match method.as_str() {
659                "get_wallet_summary" => {
660                    let _req: WalletSummaryRequest = serde_json::from_value(request)?;
661                    let wallet_summary = self.get_wallet_summary()
662                        .await
663                        .expect("Failed to fetch wallet summary");
664                    let result = serde_json::to_value(&wallet_summary)
665                        .expect("Serialization error");
666                    yield result;
667                }
668                "get_block_count_local" => {
669                    let block_count = self.get_block_count_local().await
670                        .expect("Failed to fetch block count");
671                    yield serde_json::to_value(block_count)?;
672                }
673                "peg_in" => {
674                    let req: PegInRequest = serde_json::from_value(request)?;
675                    let response = self.peg_in(req)
676                        .await
677                        .map_err(|e| anyhow::anyhow!("peg_in failed: {}", e))?;
678                    let result = serde_json::to_value(&response)?;
679                    yield result;
680                },
681                "peg_out" => {
682                    let req: PegOutRequest = serde_json::from_value(request)?;
683                    let response = self.peg_out(req)
684                        .await
685                        .map_err(|e| anyhow::anyhow!("peg_out failed: {}", e))?;
686                    let result = serde_json::to_value(&response)?;
687                    yield result;
688                },
689                "subscribe_deposit" => {
690                    let req: SubscribeDepositRequest = serde_json::from_value(request)?;
691                    for await state in self.subscribe_deposit(req.operation_id).await?.into_stream() {
692                        yield serde_json::to_value(state)?;
693                    }
694                },
695                "subscribe_withdraw" => {
696                    let req: SubscribeWithdrawRequest = serde_json::from_value(request)?;
697                    for await state in self.subscribe_withdraw_updates(req.operation_id).await?.into_stream(){
698                        yield serde_json::to_value(state)?;
699                    }
700                }
701                _ => {
702                    Err(anyhow::format_err!("Unknown method: {}", method))?;
703                }
704            }
705        })
706    }
707
708    #[cfg(feature = "cli")]
709    async fn handle_cli_command(
710        &self,
711        args: &[std::ffi::OsString],
712    ) -> anyhow::Result<serde_json::Value> {
713        cli::handle_cli_command(self, args).await
714    }
715}
716
717#[derive(Deserialize)]
718struct WalletSummaryRequest {}
719
720#[derive(Debug, Clone)]
721pub struct WalletClientContext {
722    rpc: DynBitcoindRpc,
723    wallet_descriptor: PegInDescriptor,
724    wallet_decoder: Decoder,
725    secp: Secp256k1<All>,
726    pub client_ctx: ClientContext<WalletClientModule>,
727}
728
729#[derive(Debug, Clone, Serialize, Deserialize)]
730pub struct PegInRequest {
731    pub extra_meta: serde_json::Value,
732}
733
734#[derive(Deserialize)]
735struct SubscribeDepositRequest {
736    operation_id: OperationId,
737}
738
739#[derive(Deserialize)]
740struct SubscribeWithdrawRequest {
741    operation_id: OperationId,
742}
743
744#[derive(Debug, Clone, Serialize, Deserialize)]
745pub struct PegInResponse {
746    pub deposit_address: Address<NetworkUnchecked>,
747    pub operation_id: OperationId,
748}
749
750#[derive(Debug, Clone, Serialize, Deserialize)]
751pub struct PegOutRequest {
752    pub amount_sat: u64,
753    pub destination_address: Address<NetworkUnchecked>,
754    pub extra_meta: serde_json::Value,
755}
756
757#[derive(Debug, Clone, Serialize, Deserialize)]
758pub struct PegOutResponse {
759    pub operation_id: OperationId,
760}
761
762impl Context for WalletClientContext {
763    const KIND: Option<ModuleKind> = Some(KIND);
764}
765
766impl WalletClientModule {
767    fn cfg(&self) -> &WalletClientConfig {
768        &self.data.cfg
769    }
770
771    fn get_rpc_config(cfg: &WalletClientConfig) -> BitcoinRpcConfig {
772        match BitcoinRpcConfig::get_defaults_from_env_vars() {
773            Ok(rpc_config) => {
774                // TODO: Wallet client cannot support bitcoind RPC until the bitcoin dep is
775                // updated to 0.30
776                if rpc_config.kind == "bitcoind" {
777                    cfg.default_bitcoin_rpc.clone()
778                } else {
779                    rpc_config
780                }
781            }
782            _ => cfg.default_bitcoin_rpc.clone(),
783        }
784    }
785
786    pub fn get_network(&self) -> Network {
787        self.cfg().network.0
788    }
789
790    pub fn get_finality_delay(&self) -> u32 {
791        self.cfg().finality_delay
792    }
793
794    pub fn get_fee_consensus(&self) -> FeeConsensus {
795        self.cfg().fee_consensus
796    }
797
798    async fn allocate_deposit_address_inner(
799        &self,
800        dbtx: &mut DatabaseTransaction<'_>,
801    ) -> DepositAddressInfo {
802        dbtx.ensure_isolated().expect("Must be isolated db");
803
804        let tweak_idx = get_next_peg_in_tweak_child_id(dbtx).await;
805        let (_secret_tweak_key, _, address, operation_id) =
806            self.data.derive_deposit_address(tweak_idx);
807
808        let now = fedimint_core::time::now();
809
810        dbtx.insert_new_entry(
811            &PegInTweakIndexKey(tweak_idx),
812            &PegInTweakIndexData {
813                creation_time: now,
814                next_check_time: Some(now),
815                last_check_time: None,
816                operation_id,
817                claimed: vec![],
818            },
819        )
820        .await;
821
822        DepositAddressInfo {
823            operation_id,
824            address,
825            tweak_idx,
826        }
827    }
828
829    /// Fetches the fees that would need to be paid to make the withdraw request
830    /// using [`Self::withdraw`] work *right now*.
831    ///
832    /// Note that we do not receive a guarantee that these fees will be valid in
833    /// the future, thus even the next second using these fees *may* fail.
834    /// The caller should be prepared to retry with a new fee estimate.
835    pub async fn get_withdraw_fees(
836        &self,
837        address: &bitcoin::Address,
838        amount: bitcoin::Amount,
839    ) -> anyhow::Result<PegOutFees> {
840        self.module_api
841            .fetch_peg_out_fees(address, amount)
842            .await?
843            .context("Federation didn't return peg-out fees")
844    }
845
846    /// Returns a summary of the wallet's coins
847    pub async fn get_wallet_summary(&self) -> anyhow::Result<WalletSummary> {
848        Ok(self.module_api.fetch_wallet_summary().await?)
849    }
850
851    pub async fn get_block_count_local(&self) -> anyhow::Result<u32> {
852        Ok(self.module_api.fetch_block_count_local().await?)
853    }
854
855    pub fn create_withdraw_output(
856        &self,
857        operation_id: OperationId,
858        address: bitcoin::Address,
859        amount: bitcoin::Amount,
860        fees: PegOutFees,
861    ) -> anyhow::Result<ClientOutputBundle<WalletOutput, WalletClientStates>> {
862        let output = WalletOutput::new_v0_peg_out(address, amount, fees);
863
864        let amount = output.maybe_v0_ref().expect("v0 output").amount().into();
865
866        let sm_gen = move |out_point_range: OutPointRange| {
867            assert_eq!(out_point_range.count(), 1);
868            let out_idx = out_point_range.start_idx();
869            vec![WalletClientStates::Withdraw(WithdrawStateMachine {
870                operation_id,
871                state: WithdrawStates::Created(CreatedWithdrawState {
872                    fm_outpoint: OutPoint {
873                        txid: out_point_range.txid(),
874                        out_idx,
875                    },
876                }),
877            })]
878        };
879
880        Ok(ClientOutputBundle::new(
881            vec![ClientOutput::<WalletOutput> {
882                output,
883                amounts: Amounts::new_bitcoin(amount),
884            }],
885            vec![ClientOutputSM::<WalletClientStates> {
886                state_machines: Arc::new(sm_gen),
887            }],
888        ))
889    }
890
891    pub async fn peg_in(&self, req: PegInRequest) -> anyhow::Result<PegInResponse> {
892        let deposit_address = self.safe_allocate_deposit_address(req.extra_meta).await?;
893
894        Ok(PegInResponse {
895            deposit_address: Address::from_script(
896                &deposit_address.address.script_pubkey(),
897                self.get_network(),
898            )?
899            .as_unchecked()
900            .clone(),
901            operation_id: deposit_address.operation_id,
902        })
903    }
904
905    pub async fn peg_out(&self, req: PegOutRequest) -> anyhow::Result<PegOutResponse> {
906        let amount = bitcoin::Amount::from_sat(req.amount_sat);
907        let destination = req
908            .destination_address
909            .require_network(self.get_network())?;
910
911        let fees = self.get_withdraw_fees(&destination, amount).await?;
912        let operation_id = self
913            .withdraw(&destination, amount, fees, req.extra_meta)
914            .await
915            .context("Failed to initiate withdraw")?;
916
917        Ok(PegOutResponse { operation_id })
918    }
919
920    pub fn create_rbf_withdraw_output(
921        &self,
922        operation_id: OperationId,
923        rbf: &Rbf,
924    ) -> anyhow::Result<ClientOutputBundle<WalletOutput, WalletClientStates>> {
925        let output = WalletOutput::new_v0_rbf(rbf.fees, rbf.txid);
926
927        let amount = output.maybe_v0_ref().expect("v0 output").amount().into();
928
929        let sm_gen = move |out_point_range: OutPointRange| {
930            assert_eq!(out_point_range.count(), 1);
931            let out_idx = out_point_range.start_idx();
932            vec![WalletClientStates::Withdraw(WithdrawStateMachine {
933                operation_id,
934                state: WithdrawStates::Created(CreatedWithdrawState {
935                    fm_outpoint: OutPoint {
936                        txid: out_point_range.txid(),
937                        out_idx,
938                    },
939                }),
940            })]
941        };
942
943        Ok(ClientOutputBundle::new(
944            vec![ClientOutput::<WalletOutput> {
945                output,
946                amounts: Amounts::new_bitcoin(amount),
947            }],
948            vec![ClientOutputSM::<WalletClientStates> {
949                state_machines: Arc::new(sm_gen),
950            }],
951        ))
952    }
953
954    pub async fn btc_tx_has_no_size_limit(&self) -> FederationResult<bool> {
955        Ok(self.module_api.module_consensus_version().await? >= ModuleConsensusVersion::new(2, 2))
956    }
957
958    /// Returns true if the federation's wallet module consensus version
959    /// supports processing all deposits.
960    ///
961    /// This method is safe to call offline, since it first attempts to read a
962    /// key from the db that represents the client has previously been able to
963    /// verify the wallet module consensus version. If the client has not
964    /// verified the version, it must be online to fetch the latest wallet
965    /// module consensus version.
966    pub async fn supports_safe_deposit(&self) -> bool {
967        let mut dbtx = self.db.begin_transaction().await;
968
969        let already_verified_supports_safe_deposit =
970            dbtx.get_value(&SupportsSafeDepositKey).await.is_some();
971
972        already_verified_supports_safe_deposit || {
973            match self.module_api.module_consensus_version().await {
974                Ok(module_consensus_version) => {
975                    let supported_version =
976                        SAFE_DEPOSIT_MODULE_CONSENSUS_VERSION <= module_consensus_version;
977
978                    if supported_version {
979                        dbtx.insert_new_entry(&SupportsSafeDepositKey, &()).await;
980                        dbtx.commit_tx().await;
981                    }
982
983                    supported_version
984                }
985                Err(_) => false,
986            }
987        }
988    }
989
990    /// Allocates a deposit address controlled by the federation, guaranteeing
991    /// safe handling of all deposits, including on-chain transactions exceeding
992    /// `ALEPH_BFT_UNIT_BYTE_LIMIT`.
993    ///
994    /// Returns an error if the client has never been online to verify the
995    /// federation's wallet module consensus version supports processing all
996    /// deposits.
997    pub async fn safe_allocate_deposit_address<M>(
998        &self,
999        extra_meta: M,
1000    ) -> anyhow::Result<DepositAddressInfo>
1001    where
1002        M: Serialize + MaybeSend + MaybeSync,
1003    {
1004        ensure!(
1005            self.supports_safe_deposit().await,
1006            "Wallet module consensus version doesn't support safe deposits",
1007        );
1008
1009        self.allocate_deposit_address_expert_only(extra_meta).await
1010    }
1011
1012    /// Allocates a deposit address that is controlled by the federation.
1013    ///
1014    /// This is an EXPERT ONLY method intended for power users such as Lightning
1015    /// gateways allocating liquidity, and we discourage exposing peg-in
1016    /// functionality to everyday users of a Fedimint wallet due to the
1017    /// following two limitations:
1018    ///
1019    /// The transaction sending to this address needs to be smaller than 40KB in
1020    /// order for the peg-in to be claimable. If the transaction is too large,
1021    /// funds will be lost.
1022    ///
1023    /// In the future, federations will also enforce a minimum peg-in amount to
1024    /// prevent accumulation of dust UTXOs. Peg-ins under this minimum cannot be
1025    /// claimed and funds will be lost.
1026    ///
1027    /// Everyday users should rely on Lightning to move funds into the
1028    /// federation.
1029    pub async fn allocate_deposit_address_expert_only<M>(
1030        &self,
1031        extra_meta: M,
1032    ) -> anyhow::Result<DepositAddressInfo>
1033    where
1034        M: Serialize + MaybeSend + MaybeSync,
1035    {
1036        let extra_meta_value =
1037            serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
1038        let deposit_address = self
1039            .db
1040            .autocommit(
1041                move |dbtx, _| {
1042                    let extra_meta_value_inner = extra_meta_value.clone();
1043                    Box::pin(async move {
1044                        let deposit_address = self.allocate_deposit_address_inner(dbtx).await;
1045
1046                        self.client_ctx
1047                            .manual_operation_start_dbtx(
1048                                dbtx,
1049                                deposit_address.operation_id,
1050                                WalletCommonInit::KIND.as_str(),
1051                                WalletOperationMeta {
1052                                    variant: WalletOperationMetaVariant::Deposit {
1053                                        address: deposit_address.address.clone().into_unchecked(),
1054                                        tweak_idx: Some(deposit_address.tweak_idx),
1055                                        expires_at: None,
1056                                    },
1057                                    extra_meta: extra_meta_value_inner,
1058                                },
1059                                vec![],
1060                            )
1061                            .await?;
1062
1063                        debug!(
1064                            target: LOG_CLIENT_MODULE_WALLET,
1065                            tweak_idx = %deposit_address.tweak_idx,
1066                            address = %deposit_address.address,
1067                            "Derived a new deposit address"
1068                        );
1069
1070                        // Begin watching the script address
1071                        self.rpc
1072                            .watch_script_history(&deposit_address.address.script_pubkey())
1073                            .await?;
1074
1075                        let sender = self.pegin_monitor_wakeup_sender.clone();
1076                        dbtx.on_commit(move || {
1077                            sender.send_replace(());
1078                        });
1079
1080                        Ok(deposit_address)
1081                    })
1082                },
1083                Some(100),
1084            )
1085            .await
1086            .map_err(|e| match e {
1087                AutocommitError::CommitFailed {
1088                    last_error,
1089                    attempts,
1090                } => anyhow!("Failed to commit after {attempts} attempts: {last_error}"),
1091                AutocommitError::ClosureError { error, .. } => error,
1092            })?;
1093
1094        Ok(deposit_address)
1095    }
1096
1097    /// Allocate a deposit address, bounding the gap of consecutive unused
1098    /// addresses past the last-used one, without selecting a reusable address.
1099    ///
1100    /// "Unused" here means `PegInTweakIndexData::claimed.is_empty()` — the
1101    /// `pegin_monitor` has not yet successfully claimed any deposit on this
1102    /// address. There is a small race window where a deposit has been
1103    /// observed in the mempool / awaiting confirmations but not yet claimed,
1104    /// in which case the address still looks unused. Reusing such an address
1105    /// is benign in practice: `pegin_monitor` claims any number of deposits
1106    /// per address.
1107    ///
1108    /// Let `gap` be the count of unused tweak indices strictly greater than
1109    /// the highest used index (or the count of all allocated tweaks if none
1110    /// have been used yet). Semantics:
1111    ///   - If `gap >= max_gap_size` and there is at least one unused address,
1112    ///     returns all those addresses with
1113    ///     [`MaybeNewAddress::TooManyUnusedAddresses`].
1114    ///   - Otherwise, allocates a new address (equivalent to
1115    ///     [`Self::allocate_deposit_address_expert_only`]) and returns it with
1116    ///     [`MaybeNewAddress::NewAddress`].
1117    ///
1118    /// `max_gap_size == 0` therefore means "only allocate fresh when the
1119    /// previous address was already used". The very first call still
1120    /// allocates fresh because there are no addresses to reuse.
1121    ///
1122    /// `max_gap_size == usize::MAX` makes this method behave like
1123    /// [`Self::allocate_deposit_address_expert_only`] — always allocating a new
1124    /// address.
1125    ///
1126    /// Reusable addresses are ordered by `creation_time` ascending. This lets
1127    /// callers that need a custom selection strategy choose from all available
1128    /// candidates.
1129    ///
1130    /// Caveats inherited from
1131    /// [`Self::allocate_deposit_address_expert_only`] (40KB tx limit, future
1132    /// minimum peg-in amount) apply equally to newly allocated and reused
1133    /// addresses.
1134    pub async fn allocate_deposit_address_pooled_stateless(
1135        &self,
1136        max_gap_size: usize,
1137    ) -> anyhow::Result<MaybeNewAddress> {
1138        let max_gap_size_u64 = u64::try_from(max_gap_size).unwrap_or(u64::MAX);
1139        let extra_meta_value = serde_json::Value::Null;
1140        let result = self
1141            .db
1142            .autocommit(
1143                move |dbtx, _| {
1144                    let extra_meta_value_inner = extra_meta_value.clone();
1145                    Box::pin(async move {
1146                        let unused = self.unused_pooled_deposit_addresses(dbtx).await;
1147
1148                        if max_gap_size_u64 <= unused.len() as u64 && !unused.is_empty() {
1149                            let addresses = unused
1150                                .into_iter()
1151                                .map(|(tweak_idx, data)| {
1152                                    let (_script, address, _key, operation_id) =
1153                                        self.data.derive_peg_in_script(tweak_idx);
1154
1155                                    debug_assert_eq!(operation_id, data.operation_id);
1156
1157                                    DepositAddressInfo {
1158                                        operation_id,
1159                                        address,
1160                                        tweak_idx,
1161                                    }
1162                                })
1163                                .collect();
1164
1165                            return Ok::<_, anyhow::Error>(
1166                                MaybeNewAddress::TooManyUnusedAddresses(addresses),
1167                            );
1168                        }
1169
1170                        let deposit_address = self.allocate_deposit_address_inner(dbtx).await;
1171
1172                        self.client_ctx
1173                            .manual_operation_start_dbtx(
1174                                dbtx,
1175                                deposit_address.operation_id,
1176                                WalletCommonInit::KIND.as_str(),
1177                                WalletOperationMeta {
1178                                    variant: WalletOperationMetaVariant::Deposit {
1179                                        address: deposit_address.address.clone().into_unchecked(),
1180                                        tweak_idx: Some(deposit_address.tweak_idx),
1181                                        expires_at: None,
1182                                    },
1183                                    extra_meta: extra_meta_value_inner,
1184                                },
1185                                vec![],
1186                            )
1187                            .await?;
1188
1189                        debug!(
1190                            target: LOG_CLIENT_MODULE_WALLET,
1191                            tweak_idx = %deposit_address.tweak_idx,
1192                            address = %deposit_address.address,
1193                            "Derived a new pooled deposit address"
1194                        );
1195
1196                        self.rpc
1197                            .watch_script_history(&deposit_address.address.script_pubkey())
1198                            .await?;
1199
1200                        let sender = self.pegin_monitor_wakeup_sender.clone();
1201                        dbtx.on_commit(move || {
1202                            sender.send_replace(());
1203                        });
1204
1205                        Ok(MaybeNewAddress::NewAddress(deposit_address))
1206                    })
1207                },
1208                Some(100),
1209            )
1210            .await
1211            .map_err(|e| match e {
1212                AutocommitError::CommitFailed {
1213                    last_error,
1214                    attempts,
1215                } => anyhow!("Failed to commit after {attempts} attempts: {last_error}"),
1216                AutocommitError::ClosureError { error, .. } => error,
1217            })?;
1218
1219        Ok(result)
1220    }
1221
1222    async fn unused_pooled_deposit_addresses(
1223        &self,
1224        dbtx: &mut DatabaseTransaction<'_>,
1225    ) -> Vec<(TweakIdx, PegInTweakIndexData)> {
1226        // Walk peg-in tweaks in descending order, taking unused ones
1227        // (`claimed.is_empty()`) until we hit a used one. That gives us exactly
1228        // the trailing gap — `gap + 1` reads regardless of total tweak count.
1229        // No RPC and no separately-maintained `last_used` index needed.
1230        let mut unused: Vec<(TweakIdx, PegInTweakIndexData)> = dbtx
1231            .find_by_prefix_sorted_descending(&PegInTweakIndexPrefix)
1232            .await
1233            .take_while(|(_, d)| std::future::ready(d.claimed.is_empty()))
1234            .map(|(k, v)| (k.0, v))
1235            .collect()
1236            .await;
1237
1238        // Order by `creation_time` ascending for caller selection and stable
1239        // wraparound after stateful reuse mutates `creation_time`.
1240        unused.sort_by_key(|(t, d)| (d.creation_time, *t));
1241        unused
1242    }
1243
1244    /// Allocate a deposit address, bounding the gap of consecutive unused
1245    /// addresses past the last-used one.
1246    ///
1247    /// This is a stateful wrapper over
1248    /// [`Self::allocate_deposit_address_pooled_stateless`]. If the stateless
1249    /// method returns [`MaybeNewAddress::TooManyUnusedAddresses`], this method
1250    /// picks one of those addresses with a persisted round-robin cursor and
1251    /// returns it with [`AllocateDepositOutcome::Reused`].
1252    ///
1253    /// Selection policy when reusing: round-robin over unused addresses
1254    /// ordered by `creation_time` ascending, with the cursor persisted
1255    /// per-client in the wallet module's DB. The cursor stores the next
1256    /// tweak index to consider; after picking idx X it advances to X+1, and
1257    /// when no candidate has `tweak_idx >= cursor` it wraps to the
1258    /// oldest-by-`creation_time` candidate. This gives the caller the
1259    /// predictable "cycle through the unused addresses" behavior.
1260    ///
1261    /// On reuse the `creation_time` and check-schedule fields of the
1262    /// underlying tweak entry are reset to "now", so `pegin_monitor`'s
1263    /// age-proportional polling delay restarts from zero. Without this an
1264    /// old entry would keep its long stale next-check delay and a user
1265    /// sending to it could wait a long time before the client noticed.
1266    #[allow(clippy::too_many_lines)]
1267    pub async fn allocate_deposit_address_pooled(
1268        &self,
1269        max_gap_size: usize,
1270    ) -> anyhow::Result<(DepositAddressInfo, AllocateDepositOutcome)> {
1271        let stateless = self
1272            .allocate_deposit_address_pooled_stateless(max_gap_size)
1273            .await?;
1274
1275        let reused_addresses = match stateless {
1276            MaybeNewAddress::NewAddress(deposit_address) => {
1277                return Ok((deposit_address, AllocateDepositOutcome::Fresh));
1278            }
1279            MaybeNewAddress::TooManyUnusedAddresses(addresses) => addresses,
1280        };
1281
1282        let result = self
1283            .db
1284            .autocommit(
1285                move |dbtx, _| {
1286                    let reused_addresses = reused_addresses.clone();
1287                    Box::pin(async move {
1288                        let cursor = dbtx
1289                            .get_value(&PegInPoolCursorKey)
1290                            .await
1291                            .unwrap_or(TweakIdx(0));
1292
1293                        let pick_pos = reused_addresses
1294                            .iter()
1295                            .position(|a| cursor <= a.tweak_idx)
1296                            .unwrap_or(0);
1297                        let reused_address = reused_addresses[pick_pos].clone();
1298
1299                        let existing_tweak_idx = reused_address.tweak_idx;
1300                        let existing = dbtx
1301                            .get_value(&PegInTweakIndexKey(reused_address.tweak_idx))
1302                            .await
1303                            .with_context(|| {
1304                                format!(
1305                                    "Pooled address disappeared while reusing {}",
1306                                    reused_address.tweak_idx
1307                                )
1308                            })?;
1309
1310                        ensure!(
1311                            existing.claimed.is_empty(),
1312                            "Pooled address was used while reusing {}",
1313                            reused_address.tweak_idx
1314                        );
1315
1316                        dbtx.insert_entry(&PegInPoolCursorKey, &reused_address.tweak_idx.next())
1317                            .await;
1318
1319                        // Reset the monitoring schedule so the reused address
1320                        // gets checked as aggressively as a freshly-allocated
1321                        // one. Without this, an old entry would retain its
1322                        // `age/10` next-check delay from `pegin_monitor`,
1323                        // meaning a user who sends to a long-idle pool address
1324                        // could wait hours before the client polls for the
1325                        // deposit.
1326                        let now = fedimint_core::time::now();
1327                        dbtx.insert_entry(
1328                            &PegInTweakIndexKey(reused_address.tweak_idx),
1329                            &PegInTweakIndexData {
1330                                creation_time: now,
1331                                last_check_time: None,
1332                                next_check_time: Some(now),
1333                                operation_id: existing.operation_id,
1334                                claimed: existing.claimed,
1335                            },
1336                        )
1337                        .await;
1338
1339                        let sender = self.pegin_monitor_wakeup_sender.clone();
1340                        dbtx.on_commit(move || {
1341                            sender.send_replace(());
1342                        });
1343
1344                        Ok::<_, anyhow::Error>((
1345                            reused_address,
1346                            AllocateDepositOutcome::Reused {
1347                                original_tweak_idx: existing_tweak_idx,
1348                            },
1349                        ))
1350                    })
1351                },
1352                Some(100),
1353            )
1354            .await
1355            .map_err(|e| match e {
1356                AutocommitError::CommitFailed {
1357                    last_error,
1358                    attempts,
1359                } => anyhow!("Failed to commit after {attempts} attempts: {last_error}"),
1360                AutocommitError::ClosureError { error, .. } => error,
1361            })?;
1362
1363        Ok(result)
1364    }
1365
1366    /// Returns a stream of updates about an ongoing deposit operation created
1367    /// with [`WalletClientModule::allocate_deposit_address_expert_only`].
1368    /// Returns an error for old deposit operations created prior to the 0.4
1369    /// release and not driven to completion yet. This should be rare enough
1370    /// that an indeterminate state is ok here.
1371    pub async fn subscribe_deposit(
1372        &self,
1373        operation_id: OperationId,
1374    ) -> anyhow::Result<UpdateStreamOrOutcome<DepositStateV2>> {
1375        let operation = self
1376            .client_ctx
1377            .get_operation(operation_id)
1378            .await
1379            .with_context(|| anyhow!("Operation not found: {}", operation_id.fmt_short()))?;
1380
1381        if operation.operation_module_kind() != WalletCommonInit::KIND.as_str() {
1382            bail!("Operation is not a wallet operation");
1383        }
1384
1385        let operation_meta = operation.meta::<WalletOperationMeta>();
1386
1387        let WalletOperationMetaVariant::Deposit {
1388            address, tweak_idx, ..
1389        } = operation_meta.variant
1390        else {
1391            bail!("Operation is not a deposit operation");
1392        };
1393
1394        let address = address.require_network(self.cfg().network.0)?;
1395
1396        // The old deposit operations don't have tweak_idx set
1397        let Some(tweak_idx) = tweak_idx else {
1398            // In case we are dealing with an old deposit that still uses state machines we
1399            // don't have the logic here anymore to subscribe to updates. We can still read
1400            // the final state though if it reached any.
1401            let outcome_v1 = operation
1402                .outcome::<DepositStateV1>()
1403                .context("Old pending deposit, can't subscribe to updates")?;
1404
1405            let outcome_v2 = match outcome_v1 {
1406                DepositStateV1::Claimed(tx_info) => DepositStateV2::Claimed {
1407                    btc_deposited: tx_info.btc_transaction.output[tx_info.out_idx as usize].value,
1408                    btc_out_point: bitcoin::OutPoint {
1409                        txid: tx_info.btc_transaction.compute_txid(),
1410                        vout: tx_info.out_idx,
1411                    },
1412                },
1413                DepositStateV1::Failed(error) => DepositStateV2::Failed(error),
1414                _ => bail!("Non-final outcome in operation log"),
1415            };
1416
1417            return Ok(UpdateStreamOrOutcome::Outcome(outcome_v2));
1418        };
1419
1420        Ok(self.client_ctx.outcome_or_updates(operation, operation_id, {
1421            let stream_rpc = self.rpc.clone();
1422            let stream_client_ctx = self.client_ctx.clone();
1423            let stream_script_pub_key = address.script_pubkey();
1424            move || {
1425
1426            stream! {
1427                yield DepositStateV2::WaitingForTransaction;
1428
1429                retry(
1430                    "subscribe script history",
1431                    background_backoff(),
1432                    || stream_rpc.watch_script_history(&stream_script_pub_key)
1433                ).await.expect("Will never give up");
1434                let (btc_out_point, btc_deposited) = retry(
1435                    "fetch history",
1436                    background_backoff(),
1437                    || async {
1438                        let history = stream_rpc.get_script_history(&stream_script_pub_key).await?;
1439                        history.first().and_then(|tx| {
1440                            let (out_idx, amount) = tx.output
1441                                .iter()
1442                                .enumerate()
1443                                .find_map(|(idx, output)| (output.script_pubkey == stream_script_pub_key).then_some((idx, output.value)))?;
1444                            let txid = tx.compute_txid();
1445
1446                            Some((
1447                                bitcoin::OutPoint {
1448                                    txid,
1449                                    vout: out_idx as u32,
1450                                },
1451                                amount
1452                            ))
1453                        }).context("No deposit transaction found")
1454                    }
1455                ).await.expect("Will never give up");
1456
1457                yield DepositStateV2::WaitingForConfirmation {
1458                    btc_deposited,
1459                    btc_out_point
1460                };
1461
1462                let claim_data = stream_client_ctx.module_db().wait_key_exists(&ClaimedPegInKey {
1463                    peg_in_index: tweak_idx,
1464                    btc_out_point,
1465                }).await;
1466
1467                yield DepositStateV2::Confirmed {
1468                    btc_deposited,
1469                    btc_out_point
1470                };
1471
1472                match stream_client_ctx.await_primary_module_outputs(operation_id, claim_data.change).await {
1473                    Ok(()) => yield DepositStateV2::Claimed {
1474                        btc_deposited,
1475                        btc_out_point
1476                    },
1477                    Err(e) => yield DepositStateV2::Failed(e.to_string())
1478                }
1479            }
1480        }}))
1481    }
1482
1483    pub async fn list_peg_in_tweak_idxes(&self) -> BTreeMap<TweakIdx, PegInTweakIndexData> {
1484        self.client_ctx
1485            .module_db()
1486            .clone()
1487            .begin_transaction_nc()
1488            .await
1489            .find_by_prefix(&PegInTweakIndexPrefix)
1490            .await
1491            .map(|(key, data)| (key.0, data))
1492            .collect()
1493            .await
1494    }
1495
1496    pub async fn find_tweak_idx_by_address(
1497        &self,
1498        address: bitcoin::Address<NetworkUnchecked>,
1499    ) -> anyhow::Result<TweakIdx> {
1500        let data = self.data.clone();
1501        let Some((tweak_idx, _)) = self
1502            .db
1503            .begin_transaction_nc()
1504            .await
1505            .find_by_prefix(&PegInTweakIndexPrefix)
1506            .await
1507            .filter(|(k, _)| {
1508                let (_, derived_address, _tweak_key, _) = data.derive_peg_in_script(k.0);
1509                future::ready(derived_address.into_unchecked() == address)
1510            })
1511            .next()
1512            .await
1513        else {
1514            bail!("Address not found in the list of derived keys");
1515        };
1516
1517        Ok(tweak_idx.0)
1518    }
1519    pub async fn find_tweak_idx_by_operation_id(
1520        &self,
1521        operation_id: OperationId,
1522    ) -> anyhow::Result<TweakIdx> {
1523        Ok(self
1524            .client_ctx
1525            .module_db()
1526            .clone()
1527            .begin_transaction_nc()
1528            .await
1529            .find_by_prefix(&PegInTweakIndexPrefix)
1530            .await
1531            .filter(|(_k, v)| future::ready(v.operation_id == operation_id))
1532            .next()
1533            .await
1534            .ok_or_else(|| anyhow::format_err!("OperationId not found"))?
1535            .0
1536            .0)
1537    }
1538
1539    pub async fn get_pegin_tweak_idx(
1540        &self,
1541        tweak_idx: TweakIdx,
1542    ) -> anyhow::Result<PegInTweakIndexData> {
1543        self.client_ctx
1544            .module_db()
1545            .clone()
1546            .begin_transaction_nc()
1547            .await
1548            .get_value(&PegInTweakIndexKey(tweak_idx))
1549            .await
1550            .ok_or_else(|| anyhow::format_err!("TweakIdx not found"))
1551    }
1552
1553    pub async fn get_claimed_pegins(
1554        &self,
1555        dbtx: &mut DatabaseTransaction<'_>,
1556        tweak_idx: TweakIdx,
1557    ) -> Vec<(
1558        bitcoin::OutPoint,
1559        TransactionId,
1560        Vec<fedimint_core::OutPoint>,
1561    )> {
1562        let outpoints = dbtx
1563            .get_value(&PegInTweakIndexKey(tweak_idx))
1564            .await
1565            .map(|v| v.claimed)
1566            .unwrap_or_default();
1567
1568        let mut res = vec![];
1569
1570        for outpoint in outpoints {
1571            let claimed_peg_in_data = dbtx
1572                .get_value(&ClaimedPegInKey {
1573                    peg_in_index: tweak_idx,
1574                    btc_out_point: outpoint,
1575                })
1576                .await
1577                .expect("Must have a corresponding claim record");
1578            res.push((
1579                outpoint,
1580                claimed_peg_in_data.claim_txid,
1581                claimed_peg_in_data.change,
1582            ));
1583        }
1584
1585        res
1586    }
1587
1588    /// Like [`Self::recheck_pegin_address`] but by `operation_id`
1589    pub async fn recheck_pegin_address_by_op_id(
1590        &self,
1591        operation_id: OperationId,
1592    ) -> anyhow::Result<()> {
1593        let tweak_idx = self.find_tweak_idx_by_operation_id(operation_id).await?;
1594
1595        self.recheck_pegin_address(tweak_idx).await
1596    }
1597
1598    /// Schedule given address for immediate re-check for deposits
1599    pub async fn recheck_pegin_address_by_address(
1600        &self,
1601        address: bitcoin::Address<NetworkUnchecked>,
1602    ) -> anyhow::Result<()> {
1603        self.recheck_pegin_address(self.find_tweak_idx_by_address(address).await?)
1604            .await
1605    }
1606
1607    /// Schedule given address for immediate re-check for deposits
1608    pub async fn recheck_pegin_address(&self, tweak_idx: TweakIdx) -> anyhow::Result<()> {
1609        self.db
1610            .autocommit(
1611                |dbtx, _| {
1612                    Box::pin(async {
1613                        let db_key = PegInTweakIndexKey(tweak_idx);
1614                        let db_val = dbtx
1615                            .get_value(&db_key)
1616                            .await
1617                            .ok_or_else(|| anyhow::format_err!("DBKey not found"))?;
1618
1619                        dbtx.insert_entry(
1620                            &db_key,
1621                            &PegInTweakIndexData {
1622                                next_check_time: Some(fedimint_core::time::now()),
1623                                ..db_val
1624                            },
1625                        )
1626                        .await;
1627
1628                        let sender = self.pegin_monitor_wakeup_sender.clone();
1629                        dbtx.on_commit(move || {
1630                            sender.send_replace(());
1631                        });
1632
1633                        Ok::<_, anyhow::Error>(())
1634                    })
1635                },
1636                Some(100),
1637            )
1638            .await?;
1639
1640        Ok(())
1641    }
1642
1643    /// Await for num deposit by [`OperationId`]
1644    pub async fn await_num_deposits_by_operation_id(
1645        &self,
1646        operation_id: OperationId,
1647        num_deposits: usize,
1648    ) -> anyhow::Result<()> {
1649        let tweak_idx = self.find_tweak_idx_by_operation_id(operation_id).await?;
1650        self.await_num_deposits(tweak_idx, num_deposits).await
1651    }
1652
1653    pub async fn await_num_deposits_by_address(
1654        &self,
1655        address: bitcoin::Address<NetworkUnchecked>,
1656        num_deposits: usize,
1657    ) -> anyhow::Result<()> {
1658        self.await_num_deposits(self.find_tweak_idx_by_address(address).await?, num_deposits)
1659            .await
1660    }
1661
1662    #[instrument(target = LOG_CLIENT_MODULE_WALLET, skip_all, fields(tweak_idx=?tweak_idx, num_deposists=num_deposits))]
1663    pub async fn await_num_deposits(
1664        &self,
1665        tweak_idx: TweakIdx,
1666        num_deposits: usize,
1667    ) -> anyhow::Result<()> {
1668        let operation_id = self.get_pegin_tweak_idx(tweak_idx).await?.operation_id;
1669
1670        let mut receiver = self.pegin_claimed_receiver.clone();
1671        let mut backoff = backoff_util::aggressive_backoff();
1672
1673        loop {
1674            let pegins = self
1675                .get_claimed_pegins(
1676                    &mut self.client_ctx.module_db().begin_transaction_nc().await,
1677                    tweak_idx,
1678                )
1679                .await;
1680
1681            if pegins.len() < num_deposits {
1682                debug!(target: LOG_CLIENT_MODULE_WALLET, has=pegins.len(), "Not enough deposits");
1683                self.recheck_pegin_address(tweak_idx).await?;
1684                runtime::sleep(backoff.next().unwrap_or_default()).await;
1685                receiver.changed().await?;
1686                continue;
1687            }
1688
1689            debug!(target: LOG_CLIENT_MODULE_WALLET, has=pegins.len(), "Enough deposits detected");
1690
1691            for (_outpoint, transaction_id, change) in pegins {
1692                if transaction_id == TransactionId::from_byte_array([0; 32]) && change.is_empty() {
1693                    debug!(target: LOG_CLIENT_MODULE_WALLET, "Deposited amount was too low, skipping");
1694                    continue;
1695                }
1696
1697                debug!(target: LOG_CLIENT_MODULE_WALLET, out_points=?change, "Ensuring deposists claimed");
1698                let tx_subscriber = self.client_ctx.transaction_updates(operation_id).await;
1699
1700                if let Err(e) = tx_subscriber.await_tx_accepted(transaction_id).await {
1701                    bail!("{}", e);
1702                }
1703
1704                debug!(target: LOG_CLIENT_MODULE_WALLET, out_points=?change, "Ensuring outputs claimed");
1705                self.client_ctx
1706                    .await_primary_module_outputs(operation_id, change)
1707                    .await
1708                    .expect("Cannot fail if tx was accepted and federation is honest");
1709            }
1710
1711            return Ok(());
1712        }
1713    }
1714
1715    /// Attempt to withdraw a given `amount` of Bitcoin to a destination
1716    /// `address`. The caller has to supply the fee rate to be used which can be
1717    /// fetched using [`Self::get_withdraw_fees`] and should be
1718    /// acknowledged by the user since it can be unexpectedly high.
1719    pub async fn withdraw<M: Serialize + MaybeSend + MaybeSync>(
1720        &self,
1721        address: &bitcoin::Address,
1722        amount: bitcoin::Amount,
1723        fee: PegOutFees,
1724        extra_meta: M,
1725    ) -> anyhow::Result<OperationId> {
1726        {
1727            let operation_id = OperationId(thread_rng().r#gen());
1728
1729            let withdraw_output =
1730                self.create_withdraw_output(operation_id, address.clone(), amount, fee)?;
1731            let tx_builder = TransactionBuilder::new()
1732                .with_outputs(self.client_ctx.make_client_outputs(withdraw_output));
1733
1734            let extra_meta =
1735                serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
1736            self.client_ctx
1737                .finalize_and_submit_transaction(
1738                    operation_id,
1739                    WalletCommonInit::KIND.as_str(),
1740                    {
1741                        let address = address.clone();
1742                        move |change_range: OutPointRange| WalletOperationMeta {
1743                            variant: WalletOperationMetaVariant::Withdraw {
1744                                address: address.clone().into_unchecked(),
1745                                amount,
1746                                fee,
1747                                change: change_range.into_iter().collect(),
1748                            },
1749                            extra_meta: extra_meta.clone(),
1750                        }
1751                    },
1752                    tx_builder,
1753                )
1754                .await?;
1755
1756            let mut dbtx = self.client_ctx.module_db().begin_transaction().await;
1757
1758            self.client_ctx
1759                .log_event(
1760                    &mut dbtx,
1761                    SendPaymentEvent {
1762                        operation_id,
1763                        amount: amount + fee.amount(),
1764                        fee: fee.amount(),
1765                    },
1766                )
1767                .await;
1768
1769            dbtx.commit_tx().await;
1770
1771            Ok(operation_id)
1772        }
1773    }
1774
1775    /// Attempt to increase the fee of a onchain withdraw transaction using
1776    /// replace by fee (RBF).
1777    /// This can prevent transactions from getting stuck
1778    /// in the mempool
1779    #[deprecated(
1780        since = "0.4.0",
1781        note = "RBF withdrawals are rejected by the federation"
1782    )]
1783    pub async fn rbf_withdraw<M: Serialize + MaybeSync + MaybeSend>(
1784        &self,
1785        rbf: Rbf,
1786        extra_meta: M,
1787    ) -> anyhow::Result<OperationId> {
1788        let operation_id = OperationId(thread_rng().r#gen());
1789
1790        let withdraw_output = self.create_rbf_withdraw_output(operation_id, &rbf)?;
1791        let tx_builder = TransactionBuilder::new()
1792            .with_outputs(self.client_ctx.make_client_outputs(withdraw_output));
1793
1794        let extra_meta = serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
1795        self.client_ctx
1796            .finalize_and_submit_transaction(
1797                operation_id,
1798                WalletCommonInit::KIND.as_str(),
1799                move |change_range: OutPointRange| WalletOperationMeta {
1800                    variant: WalletOperationMetaVariant::RbfWithdraw {
1801                        rbf: rbf.clone(),
1802                        change: change_range.into_iter().collect(),
1803                    },
1804                    extra_meta: extra_meta.clone(),
1805                },
1806                tx_builder,
1807            )
1808            .await?;
1809
1810        Ok(operation_id)
1811    }
1812
1813    pub async fn subscribe_withdraw_updates(
1814        &self,
1815        operation_id: OperationId,
1816    ) -> anyhow::Result<UpdateStreamOrOutcome<WithdrawState>> {
1817        let operation = self
1818            .client_ctx
1819            .get_operation(operation_id)
1820            .await
1821            .with_context(|| anyhow!("Operation not found: {}", operation_id.fmt_short()))?;
1822
1823        if operation.operation_module_kind() != WalletCommonInit::KIND.as_str() {
1824            bail!("Operation is not a wallet operation");
1825        }
1826
1827        let operation_meta = operation.meta::<WalletOperationMeta>();
1828
1829        let (WalletOperationMetaVariant::Withdraw { change, .. }
1830        | WalletOperationMetaVariant::RbfWithdraw { change, .. }) = operation_meta.variant
1831        else {
1832            bail!("Operation is not a withdraw operation");
1833        };
1834
1835        let mut operation_stream = self.notifier.subscribe(operation_id).await;
1836        let client_ctx = self.client_ctx.clone();
1837
1838        Ok(self
1839            .client_ctx
1840            .outcome_or_updates(operation, operation_id, move || {
1841                stream! {
1842                    match next_withdraw_state(&mut operation_stream).await {
1843                        Some(WithdrawStates::Created(_)) => {
1844                            yield WithdrawState::Created;
1845                        },
1846                        Some(s) => {
1847                            panic!("Unexpected state {s:?}")
1848                        },
1849                        None => return,
1850                    }
1851
1852                    // TODO: get rid of awaiting change here, there has to be a better way to make tests deterministic
1853
1854                        // Swallowing potential errors since the transaction failing  is handled by
1855                        // output outcome fetching already
1856                        let _ = client_ctx
1857                            .await_primary_module_outputs(operation_id, change)
1858                            .await;
1859
1860
1861                    match next_withdraw_state(&mut operation_stream).await {
1862                        Some(WithdrawStates::Aborted(inner)) => {
1863                            yield WithdrawState::Failed(inner.error);
1864                        },
1865                        Some(WithdrawStates::Success(inner)) => {
1866                            yield WithdrawState::Succeeded(inner.txid);
1867                        },
1868                        Some(s) => {
1869                            panic!("Unexpected state {s:?}")
1870                        },
1871                        None => {},
1872                    }
1873                }
1874            }))
1875    }
1876
1877    fn admin_auth(&self) -> anyhow::Result<ApiAuth> {
1878        self.admin_auth
1879            .clone()
1880            .ok_or_else(|| anyhow::format_err!("Admin auth not set"))
1881    }
1882
1883    pub async fn activate_consensus_version_voting(&self) -> anyhow::Result<()> {
1884        self.module_api
1885            .activate_consensus_version_voting(self.admin_auth()?)
1886            .await?;
1887
1888        Ok(())
1889    }
1890}
1891
1892/// Polls the federation checking if the activated module consensus version
1893/// supports safe deposits, saving the result in the db once it does.
1894async fn poll_supports_safe_deposit_version(db: Database, module_api: DynModuleApi) {
1895    loop {
1896        let mut dbtx = db.begin_transaction().await;
1897
1898        if dbtx.get_value(&SupportsSafeDepositKey).await.is_some() {
1899            break;
1900        }
1901
1902        module_api.wait_for_initialized_connections().await;
1903
1904        if let Ok(module_consensus_version) = module_api.module_consensus_version().await
1905            && SAFE_DEPOSIT_MODULE_CONSENSUS_VERSION <= module_consensus_version
1906        {
1907            dbtx.insert_new_entry(&SupportsSafeDepositKey, &()).await;
1908            dbtx.commit_tx().await;
1909            break;
1910        }
1911
1912        drop(dbtx);
1913
1914        if is_running_in_test_env() {
1915            // Even in tests we don't want to spam the federation with requests about it
1916            sleep(Duration::from_secs(10)).await;
1917        } else {
1918            sleep(Duration::from_hours(1)).await;
1919        }
1920    }
1921}
1922
1923/// Returns the child index to derive the next peg-in tweak key from.
1924async fn get_next_peg_in_tweak_child_id(dbtx: &mut DatabaseTransaction<'_>) -> TweakIdx {
1925    let index = dbtx
1926        .get_value(&NextPegInTweakIndexKey)
1927        .await
1928        .unwrap_or_default();
1929    dbtx.insert_entry(&NextPegInTweakIndexKey, &(index.next()))
1930        .await;
1931    index
1932}
1933
1934#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
1935pub enum WalletClientStates {
1936    Deposit(DepositStateMachine),
1937    Withdraw(WithdrawStateMachine),
1938}
1939
1940impl IntoDynInstance for WalletClientStates {
1941    type DynType = DynState;
1942
1943    fn into_dyn(self, instance_id: ModuleInstanceId) -> Self::DynType {
1944        DynState::from_typed(instance_id, self)
1945    }
1946}
1947
1948impl State for WalletClientStates {
1949    type ModuleContext = WalletClientContext;
1950
1951    fn transitions(
1952        &self,
1953        context: &Self::ModuleContext,
1954        global_context: &DynGlobalClientContext,
1955    ) -> Vec<StateTransition<Self>> {
1956        match self {
1957            WalletClientStates::Deposit(sm) => {
1958                sm_enum_variant_translation!(
1959                    sm.transitions(context, global_context),
1960                    WalletClientStates::Deposit
1961                )
1962            }
1963            WalletClientStates::Withdraw(sm) => {
1964                sm_enum_variant_translation!(
1965                    sm.transitions(context, global_context),
1966                    WalletClientStates::Withdraw
1967                )
1968            }
1969        }
1970    }
1971
1972    fn operation_id(&self) -> OperationId {
1973        match self {
1974            WalletClientStates::Deposit(sm) => sm.operation_id(),
1975            WalletClientStates::Withdraw(sm) => sm.operation_id(),
1976        }
1977    }
1978}
1979
1980#[cfg(all(test, not(target_family = "wasm")))]
1981mod tests {
1982    use std::collections::BTreeSet;
1983    use std::sync::atomic::{AtomicBool, Ordering};
1984
1985    use super::*;
1986    use crate::backup::{
1987        RECOVER_NUM_IDX_ADD_TO_LAST_USED, RecoverScanOutcome, recover_scan_idxes_for_activity,
1988    };
1989
1990    #[allow(clippy::too_many_lines)] // shut-up clippy, it's a test
1991    #[tokio::test(flavor = "multi_thread")]
1992    async fn sanity_test_recover_inner() {
1993        {
1994            let last_checked = AtomicBool::new(false);
1995            let last_checked = &last_checked;
1996            assert_eq!(
1997                recover_scan_idxes_for_activity(
1998                    TweakIdx(0),
1999                    &BTreeSet::new(),
2000                    |cur_idx| async move {
2001                        Ok(match cur_idx {
2002                            TweakIdx(9) => {
2003                                last_checked.store(true, Ordering::SeqCst);
2004                                vec![]
2005                            }
2006                            TweakIdx(10) => panic!("Shouldn't happen"),
2007                            TweakIdx(11) => {
2008                                vec![0usize] /* just for type inference */
2009                            }
2010                            _ => vec![],
2011                        })
2012                    }
2013                )
2014                .await
2015                .unwrap(),
2016                RecoverScanOutcome {
2017                    last_used_idx: None,
2018                    new_start_idx: TweakIdx(RECOVER_NUM_IDX_ADD_TO_LAST_USED),
2019                    tweak_idxes_with_pegins: BTreeSet::from([])
2020                }
2021            );
2022            assert!(last_checked.load(Ordering::SeqCst));
2023        }
2024
2025        {
2026            let last_checked = AtomicBool::new(false);
2027            let last_checked = &last_checked;
2028            assert_eq!(
2029                recover_scan_idxes_for_activity(
2030                    TweakIdx(0),
2031                    &BTreeSet::from([TweakIdx(1), TweakIdx(2)]),
2032                    |cur_idx| async move {
2033                        Ok(match cur_idx {
2034                            TweakIdx(1) => panic!("Shouldn't happen: already used (1)"),
2035                            TweakIdx(2) => panic!("Shouldn't happen: already used (2)"),
2036                            TweakIdx(11) => {
2037                                last_checked.store(true, Ordering::SeqCst);
2038                                vec![]
2039                            }
2040                            TweakIdx(12) => panic!("Shouldn't happen"),
2041                            TweakIdx(13) => {
2042                                vec![0usize] /* just for type inference */
2043                            }
2044                            _ => vec![],
2045                        })
2046                    }
2047                )
2048                .await
2049                .unwrap(),
2050                RecoverScanOutcome {
2051                    last_used_idx: Some(TweakIdx(2)),
2052                    new_start_idx: TweakIdx(2 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
2053                    tweak_idxes_with_pegins: BTreeSet::from([])
2054                }
2055            );
2056            assert!(last_checked.load(Ordering::SeqCst));
2057        }
2058
2059        {
2060            let last_checked = AtomicBool::new(false);
2061            let last_checked = &last_checked;
2062            assert_eq!(
2063                recover_scan_idxes_for_activity(
2064                    TweakIdx(10),
2065                    &BTreeSet::new(),
2066                    |cur_idx| async move {
2067                        Ok(match cur_idx {
2068                            TweakIdx(10) => vec![()],
2069                            TweakIdx(19) => {
2070                                last_checked.store(true, Ordering::SeqCst);
2071                                vec![]
2072                            }
2073                            TweakIdx(20) => panic!("Shouldn't happen"),
2074                            _ => vec![],
2075                        })
2076                    }
2077                )
2078                .await
2079                .unwrap(),
2080                RecoverScanOutcome {
2081                    last_used_idx: Some(TweakIdx(10)),
2082                    new_start_idx: TweakIdx(10 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
2083                    tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(10)])
2084                }
2085            );
2086            assert!(last_checked.load(Ordering::SeqCst));
2087        }
2088
2089        assert_eq!(
2090            recover_scan_idxes_for_activity(TweakIdx(0), &BTreeSet::new(), |cur_idx| async move {
2091                Ok(match cur_idx {
2092                    TweakIdx(6 | 15) => vec![()],
2093                    _ => vec![],
2094                })
2095            })
2096            .await
2097            .unwrap(),
2098            RecoverScanOutcome {
2099                last_used_idx: Some(TweakIdx(15)),
2100                new_start_idx: TweakIdx(15 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
2101                tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(6), TweakIdx(15)])
2102            }
2103        );
2104        assert_eq!(
2105            recover_scan_idxes_for_activity(TweakIdx(10), &BTreeSet::new(), |cur_idx| async move {
2106                Ok(match cur_idx {
2107                    TweakIdx(8) => {
2108                        vec![()] /* for type inference only */
2109                    }
2110                    TweakIdx(9) => {
2111                        panic!("Shouldn't happen")
2112                    }
2113                    _ => vec![],
2114                })
2115            })
2116            .await
2117            .unwrap(),
2118            RecoverScanOutcome {
2119                last_used_idx: None,
2120                new_start_idx: TweakIdx(9 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
2121                tweak_idxes_with_pegins: BTreeSet::from([])
2122            }
2123        );
2124        assert_eq!(
2125            recover_scan_idxes_for_activity(TweakIdx(10), &BTreeSet::new(), |cur_idx| async move {
2126                Ok(match cur_idx {
2127                    TweakIdx(9) => panic!("Shouldn't happen"),
2128                    TweakIdx(15) => vec![()],
2129                    _ => vec![],
2130                })
2131            })
2132            .await
2133            .unwrap(),
2134            RecoverScanOutcome {
2135                last_used_idx: Some(TweakIdx(15)),
2136                new_start_idx: TweakIdx(15 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
2137                tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(15)])
2138            }
2139        );
2140    }
2141}