1use std::collections::{BTreeMap, HashMap, HashSet};
9
10use anyhow::Context;
11use bitcoin::Amount;
12use bitcoin::secp256k1::{Keypair, PublicKey};
13use log::{debug, info, warn};
14
15use ark::{ProtocolEncoding, Vtxo, VtxoId};
16use ark::attestations::VtxoStatusAttestation;
17use ark::mailbox::MailboxAuthorization;
18use ark::vtxo::Full;
19use bitcoin_ext::BlockHeight;
20use server_rpc::TryFromBytes;
21use server_rpc::protos::{self, VtxoSpendState};
22use server_rpc::protos::mailbox_server::mailbox_message::Message;
23
24use crate::Wallet;
25use crate::vtxo::VtxoState;
26
27const STOP_GAP: u32 = 50;
29
30#[derive(Debug, Default, Clone)]
31pub struct RecoveryReportEntry(HashMap<VtxoId, Option<Amount>>);
32
33impl RecoveryReportEntry {
34 pub fn is_empty(&self) -> bool {
35 self.0.is_empty()
36 }
37
38 pub fn len(&self) -> usize {
39 self.0.len()
40 }
41
42 pub fn ids(&self) -> impl Iterator<Item = VtxoId> {
43 self.0.keys().cloned()
44 }
45
46 pub fn total_amount(&self) -> Amount {
47 self.0.values().filter_map(|a| *a).sum()
48 }
49
50 fn insert(&mut self, vtxo_id: VtxoId, amount: Option<Amount>) {
51 if !self.0.contains_key(&vtxo_id) {
52 self.0.insert(vtxo_id, amount);
53 }
54 }
55
56 fn remove(&mut self, vtxo_id: VtxoId) {
57 self.0.remove(&vtxo_id);
58 }
59}
60
61#[derive(Debug, Default, Clone)]
70pub struct RecoveryReport {
71 recovered: RecoveryReportEntry,
73 skipped: RecoveryReportEntry,
76 failed: RecoveryReportEntry,
79 foreign: RecoveryReportEntry,
84 exited: RecoveryReportEntry,
86}
87
88impl RecoveryReport {
89 pub fn is_complete(&self) -> bool {
95 self.failed.is_empty() && self.foreign.is_empty()
96 }
97
98 pub fn recovered(&self) -> &RecoveryReportEntry {
99 &self.recovered
100 }
101
102 pub fn push_recovered(&mut self, vtxo: &Vtxo<Full>) {
103 self.failed.remove(vtxo.id());
104 self.recovered.insert(vtxo.id(), Some(vtxo.amount()));
105 }
106
107 pub fn skipped(&self) -> &RecoveryReportEntry {
108 &self.skipped
109 }
110
111 pub fn push_skipped(&mut self, vtxo: &Vtxo<Full>) {
112 self.failed.remove(vtxo.id());
113 self.skipped.insert(vtxo.id(), Some(vtxo.amount()));
114 }
115
116 pub fn foreign(&self) -> &RecoveryReportEntry {
117 &self.foreign
118 }
119
120 pub fn push_foreign(&mut self, vtxo: &Vtxo<Full>) {
121 self.failed.remove(vtxo.id());
122 self.foreign.insert(vtxo.id(), Some(vtxo.amount()));
123 }
124
125 pub fn failed(&self) -> &RecoveryReportEntry {
126 &self.failed
127 }
128
129 pub fn push_failed(&mut self, id: VtxoId, amount: Option<Amount>) {
130 self.failed.insert(id, amount);
131 }
132
133 pub fn exited(&self) -> &RecoveryReportEntry {
134 &self.exited
135 }
136
137 pub fn push_exited(&mut self, vtxo: &Vtxo<Full>) {
138 self.failed.remove(vtxo.id());
139 self.exited.insert(vtxo.id(), Some(vtxo.amount()));
140 }
141}
142
143#[derive(Debug)]
149pub enum RecoveryStatus {
150 NotRun,
153 Failed(anyhow::Error),
156 Completed(RecoveryReport),
159}
160
161pub(crate) struct OwnedVtxo {
166 vtxo: Vtxo<Full>,
167 keypair: Keypair,
168}
169
170impl OwnedVtxo {
171 fn new(vtxo: Vtxo<Full>, keypair: Keypair) -> Self {
174 debug_assert_eq!(
175 vtxo.user_pubkey(), keypair.public_key(),
176 "OwnedVtxo keypair must match the VTXO's owner pubkey",
177 );
178 OwnedVtxo { vtxo, keypair }
179 }
180}
181
182#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
190struct ChainOrder {
191 expiry: BlockHeight,
192 depth: u16,
193}
194
195impl ChainOrder {
196 fn of(vtxo: &Vtxo<Full>) -> Self {
197 ChainOrder { expiry: vtxo.expiry_height(), depth: vtxo.exit_depth() }
198 }
199}
200
201impl Wallet {
202 fn mailbox_request(&self, checkpoint: u64) -> protos::mailbox_server::MailboxRequest {
203 let expiry = chrono::Local::now() + std::time::Duration::from_secs(60);
204 let auth = MailboxAuthorization::new(&self.recovery_mailbox_keypair(), expiry);
205 let mailbox_id = auth.mailbox();
206
207 protos::mailbox_server::MailboxRequest {
208 mailbox_id: mailbox_id.serialize(),
209 authorization: Some(auth.serialize()),
210 checkpoint,
211 }
212 }
213
214 async fn fetch_valid_owned_vtxos(&self, report: &mut RecoveryReport, ids: &[VtxoId]) ->
215 anyhow::Result<Vec<OwnedVtxo>>
216 {
217 let mut candidates = HashMap::new();
220
221 for id in ids {
222 match self.fetch_vtxo(*id).await {
223 Ok(vtxo) => {
224 candidates.insert(*id, vtxo);
225 },
226 Err(e) => {
227 warn!("Could not fetch recovery vtxo {id}: {:#}", e);
228 report.push_failed(*id, None);
229 }
230 }
231 }
232
233 let candidates = candidates.into_values().collect();
235 let owned = self.resolve_owned_vtxos(candidates, report).await?;
236
237 let mut vtxos = BTreeMap::<ChainOrder, Vec<OwnedVtxo>>::new();
239 for owned in owned {
240 vtxos.entry(ChainOrder::of(&owned.vtxo)).or_default().push(owned);
241 }
242
243 Ok(vtxos.into_values().flatten().collect())
244 }
245
246
247 async fn read_mailbox_recovery_vtxo_ids(&self) -> anyhow::Result<HashSet<VtxoId>> {
260 let (mut srv, _) = self.require_server().await?;
261
262 let mut ids = HashSet::new();
265
266 let mut iteration = 0;
267 let mut checkpoint = 0u64;
268 loop {
269 iteration += 1;
270
271 let req = self.mailbox_request(checkpoint);
272 let resp = srv.mailbox_client.read_mailbox(req).await
273 .context("error reading recovery mailbox")?.into_inner();
274
275 debug!("Recovery mailbox returned {} messages on iteration {iteration}", resp.messages.len());
276
277 let prev_checkpoint = checkpoint;
278 for msg in &resp.messages {
279 match &msg.message {
280 Some(Message::RecoveryVtxoIds(m)) => {
281 checkpoint = checkpoint.max(msg.checkpoint);
282 for raw in &m.vtxo_ids {
283 let Ok(id) = VtxoId::from_bytes(raw.clone()) else {
284 warn!("Ignoring undecodable recovery vtxo id: {raw:?}");
285 continue;
286 };
287 ids.insert(id);
288 }
289 },
290 Some(Message::Arkoor(m)) => {
291 checkpoint = checkpoint.max(msg.checkpoint);
292 for raw in &m.vtxos {
293 let Ok(vtxo) = Vtxo::<Full>::from_bytes(raw.clone()) else {
294 warn!("Ignoring undecodable vtxo: {raw:?}");
295 continue;
296 };
297 ids.insert(vtxo.id());
298 }
299 },
300 Some(Message::RoundParticipationCompleted(_)) |
301 Some(Message::IncomingLightningPayment(_)) |
302 Some(Message::LightningSendFinished(_)) => {},
303 None => {
304 warn!("Recovery mailbox returned a message with no content: {msg:?}");
305 },
306 }
307 }
308
309 if !resp.have_more {
310 break;
311 }
312
313 if checkpoint == prev_checkpoint {
317 warn!("Recovery mailbox iteration {iteration} made no progress \
318 at checkpoint {checkpoint}; stopping");
319 break;
320 }
321 }
322
323 Ok(ids)
324 }
325
326 async fn fetch_vtxo(&self, id: VtxoId) -> anyhow::Result<Vtxo<Full>> {
331 let (mut srv, _) = self.require_server().await?;
332 let resp = srv.client.get_vtxo(protos::GetVtxoRequest {
333 vtxo_id: id.to_bytes().to_vec(),
334 }).await.with_context(|| format!("error fetching vtxo {id} from server"))?.into_inner();
335
336 Vtxo::<Full>::deserialize(&resp.vtxo)
337 .with_context(|| format!("server returned an undecodable vtxo for {id}"))
338 }
339
340 async fn check_vtxo_onchain_status(&self, report: &mut RecoveryReport, vtxo: &Vtxo<Full>) -> anyhow::Result<bool> {
345 match self.inner.chain.tx_confirmed(vtxo.point().txid).await {
350 Ok(Some(height)) => {
351 self.store_vtxos(&vec![vtxo.clone()], &VtxoState::Exited).await?;
352 self.exit_mgr().start_exit_for_vtxos_including_non_standard(&vec![vtxo.to_bare()]).await?;
353 report.push_exited(vtxo);
354 debug!("Skipping recovery vtxo {}: confirmed on-chain at height {height} (exited)", vtxo.id());
355 Ok(true)
356 },
357 Ok(None) | Err(_) => Ok(false),
360 }
361 }
362
363 async fn check_vtxo_server_status(
368 &self,
369 report: &mut RecoveryReport,
370 vtxo: &Vtxo<Full>,
371 keypair: &Keypair,
372 ) -> anyhow::Result<bool> {
373 let (mut srv, _) = self.require_server().await?;
374 let vtxo_id = vtxo.id();
375 let attestation = VtxoStatusAttestation::new(vtxo_id, keypair);
376 let resp = srv.client.get_vtxo_status(protos::GetVtxoStatusRequest {
377 vtxo_id: vtxo_id.to_bytes().to_vec(),
378 attestation: attestation.serialize(),
379 }).await.with_context(|| format!("error fetching status for vtxo {vtxo_id}"))?.into_inner();
380
381 let spend_state = protos::VtxoSpendState::try_from(resp.spend_state)
382 .map_err(|_| anyhow::anyhow!(
383 "server returned unknown spend state {} for vtxo {vtxo_id}", resp.spend_state,
384 ));
385
386 match spend_state {
390 Ok(VtxoSpendState::Spendable) => return Ok(false),
391 Ok(state @ (
393 VtxoSpendState::Spent
394 | VtxoSpendState::Unclaimed
395 | VtxoSpendState::Unregistered
396 | VtxoSpendState::HtlcRecvUnclaimed
397 )) => {
398 debug!("Recovery vtxo {vtxo_id} not spendable ({state:?}), skipping");
399 report.push_skipped(vtxo);
400 },
401 Ok(VtxoSpendState::Unspecified) => {
404 warn!("Server returned an unspecified spend state for recovery vtxo {vtxo_id}");
405 report.push_failed(vtxo_id, Some(vtxo.amount()));
406 },
407 Err(e) => {
408 warn!("Could not get status for recovery vtxo {vtxo_id}: {:#}", e);
409 report.push_failed(vtxo_id, Some(vtxo.amount()));
410 },
411 }
412
413 Ok(true)
414 }
415
416 async fn resolve_owned_vtxos(
428 &self,
429 vtxos: Vec<Vtxo<Full>>,
430 report: &mut RecoveryReport,
431 ) -> anyhow::Result<Vec<OwnedVtxo>> {
432 let mut pending = HashMap::<PublicKey, Vec<Vtxo<Full>>>::new();
435 let mut matched = Vec::<(Vtxo<Full>, Keypair)>::new();
436
437 for vtxo in vtxos {
438 match self.pubkey_keypair(&vtxo.user_pubkey()).await? {
439 Some((_idx, keypair)) => matched.push((vtxo, keypair)),
440 None => pending.entry(vtxo.user_pubkey()).or_default().push(vtxo),
441 }
442 }
443
444 let start_idx = self.inner.db.get_last_vtxo_key_index().await?.map(|i| i + 1).unwrap_or(0);
446 let mut frontier = start_idx.saturating_add(STOP_GAP);
447 let mut gap = Vec::<(u32, PublicKey)>::new();
448 let mut idx = start_idx;
449 while idx <= frontier && !pending.is_empty() {
450 let keypair = self.inner.seed.derive_vtxo_keypair(idx);
451 let pubkey = keypair.public_key();
452 if let Some(owned_vtxos) = pending.remove(&pubkey) {
453 for (i, pk) in gap.drain(..) {
456 self.inner.db.store_vtxo_key(i, pk).await?;
457 }
458 self.inner.db.store_vtxo_key(idx, pubkey).await?;
459 frontier = idx.saturating_add(STOP_GAP);
460 matched.extend(owned_vtxos.into_iter().map(|v| (v, keypair)));
461 } else {
462 gap.push((idx, pubkey));
463 }
464 let Some(next_idx) = idx.checked_add(1) else { break };
467 idx = next_idx;
468 }
469
470 for vtxo in pending.into_values().flatten() {
472 report.push_foreign(&vtxo);
473 }
474
475 let mut owned = Vec::with_capacity(matched.len());
478 for (vtxo, keypair) in matched {
479 if let Err(e) = self.validate_vtxo(&vtxo).await {
480 warn!("Could not validate recovery vtxo {}: {:#}", vtxo.id(), e);
481 report.push_failed(vtxo.id(), Some(vtxo.amount()));
482 } else {
483 owned.push(OwnedVtxo::new(vtxo, keypair));
484 }
485 }
486
487 Ok(owned)
488 }
489
490 async fn inner_recover_vtxos(
491 &self,
492 report: &mut RecoveryReport,
493 ids: impl IntoIterator<Item = VtxoId>,
494 ) -> anyhow::Result<()> {
495 let ids = ids.into_iter().collect::<Vec<_>>();
496 let owned = self.fetch_valid_owned_vtxos(report, &ids).await?;
497
498 let mut spent = HashSet::<VtxoId>::new();
501
502 for o in owned {
503 let id = o.vtxo.id();
504
505 if spent.contains(&id) {
507 debug!("Skipping recovery vtxo {id}: spent into a newer recovered vtxo");
508 report.push_skipped(&o.vtxo);
509 continue;
510 }
511
512 spent.extend(o.vtxo.ancestor_ids());
514
515 if self.check_vtxo_onchain_status(report, &o.vtxo).await? {
516 continue;
517 }
518
519 if self.check_vtxo_server_status(report, &o.vtxo, &o.keypair).await? {
523 continue;
524 }
525
526 match self.store_vtxos([&o.vtxo], &VtxoState::Spendable).await {
528 Ok(()) => {
529 report.push_recovered(&o.vtxo);
530 debug!("Recovered spendable vtxo {id} ({})", o.vtxo.amount());
531 },
532 Err(e) => {
533 warn!("Failed to store recovered vtxo {id}: {:#}", e);
534 report.push_failed(id, Some(o.vtxo.amount()));
535 },
536 }
537 }
538
539 Ok(())
540 }
541
542 pub async fn recover_vtxos(&self, ids: impl IntoIterator<Item = VtxoId>)
543 -> anyhow::Result<RecoveryReport>
544 {
545 let mut report = RecoveryReport::default();
546 self.inner_recover_vtxos(&mut report, ids).await?;
547 Ok(report)
548 }
549
550 pub(crate) async fn recover_from_mailbox(&self) -> anyhow::Result<RecoveryReport> {
562 let mut report = RecoveryReport::default();
563
564 let ids = self.read_mailbox_recovery_vtxo_ids().await?;
566 debug!("Found {} distinct vtxo ids in the recovery mailbox", ids.len());
567
568 self.inner_recover_vtxos(&mut report, ids).await?;
569
570 if !report.foreign.is_empty() {
575 warn!(
576 "Recovery mailbox held {} vtxo(s) not derivable from this seed within the \
577 gap limit ({STOP_GAP}); if any are ours they were not recovered: {:?}",
578 report.foreign.len(), report.foreign,
579 );
580 }
581
582 if report.is_complete() {
583 info!(
584 "Recovered {} spendable vtxos from the recovery mailbox ({} skipped)",
585 report.recovered.len(), report.skipped.len(),
586 );
587 }
588
589 for _ in 0..3 {
591 if report.failed.is_empty() {
592 break;
593 }
594
595 let ids = report.failed.ids().collect::<Vec<_>>();
596 self.inner_recover_vtxos(&mut report, ids).await?;
597 }
598
599 if !report.failed.is_empty() {
600 warn!(
601 "Recovery incomplete: recovered {} spendable vtxos, but {} could not be \
602 checked due to errors; funds may be missing — retry recovery to recover \
603 them ({} skipped). Failed vtxos: {:?}",
604 report.recovered.len(), report.failed.len(),
605 report.skipped.len(), report.failed,
606 );
607 }
608
609 Ok(report)
610 }
611}
612
613#[cfg(test)]
614mod test {
615 use super::*;
616 use bitcoin::Amount;
617 use bitcoin::hashes::Hash;
618
619 fn dummy_id(vout: u32) -> VtxoId {
620 bitcoin::OutPoint::new(bitcoin::Txid::all_zeros(), vout).into()
621 }
622
623 #[test]
627 fn recovery_report_completeness() {
628 let clean = RecoveryReport {
629 recovered: RecoveryReportEntry(HashMap::from([(dummy_id(0), Some(Amount::from_sat(1000)))])),
630 skipped: RecoveryReportEntry(HashMap::from([(dummy_id(1), Some(Amount::from_sat(1000)))])),
631 foreign: RecoveryReportEntry(HashMap::new()),
632 failed: RecoveryReportEntry(HashMap::new()),
633 exited: RecoveryReportEntry(HashMap::new()),
634 };
635 assert!(clean.is_complete(),
636 "recovered and skipped VTXOs are clean decisions, not failures");
637
638 assert!(RecoveryReport::default().is_complete(),
639 "an empty report is trivially complete");
640 assert!(!RecoveryReport {
641 failed: RecoveryReportEntry(HashMap::from([(dummy_id(0), Some(Amount::from_sat(1000)))])),
642 ..Default::default()
643 }.is_complete(), "a failed VTXO means recovery is incomplete");
644 assert!(!RecoveryReport {
645 foreign: RecoveryReportEntry(HashMap::from([(dummy_id(2), Some(Amount::from_sat(1000)))])),
646 ..Default::default()
647 }.is_complete(), "a foreign id is likely an owned VTXO beyond the gap limit, so recovery is incomplete");
648 }
649}
650