From 7eae0271b0e4435e374e2aa8b92ff6fe2863ea96 Mon Sep 17 00:00:00 2001
From: Lucas Soriano del Pino <l.soriano.del.pino@gmail.com>
Date: Tue, 21 Sep 2021 16:13:27 +1000
Subject: [PATCH] Add TxInExt trait with find_map_signature method

---
 cfd_protocol/src/protocol.rs              |  1 +
 cfd_protocol/src/protocol/transactions.rs | 11 +++--------
 cfd_protocol/src/protocol/txin_ext.rs     | 23 +++++++++++++++++++++++
 3 files changed, 27 insertions(+), 8 deletions(-)
 create mode 100644 cfd_protocol/src/protocol/txin_ext.rs

diff --git a/cfd_protocol/src/protocol.rs b/cfd_protocol/src/protocol.rs
index c1f8b2b..b65b174 100644
--- a/cfd_protocol/src/protocol.rs
+++ b/cfd_protocol/src/protocol.rs
@@ -27,6 +27,7 @@ use std::ops::RangeInclusive;
 mod sighash_ext;
 mod transaction_ext;
 mod transactions;
+mod txin_ext;
 
 /// Static script to be used to create lock tx
 const DUMMY_2OF2_MULTISIG: &str =
diff --git a/cfd_protocol/src/protocol/transactions.rs b/cfd_protocol/src/protocol/transactions.rs
index b0f1c87..d371672 100644
--- a/cfd_protocol/src/protocol/transactions.rs
+++ b/cfd_protocol/src/protocol/transactions.rs
@@ -1,5 +1,6 @@
 use crate::protocol::sighash_ext::SigHashExt;
 use crate::protocol::transaction_ext::TransactionExt;
+use crate::protocol::txin_ext::TxInExt;
 use crate::protocol::{
     commit_descriptor, compute_adaptor_point, lock_descriptor, Payout, DUMMY_2OF2_MULTISIG,
 };
@@ -12,7 +13,7 @@ use bdk::bitcoin::{
 use bdk::descriptor::Descriptor;
 use bdk::miniscript::DescriptorTrait;
 use itertools::Itertools;
-use secp256k1_zkp::{self, schnorrsig, EcdsaAdaptorSignature, SecretKey, Signature, SECP256K1};
+use secp256k1_zkp::{self, schnorrsig, EcdsaAdaptorSignature, SecretKey, SECP256K1};
 use std::collections::HashMap;
 use std::iter::FromIterator;
 
@@ -366,13 +367,7 @@ pub fn punish_transaction(
         .context("commit transaction inputs != 1")?;
 
     let publish_them_sk = input
-        .witness
-        .iter()
-        .filter_map(|elem| {
-            let elem = elem.as_slice();
-            Signature::from_der(&elem[..elem.len() - 1]).ok()
-        })
-        .find_map(|sig| encsig.recover(SECP256K1, &sig, &pub_them_pk.key).ok())
+        .find_map_signature(|sig| encsig.recover(SECP256K1, &sig, &pub_them_pk.key).ok())
         .context("could not recover publish sk from commit tx")?;
 
     let commit_outpoint = revoked_commit_tx
diff --git a/cfd_protocol/src/protocol/txin_ext.rs b/cfd_protocol/src/protocol/txin_ext.rs
new file mode 100644
index 0000000..cc73136
--- /dev/null
+++ b/cfd_protocol/src/protocol/txin_ext.rs
@@ -0,0 +1,23 @@
+use bdk::bitcoin::secp256k1::Signature;
+use bdk::bitcoin::TxIn;
+
+pub(crate) trait TxInExt {
+    fn find_map_signature<F, R>(&self, f: F) -> Option<R>
+    where
+        F: Fn(Signature) -> Option<R>;
+}
+
+impl TxInExt for TxIn {
+    fn find_map_signature<F, R>(&self, f: F) -> Option<R>
+    where
+        F: Fn(Signature) -> Option<R>,
+    {
+        self.witness
+            .iter()
+            .filter_map(|elem| {
+                let elem = elem.as_slice();
+                Signature::from_der(&elem[..elem.len() - 1]).ok()
+            })
+            .find_map(f)
+    }
+}