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;
15mod deposit;
18pub mod events;
19mod pegin_monitor;
21mod withdraw;
22
23use std::collections::{BTreeMap, BTreeSet};
24use std::future;
25use std::sync::Arc;
26use std::time::{Duration, SystemTime};
27
28use anyhow::{Context as AnyhowContext, anyhow, bail, ensure};
29use async_stream::{stream, try_stream};
30use backup::WalletModuleBackup;
31use bitcoin::address::NetworkUnchecked;
32use bitcoin::secp256k1::{All, SECP256K1, Secp256k1};
33use bitcoin::{Address, Network, ScriptBuf};
34use client_db::{DbKeyPrefix, PegInTweakIndexKey, SupportsSafeDepositKey, TweakIdx};
35use fedimint_api_client::api::{DynModuleApi, FederationResult};
36use fedimint_bitcoind::{DynBitcoindRpc, create_esplora_rpc};
37use fedimint_client_module::module::init::{
38 ClientModuleInit, ClientModuleInitArgs, ClientModuleRecoverArgs,
39};
40use fedimint_client_module::module::{ClientContext, ClientModule, IClientModule, OutPointRange};
41use fedimint_client_module::oplog::UpdateStreamOrOutcome;
42use fedimint_client_module::sm::{Context, DynState, ModuleNotifier, State, StateTransition};
43use fedimint_client_module::transaction::{
44 ClientOutput, ClientOutputBundle, ClientOutputSM, TransactionBuilder,
45};
46use fedimint_client_module::{DynGlobalClientContext, sm_enum_variant_translation};
47use fedimint_core::core::{Decoder, IntoDynInstance, ModuleInstanceId, ModuleKind, OperationId};
48use fedimint_core::db::{
49 AutocommitError, Database, DatabaseTransaction, IDatabaseTransactionOpsCoreTyped,
50};
51use fedimint_core::encoding::{Decodable, Encodable};
52use fedimint_core::envs::{BitcoinRpcConfig, is_running_in_test_env};
53use fedimint_core::module::{
54 ApiAuth, ApiVersion, CommonModuleInit, ModuleCommon, ModuleConsensusVersion, ModuleInit,
55 MultiApiVersion,
56};
57use fedimint_core::task::{MaybeSend, MaybeSync, TaskGroup, sleep};
58use fedimint_core::util::backoff_util::background_backoff;
59use fedimint_core::util::{BoxStream, backoff_util, retry};
60use fedimint_core::{
61 Amount, OutPoint, TransactionId, apply, async_trait_maybe_send, push_db_pair_items, runtime,
62 secp256k1,
63};
64use fedimint_derive_secret::{ChildId, DerivableSecret};
65use fedimint_logging::LOG_CLIENT_MODULE_WALLET;
66use fedimint_wallet_common::config::{FeeConsensus, WalletClientConfig};
67use fedimint_wallet_common::tweakable::Tweakable;
68pub use fedimint_wallet_common::*;
69use futures::{Stream, StreamExt};
70use rand::{Rng, thread_rng};
71use secp256k1::Keypair;
72use serde::{Deserialize, Serialize};
73use strum::IntoEnumIterator;
74use tokio::sync::watch;
75use tracing::{debug, instrument};
76
77use crate::api::WalletFederationApi;
78use crate::backup::WalletRecovery;
79use crate::client_db::{
80 ClaimedPegInData, ClaimedPegInKey, ClaimedPegInPrefix, NextPegInTweakIndexKey,
81 PegInTweakIndexData, PegInTweakIndexPrefix, RecoveryFinalizedKey, SupportsSafeDepositPrefix,
82};
83use crate::deposit::DepositStateMachine;
84use crate::withdraw::{CreatedWithdrawState, WithdrawStateMachine, WithdrawStates};
85
86const WALLET_TWEAK_CHILD_ID: ChildId = ChildId(0);
87
88#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
89pub struct BitcoinTransactionData {
90 pub btc_transaction: bitcoin::Transaction,
93 pub out_idx: u32,
95}
96
97#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
98pub enum DepositStateV1 {
99 WaitingForTransaction,
100 WaitingForConfirmation(BitcoinTransactionData),
101 Confirmed(BitcoinTransactionData),
102 Claimed(BitcoinTransactionData),
103 Failed(String),
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
107pub enum DepositStateV2 {
108 WaitingForTransaction,
109 WaitingForConfirmation {
110 #[serde(with = "bitcoin::amount::serde::as_sat")]
111 btc_deposited: bitcoin::Amount,
112 btc_out_point: bitcoin::OutPoint,
113 },
114 Confirmed {
115 #[serde(with = "bitcoin::amount::serde::as_sat")]
116 btc_deposited: bitcoin::Amount,
117 btc_out_point: bitcoin::OutPoint,
118 },
119 Claimed {
120 #[serde(with = "bitcoin::amount::serde::as_sat")]
121 btc_deposited: bitcoin::Amount,
122 btc_out_point: bitcoin::OutPoint,
123 },
124 Failed(String),
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
128pub enum WithdrawState {
129 Created,
130 Succeeded(bitcoin::Txid),
131 Failed(String),
132 }
136
137async fn next_withdraw_state<S>(stream: &mut S) -> Option<WithdrawStates>
138where
139 S: Stream<Item = WalletClientStates> + Unpin,
140{
141 loop {
142 if let WalletClientStates::Withdraw(ds) = stream.next().await? {
143 return Some(ds.state);
144 }
145 tokio::task::yield_now().await;
146 }
147}
148
149#[derive(Debug, Clone, Default)]
150pub struct WalletClientInit(pub Option<DynBitcoindRpc>);
152
153impl WalletClientInit {
154 pub fn new(rpc: DynBitcoindRpc) -> Self {
155 Self(Some(rpc))
156 }
157}
158
159impl ModuleInit for WalletClientInit {
160 type Common = WalletCommonInit;
161
162 async fn dump_database(
163 &self,
164 dbtx: &mut DatabaseTransaction<'_>,
165 prefix_names: Vec<String>,
166 ) -> Box<dyn Iterator<Item = (String, Box<dyn erased_serde::Serialize + Send>)> + '_> {
167 let mut wallet_client_items: BTreeMap<String, Box<dyn erased_serde::Serialize + Send>> =
168 BTreeMap::new();
169 let filtered_prefixes = DbKeyPrefix::iter().filter(|f| {
170 prefix_names.is_empty() || prefix_names.contains(&f.to_string().to_lowercase())
171 });
172
173 for table in filtered_prefixes {
174 match table {
175 DbKeyPrefix::NextPegInTweakIndex => {
176 if let Some(index) = dbtx.get_value(&NextPegInTweakIndexKey).await {
177 wallet_client_items
178 .insert("NextPegInTweakIndex".to_string(), Box::new(index));
179 }
180 }
181 DbKeyPrefix::PegInTweakIndex => {
182 push_db_pair_items!(
183 dbtx,
184 PegInTweakIndexPrefix,
185 PegInTweakIndexKey,
186 PegInTweakIndexData,
187 wallet_client_items,
188 "Peg-In Tweak Index"
189 );
190 }
191 DbKeyPrefix::ClaimedPegIn => {
192 push_db_pair_items!(
193 dbtx,
194 ClaimedPegInPrefix,
195 ClaimedPegInKey,
196 ClaimedPegInData,
197 wallet_client_items,
198 "Claimed Peg-In"
199 );
200 }
201 DbKeyPrefix::RecoveryFinalized => {
202 if let Some(val) = dbtx.get_value(&RecoveryFinalizedKey).await {
203 wallet_client_items.insert("RecoveryFinalized".to_string(), Box::new(val));
204 }
205 }
206 DbKeyPrefix::SupportsSafeDeposit => {
207 push_db_pair_items!(
208 dbtx,
209 SupportsSafeDepositPrefix,
210 SupportsSafeDepositKey,
211 (),
212 wallet_client_items,
213 "Supports Safe Deposit"
214 );
215 }
216 DbKeyPrefix::RecoveryState
217 | DbKeyPrefix::ExternalReservedStart
218 | DbKeyPrefix::CoreInternalReservedStart
219 | DbKeyPrefix::CoreInternalReservedEnd => {}
220 }
221 }
222
223 Box::new(wallet_client_items.into_iter())
224 }
225}
226
227#[apply(async_trait_maybe_send!)]
228impl ClientModuleInit for WalletClientInit {
229 type Module = WalletClientModule;
230
231 fn supported_api_versions(&self) -> MultiApiVersion {
232 MultiApiVersion::try_from_iter([ApiVersion { major: 0, minor: 0 }])
233 .expect("no version conflicts")
234 }
235
236 async fn init(&self, args: &ClientModuleInitArgs<Self>) -> anyhow::Result<Self::Module> {
237 let data = WalletClientModuleData {
238 cfg: args.cfg().clone(),
239 module_root_secret: args.module_root_secret().clone(),
240 };
241
242 let db = args.db().clone();
243
244 let btc_rpc = self.0.clone().unwrap_or(create_esplora_rpc(
245 &WalletClientModule::get_rpc_config(args.cfg()).url,
246 )?);
247
248 let module_api = args.module_api().clone();
249
250 let (pegin_claimed_sender, pegin_claimed_receiver) = watch::channel(());
251 let (pegin_monitor_wakeup_sender, pegin_monitor_wakeup_receiver) = watch::channel(());
252
253 Ok(WalletClientModule {
254 db,
255 data,
256 module_api,
257 notifier: args.notifier().clone(),
258 rpc: btc_rpc,
259 client_ctx: args.context(),
260 pegin_monitor_wakeup_sender,
261 pegin_monitor_wakeup_receiver,
262 pegin_claimed_receiver,
263 pegin_claimed_sender,
264 task_group: args.task_group().clone(),
265 admin_auth: args.admin_auth().cloned(),
266 })
267 }
268
269 async fn recover(
277 &self,
278 args: &ClientModuleRecoverArgs<Self>,
279 snapshot: Option<&<Self::Module as ClientModule>::Backup>,
280 ) -> anyhow::Result<()> {
281 args.recover_from_history::<WalletRecovery>(self, snapshot)
282 .await
283 }
284
285 fn used_db_prefixes(&self) -> Option<BTreeSet<u8>> {
286 Some(
287 DbKeyPrefix::iter()
288 .map(|p| p as u8)
289 .chain(
290 DbKeyPrefix::ExternalReservedStart as u8
291 ..=DbKeyPrefix::CoreInternalReservedEnd as u8,
292 )
293 .collect(),
294 )
295 }
296}
297
298#[derive(Debug, Clone, Serialize, Deserialize)]
299pub struct WalletOperationMeta {
300 pub variant: WalletOperationMetaVariant,
301 pub extra_meta: serde_json::Value,
302}
303
304#[derive(Debug, Clone, Serialize, Deserialize)]
305#[serde(rename_all = "snake_case")]
306pub enum WalletOperationMetaVariant {
307 Deposit {
308 address: Address<NetworkUnchecked>,
309 #[serde(default)]
314 tweak_idx: Option<TweakIdx>,
315 #[serde(default, skip_serializing_if = "Option::is_none")]
316 expires_at: Option<SystemTime>,
317 },
318 Withdraw {
319 address: Address<NetworkUnchecked>,
320 #[serde(with = "bitcoin::amount::serde::as_sat")]
321 amount: bitcoin::Amount,
322 fee: PegOutFees,
323 change: Vec<OutPoint>,
324 },
325
326 RbfWithdraw {
327 rbf: Rbf,
328 change: Vec<OutPoint>,
329 },
330}
331
332#[derive(Debug, Clone)]
334pub struct WalletClientModuleData {
335 cfg: WalletClientConfig,
336 module_root_secret: DerivableSecret,
337}
338
339impl WalletClientModuleData {
340 fn derive_deposit_address(
341 &self,
342 idx: TweakIdx,
343 ) -> (Keypair, secp256k1::PublicKey, Address, OperationId) {
344 let idx = ChildId(idx.0);
345
346 let secret_tweak_key = self
347 .module_root_secret
348 .child_key(WALLET_TWEAK_CHILD_ID)
349 .child_key(idx)
350 .to_secp_key(fedimint_core::secp256k1::SECP256K1);
351
352 let public_tweak_key = secret_tweak_key.public_key();
353
354 let address = self
355 .cfg
356 .peg_in_descriptor
357 .tweak(&public_tweak_key, bitcoin::secp256k1::SECP256K1)
358 .address(self.cfg.network.0)
359 .unwrap();
360
361 let operation_id = OperationId(public_tweak_key.x_only_public_key().0.serialize());
363
364 (secret_tweak_key, public_tweak_key, address, operation_id)
365 }
366
367 fn derive_peg_in_script(
368 &self,
369 idx: TweakIdx,
370 ) -> (ScriptBuf, bitcoin::Address, Keypair, OperationId) {
371 let (secret_tweak_key, _, address, operation_id) = self.derive_deposit_address(idx);
372
373 (
374 self.cfg
375 .peg_in_descriptor
376 .tweak(&secret_tweak_key.public_key(), SECP256K1)
377 .script_pubkey(),
378 address,
379 secret_tweak_key,
380 operation_id,
381 )
382 }
383}
384
385#[derive(Debug)]
386pub struct WalletClientModule {
387 data: WalletClientModuleData,
388 db: Database,
389 module_api: DynModuleApi,
390 notifier: ModuleNotifier<WalletClientStates>,
391 rpc: DynBitcoindRpc,
392 client_ctx: ClientContext<Self>,
393 pegin_monitor_wakeup_sender: watch::Sender<()>,
395 pegin_monitor_wakeup_receiver: watch::Receiver<()>,
396 pegin_claimed_sender: watch::Sender<()>,
398 pegin_claimed_receiver: watch::Receiver<()>,
399 task_group: TaskGroup,
400 admin_auth: Option<ApiAuth>,
401}
402
403#[apply(async_trait_maybe_send!)]
404impl ClientModule for WalletClientModule {
405 type Init = WalletClientInit;
406 type Common = WalletModuleTypes;
407 type Backup = WalletModuleBackup;
408 type ModuleStateMachineContext = WalletClientContext;
409 type States = WalletClientStates;
410
411 fn context(&self) -> Self::ModuleStateMachineContext {
412 WalletClientContext {
413 rpc: self.rpc.clone(),
414 wallet_descriptor: self.cfg().peg_in_descriptor.clone(),
415 wallet_decoder: self.decoder(),
416 secp: Secp256k1::default(),
417 client_ctx: self.client_ctx.clone(),
418 }
419 }
420
421 async fn start(&self) {
422 self.task_group.spawn_cancellable("peg-in monitor", {
423 let client_ctx = self.client_ctx.clone();
424 let db = self.db.clone();
425 let btc_rpc = self.rpc.clone();
426 let module_api = self.module_api.clone();
427 let data = self.data.clone();
428 let pegin_claimed_sender = self.pegin_claimed_sender.clone();
429 let pegin_monitor_wakeup_receiver = self.pegin_monitor_wakeup_receiver.clone();
430 pegin_monitor::run_peg_in_monitor(
431 client_ctx,
432 db,
433 btc_rpc,
434 module_api,
435 data,
436 pegin_claimed_sender,
437 pegin_monitor_wakeup_receiver,
438 )
439 });
440
441 self.task_group
442 .spawn_cancellable("supports-safe-deposit-version", {
443 let db = self.db.clone();
444 let module_api = self.module_api.clone();
445
446 poll_supports_safe_deposit_version(db, module_api)
447 });
448 }
449
450 fn supports_backup(&self) -> bool {
451 true
452 }
453
454 async fn backup(&self) -> anyhow::Result<backup::WalletModuleBackup> {
455 let session_count = self.client_ctx.global_api().session_count().await?;
457
458 let mut dbtx = self.db.begin_transaction_nc().await;
459 let next_pegin_tweak_idx = dbtx
460 .get_value(&NextPegInTweakIndexKey)
461 .await
462 .unwrap_or_default();
463 let claimed = dbtx
464 .find_by_prefix(&PegInTweakIndexPrefix)
465 .await
466 .filter_map(|(k, v)| async move {
467 if v.claimed.is_empty() {
468 None
469 } else {
470 Some(k.0)
471 }
472 })
473 .collect()
474 .await;
475 Ok(backup::WalletModuleBackup::new_v1(
476 session_count,
477 next_pegin_tweak_idx,
478 claimed,
479 ))
480 }
481
482 fn input_fee(
483 &self,
484 _amount: Amount,
485 _input: &<Self::Common as ModuleCommon>::Input,
486 ) -> Option<Amount> {
487 Some(self.cfg().fee_consensus.peg_in_abs)
488 }
489
490 fn output_fee(
491 &self,
492 _amount: Amount,
493 _output: &<Self::Common as ModuleCommon>::Output,
494 ) -> Option<Amount> {
495 Some(self.cfg().fee_consensus.peg_out_abs)
496 }
497
498 async fn handle_rpc(
499 &self,
500 method: String,
501 request: serde_json::Value,
502 ) -> BoxStream<'_, anyhow::Result<serde_json::Value>> {
503 Box::pin(try_stream! {
504 match method.as_str() {
505 "get_wallet_summary" => {
506 let _req: WalletSummaryRequest = serde_json::from_value(request)?;
507 let wallet_summary = self.get_wallet_summary()
508 .await
509 .expect("Failed to fetch wallet summary");
510 let result = serde_json::to_value(&wallet_summary)
511 .expect("Serialization error");
512 yield result;
513 }
514 "peg_in" => {
515 let req: PegInRequest = serde_json::from_value(request)?;
516 let response = self.peg_in(req)
517 .await
518 .map_err(|e| anyhow::anyhow!("peg_in failed: {}", e))?;
519 let result = serde_json::to_value(&response)?;
520 yield result;
521 },
522 "peg_out" => {
523 let req: PegOutRequest = serde_json::from_value(request)?;
524 let response = self.peg_out(req)
525 .await
526 .map_err(|e| anyhow::anyhow!("peg_out failed: {}", e))?;
527 let result = serde_json::to_value(&response)?;
528 yield result;
529 }
530 _ => {
531 Err(anyhow::format_err!("Unknown method: {}", method))?;
532 }
533 }
534 })
535 }
536
537 #[cfg(feature = "cli")]
538 async fn handle_cli_command(
539 &self,
540 args: &[std::ffi::OsString],
541 ) -> anyhow::Result<serde_json::Value> {
542 cli::handle_cli_command(self, args).await
543 }
544}
545
546#[derive(Deserialize)]
547struct WalletSummaryRequest {}
548
549#[derive(Debug, Clone)]
550pub struct WalletClientContext {
551 rpc: DynBitcoindRpc,
552 wallet_descriptor: PegInDescriptor,
553 wallet_decoder: Decoder,
554 secp: Secp256k1<All>,
555 pub client_ctx: ClientContext<WalletClientModule>,
556}
557
558#[derive(Debug, Clone, Serialize, Deserialize)]
559pub struct PegInRequest {
560 pub extra_meta: serde_json::Value,
561}
562
563#[derive(Debug, Clone, Serialize, Deserialize)]
564pub struct PegInResponse {
565 pub deposit_address: Address<NetworkUnchecked>,
566 pub operation_id: OperationId,
567}
568
569#[derive(Debug, Clone, Serialize, Deserialize)]
570pub struct PegOutRequest {
571 pub amount_sat: u64,
572 pub destination_address: Address<NetworkUnchecked>,
573 pub extra_meta: serde_json::Value,
574}
575
576#[derive(Debug, Clone, Serialize, Deserialize)]
577pub struct PegOutResponse {
578 pub operation_id: OperationId,
579}
580
581impl Context for WalletClientContext {
582 const KIND: Option<ModuleKind> = Some(KIND);
583}
584
585impl WalletClientModule {
586 fn cfg(&self) -> &WalletClientConfig {
587 &self.data.cfg
588 }
589
590 fn get_rpc_config(cfg: &WalletClientConfig) -> BitcoinRpcConfig {
591 match BitcoinRpcConfig::get_defaults_from_env_vars() {
592 Ok(rpc_config) => {
593 if rpc_config.kind == "bitcoind" {
596 cfg.default_bitcoin_rpc.clone()
597 } else {
598 rpc_config
599 }
600 }
601 _ => cfg.default_bitcoin_rpc.clone(),
602 }
603 }
604
605 pub fn get_network(&self) -> Network {
606 self.cfg().network.0
607 }
608
609 pub fn get_fee_consensus(&self) -> FeeConsensus {
610 self.cfg().fee_consensus
611 }
612
613 async fn allocate_deposit_address_inner(
614 &self,
615 dbtx: &mut DatabaseTransaction<'_>,
616 ) -> (OperationId, Address, TweakIdx) {
617 dbtx.ensure_isolated().expect("Must be isolated db");
618
619 let tweak_idx = get_next_peg_in_tweak_child_id(dbtx).await;
620 let (_secret_tweak_key, _, address, operation_id) =
621 self.data.derive_deposit_address(tweak_idx);
622
623 let now = fedimint_core::time::now();
624
625 dbtx.insert_new_entry(
626 &PegInTweakIndexKey(tweak_idx),
627 &PegInTweakIndexData {
628 creation_time: now,
629 next_check_time: Some(now),
630 last_check_time: None,
631 operation_id,
632 claimed: vec![],
633 },
634 )
635 .await;
636
637 (operation_id, address, tweak_idx)
638 }
639
640 pub async fn get_withdraw_fees(
647 &self,
648 address: &bitcoin::Address,
649 amount: bitcoin::Amount,
650 ) -> anyhow::Result<PegOutFees> {
651 self.module_api
652 .fetch_peg_out_fees(address, amount)
653 .await?
654 .context("Federation didn't return peg-out fees")
655 }
656
657 pub async fn get_wallet_summary(&self) -> anyhow::Result<WalletSummary> {
659 Ok(self.module_api.fetch_wallet_summary().await?)
660 }
661
662 pub fn create_withdraw_output(
663 &self,
664 operation_id: OperationId,
665 address: bitcoin::Address,
666 amount: bitcoin::Amount,
667 fees: PegOutFees,
668 ) -> anyhow::Result<ClientOutputBundle<WalletOutput, WalletClientStates>> {
669 let output = WalletOutput::new_v0_peg_out(address, amount, fees);
670
671 let amount = output.maybe_v0_ref().expect("v0 output").amount().into();
672
673 let sm_gen = move |out_point_range: OutPointRange| {
674 assert_eq!(out_point_range.count(), 1);
675 let out_idx = out_point_range.start_idx();
676 vec![WalletClientStates::Withdraw(WithdrawStateMachine {
677 operation_id,
678 state: WithdrawStates::Created(CreatedWithdrawState {
679 fm_outpoint: OutPoint {
680 txid: out_point_range.txid(),
681 out_idx,
682 },
683 }),
684 })]
685 };
686
687 Ok(ClientOutputBundle::new(
688 vec![ClientOutput::<WalletOutput> { output, amount }],
689 vec![ClientOutputSM::<WalletClientStates> {
690 state_machines: Arc::new(sm_gen),
691 }],
692 ))
693 }
694
695 pub async fn peg_in(&self, req: PegInRequest) -> anyhow::Result<PegInResponse> {
696 let (operation_id, address, _) = self.safe_allocate_deposit_address(req.extra_meta).await?;
697
698 Ok(PegInResponse {
699 deposit_address: Address::from_script(&address.script_pubkey(), self.get_network())?
700 .as_unchecked()
701 .clone(),
702 operation_id,
703 })
704 }
705
706 pub async fn peg_out(&self, req: PegOutRequest) -> anyhow::Result<PegOutResponse> {
707 let amount = bitcoin::Amount::from_sat(req.amount_sat);
708 let destination = req
709 .destination_address
710 .require_network(self.get_network())?;
711
712 let fees = self.get_withdraw_fees(&destination, amount).await?;
713 let operation_id = self
714 .withdraw(&destination, amount, fees, req.extra_meta)
715 .await
716 .context("Failed to initiate withdraw")?;
717
718 Ok(PegOutResponse { operation_id })
719 }
720
721 pub fn create_rbf_withdraw_output(
722 &self,
723 operation_id: OperationId,
724 rbf: &Rbf,
725 ) -> anyhow::Result<ClientOutputBundle<WalletOutput, WalletClientStates>> {
726 let output = WalletOutput::new_v0_rbf(rbf.fees, rbf.txid);
727
728 let amount = output.maybe_v0_ref().expect("v0 output").amount().into();
729
730 let sm_gen = move |out_point_range: OutPointRange| {
731 assert_eq!(out_point_range.count(), 1);
732 let out_idx = out_point_range.start_idx();
733 vec![WalletClientStates::Withdraw(WithdrawStateMachine {
734 operation_id,
735 state: WithdrawStates::Created(CreatedWithdrawState {
736 fm_outpoint: OutPoint {
737 txid: out_point_range.txid(),
738 out_idx,
739 },
740 }),
741 })]
742 };
743
744 Ok(ClientOutputBundle::new(
745 vec![ClientOutput::<WalletOutput> { output, amount }],
746 vec![ClientOutputSM::<WalletClientStates> {
747 state_machines: Arc::new(sm_gen),
748 }],
749 ))
750 }
751
752 pub async fn btc_tx_has_no_size_limit(&self) -> FederationResult<bool> {
753 Ok(self.module_api.module_consensus_version().await? >= ModuleConsensusVersion::new(2, 2))
754 }
755
756 pub async fn supports_safe_deposit(&self) -> bool {
765 let mut dbtx = self.db.begin_transaction().await;
766
767 let already_verified_supports_safe_deposit =
768 dbtx.get_value(&SupportsSafeDepositKey).await.is_some();
769
770 already_verified_supports_safe_deposit || {
771 match self.module_api.module_consensus_version().await {
772 Ok(module_consensus_version) => {
773 let supported_version =
774 SAFE_DEPOSIT_MODULE_CONSENSUS_VERSION <= module_consensus_version;
775
776 if supported_version {
777 dbtx.insert_new_entry(&SupportsSafeDepositKey, &()).await;
778 dbtx.commit_tx().await;
779 }
780
781 supported_version
782 }
783 Err(_) => false,
784 }
785 }
786 }
787
788 pub async fn safe_allocate_deposit_address<M>(
796 &self,
797 extra_meta: M,
798 ) -> anyhow::Result<(OperationId, Address, TweakIdx)>
799 where
800 M: Serialize + MaybeSend + MaybeSync,
801 {
802 ensure!(
803 self.supports_safe_deposit().await,
804 "Wallet module consensus version doesn't support safe deposits",
805 );
806
807 self.allocate_deposit_address_expert_only(extra_meta).await
808 }
809
810 pub async fn allocate_deposit_address_expert_only<M>(
828 &self,
829 extra_meta: M,
830 ) -> anyhow::Result<(OperationId, Address, TweakIdx)>
831 where
832 M: Serialize + MaybeSend + MaybeSync,
833 {
834 let extra_meta_value =
835 serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
836 let (operation_id, address, tweak_idx) = self
837 .db
838 .autocommit(
839 move |dbtx, _| {
840 let extra_meta_value_inner = extra_meta_value.clone();
841 Box::pin(async move {
842 let (operation_id, address, tweak_idx) = self
843 .allocate_deposit_address_inner(dbtx)
844 .await;
845
846 self.client_ctx.manual_operation_start_dbtx(
847 dbtx,
848 operation_id,
849 WalletCommonInit::KIND.as_str(),
850 WalletOperationMeta {
851 variant: WalletOperationMetaVariant::Deposit {
852 address: address.clone().into_unchecked(),
853 tweak_idx: Some(tweak_idx),
854 expires_at: None,
855 },
856 extra_meta: extra_meta_value_inner,
857 },
858 vec![]
859 ).await?;
860
861 debug!(target: LOG_CLIENT_MODULE_WALLET, %tweak_idx, %address, "Derived a new deposit address");
862
863
864 let sender = self.pegin_monitor_wakeup_sender.clone();
865 dbtx.on_commit(move || {
866 sender.send_replace(());
867 });
868
869 Ok((operation_id, address, tweak_idx))
870 })
871 },
872 Some(100),
873 )
874 .await
875 .map_err(|e| match e {
876 AutocommitError::CommitFailed {
877 last_error,
878 attempts,
879 } => last_error.context(format!("Failed to commit after {attempts} attempts")),
880 AutocommitError::ClosureError { error, .. } => error,
881 })?;
882
883 Ok((operation_id, address, tweak_idx))
884 }
885
886 pub async fn subscribe_deposit(
892 &self,
893 operation_id: OperationId,
894 ) -> anyhow::Result<UpdateStreamOrOutcome<DepositStateV2>> {
895 let operation = self
896 .client_ctx
897 .get_operation(operation_id)
898 .await
899 .with_context(|| anyhow!("Operation not found: {}", operation_id.fmt_short()))?;
900
901 if operation.operation_module_kind() != WalletCommonInit::KIND.as_str() {
902 bail!("Operation is not a wallet operation");
903 }
904
905 let operation_meta = operation.meta::<WalletOperationMeta>();
906
907 let WalletOperationMetaVariant::Deposit {
908 address, tweak_idx, ..
909 } = operation_meta.variant
910 else {
911 bail!("Operation is not a deposit operation");
912 };
913
914 let address = address.require_network(self.cfg().network.0)?;
915
916 let Some(tweak_idx) = tweak_idx else {
918 let outcome_v1 = operation
922 .outcome::<DepositStateV1>()
923 .context("Old pending deposit, can't subscribe to updates")?;
924
925 let outcome_v2 = match outcome_v1 {
926 DepositStateV1::Claimed(tx_info) => DepositStateV2::Claimed {
927 btc_deposited: tx_info.btc_transaction.output[tx_info.out_idx as usize].value,
928 btc_out_point: bitcoin::OutPoint {
929 txid: tx_info.btc_transaction.compute_txid(),
930 vout: tx_info.out_idx,
931 },
932 },
933 DepositStateV1::Failed(error) => DepositStateV2::Failed(error),
934 _ => bail!("Non-final outcome in operation log"),
935 };
936
937 return Ok(UpdateStreamOrOutcome::Outcome(outcome_v2));
938 };
939
940 Ok(self.client_ctx.outcome_or_updates(operation, operation_id, {
941 let stream_rpc = self.rpc.clone();
942 let stream_client_ctx = self.client_ctx.clone();
943 let stream_script_pub_key = address.script_pubkey();
944 move || {
945
946 stream! {
947 yield DepositStateV2::WaitingForTransaction;
948
949 let (btc_out_point, btc_deposited) = retry(
950 "fetch history",
951 background_backoff(),
952 || async {
953 let history = stream_rpc.get_script_history(&stream_script_pub_key).await?;
954 history.first().and_then(|tx| {
955 let (out_idx, amount) = tx.output
956 .iter()
957 .enumerate()
958 .find_map(|(idx, output)| (output.script_pubkey == stream_script_pub_key).then_some((idx, output.value)))?;
959 let txid = tx.compute_txid();
960
961 Some((
962 bitcoin::OutPoint {
963 txid,
964 vout: out_idx as u32,
965 },
966 amount
967 ))
968 }).context("No deposit transaction found")
969 }
970 ).await.expect("Will never give up");
971
972 yield DepositStateV2::WaitingForConfirmation {
973 btc_deposited,
974 btc_out_point
975 };
976
977 let claim_data = stream_client_ctx.module_db().wait_key_exists(&ClaimedPegInKey {
978 peg_in_index: tweak_idx,
979 btc_out_point,
980 }).await;
981
982 yield DepositStateV2::Confirmed {
983 btc_deposited,
984 btc_out_point
985 };
986
987 match stream_client_ctx.await_primary_module_outputs(operation_id, claim_data.change).await {
988 Ok(()) => yield DepositStateV2::Claimed {
989 btc_deposited,
990 btc_out_point
991 },
992 Err(e) => yield DepositStateV2::Failed(e.to_string())
993 }
994 }
995 }}))
996 }
997
998 pub async fn list_peg_in_tweak_idxes(&self) -> BTreeMap<TweakIdx, PegInTweakIndexData> {
999 self.client_ctx
1000 .module_db()
1001 .clone()
1002 .begin_transaction_nc()
1003 .await
1004 .find_by_prefix(&PegInTweakIndexPrefix)
1005 .await
1006 .map(|(key, data)| (key.0, data))
1007 .collect()
1008 .await
1009 }
1010
1011 pub async fn find_tweak_idx_by_address(
1012 &self,
1013 address: bitcoin::Address<NetworkUnchecked>,
1014 ) -> anyhow::Result<TweakIdx> {
1015 let data = self.data.clone();
1016 let Some((tweak_idx, _)) = self
1017 .db
1018 .begin_transaction_nc()
1019 .await
1020 .find_by_prefix(&PegInTweakIndexPrefix)
1021 .await
1022 .filter(|(k, _)| {
1023 let (_, derived_address, _tweak_key, _) = data.derive_peg_in_script(k.0);
1024 future::ready(derived_address.into_unchecked() == address)
1025 })
1026 .next()
1027 .await
1028 else {
1029 bail!("Address not found in the list of derived keys");
1030 };
1031
1032 Ok(tweak_idx.0)
1033 }
1034 pub async fn find_tweak_idx_by_operation_id(
1035 &self,
1036 operation_id: OperationId,
1037 ) -> anyhow::Result<TweakIdx> {
1038 Ok(self
1039 .client_ctx
1040 .module_db()
1041 .clone()
1042 .begin_transaction_nc()
1043 .await
1044 .find_by_prefix(&PegInTweakIndexPrefix)
1045 .await
1046 .filter(|(_k, v)| future::ready(v.operation_id == operation_id))
1047 .next()
1048 .await
1049 .ok_or_else(|| anyhow::format_err!("OperationId not found"))?
1050 .0
1051 .0)
1052 }
1053
1054 pub async fn get_pegin_tweak_idx(
1055 &self,
1056 tweak_idx: TweakIdx,
1057 ) -> anyhow::Result<PegInTweakIndexData> {
1058 self.client_ctx
1059 .module_db()
1060 .clone()
1061 .begin_transaction_nc()
1062 .await
1063 .get_value(&PegInTweakIndexKey(tweak_idx))
1064 .await
1065 .ok_or_else(|| anyhow::format_err!("TweakIdx not found"))
1066 }
1067
1068 pub async fn get_claimed_pegins(
1069 &self,
1070 dbtx: &mut DatabaseTransaction<'_>,
1071 tweak_idx: TweakIdx,
1072 ) -> Vec<(
1073 bitcoin::OutPoint,
1074 TransactionId,
1075 Vec<fedimint_core::OutPoint>,
1076 )> {
1077 let outpoints = dbtx
1078 .get_value(&PegInTweakIndexKey(tweak_idx))
1079 .await
1080 .map(|v| v.claimed)
1081 .unwrap_or_default();
1082
1083 let mut res = vec![];
1084
1085 for outpoint in outpoints {
1086 let claimed_peg_in_data = dbtx
1087 .get_value(&ClaimedPegInKey {
1088 peg_in_index: tweak_idx,
1089 btc_out_point: outpoint,
1090 })
1091 .await
1092 .expect("Must have a corresponding claim record");
1093 res.push((
1094 outpoint,
1095 claimed_peg_in_data.claim_txid,
1096 claimed_peg_in_data.change,
1097 ));
1098 }
1099
1100 res
1101 }
1102
1103 pub async fn recheck_pegin_address_by_op_id(
1105 &self,
1106 operation_id: OperationId,
1107 ) -> anyhow::Result<()> {
1108 let tweak_idx = self.find_tweak_idx_by_operation_id(operation_id).await?;
1109
1110 self.recheck_pegin_address(tweak_idx).await
1111 }
1112
1113 pub async fn recheck_pegin_address_by_address(
1115 &self,
1116 address: bitcoin::Address<NetworkUnchecked>,
1117 ) -> anyhow::Result<()> {
1118 self.recheck_pegin_address(self.find_tweak_idx_by_address(address).await?)
1119 .await
1120 }
1121
1122 pub async fn recheck_pegin_address(&self, tweak_idx: TweakIdx) -> anyhow::Result<()> {
1124 self.db
1125 .autocommit(
1126 |dbtx, _| {
1127 Box::pin(async {
1128 let db_key = PegInTweakIndexKey(tweak_idx);
1129 let db_val = dbtx
1130 .get_value(&db_key)
1131 .await
1132 .ok_or_else(|| anyhow::format_err!("DBKey not found"))?;
1133
1134 dbtx.insert_entry(
1135 &db_key,
1136 &PegInTweakIndexData {
1137 next_check_time: Some(fedimint_core::time::now()),
1138 ..db_val
1139 },
1140 )
1141 .await;
1142
1143 let sender = self.pegin_monitor_wakeup_sender.clone();
1144 dbtx.on_commit(move || {
1145 sender.send_replace(());
1146 });
1147
1148 Ok::<_, anyhow::Error>(())
1149 })
1150 },
1151 Some(100),
1152 )
1153 .await?;
1154
1155 Ok(())
1156 }
1157
1158 pub async fn await_num_deposits_by_operation_id(
1160 &self,
1161 operation_id: OperationId,
1162 num_deposits: usize,
1163 ) -> anyhow::Result<()> {
1164 let tweak_idx = self.find_tweak_idx_by_operation_id(operation_id).await?;
1165 self.await_num_deposits(tweak_idx, num_deposits).await
1166 }
1167
1168 pub async fn await_num_deposits_by_address(
1169 &self,
1170 address: bitcoin::Address<NetworkUnchecked>,
1171 num_deposits: usize,
1172 ) -> anyhow::Result<()> {
1173 self.await_num_deposits(self.find_tweak_idx_by_address(address).await?, num_deposits)
1174 .await
1175 }
1176
1177 #[instrument(target = LOG_CLIENT_MODULE_WALLET, skip_all, fields(tweak_idx=?tweak_idx, num_deposists=num_deposits))]
1178 pub async fn await_num_deposits(
1179 &self,
1180 tweak_idx: TweakIdx,
1181 num_deposits: usize,
1182 ) -> anyhow::Result<()> {
1183 let operation_id = self.get_pegin_tweak_idx(tweak_idx).await?.operation_id;
1184
1185 let mut receiver = self.pegin_claimed_receiver.clone();
1186 let mut backoff = backoff_util::aggressive_backoff();
1187
1188 loop {
1189 let pegins = self
1190 .get_claimed_pegins(
1191 &mut self.client_ctx.module_db().begin_transaction_nc().await,
1192 tweak_idx,
1193 )
1194 .await;
1195
1196 if pegins.len() < num_deposits {
1197 debug!(target: LOG_CLIENT_MODULE_WALLET, has=pegins.len(), "Not enough deposits");
1198 self.recheck_pegin_address(tweak_idx).await?;
1199 runtime::sleep(backoff.next().unwrap_or_default()).await;
1200 receiver.changed().await?;
1201 continue;
1202 }
1203
1204 debug!(target: LOG_CLIENT_MODULE_WALLET, has=pegins.len(), "Enough deposits detected");
1205
1206 for (_outpoint, transaction_id, change) in pegins {
1207 debug!(target: LOG_CLIENT_MODULE_WALLET, out_points=?change, "Ensuring deposists claimed");
1208 let tx_subscriber = self.client_ctx.transaction_updates(operation_id).await;
1209
1210 if let Err(e) = tx_subscriber.await_tx_accepted(transaction_id).await {
1211 bail!("{}", e);
1212 }
1213
1214 debug!(target: LOG_CLIENT_MODULE_WALLET, out_points=?change, "Ensuring outputs claimed");
1215 self.client_ctx
1216 .await_primary_module_outputs(operation_id, change)
1217 .await
1218 .expect("Cannot fail if tx was accepted and federation is honest");
1219 }
1220
1221 return Ok(());
1222 }
1223 }
1224
1225 pub async fn withdraw<M: Serialize + MaybeSend + MaybeSync>(
1230 &self,
1231 address: &bitcoin::Address,
1232 amount: bitcoin::Amount,
1233 fee: PegOutFees,
1234 extra_meta: M,
1235 ) -> anyhow::Result<OperationId> {
1236 {
1237 let operation_id = OperationId(thread_rng().r#gen());
1238
1239 let withdraw_output =
1240 self.create_withdraw_output(operation_id, address.clone(), amount, fee)?;
1241 let tx_builder = TransactionBuilder::new()
1242 .with_outputs(self.client_ctx.make_client_outputs(withdraw_output));
1243
1244 let extra_meta =
1245 serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
1246 self.client_ctx
1247 .finalize_and_submit_transaction(
1248 operation_id,
1249 WalletCommonInit::KIND.as_str(),
1250 {
1251 let address = address.clone();
1252 move |change_range: OutPointRange| WalletOperationMeta {
1253 variant: WalletOperationMetaVariant::Withdraw {
1254 address: address.clone().into_unchecked(),
1255 amount,
1256 fee,
1257 change: change_range.into_iter().collect(),
1258 },
1259 extra_meta: extra_meta.clone(),
1260 }
1261 },
1262 tx_builder,
1263 )
1264 .await?;
1265
1266 Ok(operation_id)
1267 }
1268 }
1269
1270 #[deprecated(
1275 since = "0.4.0",
1276 note = "RBF withdrawals are rejected by the federation"
1277 )]
1278 pub async fn rbf_withdraw<M: Serialize + MaybeSync + MaybeSend>(
1279 &self,
1280 rbf: Rbf,
1281 extra_meta: M,
1282 ) -> anyhow::Result<OperationId> {
1283 let operation_id = OperationId(thread_rng().r#gen());
1284
1285 let withdraw_output = self.create_rbf_withdraw_output(operation_id, &rbf)?;
1286 let tx_builder = TransactionBuilder::new()
1287 .with_outputs(self.client_ctx.make_client_outputs(withdraw_output));
1288
1289 let extra_meta = serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
1290 self.client_ctx
1291 .finalize_and_submit_transaction(
1292 operation_id,
1293 WalletCommonInit::KIND.as_str(),
1294 move |change_range: OutPointRange| WalletOperationMeta {
1295 variant: WalletOperationMetaVariant::RbfWithdraw {
1296 rbf: rbf.clone(),
1297 change: change_range.into_iter().collect(),
1298 },
1299 extra_meta: extra_meta.clone(),
1300 },
1301 tx_builder,
1302 )
1303 .await?;
1304
1305 Ok(operation_id)
1306 }
1307
1308 pub async fn subscribe_withdraw_updates(
1309 &self,
1310 operation_id: OperationId,
1311 ) -> anyhow::Result<UpdateStreamOrOutcome<WithdrawState>> {
1312 let operation = self
1313 .client_ctx
1314 .get_operation(operation_id)
1315 .await
1316 .with_context(|| anyhow!("Operation not found: {}", operation_id.fmt_short()))?;
1317
1318 if operation.operation_module_kind() != WalletCommonInit::KIND.as_str() {
1319 bail!("Operation is not a wallet operation");
1320 }
1321
1322 let operation_meta = operation.meta::<WalletOperationMeta>();
1323
1324 let (WalletOperationMetaVariant::Withdraw { change, .. }
1325 | WalletOperationMetaVariant::RbfWithdraw { change, .. }) = operation_meta.variant
1326 else {
1327 bail!("Operation is not a withdraw operation");
1328 };
1329
1330 let mut operation_stream = self.notifier.subscribe(operation_id).await;
1331 let client_ctx = self.client_ctx.clone();
1332
1333 Ok(self
1334 .client_ctx
1335 .outcome_or_updates(operation, operation_id, move || {
1336 stream! {
1337 match next_withdraw_state(&mut operation_stream).await {
1338 Some(WithdrawStates::Created(_)) => {
1339 yield WithdrawState::Created;
1340 },
1341 Some(s) => {
1342 panic!("Unexpected state {s:?}")
1343 },
1344 None => return,
1345 }
1346
1347 let _ = client_ctx
1352 .await_primary_module_outputs(operation_id, change)
1353 .await;
1354
1355
1356 match next_withdraw_state(&mut operation_stream).await {
1357 Some(WithdrawStates::Aborted(inner)) => {
1358 yield WithdrawState::Failed(inner.error);
1359 },
1360 Some(WithdrawStates::Success(inner)) => {
1361 yield WithdrawState::Succeeded(inner.txid);
1362 },
1363 Some(s) => {
1364 panic!("Unexpected state {s:?}")
1365 },
1366 None => {},
1367 }
1368 }
1369 }))
1370 }
1371
1372 fn admin_auth(&self) -> anyhow::Result<ApiAuth> {
1373 self.admin_auth
1374 .clone()
1375 .ok_or_else(|| anyhow::format_err!("Admin auth not set"))
1376 }
1377
1378 pub async fn activate_consensus_version_voting(&self) -> anyhow::Result<()> {
1379 self.module_api
1380 .activate_consensus_version_voting(self.admin_auth()?)
1381 .await?;
1382
1383 Ok(())
1384 }
1385}
1386
1387async fn poll_supports_safe_deposit_version(db: Database, module_api: DynModuleApi) {
1390 loop {
1391 let mut dbtx = db.begin_transaction().await;
1392
1393 if dbtx.get_value(&SupportsSafeDepositKey).await.is_some() {
1394 break;
1395 }
1396
1397 if let Ok(module_consensus_version) = module_api.module_consensus_version().await {
1398 if SAFE_DEPOSIT_MODULE_CONSENSUS_VERSION <= module_consensus_version {
1399 dbtx.insert_new_entry(&SupportsSafeDepositKey, &()).await;
1400 dbtx.commit_tx().await;
1401 break;
1402 }
1403 }
1404
1405 drop(dbtx);
1406
1407 if is_running_in_test_env() {
1408 sleep(Duration::from_secs(10)).await;
1410 } else {
1411 sleep(Duration::from_secs(3600)).await;
1412 }
1413 }
1414}
1415
1416async fn get_next_peg_in_tweak_child_id(dbtx: &mut DatabaseTransaction<'_>) -> TweakIdx {
1418 let index = dbtx
1419 .get_value(&NextPegInTweakIndexKey)
1420 .await
1421 .unwrap_or_default();
1422 dbtx.insert_entry(&NextPegInTweakIndexKey, &(index.next()))
1423 .await;
1424 index
1425}
1426
1427#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
1428pub enum WalletClientStates {
1429 Deposit(DepositStateMachine),
1430 Withdraw(WithdrawStateMachine),
1431}
1432
1433impl IntoDynInstance for WalletClientStates {
1434 type DynType = DynState;
1435
1436 fn into_dyn(self, instance_id: ModuleInstanceId) -> Self::DynType {
1437 DynState::from_typed(instance_id, self)
1438 }
1439}
1440
1441impl State for WalletClientStates {
1442 type ModuleContext = WalletClientContext;
1443
1444 fn transitions(
1445 &self,
1446 context: &Self::ModuleContext,
1447 global_context: &DynGlobalClientContext,
1448 ) -> Vec<StateTransition<Self>> {
1449 match self {
1450 WalletClientStates::Deposit(sm) => {
1451 sm_enum_variant_translation!(
1452 sm.transitions(context, global_context),
1453 WalletClientStates::Deposit
1454 )
1455 }
1456 WalletClientStates::Withdraw(sm) => {
1457 sm_enum_variant_translation!(
1458 sm.transitions(context, global_context),
1459 WalletClientStates::Withdraw
1460 )
1461 }
1462 }
1463 }
1464
1465 fn operation_id(&self) -> OperationId {
1466 match self {
1467 WalletClientStates::Deposit(sm) => sm.operation_id(),
1468 WalletClientStates::Withdraw(sm) => sm.operation_id(),
1469 }
1470 }
1471}
1472
1473#[cfg(all(test, not(target_family = "wasm")))]
1474mod tests {
1475 use std::collections::BTreeSet;
1476 use std::sync::atomic::{AtomicBool, Ordering};
1477
1478 use super::*;
1479 use crate::backup::{
1480 RECOVER_NUM_IDX_ADD_TO_LAST_USED, RecoverScanOutcome, recover_scan_idxes_for_activity,
1481 };
1482
1483 #[allow(clippy::too_many_lines)] #[tokio::test(flavor = "multi_thread")]
1485 async fn sanity_test_recover_inner() {
1486 {
1487 let last_checked = AtomicBool::new(false);
1488 let last_checked = &last_checked;
1489 assert_eq!(
1490 recover_scan_idxes_for_activity(
1491 TweakIdx(0),
1492 &BTreeSet::new(),
1493 |cur_idx| async move {
1494 Ok(match cur_idx {
1495 TweakIdx(9) => {
1496 last_checked.store(true, Ordering::SeqCst);
1497 vec![]
1498 }
1499 TweakIdx(10) => panic!("Shouldn't happen"),
1500 TweakIdx(11) => {
1501 vec![0usize] }
1503 _ => vec![],
1504 })
1505 }
1506 )
1507 .await
1508 .unwrap(),
1509 RecoverScanOutcome {
1510 last_used_idx: None,
1511 new_start_idx: TweakIdx(RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1512 tweak_idxes_with_pegins: BTreeSet::from([])
1513 }
1514 );
1515 assert!(last_checked.load(Ordering::SeqCst));
1516 }
1517
1518 {
1519 let last_checked = AtomicBool::new(false);
1520 let last_checked = &last_checked;
1521 assert_eq!(
1522 recover_scan_idxes_for_activity(
1523 TweakIdx(0),
1524 &BTreeSet::from([TweakIdx(1), TweakIdx(2)]),
1525 |cur_idx| async move {
1526 Ok(match cur_idx {
1527 TweakIdx(1) => panic!("Shouldn't happen: already used (1)"),
1528 TweakIdx(2) => panic!("Shouldn't happen: already used (2)"),
1529 TweakIdx(11) => {
1530 last_checked.store(true, Ordering::SeqCst);
1531 vec![]
1532 }
1533 TweakIdx(12) => panic!("Shouldn't happen"),
1534 TweakIdx(13) => {
1535 vec![0usize] }
1537 _ => vec![],
1538 })
1539 }
1540 )
1541 .await
1542 .unwrap(),
1543 RecoverScanOutcome {
1544 last_used_idx: Some(TweakIdx(2)),
1545 new_start_idx: TweakIdx(2 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1546 tweak_idxes_with_pegins: BTreeSet::from([])
1547 }
1548 );
1549 assert!(last_checked.load(Ordering::SeqCst));
1550 }
1551
1552 {
1553 let last_checked = AtomicBool::new(false);
1554 let last_checked = &last_checked;
1555 assert_eq!(
1556 recover_scan_idxes_for_activity(
1557 TweakIdx(10),
1558 &BTreeSet::new(),
1559 |cur_idx| async move {
1560 Ok(match cur_idx {
1561 TweakIdx(10) => vec![()],
1562 TweakIdx(19) => {
1563 last_checked.store(true, Ordering::SeqCst);
1564 vec![]
1565 }
1566 TweakIdx(20) => panic!("Shouldn't happen"),
1567 _ => vec![],
1568 })
1569 }
1570 )
1571 .await
1572 .unwrap(),
1573 RecoverScanOutcome {
1574 last_used_idx: Some(TweakIdx(10)),
1575 new_start_idx: TweakIdx(10 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1576 tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(10)])
1577 }
1578 );
1579 assert!(last_checked.load(Ordering::SeqCst));
1580 }
1581
1582 assert_eq!(
1583 recover_scan_idxes_for_activity(TweakIdx(0), &BTreeSet::new(), |cur_idx| async move {
1584 Ok(match cur_idx {
1585 TweakIdx(6 | 15) => vec![()],
1586 _ => vec![],
1587 })
1588 })
1589 .await
1590 .unwrap(),
1591 RecoverScanOutcome {
1592 last_used_idx: Some(TweakIdx(15)),
1593 new_start_idx: TweakIdx(15 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1594 tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(6), TweakIdx(15)])
1595 }
1596 );
1597 assert_eq!(
1598 recover_scan_idxes_for_activity(TweakIdx(10), &BTreeSet::new(), |cur_idx| async move {
1599 Ok(match cur_idx {
1600 TweakIdx(8) => {
1601 vec![()] }
1603 TweakIdx(9) => {
1604 panic!("Shouldn't happen")
1605 }
1606 _ => vec![],
1607 })
1608 })
1609 .await
1610 .unwrap(),
1611 RecoverScanOutcome {
1612 last_used_idx: None,
1613 new_start_idx: TweakIdx(9 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1614 tweak_idxes_with_pegins: BTreeSet::from([])
1615 }
1616 );
1617 assert_eq!(
1618 recover_scan_idxes_for_activity(TweakIdx(10), &BTreeSet::new(), |cur_idx| async move {
1619 Ok(match cur_idx {
1620 TweakIdx(9) => panic!("Shouldn't happen"),
1621 TweakIdx(15) => vec![()],
1622 _ => vec![],
1623 })
1624 })
1625 .await
1626 .unwrap(),
1627 RecoverScanOutcome {
1628 last_used_idx: Some(TweakIdx(15)),
1629 new_start_idx: TweakIdx(15 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1630 tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(15)])
1631 }
1632 );
1633 }
1634}