From a09c414484de22f3660a35ad78a97419c6bb6256 Mon Sep 17 00:00:00 2001
From: patrini32 <171664803+patrini32@users.noreply.github.com>
Date: Sat, 8 Jun 2024 20:33:59 +0300
Subject: [PATCH] Reimplement tx_punish check without breaking database changes

---
 swap/src/cli/cancel_and_refund.rs | 16 +++++++++--
 swap/src/protocol/bob/state.rs    | 44 ++++++++++++++-----------------
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/swap/src/cli/cancel_and_refund.rs b/swap/src/cli/cancel_and_refund.rs
index b6095664..6dbff4cb 100644
--- a/swap/src/cli/cancel_and_refund.rs
+++ b/swap/src/cli/cancel_and_refund.rs
@@ -1,6 +1,6 @@
 use crate::bitcoin::{parse_rpc_error_code, RpcErrorCode, Wallet};
 use crate::protocol::bob::BobState;
-use crate::protocol::Database;
+use crate::protocol::{Database, State};
 use anyhow::{bail, Result};
 use bitcoin::Txid;
 use std::sync::Arc;
@@ -123,7 +123,19 @@ pub async fn refund(
             Ok(state)
         }
         Err(error) => {
-            if state6
+            let state3 = db
+            .get_states(swap_id)
+            .await?
+            .iter()
+            .find_map(|state| {
+                if let State::Bob(BobState::BtcLocked { state3, .. }) = state {
+                    Some(state3.clone())
+                    } else {
+                        None
+                    }
+            });
+            if state3
+                .expect("Error: weren't able to extract State3 from database.")
                 .check_for_tx_punish(bitcoin_wallet.as_ref())
                 .await
                 .is_ok()
diff --git a/swap/src/protocol/bob/state.rs b/swap/src/protocol/bob/state.rs
index 38871dd0..31417893 100644
--- a/swap/src/protocol/bob/state.rs
+++ b/swap/src/protocol/bob/state.rs
@@ -436,7 +436,6 @@ impl State3 {
             cancel_timelock: self.cancel_timelock,
             punish_timelock: self.punish_timelock,
             refund_address: self.refund_address.clone(),
-            punish_address: self.punish_address.clone(),
             tx_lock: self.tx_lock.clone(),
             tx_cancel_sig_a: self.tx_cancel_sig_a.clone(),
             tx_refund_encsig: self.tx_refund_encsig.clone(),
@@ -472,6 +471,26 @@ impl State3 {
             tx_cancel_status,
         ))
     }
+    pub async fn check_for_tx_punish(
+        &self,
+        bitcoin_wallet: &bitcoin::Wallet,
+    ) -> Result<Transaction> {
+        let tx_cancel = bitcoin::TxCancel::new(
+            &self.tx_lock,
+            self.cancel_timelock,
+            self.A,
+            self.b.public(),
+            self.tx_cancel_fee,
+        )?;
+        let tx_punish = bitcoin::TxPunish::new(
+            &tx_cancel,
+            &self.punish_address,
+            self.punish_timelock,
+            self.tx_punish_fee,
+        );
+        let tx = bitcoin_wallet.get_raw_transaction(tx_punish.txid()).await?;
+        Ok(tx)
+    }
 }
 
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
@@ -586,7 +605,6 @@ impl State4 {
             cancel_timelock: self.cancel_timelock,
             punish_timelock: self.punish_timelock,
             refund_address: self.refund_address,
-            punish_address: self.punish_address,
             tx_lock: self.tx_lock,
             tx_cancel_sig_a: self.tx_cancel_sig_a,
             tx_refund_encsig: self.tx_refund_encsig,
@@ -628,7 +646,6 @@ pub struct State6 {
     cancel_timelock: CancelTimelock,
     punish_timelock: PunishTimelock,
     refund_address: bitcoin::Address,
-    punish_address: bitcoin::Address,
     tx_lock: bitcoin::TxLock,
     tx_cancel_sig_a: Signature,
     tx_refund_encsig: bitcoin::EncryptedSignature,
@@ -681,27 +698,6 @@ impl State6 {
         Ok(tx)
     }
 
-    pub async fn check_for_tx_punish(
-        &self,
-        bitcoin_wallet: &bitcoin::Wallet,
-    ) -> Result<Transaction> {
-        let tx_cancel = bitcoin::TxCancel::new(
-            &self.tx_lock,
-            self.cancel_timelock,
-            self.A,
-            self.b.public(),
-            self.tx_cancel_fee,
-        )?;
-        let tx_punish = bitcoin::TxPunish::new(
-            &tx_cancel,
-            &self.punish_address,
-            self.punish_timelock,
-            self.tx_punish_fee,
-        );
-        let tx = bitcoin_wallet.get_raw_transaction(tx_punish.txid()).await?;
-        Ok(tx)
-    }
-
     pub async fn submit_tx_cancel(
         &self,
         bitcoin_wallet: &bitcoin::Wallet,