From 98e779a102627add12a02176c90544c027afde09 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Thu, 1 Oct 2020 14:55:40 +0200
Subject: [PATCH] wallet: use constant input source for change calculation

To fix a bug where specifying multiple UTXOs that are by themselves
large enough to satisfy the output amount would lead to the rest of them
being added to fees, we need to provide the transaction author with a
constant list of UTXOs. If we didn't, the author would only consider one
input and calculate the change based on that alone. But since we'd add
all inputs to the PSBT, the rest of the amounts would go to fees.
---
 wallet/psbt.go      | 38 ++++++++++++++++++++++++++++++++++++--
 wallet/psbt_test.go | 20 ++++++++++++++------
 2 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/wallet/psbt.go b/wallet/psbt.go
index e2d4fe8..1a68a71 100644
--- a/wallet/psbt.go
+++ b/wallet/psbt.go
@@ -134,7 +134,14 @@ func (w *Wallet) FundPsbt(packet *psbt.Packet, account uint32,
 		}
 
 		// We can leverage the fee calculation of the txauthor package
-		// if we provide the selected UTXOs as a coin source.
+		// if we provide the selected UTXOs as a coin source. We just
+		// need to make sure we always return the full list of user-
+		// selected UTXOs rather than a subset, otherwise our change
+		// amount will be off (in case the user selected multiple UTXOs
+		// that are large enough on their own). That's why we use our
+		// own static input source creator instead of the more generic
+		// makeInputSource() that selects a subset that is "large
+		// enough".
 		credits := make([]wtxmgr.Credit, len(txIn))
 		for idx, in := range txIn {
 			utxo := packet.Inputs[idx].WitnessUtxo
@@ -144,7 +151,7 @@ func (w *Wallet) FundPsbt(packet *psbt.Packet, account uint32,
 				PkScript: utxo.PkScript,
 			}
 		}
-		inputSource := makeInputSource(credits)
+		inputSource := constantInputSource(credits)
 
 		// We also need a change source which needs to be able to insert
 		// a new change addresse into the database.
@@ -319,3 +326,30 @@ func (w *Wallet) FinalizePsbt(packet *psbt.Packet) error {
 
 	return nil
 }
+
+// constantInputSource creates an input source function that always returns the
+// static set of user-selected UTXOs.
+func constantInputSource(eligible []wtxmgr.Credit) txauthor.InputSource {
+	// Current inputs and their total value. These won't change over
+	// different invocations as we want our inputs to remain static since
+	// they're selected by the user.
+	currentTotal := btcutil.Amount(0)
+	currentInputs := make([]*wire.TxIn, 0, len(eligible))
+	currentScripts := make([][]byte, 0, len(eligible))
+	currentInputValues := make([]btcutil.Amount, 0, len(eligible))
+
+	for _, credit := range eligible {
+		nextInput := wire.NewTxIn(&credit.OutPoint, nil, nil)
+		currentTotal += credit.Amount
+		currentInputs = append(currentInputs, nextInput)
+		currentScripts = append(currentScripts, credit.PkScript)
+		currentInputValues = append(currentInputValues, credit.Amount)
+	}
+
+	return func(target btcutil.Amount) (btcutil.Amount, []*wire.TxIn,
+		[]btcutil.Amount, [][]byte, error) {
+
+		return currentTotal, currentInputs, currentInputValues,
+			currentScripts, nil
+	}
+}
diff --git a/wallet/psbt_test.go b/wallet/psbt_test.go
index dfb67cd..b9747b2 100644
--- a/wallet/psbt_test.go
+++ b/wallet/psbt_test.go
@@ -68,6 +68,8 @@ func TestFundPsbt(t *testing.T) {
 		feeRateSatPerKB   btcutil.Amount
 		expectedErr       string
 		validatePackage   bool
+		expectedFee       int64
+		expectedChange    int64
 		numExpectedInputs int
 	}{{
 		name: "no outputs provided",
@@ -106,6 +108,8 @@ func TestFundPsbt(t *testing.T) {
 		feeRateSatPerKB:   2000, // 2 sat/byte
 		expectedErr:       "",
 		validatePackage:   true,
+		expectedChange:    1000000 - 150000 - 368,
+		expectedFee:       368,
 		numExpectedInputs: 1,
 	}, {
 		name: "two outputs, two inputs",
@@ -136,6 +140,8 @@ func TestFundPsbt(t *testing.T) {
 		feeRateSatPerKB:   2000, // 2 sat/byte
 		expectedErr:       "",
 		validatePackage:   true,
+		expectedFee:       506,
+		expectedChange:    2000000 - 150000 - 506,
 		numExpectedInputs: 2,
 	}}
 
@@ -296,15 +302,17 @@ func TestFundPsbt(t *testing.T) {
 					txOuts[p2wshIndex].PkScript)
 			}
 
-			// Finally, check the change output size and that it
-			// belongs to the wallet.
-			expectedFee := int64(368)
-			expectedChange := 1000000 - 150000 - expectedFee
-			if txOuts[changeIndex].Value != expectedChange {
+			// Finally, check the change output size and fee.
+			fee := totalIn - totalOut
+			if fee != tc.expectedFee {
+				t.Fatalf("unexpected fee, got %d wanted %d",
+					fee, tc.expectedFee)
+			}
+			if txOuts[changeIndex].Value != tc.expectedChange {
 				t.Fatalf("unexpected change output size, got "+
 					"%d wanted %d",
 					txOuts[changeIndex].Value,
-					expectedChange)
+					tc.expectedChange)
 			}
 		})
 	}