Introduce Watchable abstraction for Bitcoin wallet

We have a repeated pattern where we construct one of our
Tx{Cancel,Redeem,Punish,Refund,Lock} transactions and wait until
the status of this transaction changes. We can make this more
ergonomic by creating and implementing a `Watchable` trait that
gives access to the TxId and relevant script for this transaction.
This allows us to remove a parameter from the `watch_until_status`
function.

Additionally, there is a 2nd pattern: "Completing" one of these
transaction and waiting until they are confirmed with the configured
number of blocks for finality. We can make this more ergonomic by
returning a future from `broadcast` that callers can await in case
they want to wait for the broadcasted transaction to reach finality.

swap/src/bitcoin/cancel.rs

@@ -1,3 +1,4 @@
+use crate::bitcoin::wallet::Watchable;
 use crate::bitcoin::{
     build_shared_output_descriptor, Address, Amount, BlockHeight, PublicKey, Transaction, TxLock,
     TX_FEE,
@@ -81,7 +82,7 @@ impl PartialEq<PunishTimelock> for u32 {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct TxCancel {
     inner: Transaction,
     digest: SigHash,
@@ -148,16 +149,6 @@ impl TxCancel {
         OutPoint::new(self.inner.txid(), 0)
     }
 
-    /// Return the relevant script_pubkey of this transaction.
-    ///
-    /// Even though a transaction can have multiple outputs, the nature of our
-    /// protocol is that there is only one relevant output within this one.
-    /// As such, subscribing or inquiring the status of this script allows us to
-    /// check the status of the whole transaction.
-    pub fn script_pubkey(&self) -> Script {
-        self.output_descriptor.script_pubkey()
-    }
-
     pub fn add_signatures(
         self,
         (A, sig_a): (PublicKey, Signature),
@@ -216,3 +207,13 @@ impl TxCancel {
         }
     }
 }
+
+impl Watchable for TxCancel {
+    fn id(&self) -> Txid {
+        self.txid()
+    }
+
+    fn script(&self) -> Script {
+        self.output_descriptor.script_pubkey()
+    }
+}

swap/src/bitcoin/lock.rs

@@ -1,3 +1,4 @@
+use crate::bitcoin::wallet::Watchable;
 use crate::bitcoin::{
     build_shared_output_descriptor, Address, Amount, PublicKey, Transaction, Wallet, TX_FEE,
 };
@@ -105,3 +106,13 @@ impl From<TxLock> for PartiallySignedTransaction {
         from.inner
     }
 }
+
+impl Watchable for TxLock {
+    fn id(&self) -> Txid {
+        self.txid()
+    }
+
+    fn script(&self) -> Script {
+        self.output_descriptor.script_pubkey()
+    }
+}

swap/src/bitcoin/punish.rs

@@ -1,15 +1,18 @@
-use crate::bitcoin::{self, Address, PunishTimelock, Transaction, TxCancel};
+use crate::bitcoin::wallet::Watchable;
+use crate::bitcoin::{self, Address, PunishTimelock, Transaction, TxCancel, Txid};
 use ::bitcoin::util::bip143::SigHashCache;
 use ::bitcoin::{SigHash, SigHashType};
 use anyhow::{Context, Result};
+use bdk::bitcoin::Script;
 use miniscript::{Descriptor, DescriptorTrait};
 use std::collections::HashMap;
 
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 pub struct TxPunish {
     inner: Transaction,
     digest: SigHash,
     cancel_output_descriptor: Descriptor<::bitcoin::PublicKey>,
+    watch_script: Script,
 }
 
 impl TxPunish {
@@ -31,6 +34,7 @@ impl TxPunish {
             inner: tx_punish,
             digest,
             cancel_output_descriptor: tx_cancel.output_descriptor.clone(),
+            watch_script: punish_address.script_pubkey(),
         }
     }
 
@@ -68,3 +72,13 @@ impl TxPunish {
         Ok(tx_punish)
     }
 }
+
+impl Watchable for TxPunish {
+    fn id(&self) -> Txid {
+        self.inner.txid()
+    }
+
+    fn script(&self) -> Script {
+        self.watch_script.clone()
+    }
+}

swap/src/bitcoin/redeem.rs

