btcstaking: compose witness of a slashing tx (#43)

btcstaking/staking.go

@@ -404,45 +404,50 @@ func BuildStakingOutput(
 	return wire.NewTxOut(int64(stAmount), pkScript), script, nil
 }
 
-// BuildWitnessToSpendStakingOutput builds witness for spending staking as single staker
-// Current assumptions:
-// - staking output is the only input to the transaction
-// - staking output is valid staking output
-func BuildWitnessToSpendStakingOutput(
-	tx *wire.MsgTx,
-	stakingOutput *wire.TxOut,
+// NewWitnessFromStakingScriptAndSignature creates witness for spending
+// staking from the given staking script and the given signature of
+// a single party in the multisig
+func NewWitnessFromStakingScriptAndSignature(
 	stakingScript []byte,
-	privKey *btcec.PrivateKey) (wire.TxWitness, error) {
-
+	sig *schnorr.Signature,
+) (wire.TxWitness, error) {
+	// get ctrlBlockBytes
 	internalPubKey := UnspendableKeyPathInternalPubKey()
-
 	tapLeaf := txscript.NewBaseTapLeaf(stakingScript)
-
 	tapScriptTree := txscript.AssembleTaprootScriptTree(tapLeaf)
-
-	ctrlBlock := tapScriptTree.LeafMerkleProofs[0].ToControlBlock(
-		internalPubKey,
-	)
-
+	ctrlBlock := tapScriptTree.LeafMerkleProofs[0].ToControlBlock(internalPubKey)
 	ctrlBlockBytes, err := ctrlBlock.ToBytes()
-
-	if err != nil {
-		return nil, err
-	}
-
-	sig, err := SignTxWithOneScriptSpendInputFromTapLeaf(tx, stakingOutput, privKey, tapLeaf)
-
 	if err != nil {
 		return nil, err
 	}
 
+	// compose witness stack
 	witnessStack := wire.TxWitness(make([][]byte, 3))
 	witnessStack[0] = sig.Serialize()
 	witnessStack[1] = stakingScript
 	witnessStack[2] = ctrlBlockBytes
 	return witnessStack, nil
 }
 
+// BuildWitnessToSpendStakingOutput builds witness for spending staking as single staker
+// Current assumptions:
+// - staking output is the only input to the transaction
+// - staking output is valid staking output
+func BuildWitnessToSpendStakingOutput(
+	slashingMsgTx *wire.MsgTx, // slashing tx
+	stakingOutput *wire.TxOut,
+	stakingScript []byte,
+	privKey *btcec.PrivateKey,
+) (wire.TxWitness, error) {
+	tapLeaf := txscript.NewBaseTapLeaf(stakingScript)
+	sig, err := SignTxWithOneScriptSpendInputFromTapLeaf(slashingMsgTx, stakingOutput, privKey, tapLeaf)
+	if err != nil {
+		return nil, err
+	}
+
+	return NewWitnessFromStakingScriptAndSignature(stakingScript, sig)
+}
+
 // ValidateStakingOutputPkScript validates that:
 // - provided output commits to the given script with unspendable spending key path
 // - provided script is valid staking script

btcstaking/staking_test.go

@@ -1,14 +1,13 @@
 package btcstaking_test
 
 import (
-	"bytes"
-	"fmt"
 	"math"
 	"math/rand"
 	"testing"
 	"time"
 
 	"github.com/babylonchain/babylon/btcstaking"
+	btctest "github.com/babylonchain/babylon/testutil/bitcoin"
 	"github.com/babylonchain/babylon/testutil/datagen"
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
@@ -160,68 +159,6 @@ func FuzzGeneratingSignatureValidation(f *testing.F) {
 	})
 }
 
-// Help function to assert the execution of a script engine. Copied from:
-// https://github.com/lightningnetwork/lnd/blob/master/input/script_utils_test.go#L24
-func assertEngineExecution(t *testing.T, testNum int, valid bool,
-	newEngine func() (*txscript.Engine, error)) {
-
-	t.Helper()
-
-	// Get a new VM to execute.
-	vm, err := newEngine()
-	require.NoError(t, err, "unable to create engine")
-
-	// Execute the VM, only go on to the step-by-step execution if
-	// it doesn't validate as expected.
-	vmErr := vm.Execute()
-	if valid == (vmErr == nil) {
-		return
-	}
-
-	// Now that the execution didn't match what we expected, fetch a new VM
-	// to step through.
-	vm, err = newEngine()
-	require.NoError(t, err, "unable to create engine")
-
-	// This buffer will trace execution of the Script, dumping out
-	// to stdout.
-	var debugBuf bytes.Buffer
-
-	done := false
-	for !done {
-		dis, err := vm.DisasmPC()
-		if err != nil {
-			t.Fatalf("stepping (%v)\n", err)
-		}
-		debugBuf.WriteString(fmt.Sprintf("stepping %v\n", dis))
-
-		done, err = vm.Step()
-		if err != nil && valid {
-			t.Log(debugBuf.String())
-			t.Fatalf("spend test case #%v failed, spend "+
-				"should be valid: %v", testNum, err)
-		} else if err == nil && !valid && done {
-			t.Log(debugBuf.String())
-			t.Fatalf("spend test case #%v succeed, spend "+
-				"should be invalid: %v", testNum, err)
-		}
-
-		debugBuf.WriteString(fmt.Sprintf("Stack: %v", vm.GetStack()))
-		debugBuf.WriteString(fmt.Sprintf("AltStack: %v", vm.GetAltStack()))
-	}
-
-	// If we get to this point the unexpected case was not reached
-	// during step execution, which happens for some checks, like
-	// the clean-stack rule.
-	validity := "invalid"
-	if valid {
-		validity = "valid"
-	}
-
-	t.Log(debugBuf.String())
-	t.Fatalf("%v spend test case #%v execution ended with: %v", validity, testNum, vmErr)
-}
-
 func TestStakingScriptExecutionSingleStaker(t *testing.T) {
 	const (
 		stakingValue      = btcutil.Amount(2 * 10e8)
@@ -295,7 +232,7 @@ func TestStakingScriptExecutionSingleStaker(t *testing.T) {
 			prevOutputFetcher,
 		)
 	}
-	assertEngineExecution(t, 0, true, newEngine)
+	btctest.AssertEngineExecution(t, 0, true, newEngine)
 }
 
 func TestStakingScriptExecutionMulitSig(t *testing.T) {
@@ -395,5 +332,5 @@ func TestStakingScriptExecutionMulitSig(t *testing.T) {
 			prevOutputFetcher,
 		)
 	}
-	assertEngineExecution(t, 0, true, newEngine)
+	btctest.AssertEngineExecution(t, 0, true, newEngine)
 }

