Implement swap recover function for Alice

This introduces a lot of duplication between the binary and the
library, but it's okay because this module should only be a temporary
measure until we allow recovery to be handled by the original state
machine.

Also, fix a bug in `xmr_btc::alice::action_generator` caused by the
incorrect assumption that Alice's ability to punish Bob could be
determined before the cancel transaction hits the blockchain.

swap/src/lib.rs

@@ -6,6 +6,7 @@ pub mod bitcoin;
 pub mod bob;
 pub mod monero;
 pub mod network;
+pub mod recover;
 pub mod state;
 pub mod storage;
 pub mod tor;

swap/src/main.rs

@@ -16,7 +16,6 @@ use anyhow::Result;
 use futures::{channel::mpsc, StreamExt};
 use libp2p::Multiaddr;
 use log::LevelFilter;
-use prettytable::{row, Table};
 use std::{io, io::Write, process, sync::Arc};
 use structopt::StructOpt;
 use swap::{
@@ -31,13 +30,11 @@ use swap::{
 use tempfile::tempdir;
 use tracing::info;
 
-#[macro_use]
-extern crate prettytable;
-
 mod cli;
 mod trace;
 
 use cli::Options;
+use swap::storage::Database;
 
 // TODO: Add root seed file instead of generating new seed each run.
 
@@ -47,6 +44,9 @@ async fn main() -> Result<()> {
 
     trace::init_tracing(LevelFilter::Debug)?;
 
+    let db_dir = tempdir()?;
+    let db = Database::open(db_dir.path()).unwrap();
+
     match opt {
         Options::Alice {
             bitcoind_url,
@@ -85,7 +85,11 @@ async fn main() -> Result<()> {
 
             let monero_wallet = Arc::new(monero::Wallet::new(monerod_url));
 
-            swap_as_alice(bitcoin_wallet, monero_wallet, dblisten_addr,
+            swap_as_alice(
+                bitcoin_wallet,
+                monero_wallet,
+                db,
+                listen_addr,
                 transport,
                 behaviour,
             )
@@ -115,12 +119,10 @@ async fn main() -> Result<()> {
 
             let monero_wallet = Arc::new(monero::Wallet::new(monerod_url));
 
-            let db = Database::open(db_dir.path()).unwrap();
-
             swap_as_bob(
                 bitcoin_wallet,
                 monero_wallet,
-                db
+                db,
                 satoshis,
                 alice_addr,
                 transport,
@@ -159,7 +161,15 @@ async fn swap_as_alice(
     transport: SwapTransport,
     behaviour: Alice,
 ) -> Result<()> {
-    alice::swap(bitcoin_wallet, monero_wallet, addr, transport, behaviour).await
+    alice::swap(
+        bitcoin_wallet,
+        monero_wallet,
+        db,
+        addr,
+        transport,
+        behaviour,
+    )
+    .await
 }
 
 async fn swap_as_bob(

swap/src/recover.rs

@@ -0,0 +1,187 @@
+use crate::{
+    monero::CreateWalletForOutput,
+    state::{Alice, Bob, Swap},
+};
+use anyhow::Result;
+use futures::{
+    future::{select, Either},
+    pin_mut,
+};
+use xmr_btc::bitcoin::{
+    poll_until_block_height_is_gte, BroadcastSignedTransaction, TransactionBlockHeight, TxCancel,
+    TxPunish, TxRefund, WatchForRawTransaction,
+};
+
+pub async fn recover(
+    bitcoin_wallet: crate::bitcoin::Wallet,
+    monero_wallet: crate::monero::Wallet,
+    state: Swap,
+) -> Result<()> {
+    match state {
+        Swap::Alice(state) => alice_recover(bitcoin_wallet, monero_wallet, state).await,
+        Swap::Bob(state) => bob_recover(bitcoin_wallet, monero_wallet, state).await,
+    }
+}
+
+pub async fn alice_recover(
+    bitcoin_wallet: crate::bitcoin::Wallet,
+    monero_wallet: crate::monero::Wallet,
+    state: Alice,
+) -> Result<()> {
+    match state {
+        Alice::Handshaken(_) | Alice::BtcLocked(_) | Alice::SwapComplete => Ok(()),
+        Alice::XmrLocked(state) => {
+            let tx_cancel = TxCancel::new(
+                &state.tx_lock,
+                state.refund_timelock,
+                state.a.public(),
+                state.B.clone(),
+            );
+
+            // Ensure that TxCancel is on the blockchain
+            if bitcoin_wallet
+                .0
+                .get_raw_transaction(tx_cancel.txid())
+                .await
+                .is_err()
+            {
+                let tx_lock_height = bitcoin_wallet
+                    .transaction_block_height(state.tx_lock.txid())
+                    .await;
+                poll_until_block_height_is_gte(
+                    &bitcoin_wallet,
+                    tx_lock_height + state.refund_timelock,
+                )
+                .await;
+
+                let sig_a = state.a.sign(tx_cancel.digest());
+                let sig_b = state.tx_cancel_sig_bob.clone();
+
+                let tx_cancel = tx_cancel
+                    .clone()
+                    .add_signatures(
+                        &state.tx_lock,
+                        (state.a.public(), sig_a),
+                        (state.B.clone(), sig_b),
+                    )
+                    .expect("sig_{a,b} to be valid signatures for tx_cancel");
+
+                // TODO: We should not fail if the transaction is already on the blockchain
+                bitcoin_wallet
+                    .broadcast_signed_transaction(tx_cancel)
+                    .await?;
+            }
+
+            let tx_cancel_height = bitcoin_wallet
+                .transaction_block_height(tx_cancel.txid())
+                .await;
+            let poll_until_bob_can_be_punished = poll_until_block_height_is_gte(
+                &bitcoin_wallet,
+                tx_cancel_height + state.punish_timelock,
+            );
+            pin_mut!(poll_until_bob_can_be_punished);
+
+            let tx_refund = TxRefund::new(&tx_cancel, &state.refund_address);
+
+            match select(
+                bitcoin_wallet.watch_for_raw_transaction(tx_refund.txid()),
+                poll_until_bob_can_be_punished,
+            )
+            .await
+            {
+                Either::Left((tx_refund_published, ..)) => {
+                    let tx_refund_sig = tx_refund
+                        .extract_signature_by_key(tx_refund_published, state.a.public())?;
+                    let tx_refund_encsig = state
+                        .a
+                        .encsign(state.S_b_bitcoin.clone(), tx_refund.digest());
+
+                    let s_b = xmr_btc::bitcoin::recover(
+                        state.S_b_bitcoin,
+                        tx_refund_sig,
+                        tx_refund_encsig,
+                    )?;
+                    let s_b = monero::PrivateKey::from_scalar(
+                        xmr_btc::monero::Scalar::from_bytes_mod_order(s_b.to_bytes()),
+                    );
+
+                    let s_a = monero::PrivateKey {
+                        scalar: state.s_a.into_ed25519(),
+                    };
+
+                    monero_wallet
+                        .create_and_load_wallet_for_output(s_a + s_b, state.v)
+                        .await?;
+                }
+                Either::Right(_) => {
+                    let tx_punish =
+                        TxPunish::new(&tx_cancel, &state.punish_address, state.punish_timelock);
+
+                    let sig_a = state.a.sign(tx_punish.digest());
+                    let sig_b = state.tx_cancel_sig_bob.clone();
+
+                    let sig_tx_punish = tx_punish.add_signatures(
+                        &tx_cancel,
+                        (state.a.public(), sig_a),
+                        (state.B.clone(), sig_b),
+                    )?;
+
+                    bitcoin_wallet
+                        .broadcast_signed_transaction(sig_tx_punish)
+                        .await?;
+                }
+            };
+
+            Ok(())
+        }
+        Alice::BtcRedeemable { redeem_tx, .. } => {
+            bitcoin_wallet
+                .broadcast_signed_transaction(redeem_tx)
+                .await?;
+            Ok(())
+        }
+        Alice::BtcPunishable(state) => {
+            let tx_cancel = TxCancel::new(
+                &state.tx_lock,
+                state.refund_timelock,
+                state.a.public(),
+                state.B.clone(),
+            );
+
+            let tx_punish = TxPunish::new(&tx_cancel, &state.punish_address, state.punish_timelock);
+
+            let sig_a = state.a.sign(tx_punish.digest());
+            let sig_b = state.tx_cancel_sig_bob.clone();
+
+            let sig_tx_punish = tx_punish.add_signatures(
+                &tx_cancel,
+                (state.a.public(), sig_a),
+                (state.B.clone(), sig_b),
+            )?;
+
+            bitcoin_wallet
+                .broadcast_signed_transaction(sig_tx_punish)
+                .await?;
+
+            Ok(())
+        }
+        Alice::BtcRefunded {
+            view_key,
+            spend_key,
+            ..
+        } => {
+            monero_wallet
+                .create_and_load_wallet_for_output(spend_key, view_key)
+                .await?;
+            Ok(())
+        }
+    }
+}
+
+pub async fn bob_recover(
+    _bitcoin_wallet: crate::bitcoin::Wallet,
+    _monero_wallet: crate::monero::Wallet,
+    _state: Bob,
+) -> Result<()> {
+    todo!()
+}

swap/src/storage.rs

@@ -33,7 +33,7 @@ impl Database {
             .context("Could not flush db")
     }
 
-    pub fn get_latest_state(&self, swap_id: Uuid) -> anyhow::Result<Swap> {
+    pub fn get_state(&self, swap_id: Uuid) -> anyhow::Result<Swap> {
         let key = serialize(&swap_id)?;
 
         let encoded = self
@@ -103,11 +103,11 @@ mod tests {
             .expect("Failed to save first state");
 
         let recovered_1 = db
-            .get_latest_state(swap_id_1)
+            .get_state(swap_id_1)
             .expect("Failed to recover first state");
 
         let recovered_2 = db
-            .get_latest_state(swap_id_2)
+            .get_state(swap_id_2)
             .expect("Failed to recover second state");
 
         assert_eq!(recovered_1, state_1);
@@ -126,7 +126,7 @@ mod tests {
             .await
             .expect("Failed to save state the first time");
         let recovered = db
-            .get_latest_state(swap_id)
+            .get_state(swap_id)
             .expect("Failed to recover state the first time");
 
         // We insert and recover twice to ensure database implementation allows the
@@ -135,7 +135,7 @@ mod tests {
             .await
             .expect("Failed to save state the second time");
         let recovered = db
-            .get_latest_state(swap_id)
+            .get_state(swap_id)
             .expect("Failed to recover state the second time");
 
         assert_eq!(recovered, state);