From b34531d2d3a4952a5fa76843680e2c037203956c Mon Sep 17 00:00:00 2001
From: KonradStaniec <konrad.staniec@gmail.com>
Date: Wed, 16 Oct 2024 13:27:53 +0200
Subject: [PATCH] Add filtering of utxos used in transaction building

---
 itest/e2e_test.go             |  4 ++++
 staker/stakerapp.go           | 21 ++++++++++++++++++++-
 walletcontroller/client.go    | 22 ++++++++++++++++++----
 walletcontroller/interface.go | 12 +++++++++++-
 4 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/itest/e2e_test.go b/itest/e2e_test.go
index 3b9fc3e..4ffec84 100644
--- a/itest/e2e_test.go
+++ b/itest/e2e_test.go
@@ -544,6 +544,7 @@ func (tm *TestManager) FinalizeUntilEpoch(t *testing.T, epoch uint64) {
 			},
 			2000,
 			tm.MinerAddr,
+			nil,
 		)
 		require.NoError(t, err)
 		_, err = tm.Sa.Wallet().SendRawTransaction(tx1, true)
@@ -557,6 +558,7 @@ func (tm *TestManager) FinalizeUntilEpoch(t *testing.T, epoch uint64) {
 			},
 			2000,
 			tm.MinerAddr,
+			nil,
 		)
 		require.NoError(t, err)
 		_, err = tm.Sa.Wallet().SendRawTransaction(tx2, true)
@@ -804,6 +806,7 @@ func (tm *TestManager) sendWatchedStakingTx(
 		[]*wire.TxOut{stakingInfo.StakingOutput},
 		2000,
 		tm.MinerAddr,
+		nil,
 	)
 	require.NoError(t, err)
 	txHash := tx.TxHash()
@@ -1677,6 +1680,7 @@ func TestBitcoindWalletRpcApi(t *testing.T) {
 		[]*wire.TxOut{newOutput},
 		btcutil.Amount(2000),
 		walletAddress,
+		nil,
 	)
 	require.NoError(t, err)
 
diff --git a/staker/stakerapp.go b/staker/stakerapp.go
index ba729e0..d39b8fc 100644
--- a/staker/stakerapp.go
+++ b/staker/stakerapp.go
@@ -1705,6 +1705,20 @@ func (app *StakerApp) WatchStaking(
 	}
 }
 
+func (app *StakerApp) filteUtxoFnGen() walletcontroller.UseUtxoFn {
+	return func(utxo walletcontroller.Utxo) bool {
+		outpoint := utxo.OutPoint
+
+		used, err := app.txTracker.OutpointUsed(&outpoint)
+
+		if err != nil {
+			return false
+		}
+
+		return !used
+	}
+}
+
 func (app *StakerApp) StakeFunds(
 	stakerAddress btcutil.Address,
 	stakingAmount btcutil.Amount,
@@ -1809,7 +1823,12 @@ func (app *StakerApp) StakeFunds(
 
 	// Create unsigned transaction by wallet without signing. Signing will happen
 	// in next steps
-	tx, err := app.wc.CreateTransaction([]*wire.TxOut{stakingInfo.StakingOutput}, btcutil.Amount(feeRate), stakerAddress)
+	tx, err := app.wc.CreateTransaction(
+		[]*wire.TxOut{stakingInfo.StakingOutput},
+		btcutil.Amount(feeRate),
+		stakerAddress,
+		app.filteUtxoFnGen(),
+	)
 
 	if err != nil {
 		return nil, err
diff --git a/walletcontroller/client.go b/walletcontroller/client.go
index 65770f0..ce48e1a 100644
--- a/walletcontroller/client.go
+++ b/walletcontroller/client.go
@@ -140,7 +140,9 @@ func (w *RpcWalletController) NetworkName() string {
 func (w *RpcWalletController) CreateTransaction(
 	outputs []*wire.TxOut,
 	feeRatePerKb btcutil.Amount,
-	changeAddres btcutil.Address) (*wire.MsgTx, error) {
+	changeAddres btcutil.Address,
+	useUtxoFn UseUtxoFn,
+) (*wire.MsgTx, error) {
 
 	utxoResults, err := w.ListUnspent()
 
@@ -154,9 +156,20 @@ func (w *RpcWalletController) CreateTransaction(
 		return nil, err
 	}
 
+	var utxosToUse []Utxo
+	if useUtxoFn != nil {
+		for _, u := range utxos {
+			if useUtxoFn(u) {
+				utxosToUse = append(utxosToUse, u)
+			}
+		}
+	} else {
+		utxosToUse = utxos
+	}
+
 	// sort utxos by amount from highest to lowest, this is effectively strategy of using
 	// largest inputs first
-	sort.Sort(sort.Reverse(byAmount(utxos)))
+	sort.Sort(sort.Reverse(byAmount(utxosToUse)))
 
 	changeScript, err := txscript.PayToAddrScript(changeAddres)
 
@@ -164,7 +177,7 @@ func (w *RpcWalletController) CreateTransaction(
 		return nil, err
 	}
 
-	tx, err := buildTxFromOutputs(utxos, outputs, feeRatePerKb, changeScript)
+	tx, err := buildTxFromOutputs(utxosToUse, outputs, feeRatePerKb, changeScript)
 
 	if err != nil {
 		return nil, err
@@ -177,8 +190,9 @@ func (w *RpcWalletController) CreateAndSignTx(
 	outputs []*wire.TxOut,
 	feeRatePerKb btcutil.Amount,
 	changeAddress btcutil.Address,
+	useUtxoFn UseUtxoFn,
 ) (*wire.MsgTx, error) {
-	tx, err := w.CreateTransaction(outputs, feeRatePerKb, changeAddress)
+	tx, err := w.CreateTransaction(outputs, feeRatePerKb, changeAddress, useUtxoFn)
 
 	if err != nil {
 		return nil, err
diff --git a/walletcontroller/interface.go b/walletcontroller/interface.go
index 2a776ba..d0e2b2d 100644
--- a/walletcontroller/interface.go
+++ b/walletcontroller/interface.go
@@ -37,21 +37,31 @@ type TaprootSigningResult struct {
 	FullInputWitness wire.TxWitness
 }
 
+// Function to filer utxos that should be used in transaction creation
+type UseUtxoFn func(utxo Utxo) bool
+
 type WalletController interface {
 	UnlockWallet(timeoutSecs int64) error
 	AddressPublicKey(address btcutil.Address) (*btcec.PublicKey, error)
 	ImportPrivKey(privKeyWIF *btcutil.WIF) error
 	NetworkName() string
+	// passning nil usedUtxoFilter will use all possible spendable utxos to choose
+	// inputs
 	CreateTransaction(
 		outputs []*wire.TxOut,
 		feeRatePerKb btcutil.Amount,
-		changeScript btcutil.Address) (*wire.MsgTx, error)
+		changeScript btcutil.Address,
+		usedUtxoFilter UseUtxoFn,
+	) (*wire.MsgTx, error)
 	SignRawTransaction(tx *wire.MsgTx) (*wire.MsgTx, bool, error)
 	// requires wallet to be unlocked
+	// passning nil usedUtxoFilter will use all possible spendable utxos to choose
+	// inputs
 	CreateAndSignTx(
 		outputs []*wire.TxOut,
 		feeRatePerKb btcutil.Amount,
 		changeAddress btcutil.Address,
+		usedUtxoFilter UseUtxoFn,
 	) (*wire.MsgTx, error)
 	SendRawTransaction(tx *wire.MsgTx, allowHighFees bool) (*chainhash.Hash, error)
 	ListOutputs(onlySpendable bool) ([]Utxo, error)