testutil/bitcoin/utils.go

@@ -0,0 +1,72 @@
+package bitcoin
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+
+	"github.com/btcsuite/btcd/txscript"
+	"github.com/stretchr/testify/require"
+)
+
+// Help function to assert the execution of a script engine. Copied from:
+// https://github.com/lightningnetwork/lnd/blob/master/input/script_utils_test.go#L24
+func AssertEngineExecution(t *testing.T, testNum int, valid bool,
+	newEngine func() (*txscript.Engine, error)) {
+
+	t.Helper()
+
+	// Get a new VM to execute.
+	vm, err := newEngine()
+	require.NoError(t, err, "unable to create engine")
+
+	// Execute the VM, only go on to the step-by-step execution if
+	// it doesn't validate as expected.
+	vmErr := vm.Execute()
+	if valid == (vmErr == nil) {
+		return
+	}
+
+	// Now that the execution didn't match what we expected, fetch a new VM
+	// to step through.
+	vm, err = newEngine()
+	require.NoError(t, err, "unable to create engine")
+
+	// This buffer will trace execution of the Script, dumping out
+	// to stdout.
+	var debugBuf bytes.Buffer
+
+	done := false
+	for !done {
+		dis, err := vm.DisasmPC()
+		if err != nil {
+			t.Fatalf("stepping (%v)\n", err)
+		}
+		debugBuf.WriteString(fmt.Sprintf("stepping %v\n", dis))
+
+		done, err = vm.Step()
+		if err != nil && valid {
+			t.Log(debugBuf.String())
+			t.Fatalf("spend test case #%v failed, spend "+
+				"should be valid: %v", testNum, err)
+		} else if err == nil && !valid && done {
+			t.Log(debugBuf.String())
+			t.Fatalf("spend test case #%v succeed, spend "+
+				"should be invalid: %v", testNum, err)
+		}
+
+		debugBuf.WriteString(fmt.Sprintf("Stack: %v", vm.GetStack()))
+		debugBuf.WriteString(fmt.Sprintf("AltStack: %v", vm.GetAltStack()))
+	}
+
+	// If we get to this point the unexpected case was not reached
+	// during step execution, which happens for some checks, like
+	// the clean-stack rule.
+	validity := "invalid"
+	if valid {
+		validity = "valid"
+	}
+
+	t.Log(debugBuf.String())
+	t.Fatalf("%v spend test case #%v execution ended with: %v", validity, testNum, vmErr)
+}

