wallet: add ComputeInputScript

wallet/signer.go

@@ -0,0 +1,102 @@
+// Copyright (c) 2020 The btcsuite developers
+// Use of this source code is governed by an ISC
+// license that can be found in the LICENSE file.
+
+package wallet
+
+import (
+	"github.com/btcsuite/btcd/btcec"
+	"github.com/btcsuite/btcd/txscript"
+	"github.com/btcsuite/btcd/wire"
+	"github.com/btcsuite/btcutil"
+	"github.com/btcsuite/btcwallet/waddrmgr"
+)
+
+// PrivKeyTweaker is a function type that can be used to pass in a callback for
+// tweaking a private key before it's used to sign an input.
+type PrivKeyTweaker func(*btcec.PrivateKey) (*btcec.PrivateKey, error)
+
+// ComputeInputScript generates a complete InputScript for the passed
+// transaction with the signature as defined within the passed SignDescriptor.
+// This method is capable of generating the proper input script for both
+// regular p2wkh output and p2wkh outputs nested within a regular p2sh output.
+func (w *Wallet) ComputeInputScript(tx *wire.MsgTx, output *wire.TxOut,
+	inputIndex int, sigHashes *txscript.TxSigHashes,
+	hashType txscript.SigHashType, tweaker PrivKeyTweaker) (wire.TxWitness,
+	[]byte, error) {
+
+	// First make sure we can sign for the input by making sure the script
+	// in the UTXO belongs to our wallet and we have the private key for it.
+	walletAddr, err := w.fetchOutputAddr(output.PkScript)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	pka := walletAddr.(waddrmgr.ManagedPubKeyAddress)
+	privKey, err := pka.PrivKey()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	var (
+		witnessProgram []byte
+		sigScript      []byte
+	)
+
+	switch {
+	// If we're spending p2wkh output nested within a p2sh output, then
+	// we'll need to attach a sigScript in addition to witness data.
+	case pka.AddrType() == waddrmgr.NestedWitnessPubKey:
+		pubKey := privKey.PubKey()
+		pubKeyHash := btcutil.Hash160(pubKey.SerializeCompressed())
+
+		// Next, we'll generate a valid sigScript that will allow us to
+		// spend the p2sh output. The sigScript will contain only a
+		// single push of the p2wkh witness program corresponding to
+		// the matching public key of this address.
+		p2wkhAddr, err := btcutil.NewAddressWitnessPubKeyHash(
+			pubKeyHash, w.chainParams,
+		)
+		if err != nil {
+			return nil, nil, err
+		}
+		witnessProgram, err = txscript.PayToAddrScript(p2wkhAddr)
+		if err != nil {
+			return nil, nil, err
+		}
+
+		bldr := txscript.NewScriptBuilder()
+		bldr.AddData(witnessProgram)
+		sigScript, err = bldr.Script()
+		if err != nil {
+			return nil, nil, err
+		}
+
+	// Otherwise, this is a regular p2wkh output, so we include the
+	// witness program itself as the subscript to generate the proper
+	// sighash digest. As part of the new sighash digest algorithm, the
+	// p2wkh witness program will be expanded into a regular p2kh
+	// script.
+	default:
+		witnessProgram = output.PkScript
+	}
+
+	// If we need to maybe tweak our private key, do it now.
+	if tweaker != nil {
+		privKey, err = tweaker(privKey)
+		if err != nil {
+			return nil, nil, err
+		}
+	}
+
+	// Generate a valid witness stack for the input.
+	witnessScript, err := txscript.WitnessSignature(
+		tx, sigHashes, inputIndex, output.Value, witnessProgram,
+		hashType, privKey, true,
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return witnessScript, sigScript, nil
+}

wallet/signer_test.go

@@ -0,0 +1,103 @@
+// Copyright (c) 2020 The btcsuite developers
+// Use of this source code is governed by an ISC
+// license that can be found in the LICENSE file.
+
+package wallet
+
+import (
+	"testing"
+
+	"github.com/btcsuite/btcd/txscript"
+	"github.com/btcsuite/btcd/wire"
+	"github.com/btcsuite/btcutil"
+	"github.com/btcsuite/btcwallet/waddrmgr"
+)
+
+// TestComputeInputScript checks that the wallet can create the full
+// witness script for a witness output.
+func TestComputeInputScript(t *testing.T) {
+	testCases := []struct {
+		name              string
+		scope             waddrmgr.KeyScope
+		expectedScriptLen int
+	}{{
+		name:              "BIP084 P2WKH",
+		scope:             waddrmgr.KeyScopeBIP0084,
+		expectedScriptLen: 0,
+	}, {
+		name:              "BIP049 nested P2WKH",
+		scope:             waddrmgr.KeyScopeBIP0049Plus,
+		expectedScriptLen: 23,
+	}}
+
+	w, cleanup := testWallet(t)
+	defer cleanup()
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			runTestCase(t, w, tc.scope, tc.expectedScriptLen)
+		})
+	}
+}
+
+func runTestCase(t *testing.T, w *Wallet, scope waddrmgr.KeyScope,
+	scriptLen int) {
+
+	// Create an address we can use to send some coins to.
+	addr, err := w.CurrentAddress(0, scope)
+	if err != nil {
+		t.Fatalf("unable to get current address: %v", addr)
+	}
+	p2shAddr, err := txscript.PayToAddrScript(addr)
+	if err != nil {
+		t.Fatalf("unable to convert wallet address to p2sh: %v", err)
+	}
+
+	// Add an output paying to the wallet's address to the database.
+	utxOut := wire.NewTxOut(100000, p2shAddr)
+	incomingTx := &wire.MsgTx{
+		TxIn:  []*wire.TxIn{{}},
+		TxOut: []*wire.TxOut{utxOut},
+	}
+	addUtxo(t, w, incomingTx)
+
+	// Create a transaction that spends the UTXO created above and spends to
+	// the same address again.
+	prevOut := wire.OutPoint{
+		Hash:  incomingTx.TxHash(),
+		Index: 0,
+	}
+	outgoingTx := &wire.MsgTx{
+		TxIn: []*wire.TxIn{{
+			PreviousOutPoint: prevOut,
+		}},
+		TxOut: []*wire.TxOut{utxOut},
+	}
+	sigHashes := txscript.NewTxSigHashes(outgoingTx)
+
+	// Compute the input script to spend the UTXO now.
+	witness, script, err := w.ComputeInputScript(
+		outgoingTx, utxOut, 0, sigHashes, txscript.SigHashAll, nil,
+	)
+	if err != nil {
+		t.Fatalf("error computing input script: %v", err)
+	}
+	if len(script) != scriptLen {
+		t.Fatalf("unexpected script length, got %d wanted %d",
+			len(script), scriptLen)
+	}
+	if len(witness) != 2 {
+		t.Fatalf("unexpected witness stack length, got %d, wanted %d",
+			len(witness), 2)
+	}
+
+	// Finally verify that the created witness is valid.
+	outgoingTx.TxIn[0].Witness = witness
+	outgoingTx.TxIn[0].SignatureScript = script
+	err = validateMsgTx(
+		outgoingTx, [][]byte{utxOut.PkScript}, []btcutil.Amount{100000},
+	)
+	if err != nil {
+		t.Fatalf("error validating tx: %v", err)
+	}
+}