feat: Add support for stake expansion (#127)

* chore: initial add btc stk exp sig

* chore: add stk exp to the parsed sig request

* fix: add missing funding output to stake expansion

* chore: add TODO verifier

* chore: add missing stk expansion fields

* chore: add details for stk exp signing err

* update based on latest changes

* address TODO

* update validations

* fix build

* add chlog

* update tools pkg

* update Dockerfiles go version

* update test

* fix bbn node startup

* fix tst

* update babylon dep to v3 snapshot

* update babylon dep to v3 snapshot in tools

* changes based on review comments

* add basic validations

* update bbn snapshot

* update stake expansion sig

* add e2e test

* update bbn snapshot

* go mod tidy

* changes based on review comments

* revert go mod change

---------

Co-authored-by: tom <[email protected]>

Author: RafilxTenfen

CHANGELOG.md

@@ -40,6 +40,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ### Improvements
 
 * [#125](https://github.com/babylonlabs-io/covenant-emulator/pull/125) Update config defaults for gas adjustment and cov batch sigs.
+* [#127](https://github.com/babylonlabs-io/covenant-emulator/pull/127) Add Stake Expansion support
 
 ## v0.15.1
 

Dockerfile

@@ -1,4 +1,4 @@
-FROM golang:1.23 AS builder
+FROM golang:1.23.10 AS builder
 
 # hadolint ignore=DL3008
 RUN apt-get update && apt-get install --no-install-recommends -y ca-certificates make git bash gcc curl jq && rm -rf /var/lib/apt/lists/*

Makefile

@@ -1,7 +1,7 @@
 BUILDDIR ?= $(CURDIR)/build
 TOOLS_DIR := tools
 
-BABYLON_PKG := github.com/babylonlabs-io/babylon/cmd/babylond
+BABYLON_PKG := github.com/babylonlabs-io/babylon/v3/cmd/babylond
 
 GO_BIN := ${GOPATH}/bin
 

clientcontroller/babylon.go

@@ -10,12 +10,12 @@ import (
 	"github.com/btcsuite/btcd/btcec/v2"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
-	"github.com/babylonlabs-io/babylon/client/babylonclient"
-	bbnclient "github.com/babylonlabs-io/babylon/client/client"
-	bbntypes "github.com/babylonlabs-io/babylon/types"
-	btcctypes "github.com/babylonlabs-io/babylon/x/btccheckpoint/types"
-	btclctypes "github.com/babylonlabs-io/babylon/x/btclightclient/types"
-	btcstakingtypes "github.com/babylonlabs-io/babylon/x/btcstaking/types"
+	"github.com/babylonlabs-io/babylon/v3/client/babylonclient"
+	bbnclient "github.com/babylonlabs-io/babylon/v3/client/client"
+	bbntypes "github.com/babylonlabs-io/babylon/v3/types"
+	btcctypes "github.com/babylonlabs-io/babylon/v3/x/btccheckpoint/types"
+	btclctypes "github.com/babylonlabs-io/babylon/v3/x/btclightclient/types"
+	btcstakingtypes "github.com/babylonlabs-io/babylon/v3/x/btcstaking/types"
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg"
 	sdkclient "github.com/cosmos/cosmos-sdk/client"
@@ -165,14 +165,22 @@ func (bc *BabylonController) SubmitCovenantSigs(covSigs []*types.CovenantSigs) (
 	msgs := make([]sdk.Msg, 0, len(covSigs))
 	for _, covSig := range covSigs {
 		bip340UnbondingSig := bbntypes.NewBIP340SignatureFromBTCSig(covSig.UnbondingSig)
-		msgs = append(msgs, &btcstakingtypes.MsgAddCovenantSigs{
+		msg := &btcstakingtypes.MsgAddCovenantSigs{
 			Signer:                  bc.mustGetTxSigner(),
 			Pk:                      bbntypes.NewBIP340PubKeyFromBTCPK(covSig.PublicKey),
 			StakingTxHash:           covSig.StakingTxHash.String(),
 			SlashingTxSigs:          covSig.SlashingSigs,
 			UnbondingTxSig:          bip340UnbondingSig,
 			SlashingUnbondingTxSigs: covSig.SlashingUnbondingSigs,
-		})
+			StakeExpansionTxSig:     nil,
+		}
+
+		if covSig.StkExpSig != nil {
+			stkExpSig := bbntypes.NewBIP340SignatureFromBTCSig(covSig.StkExpSig)
+			msg.StakeExpansionTxSig = stkExpSig
+		}
+
+		msgs = append(msgs, msg)
 	}
 	res, err := bc.reliablySendMsgs(msgs)
 	if err != nil {
@@ -198,6 +206,15 @@ func (bc *BabylonController) QueryVerifiedDelegations(limit uint64) ([]*types.De
 	return bc.queryDelegationsWithStatus(btcstakingtypes.BTCDelegationStatus_VERIFIED, limit, nil)
 }
 
+// QueryBTCDelegation queries the BTC delegation by the tx hash
+func (bc *BabylonController) QueryBTCDelegation(stakingTxHashHex string) (*types.Delegation, error) {
+	resp, err := bc.bbnClient.QueryClient.BTCDelegation(stakingTxHashHex)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query BTC delegation %s: %v", stakingTxHashHex, err)
+	}
+	return DelegationRespToDelegation(resp.BtcDelegation)
+}
+
 // queryDelegationsWithStatus queries BTC delegations that need a Covenant signature
 // with the given status (either pending or unbonding)
 // it is only used when the program is running in Covenant mode
@@ -308,7 +325,7 @@ func DelegationRespToDelegation(del *btcstakingtypes.BTCDelegationResponse) (*ty
 		return nil, fmt.Errorf("total sat (%d) is larger than the maximum int64", del.TotalSat)
 	}
 
-	return &types.Delegation{
+	respDel := &types.Delegation{
 		BtcPk:            del.BtcPk.MustToBTCPK(),
 		FpBtcPks:         fpBtcPks,
 		TotalSat:         btcutil.Amount(del.TotalSat),
@@ -322,7 +339,17 @@ func DelegationRespToDelegation(del *btcstakingtypes.BTCDelegationResponse) (*ty
 		UnbondingTime:    uint16(del.UnbondingTime),
 		BtcUndelegation:  undelegation,
 		ParamsVersion:    del.ParamsVersion,
-	}, nil
+		StakeExpansion:   nil,
+	}
+
+	if del.StkExp != nil {
+		respDel.StakeExpansion = &types.DelegationStakeExpansion{
+			PreviousStakingTxHashHex: del.StkExp.PreviousStakingTxHashHex,
+			OtherFundingTxOutHex:     del.StkExp.OtherFundingTxOutHex,
+		}
+	}
+
+	return respDel, nil
 }
 
 func UndelegationRespToUndelegation(undel *btcstakingtypes.BTCUndelegationResponse) (*types.Undelegation, error) {
@@ -426,6 +453,57 @@ func (bc *BabylonController) CreateBTCDelegation(
 	return &types.TxResponse{TxHash: res.TxHash}, nil
 }
 
+// CreateStakeExpansionDelegation creates a BTC stake expansion delegation using MsgBtcStakeExpand
+// Currently this is only used for e2e tests, probably does not need to add it into the interface
+func (bc *BabylonController) CreateStakeExpansionDelegation(
+	delBtcPk *bbntypes.BIP340PubKey,
+	fpPks []*btcec.PublicKey,
+	pop *btcstakingtypes.ProofOfPossessionBTC,
+	stakingTime uint32,
+	stakingValue int64,
+	stakingTxInfo *btcctypes.TransactionInfo,
+	slashingTx *btcstakingtypes.BTCSlashingTx,
+	delSlashingSig *bbntypes.BIP340Signature,
+	unbondingTx []byte,
+	unbondingTime uint32,
+	unbondingValue int64,
+	unbondingSlashingTx *btcstakingtypes.BTCSlashingTx,
+	delUnbondingSlashingSig *bbntypes.BIP340Signature,
+	previousStakingTxHash string,
+	fundingTx []byte,
+) (*types.TxResponse, error) {
+	fpBtcPks := make([]bbntypes.BIP340PubKey, 0, len(fpPks))
+	for _, v := range fpPks {
+		fpBtcPks = append(fpBtcPks, *bbntypes.NewBIP340PubKeyFromBTCPK(v))
+	}
+
+	msg := &btcstakingtypes.MsgBtcStakeExpand{
+		StakerAddr:                    bc.mustGetTxSigner(),
+		Pop:                           pop,
+		BtcPk:                         delBtcPk,
+		FpBtcPkList:                   fpBtcPks,
+		StakingTime:                   stakingTime,
+		StakingValue:                  stakingValue,
+		StakingTx:                     stakingTxInfo.Transaction,
+		SlashingTx:                    slashingTx,
+		DelegatorSlashingSig:          delSlashingSig,
+		UnbondingTx:                   unbondingTx,
+		UnbondingTime:                 unbondingTime,
+		UnbondingValue:                unbondingValue,
+		UnbondingSlashingTx:           unbondingSlashingTx,
+		DelegatorUnbondingSlashingSig: delUnbondingSlashingSig,
+		PreviousStakingTxHash:         previousStakingTxHash,
+		FundingTx:                     fundingTx,
+	}
+
+	res, err := bc.reliablySendMsg(msg)
+	if err != nil {
+		return nil, err
+	}
+
+	return &types.TxResponse{TxHash: res.TxHash}, nil
+}
+
 // Register a finality provider to Babylon
 // Currently this is only used for e2e tests, probably does not need to add it into the interface
 func (bc *BabylonController) RegisterFinalityProvider(

clientcontroller/babylon_test.go

@@ -4,9 +4,9 @@ import (
 	"math/rand"
 	"testing"
 
-	"github.com/babylonlabs-io/babylon/testutil/datagen"
-	bbntypes "github.com/babylonlabs-io/babylon/types"
-	"github.com/babylonlabs-io/babylon/x/btcstaking/types"
+	"github.com/babylonlabs-io/babylon/v3/testutil/datagen"
+	bbntypes "github.com/babylonlabs-io/babylon/v3/types"
+	"github.com/babylonlabs-io/babylon/v3/x/btcstaking/types"
 	"github.com/babylonlabs-io/covenant-emulator/clientcontroller"
 	"github.com/babylonlabs-io/covenant-emulator/testutil"
 	"github.com/stretchr/testify/require"

clientcontroller/interface.go

@@ -25,6 +25,9 @@ type (
 		// QueryPendingDelegations queries BTC delegations that are in status of pending
 		QueryPendingDelegations(limit uint64, filter FilterFn) ([]*types.Delegation, error)
 
+		// QueryBTCDelegation queries the BTC delegation by the tx hash
+		QueryBTCDelegation(stakingTxHashHex string) (*types.Delegation, error)
+
 		QueryStakingParamsByVersion(version uint32) (*types.StakingParams, error)
 
 		Close() error

clientcontroller/retry_utils.go

@@ -2,7 +2,7 @@ package clientcontroller
 
 import (
 	sdkErr "cosmossdk.io/errors"
-	btcstakingtypes "github.com/babylonlabs-io/babylon/x/btcstaking/types"
+	btcstakingtypes "github.com/babylonlabs-io/babylon/v3/x/btcstaking/types"
 )
 
 // these errors are considered unrecoverable because these indicate

cmd/covd/key.go

@@ -4,7 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 
-	"github.com/babylonlabs-io/babylon/types"
+	"github.com/babylonlabs-io/babylon/v3/types"
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/jessevdk/go-flags"
 

config/babylon.go

@@ -3,7 +3,7 @@ package config
 import (
 	"time"
 
-	bbncfg "github.com/babylonlabs-io/babylon/client/config"
+	bbncfg "github.com/babylonlabs-io/babylon/v3/client/config"
 )
 
 type BBNConfig struct {

covenant-signer/Dockerfile

@@ -1,4 +1,4 @@
-FROM golang:1.23.1-alpine AS builder
+FROM golang:1.23.10-alpine AS builder
 
 # Use muslc for static libs
 ARG BUILD_TAGS="muslc"