testutil/datagen/btcstaking.go

@@ -87,7 +87,15 @@ func GenRandomBTCDelegation(r *rand.Rand, valBTCPK *bbn.BIP340PubKey, startHeigh
 	}, nil
 }
 
-func GenBTCStakingSlashingTx(r *rand.Rand, stakerSK *btcec.PrivateKey, validatorPK, juryPK *btcec.PublicKey, stakingTimeBlocks uint16, stakingValue int64, slashingAddress string) (*bstypes.StakingTx, *bstypes.BTCSlashingTx, error) {
+func GenBTCStakingSlashingTx(
+	r *rand.Rand,
+	stakerSK *btcec.PrivateKey,
+	validatorPK *btcec.PublicKey,
+	juryPK *btcec.PublicKey,
+	stakingTimeBlocks uint16,
+	stakingValue int64,
+	slashingAddress string,
+) (*bstypes.StakingTx, *bstypes.BTCSlashingTx, error) {
 	btcNet := &chaincfg.SimNetParams
 
 	stakingOutput, stakingScript, err := btcstaking.BuildStakingOutput(

x/btcstaking/types/btc_slashing_tx.go

@@ -3,6 +3,7 @@ package types
 import (
 	"bytes"
 	"encoding/hex"
+	"fmt"
 
 	"github.com/babylonchain/babylon/btcstaking"
 	bbn "github.com/babylonchain/babylon/types"
@@ -134,3 +135,59 @@ func (tx *BTCSlashingTx) VerifySignature(stakingPkScript []byte, stakingAmount i
 		*sig,
 	)
 }
+
+// ToMsgTxWithWitness generates a BTC slashing tx with witness from
+// - the staking tx
+// - validator signature
+// - delegator signature
+// - jury signature
+func (tx *BTCSlashingTx) ToMsgTxWithWitness(stakingTx *StakingTx, valSig, delSig, jurySig *bbn.BIP340Signature) (*wire.MsgTx, error) {
+	// get staking script
+	stakingScript := stakingTx.StakingScript
+
+	// get Schnorr signatures
+	valSchnorrSig, err := valSig.ToBTCSig()
+	if err != nil {
+		return nil, fmt.Errorf("failed to convert BTC validator signature to Schnorr signature format: %w", err)
+	}
+	delSchnorrSig, err := delSig.ToBTCSig()
+	if err != nil {
+		return nil, fmt.Errorf("failed to convert BTC delegator signature to Schnorr signature format: %w", err)
+	}
+	jurySchnorrSig, err := jurySig.ToBTCSig()
+	if err != nil {
+		return nil, fmt.Errorf("failed to convert jury signature to Schnorr signature format: %w", err)
+	}
+
+	// build witness from each signature
+	valWitness, err := btcstaking.NewWitnessFromStakingScriptAndSignature(stakingScript, valSchnorrSig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to build witness for BTC validator: %w", err)
+	}
+	delWitness, err := btcstaking.NewWitnessFromStakingScriptAndSignature(stakingScript, delSchnorrSig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to build witness for BTC delegator: %w", err)
+	}
+	juryWitness, err := btcstaking.NewWitnessFromStakingScriptAndSignature(stakingScript, jurySchnorrSig)
+	if err != nil {
+		return nil, fmt.Errorf("failed to build witness for jury: %w", err)
+	}
+
+	// To Construct valid witness, for multisig case we need:
+	// - jury signature - witnessJury[0]
+	// - validator signature - witnessValidator[0]
+	// - staker signature - witnessStaker[0]
+	// - empty signature - which is just an empty byte array which signals we are going to use multisig.
+	// 	 This must be signature on top of the stack.
+	// - whole script - witnessStaker[1] (any other witness[1] will work as well)
+	// - control block - witnessStaker[2] (any other witness[2] will work as well)
+	slashingMsgTx, err := tx.ToMsgTx()
+	if err != nil {
+		return nil, fmt.Errorf("failed to convert slashing tx to Bitcoin format: %w", err)
+	}
+	slashingMsgTx.TxIn[0].Witness = [][]byte{
+		juryWitness[0], valWitness[0], delWitness[0], []byte{}, delWitness[1], delWitness[2],
+	}
+
+	return slashingMsgTx, nil
+}