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, BitcoinHash, OutPoint, TransactionId, apply, async_trait_maybe_send,
62 push_db_pair_items, runtime, 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 "subscribe_deposit" => {
531 let req: SubscribeDepositRequest = serde_json::from_value(request)?;
532 for await state in self.subscribe_deposit(req.operation_id).await?.into_stream() {
533 yield serde_json::to_value(state)?;
534 }
535 }
536 _ => {
537 Err(anyhow::format_err!("Unknown method: {}", method))?;
538 }
539 }
540 })
541 }
542
543 #[cfg(feature = "cli")]
544 async fn handle_cli_command(
545 &self,
546 args: &[std::ffi::OsString],
547 ) -> anyhow::Result<serde_json::Value> {
548 cli::handle_cli_command(self, args).await
549 }
550}
551
552#[derive(Deserialize)]
553struct WalletSummaryRequest {}
554
555#[derive(Debug, Clone)]
556pub struct WalletClientContext {
557 rpc: DynBitcoindRpc,
558 wallet_descriptor: PegInDescriptor,
559 wallet_decoder: Decoder,
560 secp: Secp256k1<All>,
561 pub client_ctx: ClientContext<WalletClientModule>,
562}
563
564#[derive(Debug, Clone, Serialize, Deserialize)]
565pub struct PegInRequest {
566 pub extra_meta: serde_json::Value,
567}
568#[derive(Deserialize)]
569struct SubscribeDepositRequest {
570 operation_id: OperationId,
571}
572
573#[derive(Debug, Clone, Serialize, Deserialize)]
574pub struct PegInResponse {
575 pub deposit_address: Address<NetworkUnchecked>,
576 pub operation_id: OperationId,
577}
578
579#[derive(Debug, Clone, Serialize, Deserialize)]
580pub struct PegOutRequest {
581 pub amount_sat: u64,
582 pub destination_address: Address<NetworkUnchecked>,
583 pub extra_meta: serde_json::Value,
584}
585
586#[derive(Debug, Clone, Serialize, Deserialize)]
587pub struct PegOutResponse {
588 pub operation_id: OperationId,
589}
590
591impl Context for WalletClientContext {
592 const KIND: Option<ModuleKind> = Some(KIND);
593}
594
595impl WalletClientModule {
596 fn cfg(&self) -> &WalletClientConfig {
597 &self.data.cfg
598 }
599
600 fn get_rpc_config(cfg: &WalletClientConfig) -> BitcoinRpcConfig {
601 match BitcoinRpcConfig::get_defaults_from_env_vars() {
602 Ok(rpc_config) => {
603 if rpc_config.kind == "bitcoind" {
606 cfg.default_bitcoin_rpc.clone()
607 } else {
608 rpc_config
609 }
610 }
611 _ => cfg.default_bitcoin_rpc.clone(),
612 }
613 }
614
615 pub fn get_network(&self) -> Network {
616 self.cfg().network.0
617 }
618
619 pub fn get_fee_consensus(&self) -> FeeConsensus {
620 self.cfg().fee_consensus
621 }
622
623 async fn allocate_deposit_address_inner(
624 &self,
625 dbtx: &mut DatabaseTransaction<'_>,
626 ) -> (OperationId, Address, TweakIdx) {
627 dbtx.ensure_isolated().expect("Must be isolated db");
628
629 let tweak_idx = get_next_peg_in_tweak_child_id(dbtx).await;
630 let (_secret_tweak_key, _, address, operation_id) =
631 self.data.derive_deposit_address(tweak_idx);
632
633 let now = fedimint_core::time::now();
634
635 dbtx.insert_new_entry(
636 &PegInTweakIndexKey(tweak_idx),
637 &PegInTweakIndexData {
638 creation_time: now,
639 next_check_time: Some(now),
640 last_check_time: None,
641 operation_id,
642 claimed: vec![],
643 },
644 )
645 .await;
646
647 (operation_id, address, tweak_idx)
648 }
649
650 pub async fn get_withdraw_fees(
657 &self,
658 address: &bitcoin::Address,
659 amount: bitcoin::Amount,
660 ) -> anyhow::Result<PegOutFees> {
661 self.module_api
662 .fetch_peg_out_fees(address, amount)
663 .await?
664 .context("Federation didn't return peg-out fees")
665 }
666
667 pub async fn get_wallet_summary(&self) -> anyhow::Result<WalletSummary> {
669 Ok(self.module_api.fetch_wallet_summary().await?)
670 }
671
672 pub fn create_withdraw_output(
673 &self,
674 operation_id: OperationId,
675 address: bitcoin::Address,
676 amount: bitcoin::Amount,
677 fees: PegOutFees,
678 ) -> anyhow::Result<ClientOutputBundle<WalletOutput, WalletClientStates>> {
679 let output = WalletOutput::new_v0_peg_out(address, amount, fees);
680
681 let amount = output.maybe_v0_ref().expect("v0 output").amount().into();
682
683 let sm_gen = move |out_point_range: OutPointRange| {
684 assert_eq!(out_point_range.count(), 1);
685 let out_idx = out_point_range.start_idx();
686 vec![WalletClientStates::Withdraw(WithdrawStateMachine {
687 operation_id,
688 state: WithdrawStates::Created(CreatedWithdrawState {
689 fm_outpoint: OutPoint {
690 txid: out_point_range.txid(),
691 out_idx,
692 },
693 }),
694 })]
695 };
696
697 Ok(ClientOutputBundle::new(
698 vec![ClientOutput::<WalletOutput> { output, amount }],
699 vec![ClientOutputSM::<WalletClientStates> {
700 state_machines: Arc::new(sm_gen),
701 }],
702 ))
703 }
704
705 pub async fn peg_in(&self, req: PegInRequest) -> anyhow::Result<PegInResponse> {
706 let (operation_id, address, _) = self.safe_allocate_deposit_address(req.extra_meta).await?;
707
708 Ok(PegInResponse {
709 deposit_address: Address::from_script(&address.script_pubkey(), self.get_network())?
710 .as_unchecked()
711 .clone(),
712 operation_id,
713 })
714 }
715
716 pub async fn peg_out(&self, req: PegOutRequest) -> anyhow::Result<PegOutResponse> {
717 let amount = bitcoin::Amount::from_sat(req.amount_sat);
718 let destination = req
719 .destination_address
720 .require_network(self.get_network())?;
721
722 let fees = self.get_withdraw_fees(&destination, amount).await?;
723 let operation_id = self
724 .withdraw(&destination, amount, fees, req.extra_meta)
725 .await
726 .context("Failed to initiate withdraw")?;
727
728 Ok(PegOutResponse { operation_id })
729 }
730
731 pub fn create_rbf_withdraw_output(
732 &self,
733 operation_id: OperationId,
734 rbf: &Rbf,
735 ) -> anyhow::Result<ClientOutputBundle<WalletOutput, WalletClientStates>> {
736 let output = WalletOutput::new_v0_rbf(rbf.fees, rbf.txid);
737
738 let amount = output.maybe_v0_ref().expect("v0 output").amount().into();
739
740 let sm_gen = move |out_point_range: OutPointRange| {
741 assert_eq!(out_point_range.count(), 1);
742 let out_idx = out_point_range.start_idx();
743 vec![WalletClientStates::Withdraw(WithdrawStateMachine {
744 operation_id,
745 state: WithdrawStates::Created(CreatedWithdrawState {
746 fm_outpoint: OutPoint {
747 txid: out_point_range.txid(),
748 out_idx,
749 },
750 }),
751 })]
752 };
753
754 Ok(ClientOutputBundle::new(
755 vec![ClientOutput::<WalletOutput> { output, amount }],
756 vec![ClientOutputSM::<WalletClientStates> {
757 state_machines: Arc::new(sm_gen),
758 }],
759 ))
760 }
761
762 pub async fn btc_tx_has_no_size_limit(&self) -> FederationResult<bool> {
763 Ok(self.module_api.module_consensus_version().await? >= ModuleConsensusVersion::new(2, 2))
764 }
765
766 pub async fn supports_safe_deposit(&self) -> bool {
775 let mut dbtx = self.db.begin_transaction().await;
776
777 let already_verified_supports_safe_deposit =
778 dbtx.get_value(&SupportsSafeDepositKey).await.is_some();
779
780 already_verified_supports_safe_deposit || {
781 match self.module_api.module_consensus_version().await {
782 Ok(module_consensus_version) => {
783 let supported_version =
784 SAFE_DEPOSIT_MODULE_CONSENSUS_VERSION <= module_consensus_version;
785
786 if supported_version {
787 dbtx.insert_new_entry(&SupportsSafeDepositKey, &()).await;
788 dbtx.commit_tx().await;
789 }
790
791 supported_version
792 }
793 Err(_) => false,
794 }
795 }
796 }
797
798 pub async fn safe_allocate_deposit_address<M>(
806 &self,
807 extra_meta: M,
808 ) -> anyhow::Result<(OperationId, Address, TweakIdx)>
809 where
810 M: Serialize + MaybeSend + MaybeSync,
811 {
812 ensure!(
813 self.supports_safe_deposit().await,
814 "Wallet module consensus version doesn't support safe deposits",
815 );
816
817 self.allocate_deposit_address_expert_only(extra_meta).await
818 }
819
820 pub async fn allocate_deposit_address_expert_only<M>(
838 &self,
839 extra_meta: M,
840 ) -> anyhow::Result<(OperationId, Address, TweakIdx)>
841 where
842 M: Serialize + MaybeSend + MaybeSync,
843 {
844 let extra_meta_value =
845 serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
846 let (operation_id, address, tweak_idx) = self
847 .db
848 .autocommit(
849 move |dbtx, _| {
850 let extra_meta_value_inner = extra_meta_value.clone();
851 Box::pin(async move {
852 let (operation_id, address, tweak_idx) = self
853 .allocate_deposit_address_inner(dbtx)
854 .await;
855
856 self.client_ctx.manual_operation_start_dbtx(
857 dbtx,
858 operation_id,
859 WalletCommonInit::KIND.as_str(),
860 WalletOperationMeta {
861 variant: WalletOperationMetaVariant::Deposit {
862 address: address.clone().into_unchecked(),
863 tweak_idx: Some(tweak_idx),
864 expires_at: None,
865 },
866 extra_meta: extra_meta_value_inner,
867 },
868 vec![]
869 ).await?;
870
871 debug!(target: LOG_CLIENT_MODULE_WALLET, %tweak_idx, %address, "Derived a new deposit address");
872
873
874 let sender = self.pegin_monitor_wakeup_sender.clone();
875 dbtx.on_commit(move || {
876 sender.send_replace(());
877 });
878
879 Ok((operation_id, address, tweak_idx))
880 })
881 },
882 Some(100),
883 )
884 .await
885 .map_err(|e| match e {
886 AutocommitError::CommitFailed {
887 last_error,
888 attempts,
889 } => last_error.context(format!("Failed to commit after {attempts} attempts")),
890 AutocommitError::ClosureError { error, .. } => error,
891 })?;
892
893 Ok((operation_id, address, tweak_idx))
894 }
895
896 pub async fn subscribe_deposit(
902 &self,
903 operation_id: OperationId,
904 ) -> anyhow::Result<UpdateStreamOrOutcome<DepositStateV2>> {
905 let operation = self
906 .client_ctx
907 .get_operation(operation_id)
908 .await
909 .with_context(|| anyhow!("Operation not found: {}", operation_id.fmt_short()))?;
910
911 if operation.operation_module_kind() != WalletCommonInit::KIND.as_str() {
912 bail!("Operation is not a wallet operation");
913 }
914
915 let operation_meta = operation.meta::<WalletOperationMeta>();
916
917 let WalletOperationMetaVariant::Deposit {
918 address, tweak_idx, ..
919 } = operation_meta.variant
920 else {
921 bail!("Operation is not a deposit operation");
922 };
923
924 let address = address.require_network(self.cfg().network.0)?;
925
926 let Some(tweak_idx) = tweak_idx else {
928 let outcome_v1 = operation
932 .outcome::<DepositStateV1>()
933 .context("Old pending deposit, can't subscribe to updates")?;
934
935 let outcome_v2 = match outcome_v1 {
936 DepositStateV1::Claimed(tx_info) => DepositStateV2::Claimed {
937 btc_deposited: tx_info.btc_transaction.output[tx_info.out_idx as usize].value,
938 btc_out_point: bitcoin::OutPoint {
939 txid: tx_info.btc_transaction.compute_txid(),
940 vout: tx_info.out_idx,
941 },
942 },
943 DepositStateV1::Failed(error) => DepositStateV2::Failed(error),
944 _ => bail!("Non-final outcome in operation log"),
945 };
946
947 return Ok(UpdateStreamOrOutcome::Outcome(outcome_v2));
948 };
949
950 Ok(self.client_ctx.outcome_or_updates(operation, operation_id, {
951 let stream_rpc = self.rpc.clone();
952 let stream_client_ctx = self.client_ctx.clone();
953 let stream_script_pub_key = address.script_pubkey();
954 move || {
955
956 stream! {
957 yield DepositStateV2::WaitingForTransaction;
958
959 let (btc_out_point, btc_deposited) = retry(
960 "fetch history",
961 background_backoff(),
962 || async {
963 let history = stream_rpc.get_script_history(&stream_script_pub_key).await?;
964 history.first().and_then(|tx| {
965 let (out_idx, amount) = tx.output
966 .iter()
967 .enumerate()
968 .find_map(|(idx, output)| (output.script_pubkey == stream_script_pub_key).then_some((idx, output.value)))?;
969 let txid = tx.compute_txid();
970
971 Some((
972 bitcoin::OutPoint {
973 txid,
974 vout: out_idx as u32,
975 },
976 amount
977 ))
978 }).context("No deposit transaction found")
979 }
980 ).await.expect("Will never give up");
981
982 yield DepositStateV2::WaitingForConfirmation {
983 btc_deposited,
984 btc_out_point
985 };
986
987 let claim_data = stream_client_ctx.module_db().wait_key_exists(&ClaimedPegInKey {
988 peg_in_index: tweak_idx,
989 btc_out_point,
990 }).await;
991
992 yield DepositStateV2::Confirmed {
993 btc_deposited,
994 btc_out_point
995 };
996
997 match stream_client_ctx.await_primary_module_outputs(operation_id, claim_data.change).await {
998 Ok(()) => yield DepositStateV2::Claimed {
999 btc_deposited,
1000 btc_out_point
1001 },
1002 Err(e) => yield DepositStateV2::Failed(e.to_string())
1003 }
1004 }
1005 }}))
1006 }
1007
1008 pub async fn list_peg_in_tweak_idxes(&self) -> BTreeMap<TweakIdx, PegInTweakIndexData> {
1009 self.client_ctx
1010 .module_db()
1011 .clone()
1012 .begin_transaction_nc()
1013 .await
1014 .find_by_prefix(&PegInTweakIndexPrefix)
1015 .await
1016 .map(|(key, data)| (key.0, data))
1017 .collect()
1018 .await
1019 }
1020
1021 pub async fn find_tweak_idx_by_address(
1022 &self,
1023 address: bitcoin::Address<NetworkUnchecked>,
1024 ) -> anyhow::Result<TweakIdx> {
1025 let data = self.data.clone();
1026 let Some((tweak_idx, _)) = self
1027 .db
1028 .begin_transaction_nc()
1029 .await
1030 .find_by_prefix(&PegInTweakIndexPrefix)
1031 .await
1032 .filter(|(k, _)| {
1033 let (_, derived_address, _tweak_key, _) = data.derive_peg_in_script(k.0);
1034 future::ready(derived_address.into_unchecked() == address)
1035 })
1036 .next()
1037 .await
1038 else {
1039 bail!("Address not found in the list of derived keys");
1040 };
1041
1042 Ok(tweak_idx.0)
1043 }
1044 pub async fn find_tweak_idx_by_operation_id(
1045 &self,
1046 operation_id: OperationId,
1047 ) -> anyhow::Result<TweakIdx> {
1048 Ok(self
1049 .client_ctx
1050 .module_db()
1051 .clone()
1052 .begin_transaction_nc()
1053 .await
1054 .find_by_prefix(&PegInTweakIndexPrefix)
1055 .await
1056 .filter(|(_k, v)| future::ready(v.operation_id == operation_id))
1057 .next()
1058 .await
1059 .ok_or_else(|| anyhow::format_err!("OperationId not found"))?
1060 .0
1061 .0)
1062 }
1063
1064 pub async fn get_pegin_tweak_idx(
1065 &self,
1066 tweak_idx: TweakIdx,
1067 ) -> anyhow::Result<PegInTweakIndexData> {
1068 self.client_ctx
1069 .module_db()
1070 .clone()
1071 .begin_transaction_nc()
1072 .await
1073 .get_value(&PegInTweakIndexKey(tweak_idx))
1074 .await
1075 .ok_or_else(|| anyhow::format_err!("TweakIdx not found"))
1076 }
1077
1078 pub async fn get_claimed_pegins(
1079 &self,
1080 dbtx: &mut DatabaseTransaction<'_>,
1081 tweak_idx: TweakIdx,
1082 ) -> Vec<(
1083 bitcoin::OutPoint,
1084 TransactionId,
1085 Vec<fedimint_core::OutPoint>,
1086 )> {
1087 let outpoints = dbtx
1088 .get_value(&PegInTweakIndexKey(tweak_idx))
1089 .await
1090 .map(|v| v.claimed)
1091 .unwrap_or_default();
1092
1093 let mut res = vec![];
1094
1095 for outpoint in outpoints {
1096 let claimed_peg_in_data = dbtx
1097 .get_value(&ClaimedPegInKey {
1098 peg_in_index: tweak_idx,
1099 btc_out_point: outpoint,
1100 })
1101 .await
1102 .expect("Must have a corresponding claim record");
1103 res.push((
1104 outpoint,
1105 claimed_peg_in_data.claim_txid,
1106 claimed_peg_in_data.change,
1107 ));
1108 }
1109
1110 res
1111 }
1112
1113 pub async fn recheck_pegin_address_by_op_id(
1115 &self,
1116 operation_id: OperationId,
1117 ) -> anyhow::Result<()> {
1118 let tweak_idx = self.find_tweak_idx_by_operation_id(operation_id).await?;
1119
1120 self.recheck_pegin_address(tweak_idx).await
1121 }
1122
1123 pub async fn recheck_pegin_address_by_address(
1125 &self,
1126 address: bitcoin::Address<NetworkUnchecked>,
1127 ) -> anyhow::Result<()> {
1128 self.recheck_pegin_address(self.find_tweak_idx_by_address(address).await?)
1129 .await
1130 }
1131
1132 pub async fn recheck_pegin_address(&self, tweak_idx: TweakIdx) -> anyhow::Result<()> {
1134 self.db
1135 .autocommit(
1136 |dbtx, _| {
1137 Box::pin(async {
1138 let db_key = PegInTweakIndexKey(tweak_idx);
1139 let db_val = dbtx
1140 .get_value(&db_key)
1141 .await
1142 .ok_or_else(|| anyhow::format_err!("DBKey not found"))?;
1143
1144 dbtx.insert_entry(
1145 &db_key,
1146 &PegInTweakIndexData {
1147 next_check_time: Some(fedimint_core::time::now()),
1148 ..db_val
1149 },
1150 )
1151 .await;
1152
1153 let sender = self.pegin_monitor_wakeup_sender.clone();
1154 dbtx.on_commit(move || {
1155 sender.send_replace(());
1156 });
1157
1158 Ok::<_, anyhow::Error>(())
1159 })
1160 },
1161 Some(100),
1162 )
1163 .await?;
1164
1165 Ok(())
1166 }
1167
1168 pub async fn await_num_deposits_by_operation_id(
1170 &self,
1171 operation_id: OperationId,
1172 num_deposits: usize,
1173 ) -> anyhow::Result<()> {
1174 let tweak_idx = self.find_tweak_idx_by_operation_id(operation_id).await?;
1175 self.await_num_deposits(tweak_idx, num_deposits).await
1176 }
1177
1178 pub async fn await_num_deposits_by_address(
1179 &self,
1180 address: bitcoin::Address<NetworkUnchecked>,
1181 num_deposits: usize,
1182 ) -> anyhow::Result<()> {
1183 self.await_num_deposits(self.find_tweak_idx_by_address(address).await?, num_deposits)
1184 .await
1185 }
1186
1187 #[instrument(target = LOG_CLIENT_MODULE_WALLET, skip_all, fields(tweak_idx=?tweak_idx, num_deposists=num_deposits))]
1188 pub async fn await_num_deposits(
1189 &self,
1190 tweak_idx: TweakIdx,
1191 num_deposits: usize,
1192 ) -> anyhow::Result<()> {
1193 let operation_id = self.get_pegin_tweak_idx(tweak_idx).await?.operation_id;
1194
1195 let mut receiver = self.pegin_claimed_receiver.clone();
1196 let mut backoff = backoff_util::aggressive_backoff();
1197
1198 loop {
1199 let pegins = self
1200 .get_claimed_pegins(
1201 &mut self.client_ctx.module_db().begin_transaction_nc().await,
1202 tweak_idx,
1203 )
1204 .await;
1205
1206 if pegins.len() < num_deposits {
1207 debug!(target: LOG_CLIENT_MODULE_WALLET, has=pegins.len(), "Not enough deposits");
1208 self.recheck_pegin_address(tweak_idx).await?;
1209 runtime::sleep(backoff.next().unwrap_or_default()).await;
1210 receiver.changed().await?;
1211 continue;
1212 }
1213
1214 debug!(target: LOG_CLIENT_MODULE_WALLET, has=pegins.len(), "Enough deposits detected");
1215
1216 for (_outpoint, transaction_id, change) in pegins {
1217 if transaction_id == TransactionId::from_byte_array([0; 32]) && change.is_empty() {
1218 debug!(target: LOG_CLIENT_MODULE_WALLET, "Deposited amount was too low, skipping");
1219 continue;
1220 }
1221
1222 debug!(target: LOG_CLIENT_MODULE_WALLET, out_points=?change, "Ensuring deposists claimed");
1223 let tx_subscriber = self.client_ctx.transaction_updates(operation_id).await;
1224
1225 if let Err(e) = tx_subscriber.await_tx_accepted(transaction_id).await {
1226 bail!("{}", e);
1227 }
1228
1229 debug!(target: LOG_CLIENT_MODULE_WALLET, out_points=?change, "Ensuring outputs claimed");
1230 self.client_ctx
1231 .await_primary_module_outputs(operation_id, change)
1232 .await
1233 .expect("Cannot fail if tx was accepted and federation is honest");
1234 }
1235
1236 return Ok(());
1237 }
1238 }
1239
1240 pub async fn withdraw<M: Serialize + MaybeSend + MaybeSync>(
1245 &self,
1246 address: &bitcoin::Address,
1247 amount: bitcoin::Amount,
1248 fee: PegOutFees,
1249 extra_meta: M,
1250 ) -> anyhow::Result<OperationId> {
1251 {
1252 let operation_id = OperationId(thread_rng().r#gen());
1253
1254 let withdraw_output =
1255 self.create_withdraw_output(operation_id, address.clone(), amount, fee)?;
1256 let tx_builder = TransactionBuilder::new()
1257 .with_outputs(self.client_ctx.make_client_outputs(withdraw_output));
1258
1259 let extra_meta =
1260 serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
1261 self.client_ctx
1262 .finalize_and_submit_transaction(
1263 operation_id,
1264 WalletCommonInit::KIND.as_str(),
1265 {
1266 let address = address.clone();
1267 move |change_range: OutPointRange| WalletOperationMeta {
1268 variant: WalletOperationMetaVariant::Withdraw {
1269 address: address.clone().into_unchecked(),
1270 amount,
1271 fee,
1272 change: change_range.into_iter().collect(),
1273 },
1274 extra_meta: extra_meta.clone(),
1275 }
1276 },
1277 tx_builder,
1278 )
1279 .await?;
1280
1281 Ok(operation_id)
1282 }
1283 }
1284
1285 #[deprecated(
1290 since = "0.4.0",
1291 note = "RBF withdrawals are rejected by the federation"
1292 )]
1293 pub async fn rbf_withdraw<M: Serialize + MaybeSync + MaybeSend>(
1294 &self,
1295 rbf: Rbf,
1296 extra_meta: M,
1297 ) -> anyhow::Result<OperationId> {
1298 let operation_id = OperationId(thread_rng().r#gen());
1299
1300 let withdraw_output = self.create_rbf_withdraw_output(operation_id, &rbf)?;
1301 let tx_builder = TransactionBuilder::new()
1302 .with_outputs(self.client_ctx.make_client_outputs(withdraw_output));
1303
1304 let extra_meta = serde_json::to_value(extra_meta).expect("Failed to serialize extra meta");
1305 self.client_ctx
1306 .finalize_and_submit_transaction(
1307 operation_id,
1308 WalletCommonInit::KIND.as_str(),
1309 move |change_range: OutPointRange| WalletOperationMeta {
1310 variant: WalletOperationMetaVariant::RbfWithdraw {
1311 rbf: rbf.clone(),
1312 change: change_range.into_iter().collect(),
1313 },
1314 extra_meta: extra_meta.clone(),
1315 },
1316 tx_builder,
1317 )
1318 .await?;
1319
1320 Ok(operation_id)
1321 }
1322
1323 pub async fn subscribe_withdraw_updates(
1324 &self,
1325 operation_id: OperationId,
1326 ) -> anyhow::Result<UpdateStreamOrOutcome<WithdrawState>> {
1327 let operation = self
1328 .client_ctx
1329 .get_operation(operation_id)
1330 .await
1331 .with_context(|| anyhow!("Operation not found: {}", operation_id.fmt_short()))?;
1332
1333 if operation.operation_module_kind() != WalletCommonInit::KIND.as_str() {
1334 bail!("Operation is not a wallet operation");
1335 }
1336
1337 let operation_meta = operation.meta::<WalletOperationMeta>();
1338
1339 let (WalletOperationMetaVariant::Withdraw { change, .. }
1340 | WalletOperationMetaVariant::RbfWithdraw { change, .. }) = operation_meta.variant
1341 else {
1342 bail!("Operation is not a withdraw operation");
1343 };
1344
1345 let mut operation_stream = self.notifier.subscribe(operation_id).await;
1346 let client_ctx = self.client_ctx.clone();
1347
1348 Ok(self
1349 .client_ctx
1350 .outcome_or_updates(operation, operation_id, move || {
1351 stream! {
1352 match next_withdraw_state(&mut operation_stream).await {
1353 Some(WithdrawStates::Created(_)) => {
1354 yield WithdrawState::Created;
1355 },
1356 Some(s) => {
1357 panic!("Unexpected state {s:?}")
1358 },
1359 None => return,
1360 }
1361
1362 let _ = client_ctx
1367 .await_primary_module_outputs(operation_id, change)
1368 .await;
1369
1370
1371 match next_withdraw_state(&mut operation_stream).await {
1372 Some(WithdrawStates::Aborted(inner)) => {
1373 yield WithdrawState::Failed(inner.error);
1374 },
1375 Some(WithdrawStates::Success(inner)) => {
1376 yield WithdrawState::Succeeded(inner.txid);
1377 },
1378 Some(s) => {
1379 panic!("Unexpected state {s:?}")
1380 },
1381 None => {},
1382 }
1383 }
1384 }))
1385 }
1386
1387 fn admin_auth(&self) -> anyhow::Result<ApiAuth> {
1388 self.admin_auth
1389 .clone()
1390 .ok_or_else(|| anyhow::format_err!("Admin auth not set"))
1391 }
1392
1393 pub async fn activate_consensus_version_voting(&self) -> anyhow::Result<()> {
1394 self.module_api
1395 .activate_consensus_version_voting(self.admin_auth()?)
1396 .await?;
1397
1398 Ok(())
1399 }
1400}
1401
1402async fn poll_supports_safe_deposit_version(db: Database, module_api: DynModuleApi) {
1405 loop {
1406 let mut dbtx = db.begin_transaction().await;
1407
1408 if dbtx.get_value(&SupportsSafeDepositKey).await.is_some() {
1409 break;
1410 }
1411
1412 if let Ok(module_consensus_version) = module_api.module_consensus_version().await {
1413 if SAFE_DEPOSIT_MODULE_CONSENSUS_VERSION <= module_consensus_version {
1414 dbtx.insert_new_entry(&SupportsSafeDepositKey, &()).await;
1415 dbtx.commit_tx().await;
1416 break;
1417 }
1418 }
1419
1420 drop(dbtx);
1421
1422 if is_running_in_test_env() {
1423 sleep(Duration::from_secs(10)).await;
1425 } else {
1426 sleep(Duration::from_secs(3600)).await;
1427 }
1428 }
1429}
1430
1431async fn get_next_peg_in_tweak_child_id(dbtx: &mut DatabaseTransaction<'_>) -> TweakIdx {
1433 let index = dbtx
1434 .get_value(&NextPegInTweakIndexKey)
1435 .await
1436 .unwrap_or_default();
1437 dbtx.insert_entry(&NextPegInTweakIndexKey, &(index.next()))
1438 .await;
1439 index
1440}
1441
1442#[derive(Debug, Clone, Eq, PartialEq, Hash, Decodable, Encodable)]
1443pub enum WalletClientStates {
1444 Deposit(DepositStateMachine),
1445 Withdraw(WithdrawStateMachine),
1446}
1447
1448impl IntoDynInstance for WalletClientStates {
1449 type DynType = DynState;
1450
1451 fn into_dyn(self, instance_id: ModuleInstanceId) -> Self::DynType {
1452 DynState::from_typed(instance_id, self)
1453 }
1454}
1455
1456impl State for WalletClientStates {
1457 type ModuleContext = WalletClientContext;
1458
1459 fn transitions(
1460 &self,
1461 context: &Self::ModuleContext,
1462 global_context: &DynGlobalClientContext,
1463 ) -> Vec<StateTransition<Self>> {
1464 match self {
1465 WalletClientStates::Deposit(sm) => {
1466 sm_enum_variant_translation!(
1467 sm.transitions(context, global_context),
1468 WalletClientStates::Deposit
1469 )
1470 }
1471 WalletClientStates::Withdraw(sm) => {
1472 sm_enum_variant_translation!(
1473 sm.transitions(context, global_context),
1474 WalletClientStates::Withdraw
1475 )
1476 }
1477 }
1478 }
1479
1480 fn operation_id(&self) -> OperationId {
1481 match self {
1482 WalletClientStates::Deposit(sm) => sm.operation_id(),
1483 WalletClientStates::Withdraw(sm) => sm.operation_id(),
1484 }
1485 }
1486}
1487
1488#[cfg(all(test, not(target_family = "wasm")))]
1489mod tests {
1490 use std::collections::BTreeSet;
1491 use std::sync::atomic::{AtomicBool, Ordering};
1492
1493 use super::*;
1494 use crate::backup::{
1495 RECOVER_NUM_IDX_ADD_TO_LAST_USED, RecoverScanOutcome, recover_scan_idxes_for_activity,
1496 };
1497
1498 #[allow(clippy::too_many_lines)] #[tokio::test(flavor = "multi_thread")]
1500 async fn sanity_test_recover_inner() {
1501 {
1502 let last_checked = AtomicBool::new(false);
1503 let last_checked = &last_checked;
1504 assert_eq!(
1505 recover_scan_idxes_for_activity(
1506 TweakIdx(0),
1507 &BTreeSet::new(),
1508 |cur_idx| async move {
1509 Ok(match cur_idx {
1510 TweakIdx(9) => {
1511 last_checked.store(true, Ordering::SeqCst);
1512 vec![]
1513 }
1514 TweakIdx(10) => panic!("Shouldn't happen"),
1515 TweakIdx(11) => {
1516 vec![0usize] }
1518 _ => vec![],
1519 })
1520 }
1521 )
1522 .await
1523 .unwrap(),
1524 RecoverScanOutcome {
1525 last_used_idx: None,
1526 new_start_idx: TweakIdx(RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1527 tweak_idxes_with_pegins: BTreeSet::from([])
1528 }
1529 );
1530 assert!(last_checked.load(Ordering::SeqCst));
1531 }
1532
1533 {
1534 let last_checked = AtomicBool::new(false);
1535 let last_checked = &last_checked;
1536 assert_eq!(
1537 recover_scan_idxes_for_activity(
1538 TweakIdx(0),
1539 &BTreeSet::from([TweakIdx(1), TweakIdx(2)]),
1540 |cur_idx| async move {
1541 Ok(match cur_idx {
1542 TweakIdx(1) => panic!("Shouldn't happen: already used (1)"),
1543 TweakIdx(2) => panic!("Shouldn't happen: already used (2)"),
1544 TweakIdx(11) => {
1545 last_checked.store(true, Ordering::SeqCst);
1546 vec![]
1547 }
1548 TweakIdx(12) => panic!("Shouldn't happen"),
1549 TweakIdx(13) => {
1550 vec![0usize] }
1552 _ => vec![],
1553 })
1554 }
1555 )
1556 .await
1557 .unwrap(),
1558 RecoverScanOutcome {
1559 last_used_idx: Some(TweakIdx(2)),
1560 new_start_idx: TweakIdx(2 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1561 tweak_idxes_with_pegins: BTreeSet::from([])
1562 }
1563 );
1564 assert!(last_checked.load(Ordering::SeqCst));
1565 }
1566
1567 {
1568 let last_checked = AtomicBool::new(false);
1569 let last_checked = &last_checked;
1570 assert_eq!(
1571 recover_scan_idxes_for_activity(
1572 TweakIdx(10),
1573 &BTreeSet::new(),
1574 |cur_idx| async move {
1575 Ok(match cur_idx {
1576 TweakIdx(10) => vec![()],
1577 TweakIdx(19) => {
1578 last_checked.store(true, Ordering::SeqCst);
1579 vec![]
1580 }
1581 TweakIdx(20) => panic!("Shouldn't happen"),
1582 _ => vec![],
1583 })
1584 }
1585 )
1586 .await
1587 .unwrap(),
1588 RecoverScanOutcome {
1589 last_used_idx: Some(TweakIdx(10)),
1590 new_start_idx: TweakIdx(10 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1591 tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(10)])
1592 }
1593 );
1594 assert!(last_checked.load(Ordering::SeqCst));
1595 }
1596
1597 assert_eq!(
1598 recover_scan_idxes_for_activity(TweakIdx(0), &BTreeSet::new(), |cur_idx| async move {
1599 Ok(match cur_idx {
1600 TweakIdx(6 | 15) => vec![()],
1601 _ => vec![],
1602 })
1603 })
1604 .await
1605 .unwrap(),
1606 RecoverScanOutcome {
1607 last_used_idx: Some(TweakIdx(15)),
1608 new_start_idx: TweakIdx(15 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1609 tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(6), TweakIdx(15)])
1610 }
1611 );
1612 assert_eq!(
1613 recover_scan_idxes_for_activity(TweakIdx(10), &BTreeSet::new(), |cur_idx| async move {
1614 Ok(match cur_idx {
1615 TweakIdx(8) => {
1616 vec![()] }
1618 TweakIdx(9) => {
1619 panic!("Shouldn't happen")
1620 }
1621 _ => vec![],
1622 })
1623 })
1624 .await
1625 .unwrap(),
1626 RecoverScanOutcome {
1627 last_used_idx: None,
1628 new_start_idx: TweakIdx(9 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1629 tweak_idxes_with_pegins: BTreeSet::from([])
1630 }
1631 );
1632 assert_eq!(
1633 recover_scan_idxes_for_activity(TweakIdx(10), &BTreeSet::new(), |cur_idx| async move {
1634 Ok(match cur_idx {
1635 TweakIdx(9) => panic!("Shouldn't happen"),
1636 TweakIdx(15) => vec![()],
1637 _ => vec![],
1638 })
1639 })
1640 .await
1641 .unwrap(),
1642 RecoverScanOutcome {
1643 last_used_idx: Some(TweakIdx(15)),
1644 new_start_idx: TweakIdx(15 + RECOVER_NUM_IDX_ADD_TO_LAST_USED),
1645 tweak_idxes_with_pegins: BTreeSet::from([TweakIdx(15)])
1646 }
1647 );
1648 }
1649}