From 78d8c81e0aa9bbea8382b1770fe6c2f1893a446b Mon Sep 17 00:00:00 2001 From: Oliver Gugger Date: Thu, 27 Aug 2020 21:14:59 +0200 Subject: [PATCH] wallet: add ComputeInputScript --- wallet/signer.go | 102 +++++++++++++++++++++++++++++++++++++++++ wallet/signer_test.go | 103 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 205 insertions(+) create mode 100644 wallet/signer.go create mode 100644 wallet/signer_test.go diff --git a/wallet/signer.go b/wallet/signer.go new file mode 100644 index 0000000..390022b --- /dev/null +++ b/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 +} diff --git a/wallet/signer_test.go b/wallet/signer_test.go new file mode 100644 index 0000000..0c0cd00 --- /dev/null +++ b/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) + } +}