1use std::collections::{BTreeMap, HashMap, HashSet};
9
10use anyhow::Context;
11use bitcoin::Amount;
12use bitcoin::secp256k1::Keypair;
13use log::{debug, info, warn};
14
15use ark::{ProtocolEncoding, Vtxo, VtxoId};
16use ark::mailbox::MailboxAuthorization;
17use ark::vtxo::Full;
18use bitcoin_ext::BlockHeight;
19use server_rpc::TryFromBytes;
20use server_rpc::protos::{self, VtxoSpendState};
21use server_rpc::protos::mailbox_server::mailbox_message::Message;
22
23use crate::Wallet;
24use crate::vtxo::VtxoState;
25
26#[derive(Debug, Default, Clone)]
27pub struct RecoveryReportEntry(HashMap<VtxoId, Option<Amount>>);
28
29impl RecoveryReportEntry {
30 pub fn is_empty(&self) -> bool {
31 self.0.is_empty()
32 }
33
34 pub fn len(&self) -> usize {
35 self.0.len()
36 }
37
38 pub fn ids(&self) -> impl Iterator<Item = VtxoId> {
39 self.0.keys().cloned()
40 }
41
42 pub fn total_amount(&self) -> Amount {
43 self.0.values().filter_map(|a| *a).sum()
44 }
45
46 fn insert(&mut self, vtxo_id: VtxoId, amount: Option<Amount>) {
47 if !self.0.contains_key(&vtxo_id) {
48 self.0.insert(vtxo_id, amount);
49 }
50 }
51
52 fn remove(&mut self, vtxo_id: VtxoId) {
53 self.0.remove(&vtxo_id);
54 }
55}
56
57#[derive(Debug, Default, Clone)]
66pub struct RecoveryReport {
67 recovered: RecoveryReportEntry,
69 skipped: RecoveryReportEntry,
72 failed: RecoveryReportEntry,
75 foreign: RecoveryReportEntry,
80 exited: RecoveryReportEntry,
82}
83
84impl RecoveryReport {
85 pub fn is_complete(&self) -> bool {
91 self.failed.is_empty() && self.foreign.is_empty()
92 }
93
94 pub fn recovered(&self) -> &RecoveryReportEntry {
95 &self.recovered
96 }
97
98 pub fn push_recovered(&mut self, vtxo: &Vtxo<Full>) {
99 self.failed.remove(vtxo.id());
100 self.recovered.insert(vtxo.id(), Some(vtxo.amount()));
101 }
102
103 pub fn skipped(&self) -> &RecoveryReportEntry {
104 &self.skipped
105 }
106
107 pub fn push_skipped(&mut self, vtxo: &Vtxo<Full>) {
108 self.failed.remove(vtxo.id());
109 self.skipped.insert(vtxo.id(), Some(vtxo.amount()));
110 }
111
112 pub fn foreign(&self) -> &RecoveryReportEntry {
113 &self.foreign
114 }
115
116 pub fn push_foreign(&mut self, vtxo: &Vtxo<Full>) {
117 self.failed.remove(vtxo.id());
118 self.foreign.insert(vtxo.id(), Some(vtxo.amount()));
119 }
120
121 pub fn failed(&self) -> &RecoveryReportEntry {
122 &self.failed
123 }
124
125 pub fn push_failed(&mut self, id: VtxoId, amount: Option<Amount>) {
126 self.failed.insert(id, amount);
127 }
128
129 pub fn exited(&self) -> &RecoveryReportEntry {
130 &self.exited
131 }
132
133 pub fn push_exited(&mut self, vtxo: &Vtxo<Full>) {
134 self.failed.remove(vtxo.id());
135 self.exited.insert(vtxo.id(), Some(vtxo.amount()));
136 }
137}
138
139#[derive(Debug)]
145pub enum RecoveryStatus {
146 NotRun,
149 Failed(anyhow::Error),
152 Completed(RecoveryReport),
155}
156
157pub(crate) struct OwnedVtxo {
162 vtxo: Vtxo<Full>,
163 keypair: Keypair,
164}
165
166impl OwnedVtxo {
167 fn new(vtxo: Vtxo<Full>, keypair: Keypair) -> Self {
170 debug_assert_eq!(
171 vtxo.user_pubkey(), keypair.public_key(),
172 "OwnedVtxo keypair must match the VTXO's owner pubkey",
173 );
174 OwnedVtxo { vtxo, keypair }
175 }
176}
177
178#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
186struct ChainOrder {
187 expiry: BlockHeight,
188 depth: u16,
189}
190
191impl ChainOrder {
192 fn of(vtxo: &Vtxo<Full>) -> Self {
193 ChainOrder { expiry: vtxo.expiry_height(), depth: vtxo.exit_depth() }
194 }
195}
196
197impl Wallet {
198 fn mailbox_request(&self, checkpoint: u64) -> protos::mailbox_server::MailboxRequest {
199 let expiry = chrono::Local::now() + std::time::Duration::from_secs(60);
200 let auth = MailboxAuthorization::new(&self.recovery_mailbox_keypair(), expiry);
201 let mailbox_id = auth.mailbox();
202
203 protos::mailbox_server::MailboxRequest {
204 mailbox_id: mailbox_id.serialize(),
205 authorization: Some(auth.serialize()),
206 checkpoint,
207 }
208 }
209
210 async fn fetch_valid_owned_vtxos(
211 &self,
212 report: &mut RecoveryReport,
213 ids: &[VtxoId],
214 gap_limit: u32,
215 ) -> anyhow::Result<Vec<OwnedVtxo>> {
216 let mut candidates = HashMap::new();
219
220 for id in ids {
221 match self.fetch_vtxo(*id).await {
222 Ok(vtxo) => {
223 candidates.insert(*id, vtxo);
224 },
225 Err(e) => {
226 warn!("Could not fetch recovery vtxo {id}: {:#}", e);
227 report.push_failed(*id, None);
228 }
229 }
230 }
231
232 let candidates = candidates.into_values().collect();
234 let owned = self.resolve_owned_vtxos(candidates, gap_limit, report).await?;
235
236 let mut vtxos = BTreeMap::<ChainOrder, Vec<OwnedVtxo>>::new();
238 for owned in owned {
239 vtxos.entry(ChainOrder::of(&owned.vtxo)).or_default().push(owned);
240 }
241
242 Ok(vtxos.into_values().flatten().collect())
243 }
244
245
246 async fn read_mailbox_recovery_vtxo_ids(&self) -> anyhow::Result<HashSet<VtxoId>> {
259 let (mut srv, _) = self.require_server().await?;
260
261 let mut ids = HashSet::new();
264
265 let mut iteration = 0;
266 let mut checkpoint = 0u64;
267 loop {
268 iteration += 1;
269
270 let req = self.mailbox_request(checkpoint);
271 let resp = srv.mailbox_client.read_mailbox(req).await
272 .context("error reading recovery mailbox")?.into_inner();
273
274 debug!("Recovery mailbox returned {} messages on iteration {iteration}", resp.messages.len());
275
276 let prev_checkpoint = checkpoint;
277 for msg in &resp.messages {
278 match &msg.message {
279 Some(Message::RecoveryVtxoIds(m)) => {
280 checkpoint = checkpoint.max(msg.checkpoint);
281 for raw in &m.vtxo_ids {
282 let Ok(id) = VtxoId::from_bytes(raw.clone()) else {
283 warn!("Ignoring undecodable recovery vtxo id: {raw:?}");
284 continue;
285 };
286 ids.insert(id);
287 }
288 },
289 Some(Message::Arkoor(m)) => {
290 checkpoint = checkpoint.max(msg.checkpoint);
291 for raw in &m.vtxos {
292 let Ok(vtxo) = Vtxo::<Full>::from_bytes(raw.clone()) else {
293 warn!("Ignoring undecodable vtxo: {raw:?}");
294 continue;
295 };
296 ids.insert(vtxo.id());
297 }
298 },
299 Some(Message::RoundParticipationCompleted(_)) |
300 Some(Message::IncomingLightningPayment(_)) |
301 Some(Message::LightningSendFinished(_)) => {},
302 None => {
303 warn!("Recovery mailbox returned a message with no content: {msg:?}");
304 },
305 }
306 }
307
308 if !resp.have_more {
309 break;
310 }
311
312 if checkpoint == prev_checkpoint {
316 warn!("Recovery mailbox iteration {iteration} made no progress \
317 at checkpoint {checkpoint}; stopping");
318 break;
319 }
320 }
321
322 Ok(ids)
323 }
324
325 async fn fetch_vtxo(&self, id: VtxoId) -> anyhow::Result<Vtxo<Full>> {
330 let (mut srv, _) = self.require_server().await?;
331 let resp = srv.client.get_vtxo(protos::GetVtxoRequest {
332 vtxo_id: id.to_bytes().to_vec(),
333 }).await.with_context(|| format!("error fetching vtxo {id} from server"))?.into_inner();
334
335 Vtxo::<Full>::deserialize(&resp.vtxo)
336 .with_context(|| format!("server returned an undecodable vtxo for {id}"))
337 }
338
339 async fn check_vtxo_onchain_status(&self, report: &mut RecoveryReport, vtxo: &Vtxo<Full>) -> anyhow::Result<bool> {
344 match self.inner.chain.tx_confirmed(vtxo.point().txid).await {
349 Ok(Some(height)) => {
350 self.store_vtxos(&vec![vtxo.clone()], &VtxoState::Exited).await?;
351 self.exit_mgr().start_exit_for_vtxos_including_non_standard(&vec![vtxo.to_bare()]).await?;
352 report.push_exited(vtxo);
353 debug!("Skipping recovery vtxo {}: confirmed on-chain at height {height} (exited)", vtxo.id());
354 Ok(true)
355 },
356 Ok(None) | Err(_) => Ok(false),
359 }
360 }
361
362 async fn check_vtxo_server_status(
367 &self,
368 report: &mut RecoveryReport,
369 vtxo: &Vtxo<Full>,
370 keypair: &Keypair,
371 ) -> anyhow::Result<bool> {
372 let vtxo_id = vtxo.id();
373 let spend_state = self.fetch_vtxo_spend_state(vtxo_id, keypair).await;
376
377 match spend_state {
381 Ok(VtxoSpendState::Spendable) => return Ok(false),
382 Ok(VtxoSpendState::Spent) => {
385 debug!("Recovery vtxo {vtxo_id} already spent, skipping");
386 self.store_vtxos([vtxo], &VtxoState::Spent).await
387 .context("Failed to record spent recovery vtxo")?;
388 report.push_skipped(vtxo);
389 },
390 Ok(state @ (
395 VtxoSpendState::Unclaimed
396 | VtxoSpendState::Unregistered
397 | VtxoSpendState::HtlcRecvUnclaimed
398 )) => {
399 debug!("Recovery vtxo {vtxo_id} not spendable ({state:?}), skipping");
400 report.push_skipped(vtxo);
401 },
402 Ok(VtxoSpendState::Unspecified) => {
405 warn!("Server returned an unspecified spend state for recovery vtxo {vtxo_id}");
406 report.push_failed(vtxo_id, Some(vtxo.amount()));
407 },
408 Err(e) => {
409 warn!("Could not get status for recovery vtxo {vtxo_id}: {:#}", e);
410 report.push_failed(vtxo_id, Some(vtxo.amount()));
411 },
412 }
413
414 Ok(true)
415 }
416
417 async fn resolve_owned_vtxos(
427 &self,
428 vtxos: Vec<Vtxo<Full>>,
429 gap_limit: u32,
430 report: &mut RecoveryReport,
431 ) -> anyhow::Result<Vec<OwnedVtxo>> {
432 let user_pubkeys = vtxos.iter().map(|v| v.user_pubkey()).collect::<Vec<_>>();
435 let keypairs = self.find_vtxo_keypairs(user_pubkeys, gap_limit).await?;
436
437 let mut matched = Vec::<(Vtxo<Full>, Keypair)>::new();
440 for vtxo in vtxos {
441 match keypairs.get(&vtxo.user_pubkey()) {
442 Some(keypair) => matched.push((vtxo, *keypair)),
443 None => report.push_foreign(&vtxo),
444 }
445 }
446
447 let mut owned = Vec::with_capacity(matched.len());
450 for (vtxo, keypair) in matched {
451 if let Err(e) = self.validate_vtxo(&vtxo).await {
452 warn!("Could not validate recovery vtxo {}: {:#}", vtxo.id(), e);
453 report.push_failed(vtxo.id(), Some(vtxo.amount()));
454 } else {
455 owned.push(OwnedVtxo::new(vtxo, keypair));
456 }
457 }
458
459 Ok(owned)
460 }
461
462 async fn inner_recover_vtxos(
463 &self,
464 report: &mut RecoveryReport,
465 ids: impl IntoIterator<Item = VtxoId>,
466 gap_limit: u32,
467 ) -> anyhow::Result<()> {
468 let ids = ids.into_iter().collect::<Vec<_>>();
469 let owned = self.fetch_valid_owned_vtxos(report, &ids, gap_limit).await?;
470
471 let mut spent = HashSet::<VtxoId>::new();
474
475 for o in owned {
476 let id = o.vtxo.id();
477
478 if spent.contains(&id) {
480 debug!("Skipping recovery vtxo {id}: spent into a newer recovered vtxo");
481 self.store_vtxos([&o.vtxo], &VtxoState::Spent).await
482 .context("Failed to record previously spent vtxo")?;
483 report.push_skipped(&o.vtxo);
484 continue;
485 }
486
487 spent.extend(o.vtxo.ancestor_ids());
489
490 if self.check_vtxo_onchain_status(report, &o.vtxo).await? {
491 continue;
492 }
493
494 if self.check_vtxo_server_status(report, &o.vtxo, &o.keypair).await? {
498 continue;
499 }
500
501 match self.store_vtxos([&o.vtxo], &VtxoState::Spendable).await {
503 Ok(()) => {
504 report.push_recovered(&o.vtxo);
505 debug!("Recovered spendable vtxo {id} ({})", o.vtxo.amount());
506 },
507 Err(e) => {
508 warn!("Failed to store recovered vtxo {id}: {:#}", e);
509 report.push_failed(id, Some(o.vtxo.amount()));
510 },
511 }
512 }
513
514 Ok(())
515 }
516
517 pub async fn recover_vtxos(
522 &self,
523 ids: impl IntoIterator<Item = VtxoId>,
524 gap_limit: Option<u32>,
525 ) -> anyhow::Result<RecoveryReport> {
526 let mut report = RecoveryReport::default();
527 let gap_limit = gap_limit.unwrap_or(self.inner.config.vtxo_key_gap_limit);
528 self.inner_recover_vtxos(&mut report, ids, gap_limit).await?;
529 Ok(report)
530 }
531
532 pub(crate) async fn recover_from_mailbox(&self) -> anyhow::Result<RecoveryReport> {
544 let mut report = RecoveryReport::default();
545 let gap_limit = self.inner.config.vtxo_key_gap_limit;
546
547 let ids = self.read_mailbox_recovery_vtxo_ids().await?;
549 debug!("Found {} distinct vtxo ids in the recovery mailbox", ids.len());
550
551 self.inner_recover_vtxos(&mut report, ids, gap_limit).await?;
552
553 if !report.foreign.is_empty() {
558 warn!(
559 "Recovery mailbox held {} vtxo(s) not derivable from this seed within the \
560 gap limit ({gap_limit}); if any are ours they were not recovered: {:?}",
561 report.foreign.len(), report.foreign,
562 );
563 }
564
565 if report.is_complete() {
566 info!(
567 "Recovered {} spendable vtxos from the recovery mailbox ({} skipped)",
568 report.recovered.len(), report.skipped.len(),
569 );
570 }
571
572 for _ in 0..3 {
574 if report.failed.is_empty() {
575 break;
576 }
577
578 let ids = report.failed.ids().collect::<Vec<_>>();
579 self.inner_recover_vtxos(&mut report, ids, gap_limit).await?;
580 }
581
582 if !report.failed.is_empty() {
583 warn!(
584 "Recovery incomplete: recovered {} spendable vtxos, but {} could not be \
585 checked due to errors; funds may be missing — retry recovery to recover \
586 them ({} skipped). Failed vtxos: {:?}",
587 report.recovered.len(), report.failed.len(),
588 report.skipped.len(), report.failed,
589 );
590 }
591
592 Ok(report)
593 }
594}
595
596#[cfg(test)]
597mod test {
598 use super::*;
599 use bitcoin::Amount;
600 use bitcoin::hashes::Hash;
601
602 fn dummy_id(vout: u32) -> VtxoId {
603 bitcoin::OutPoint::new(bitcoin::Txid::all_zeros(), vout).into()
604 }
605
606 #[test]
610 fn recovery_report_completeness() {
611 let clean = RecoveryReport {
612 recovered: RecoveryReportEntry(HashMap::from([(dummy_id(0), Some(Amount::from_sat(1000)))])),
613 skipped: RecoveryReportEntry(HashMap::from([(dummy_id(1), Some(Amount::from_sat(1000)))])),
614 foreign: RecoveryReportEntry(HashMap::new()),
615 failed: RecoveryReportEntry(HashMap::new()),
616 exited: RecoveryReportEntry(HashMap::new()),
617 };
618 assert!(clean.is_complete(),
619 "recovered and skipped VTXOs are clean decisions, not failures");
620
621 assert!(RecoveryReport::default().is_complete(),
622 "an empty report is trivially complete");
623 assert!(!RecoveryReport {
624 failed: RecoveryReportEntry(HashMap::from([(dummy_id(0), Some(Amount::from_sat(1000)))])),
625 ..Default::default()
626 }.is_complete(), "a failed VTXO means recovery is incomplete");
627 assert!(!RecoveryReport {
628 foreign: RecoveryReportEntry(HashMap::from([(dummy_id(2), Some(Amount::from_sat(1000)))])),
629 ..Default::default()
630 }.is_complete(), "a foreign id is likely an owned VTXO beyond the gap limit, so recovery is incomplete");
631 }
632}
633