@@ -1,3 +1,4 @@
+use crate::bitcoin::wallet::Watchable;
 use crate::bitcoin::{
     verify_encsig, verify_sig, Address, EmptyWitnessStack, EncryptedSignature, NoInputs,
     NotThreeWitnesses, PublicKey, SecretKey, TooManyInputs, Transaction, TxLock,
@@ -5,6 +6,7 @@ use crate::bitcoin::{
 use ::bitcoin::util::bip143::SigHashCache;
 use ::bitcoin::{SigHash, SigHashType, Txid};
 use anyhow::{bail, Context, Result};
+use bitcoin::Script;
 use ecdsa_fun::adaptor::{Adaptor, HashTranscript};
 use ecdsa_fun::fun::Scalar;
 use ecdsa_fun::nonce::Deterministic;
@@ -13,11 +15,12 @@ use miniscript::{Descriptor, DescriptorTrait};
 use sha2::Sha256;
 use std::collections::HashMap;
 
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct TxRedeem {
     inner: Transaction,
     digest: SigHash,
     lock_output_descriptor: Descriptor<::bitcoin::PublicKey>,
+    watch_script: Script,
 }
 
 impl TxRedeem {
@@ -37,6 +40,7 @@ impl TxRedeem {
             inner: tx_redeem,
             digest,
             lock_output_descriptor: tx_lock.output_descriptor.clone(),
+            watch_script: redeem_address.script_pubkey(),
         }
     }
 
@@ -130,3 +134,13 @@ impl TxRedeem {
         Ok(sig)
     }
 }
+
+impl Watchable for TxRedeem {
+    fn id(&self) -> Txid {
+        self.txid()
+    }
+
+    fn script(&self) -> Script {
+        self.watch_script.clone()
+    }
+}

swap/src/bitcoin/refund.rs

@@ -1,3 +1,4 @@
+use crate::bitcoin::wallet::Watchable;
 use crate::bitcoin::{
     verify_sig, Address, EmptyWitnessStack, NoInputs, NotThreeWitnesses, PublicKey, TooManyInputs,
     Transaction, TxCancel,
@@ -5,6 +6,7 @@ use crate::bitcoin::{
 use ::bitcoin::util::bip143::SigHashCache;
 use ::bitcoin::{SigHash, SigHashType, Txid};
 use anyhow::{bail, Context, Result};
+use bitcoin::Script;
 use ecdsa_fun::Signature;
 use miniscript::{Descriptor, DescriptorTrait};
 use std::collections::HashMap;
@@ -14,6 +16,7 @@ pub struct TxRefund {
     inner: Transaction,
     digest: SigHash,
     cancel_output_descriptor: Descriptor<::bitcoin::PublicKey>,
+    watch_script: Script,
 }
 
 impl TxRefund {
@@ -31,6 +34,7 @@ impl TxRefund {
             inner: tx_punish,
             digest,
             cancel_output_descriptor: tx_cancel.output_descriptor.clone(),
+            watch_script: refund_address.script_pubkey(),
         }
     }
 
@@ -110,3 +114,13 @@ impl TxRefund {
         Ok(sig)
     }
 }
+
+impl Watchable for TxRefund {
+    fn id(&self) -> Txid {
+        self.txid()
+    }
+
+    fn script(&self) -> Script {
+        self.watch_script.clone()
+    }
+}

swap/src/bitcoin/wallet.rs

@@ -13,6 +13,7 @@ use reqwest::Url;
 use std::collections::BTreeMap;
 use std::convert::TryFrom;
 use std::fmt;
+use std::future::Future;
 use std::path::Path;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
@@ -159,9 +160,22 @@ impl Wallet {
 
     /// Broadcast the given transaction to the network and emit a log statement
     /// if done so successfully.
-    pub async fn broadcast(&self, transaction: Transaction, kind: &str) -> Result<Txid> {
+    ///
+    /// Returns the transaction ID and a future for when the transaction meets
+    /// the configured finality confirmations.
+    pub async fn broadcast(
+        &self,
+        transaction: Transaction,
+        kind: &str,
+    ) -> Result<(Txid, impl Future<Output = Result<()>> + '_)> {
         let txid = transaction.txid();
 
+        // to watch for confirmations, watching a single output is enough
+        let watcher = self.wait_for_transaction_finality(
+            (txid, transaction.output[0].script_pubkey.clone()),
+            kind.to_owned(),
+        );
+
         self.wallet
             .lock()
             .await
@@ -172,7 +186,7 @@ impl Wallet {
 
         tracing::info!(%txid, "Published Bitcoin {} transaction", kind);
 
-        Ok(txid)
+        Ok((txid, watcher))
     }
 
     pub async fn sign_and_finalize(&self, psbt: PartiallySignedTransaction) -> Result<Transaction> {
@@ -193,20 +207,27 @@ impl Wallet {
             .ok_or_else(|| anyhow!("Could not get raw tx with id: {}", txid))
     }
 
-    pub async fn status_of_script(&self, script: &Script, txid: &Txid) -> Result<ScriptStatus> {
-        self.client.lock().await.status_of_script(script, txid)
+    pub async fn status_of_script<T>(&self, tx: &T) -> Result<ScriptStatus>
+    where
+        T: Watchable,
+    {
+        self.client.lock().await.status_of_script(tx)
     }
 
-    pub async fn watch_until_status(
+    pub async fn watch_until_status<T>(
         &self,
-        txid: Txid,
-        script: Script,
+        tx: &T,
         mut status_fn: impl FnMut(ScriptStatus) -> bool,
-    ) -> Result<()> {
+    ) -> Result<()>
+    where
+        T: Watchable,
+    {
+        let txid = tx.id();
+
         let mut last_status = None;
 
         loop {
-            let new_status = self.client.lock().await.status_of_script(&script, &txid)?;
+            let new_status = self.client.lock().await.status_of_script(tx)?;
 
             if Some(new_status) != last_status {
                 tracing::debug!(%txid, "Transaction is {}", new_status);
@@ -223,23 +244,23 @@ impl Wallet {
         Ok(())
     }
 
-    pub async fn wait_for_transaction_finality(
-        &self,
-        txid: Txid,
-        script_to_watch: Script,
-    ) -> Result<()> {
+    async fn wait_for_transaction_finality<T>(&self, tx: T, kind: String) -> Result<()>
+    where
+        T: Watchable,
+    {
         let conf_target = self.bitcoin_finality_confirmations;
+        let txid = tx.id();
 
-        tracing::info!(%txid, "Waiting for {} confirmation{} of Bitcoin transaction", conf_target, if conf_target > 1 { "s" } else { "" });
+        tracing::info!(%txid, "Waiting for {} confirmation{} of Bitcoin {} transaction", conf_target, if conf_target > 1 { "s" } else { "" }, kind);
 
         let mut seen_confirmations = 0;
 
-        self.watch_until_status(txid, script_to_watch, |status| match status {
+        self.watch_until_status(&tx, |status| match status {
             ScriptStatus::Confirmed(inner) => {
                 let confirmations = inner.confirmations();
 
                 if confirmations > seen_confirmations {
-                    tracing::info!(%txid, "Bitcoin tx has {} out of {} confirmation{}", confirmations, conf_target, if conf_target > 1 { "s" } else { "" });
+                    tracing::info!(%txid, "Bitcoin {} tx has {} out of {} confirmation{}", kind, confirmations, conf_target, if conf_target > 1 { "s" } else { "" });
                     seen_confirmations = confirmations;
                 }
 
@@ -260,6 +281,27 @@ impl Wallet {
     }
 }
 
+/// Defines a watchable transaction.
+///
+/// For a transaction to be watchable, we need to know two things: Its
+/// transaction ID and the specific output script that is going to change.
+/// A transaction can obviously have multiple outputs but our protocol purposes,
+/// we are usually interested in a specific one.
+pub trait Watchable {
+    fn id(&self) -> Txid;
+    fn script(&self) -> Script;
+}
+
+impl Watchable for (Txid, Script) {
+    fn id(&self) -> Txid {
+        self.0
+    }
+
+    fn script(&self) -> Script {
+        self.1.clone()
+    }
+}
+
 struct Client {
     electrum: bdk::electrum_client::Client,
     latest_block: BlockHeight,
@@ -321,18 +363,24 @@ impl Client {
         Ok(())
     }
 
-    fn status_of_script(&mut self, script: &Script, txid: &Txid) -> Result<ScriptStatus> {
-        if !self.script_history.contains_key(script) {
+    fn status_of_script<T>(&mut self, tx: &T) -> Result<ScriptStatus>
+    where
+        T: Watchable,
+    {
+        let txid = tx.id();
+        let script = tx.script();
+
+        if !self.script_history.contains_key(&script) {
             self.script_history.insert(script.clone(), vec![]);
         }
 
         self.drain_notifications()?;
 
-        let history = self.script_history.entry(script.clone()).or_default();
+        let history = self.script_history.entry(script).or_default();
 
         let history_of_tx = history
             .iter()
-            .filter(|entry| &entry.tx_hash == txid)
+            .filter(|entry| entry.tx_hash == txid)
             .collect::<Vec<_>>();
 
         match history_of_tx.as_slice() {