feat: add sign tx util functions

2026-01-12 22:43:22 +08:00 · 2024-09-06 21:58:02 +07:00
parent 6e8a846c27
commit 3d5f3b414c
2 changed files with 240 additions and 0 deletions
--- a/pkg/btcutils/signature.go
+++ b/pkg/btcutils/signature.go
@@ -3,8 +3,12 @@ package btcutils
 import (
 	"github.com/Cleverse/go-utilities/utils"
 	verifier "github.com/bitonicnl/verify-signed-message/pkg"
+	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/chaincfg"
+	"github.com/btcsuite/btcd/txscript"
+	"github.com/btcsuite/btcd/wire"
 	"github.com/cockroachdb/errors"
+	"github.com/gaze-network/indexer-network/common/errs"
 )

 func VerifySignature(address string, message string, sigBase64 string, defaultNet ...*chaincfg.Params) error {
@@ -19,3 +23,121 @@ func VerifySignature(address string, message string, sigBase64 string, defaultNe
 	}
 	return nil
 }
+
+func SignTxInput(tx *wire.MsgTx, privateKey *btcec.PrivateKey, prevTxOut *wire.TxOut, inputIndex int) (*wire.MsgTx, error) {
+	if privateKey == nil {
+		return nil, errors.Wrap(errs.InvalidArgument, "PrivateKey is required")
+	}
+	if tx == nil {
+		return nil, errors.Wrap(errs.InvalidArgument, "Tx is required")
+	}
+	if prevTxOut == nil {
+		return nil, errors.Wrap(errs.InvalidArgument, "PrevTxOut is required")
+	}
+
+	prevOutFetcher := txscript.NewCannedPrevOutputFetcher(prevTxOut.PkScript, prevTxOut.Value)
+	sigHashes := txscript.NewTxSigHashes(tx, prevOutFetcher)
+	if len(tx.TxIn) <= inputIndex {
+		return nil, errors.Errorf("input to sign (%d) is out of range", inputIndex)
+	}
+	address, err := ExtractAddressFromPkScript(prevTxOut.PkScript)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to extract address")
+	}
+
+	switch address.Type() {
+	case AddressP2TR:
+		witness, err := txscript.TaprootWitnessSignature(
+			tx,
+			sigHashes,
+			inputIndex,
+			prevTxOut.Value,
+			prevTxOut.PkScript,
+			txscript.SigHashAll|txscript.SigHashAnyOneCanPay,
+			privateKey)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to sign")
+		}
+		tx.TxIn[inputIndex].Witness = witness
+	case AddressP2WPKH:
+		witness, err := txscript.WitnessSignature(
+			tx,
+			sigHashes,
+			inputIndex,
+			prevTxOut.Value,
+			prevTxOut.PkScript,
+			txscript.SigHashAll|txscript.SigHashAnyOneCanPay,
+			privateKey,
+			true,
+		)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to sign")
+		}
+		tx.TxIn[inputIndex].Witness = witness
+	case AddressP2PKH:
+		sigScript, err := txscript.SignatureScript(
+			tx,
+			inputIndex,
+			prevTxOut.PkScript,
+			txscript.SigHashAll|txscript.SigHashAnyOneCanPay,
+			privateKey,
+			true,
+		)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to sign")
+		}
+		tx.TxIn[inputIndex].SignatureScript = sigScript
+	default:
+		return nil, errors.Wrapf(errs.NotSupported, "unsupported input address type %s", address.Type())
+	}
+
+	return tx, nil
+}
+
+func SignTxInputTapScript(tx *wire.MsgTx, privateKey *btcec.PrivateKey, prevTxOut *wire.TxOut, inputIndex int) (*wire.MsgTx, error) {
+	if privateKey == nil {
+		return nil, errors.Wrap(errs.InvalidArgument, "PrivateKey is required")
+	}
+	if tx == nil {
+		return nil, errors.Wrap(errs.InvalidArgument, "Tx is required")
+	}
+	if prevTxOut == nil {
+		return nil, errors.Wrap(errs.InvalidArgument, "PrevTxOut is required")
+	}
+
+	prevOutFetcher := txscript.NewCannedPrevOutputFetcher(prevTxOut.PkScript, prevTxOut.Value)
+	sigHashes := txscript.NewTxSigHashes(tx, prevOutFetcher)
+	if len(tx.TxIn) <= inputIndex {
+		return nil, errors.Errorf("input to sign (%d) is out of range", inputIndex)
+	}
+	address, err := ExtractAddressFromPkScript(prevTxOut.PkScript)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to extract address")
+	}
+
+	if address.Type() != AddressTaproot {
+		return nil, errors.Errorf("input type must be %s", AddressTaproot)
+	}
+
+	witness := tx.TxIn[inputIndex].Witness
+	if len(witness) != 3 {
+		return nil, errors.Wrapf(errs.InvalidArgument, "invalid witness length: expected 3, got %d", len(witness))
+	}
+
+	tapLeaf := txscript.NewBaseTapLeaf(witness[1])
+	signature, err := txscript.RawTxInTapscriptSignature(
+		tx,
+		sigHashes,
+		inputIndex,
+		prevTxOut.Value,
+		prevTxOut.PkScript,
+		tapLeaf,
+		txscript.SigHashAll|txscript.SigHashAnyOneCanPay,
+		privateKey)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to sign")
+	}
+	tx.TxIn[inputIndex].Witness[0] = signature
+
+	return tx, nil
+}
--- a/pkg/btcutils/signature_test.go
+++ b/pkg/btcutils/signature_test.go
@@ -3,8 +3,14 @@ package btcutils
 import (
 	"testing"

+	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/btcsuite/btcd/btcec/v2/schnorr"
+	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg"
+	"github.com/btcsuite/btcd/txscript"
+	"github.com/btcsuite/btcd/wire"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )

 func TestVerifySignature(t *testing.T) {
@@ -67,3 +73,115 @@ func TestVerifySignature(t *testing.T) {
 		assert.Error(t, err)
 	}
 }
