Merge pull request #842 from Torakushi/fund-psbt

wallet: add ability to use custom change scope for FundPsbt

wallet/psbt.go

@@ -23,8 +23,10 @@ import (
 // FundPsbt creates a fully populated PSBT packet that contains enough inputs to
 // fund the outputs specified in the passed in packet with the specified fee
 // rate. If there is change left, a change output from the wallet is added and
-// the index of the change output is returned. Otherwise no additional output
-// is created and the index -1 is returned.
+// the index of the change output is returned. If no custom change scope is
+// specified, we will use the coin selection scope (if not nil) or the BIP0086
+// scope by default. Otherwise, no additional output is created and the
+// index -1 is returned.
 //
 // NOTE: If the packet doesn't contain any inputs, coin selection is performed
 // automatically, only selecting inputs from the account based on the given key
@@ -43,7 +45,8 @@ import (
 // responsibility to lock the inputs before handing the partial transaction out.
 func (w *Wallet) FundPsbt(packet *psbt.Packet, keyScope *waddrmgr.KeyScope,
 	minConfs int32, account uint32, feeSatPerKB btcutil.Amount,
-	coinSelectionStrategy CoinSelectionStrategy) (int32, error) {
+	coinSelectionStrategy CoinSelectionStrategy,
+	optFuncs ...TxCreateOption) (int32, error) {
 
 	// Make sure the packet is well formed. We only require there to be at
 	// least one input or output.
@@ -131,6 +134,7 @@ func (w *Wallet) FundPsbt(packet *psbt.Packet, keyScope *waddrmgr.KeyScope,
 		tx, err = w.CreateSimpleTx(
 			keyScope, account, packet.UnsignedTx.TxOut, minConfs,
 			feeSatPerKB, coinSelectionStrategy, false,
+			optFuncs...,
 		)
 		if err != nil {
 			return 0, fmt.Errorf("error creating funding TX: %v",
@@ -176,11 +180,21 @@ func (w *Wallet) FundPsbt(packet *psbt.Packet, keyScope *waddrmgr.KeyScope,
 		}
 		inputSource := constantInputSource(credits)
 
+		// Build the TxCreateOption to retrieve the change scope.
+		opts := defaultTxCreateOptions()
+		for _, optFunc := range optFuncs {
+			optFunc(opts)
+		}
+
+		if opts.changeKeyScope == nil {
+			opts.changeKeyScope = keyScope
+		}
+
 		// We also need a change source which needs to be able to insert
 		// a new change address into the database.
 		err = walletdb.Update(w.db, func(dbtx walletdb.ReadWriteTx) error {
 			_, changeSource, err := w.addrMgrWithChangeSource(
-				dbtx, keyScope, account,
+				dbtx, opts.changeKeyScope, account,
 			)
 			if err != nil {
 				return err

wallet/psbt_test.go

@@ -76,6 +76,7 @@ func TestFundPsbt(t *testing.T) {
 		name                    string
 		packet                  *psbt.Packet
 		feeRateSatPerKB         btcutil.Amount
+		changeKeyScope          *waddrmgr.KeyScope
 		expectedErr             string
 		validatePackage         bool
 		expectedChangeBeforeFee int64
@@ -217,6 +218,41 @@ func TestFundPsbt(t *testing.T) {
 				"index after sorting",
 			)
 		},
+	}, {
+		name: "one input and a custom change scope: BIP0084",
+		packet: &psbt.Packet{
+			UnsignedTx: &wire.MsgTx{
+				TxIn: []*wire.TxIn{{
+					PreviousOutPoint: utxo1,
+				}},
+			},
+			Inputs: []psbt.PInput{{}},
+		},
+		feeRateSatPerKB:         20000,
+		validatePackage:         true,
+		changeKeyScope:          &waddrmgr.KeyScopeBIP0084,
+		expectedInputs:          []wire.OutPoint{utxo1},
+		expectedChangeBeforeFee: utxo1Amount,
+	}, {
+		name: "no inputs and a custom change scope: BIP0084",
+		packet: &psbt.Packet{
+			UnsignedTx: &wire.MsgTx{
+				TxOut: []*wire.TxOut{{
+					PkScript: testScriptP2WSH,
+					Value:    100000,
+				}, {
+					PkScript: testScriptP2WKH,
+					Value:    50000,
+				}},
+			},
+			Outputs: []psbt.POutput{{}, {}},
+		},
+		feeRateSatPerKB:         2000, // 2 sat/byte
+		expectedErr:             "",
+		validatePackage:         true,
+		changeKeyScope:          &waddrmgr.KeyScopeBIP0084,
+		expectedChangeBeforeFee: utxo1Amount - 150000,
+		expectedInputs:          []wire.OutPoint{utxo1},
 	}}
 
 	calcFee := func(feeRateSatPerKB btcutil.Amount,
@@ -244,8 +280,9 @@ func TestFundPsbt(t *testing.T) {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			changeIndex, err := w.FundPsbt(
-				tc.packet, nil, 1, 0, tc.feeRateSatPerKB,
-				CoinSelectionLargest,
+				tc.packet, nil, 1, 0,
+				tc.feeRateSatPerKB, CoinSelectionLargest,
+				WithCustomChangeScope(tc.changeKeyScope),
 			)
 
 			// In any case, unlock the UTXO before continuing, we
@@ -306,6 +343,10 @@ func TestFundPsbt(t *testing.T) {
 			// to a change output.
 			require.EqualValues(t, 1, b32d.Bip32Path[3])
 
+			assertChangeOutputScope(
+				t, changeTxOut.PkScript, tc.changeKeyScope,
+			)
+
 			if txscript.IsPayToTaproot(changeTxOut.PkScript) {
 				require.NotEmpty(
 					t, changeOutput.TaprootInternalKey,
@@ -354,6 +395,28 @@ func assertTxInputs(t *testing.T, packet *psbt.Packet,
 	}
 }
 
+// assertChangeOutputScope checks if the pkScript has the right type.
+func assertChangeOutputScope(t *testing.T, pkScript []byte,
+	changeScope *waddrmgr.KeyScope) {
+
+	// By default (changeScope == nil), the script should
+	// be a pay-to-taproot one.
+	switch changeScope {
+	case nil, &waddrmgr.KeyScopeBIP0086:
+		require.True(t, txscript.IsPayToTaproot(pkScript))
+
+	case &waddrmgr.KeyScopeBIP0049Plus, &waddrmgr.KeyScopeBIP0084:
+		require.True(t, txscript.IsPayToWitnessPubKeyHash(pkScript))
+
+	case &waddrmgr.KeyScopeBIP0044:
+		require.True(t, txscript.IsPayToPubKeyHash(pkScript))
+
+	default:
+		require.Fail(t, "assertChangeOutputScope error",
+			"change scope: %s", changeScope.String())
+	}
+}
+
 func containsUtxo(list []wire.OutPoint, candidate wire.OutPoint) bool {
 	for _, utxo := range list {
 		if utxo == candidate {