From 783bae0f857510133815b2c80cd2c87b13dc28ba Mon Sep 17 00:00:00 2001
From: patrini32 <171664803+patrini32@users.noreply.github.com>
Date: Sat, 8 Jun 2024 21:32:41 +0300
Subject: [PATCH] Revert "Reimplement tx_punish check without breaking database
 changes"

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

diff --git a/swap/src/cli/cancel_and_refund.rs b/swap/src/cli/cancel_and_refund.rs
index 6dbff4cb..b6095664 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, State};
+use crate::protocol::Database;
 use anyhow::{bail, Result};
 use bitcoin::Txid;
 use std::sync::Arc;
@@ -123,19 +123,7 @@ pub async fn refund(
             Ok(state)
         }
         Err(error) => {
-            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.")
+            if state6
                 .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 31417893..38871dd0 100644
--- a/swap/src/protocol/bob/state.rs
+++ b/swap/src/protocol/bob/state.rs
@@ -436,6 +436,7 @@ 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(),
@@ -471,26 +472,6 @@ 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)]
@@ -605,6 +586,7 @@ 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,
@@ -646,6 +628,7 @@ 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,
@@ -698,6 +681,27 @@ 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,