+
+func TestSignTxInput(t *testing.T) {
+	generateTxAndPrevTxOutFromPkScript := func(pkScript []byte) (*wire.MsgTx, *wire.TxOut) {
+		tx := wire.NewMsgTx(wire.TxVersion)
+		tx.AddTxIn(&wire.TxIn{
+			PreviousOutPoint: wire.OutPoint{
+				Index: 1,
+			},
+		})
+		txOut := &wire.TxOut{
+			Value: 1e8, PkScript: pkScript,
+		}
+		tx.AddTxOut(txOut)
+		// using same value and pkScript as input for simplicity
+		return tx, txOut
+	}
+	verifySignedTx := func(t *testing.T, signedTx *wire.MsgTx, prevTxOut *wire.TxOut) {
+		t.Helper()
+
+		prevOutFetcher := txscript.NewCannedPrevOutputFetcher(prevTxOut.PkScript, prevTxOut.Value)
+		sigHashes := txscript.NewTxSigHashes(signedTx, prevOutFetcher)
+		vm, err := txscript.NewEngine(
+			prevTxOut.PkScript, signedTx, 0, txscript.StandardVerifyFlags,
+			nil, sigHashes, prevTxOut.Value, prevOutFetcher,
+		)
+		require.NoError(t, err)
+		require.NoError(t, vm.Execute(), "error during signature verification") // no error means success
+	}
+
+	privKey, _ := btcec.NewPrivateKey()
+	t.Run("P2TR input", func(t *testing.T) {
+		taprootKey := txscript.ComputeTaprootKeyNoScript(privKey.PubKey())
+
+		pkScript, err := txscript.PayToTaprootScript(taprootKey)
+		require.NoError(t, err)
+
+		tx, prevTxOut := generateTxAndPrevTxOutFromPkScript(pkScript)
+		signedTx, err := SignTxInput(
+			tx, privKey, prevTxOut, 0,
+		)
+		require.NoError(t, err)
+		verifySignedTx(t, signedTx, prevTxOut)
+	})
+	t.Run("tapscript input", func(t *testing.T) {
+		internalKey := privKey.PubKey()
+
+		// Our script will be a simple OP_CHECKSIG as the sole leaf of a
+		// tapscript tree.
+		builder := txscript.NewScriptBuilder()
+		builder.AddData(schnorr.SerializePubKey(internalKey))
+		builder.AddOp(txscript.OP_CHECKSIG)
+		tapScript, err := builder.Script()
+		require.NoError(t, err)
+
+		tapLeaf := txscript.NewBaseTapLeaf(tapScript)
+		tapScriptTree := txscript.AssembleTaprootScriptTree(tapLeaf)
+
+		controlBlock := tapScriptTree.LeafMerkleProofs[0].ToControlBlock(
+			internalKey,
+		)
+		controlBlockBytes, err := controlBlock.ToBytes()
+		require.NoError(t, err)
+
+		tapScriptRootHash := tapScriptTree.RootNode.TapHash()
+		outputKey := txscript.ComputeTaprootOutputKey(
+			internalKey, tapScriptRootHash[:],
+		)
+		p2trScript, err := txscript.PayToTaprootScript(outputKey)
+		require.NoError(t, err)
+
+		tx, prevTxOut := generateTxAndPrevTxOutFromPkScript(p2trScript)
+		tx.TxIn[0].Witness = wire.TxWitness{
+			{},
+			tapScript,
+			controlBlockBytes,
+		}
+		signedTx, err := SignTxInputTapScript(
+			tx, privKey, prevTxOut, 0,
+		)
+		require.NoError(t, err)
+		verifySignedTx(t, signedTx, prevTxOut)
+	})
+	t.Run("P2WPKH input", func(t *testing.T) {
+		pubKey := privKey.PubKey()
+		pubKeyHash := btcutil.Hash160(pubKey.SerializeCompressed())
+
+		pkScript, err := txscript.NewScriptBuilder().
+			AddOp(txscript.OP_0).
+			AddData(pubKeyHash).
+			Script()
+
+		tx, prevTxOut := generateTxAndPrevTxOutFromPkScript(pkScript)
+		signedTx, err := SignTxInput(
+			tx, privKey, prevTxOut, 0,
+		)
+		require.NoError(t, err)
+		verifySignedTx(t, signedTx, prevTxOut)
+	})
+	t.Run("P2PKH input", func(t *testing.T) {
+		pubKey := privKey.PubKey()
+		pubKeyHash := btcutil.Hash160(pubKey.SerializeCompressed())
+		address, err := btcutil.NewAddressPubKeyHash(pubKeyHash, &chaincfg.MainNetParams)
+		pkScript, err := txscript.PayToAddrScript(address)
+
+		tx, prevTxOut := generateTxAndPrevTxOutFromPkScript(pkScript)
+		signedTx, err := SignTxInput(
+			tx, privKey, prevTxOut, 0,
+		)
+		require.NoError(t, err)
+		verifySignedTx(t, signedTx, prevTxOut)
+	})
+}