feat: add check if the covenant is the covenant committee

CHANGELOG.md

@@ -37,6 +37,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## Unreleased
 
+### Improvements
+
+* [#91](https://github.com/babylonlabs-io/covenant-emulator/pull/91) Add verification
+if the covenant public key was present in the params of the versioned BTC delegation
+
 ## v0.11.2
 
 ### Bug fixes

covenant/cache_params.go

@@ -0,0 +1,80 @@
+package covenant
+
+import (
+	"sync"
+
+	"github.com/avast/retry-go/v4"
+	"github.com/babylonlabs-io/covenant-emulator/clientcontroller"
+	"github.com/babylonlabs-io/covenant-emulator/types"
+	"go.uber.org/zap"
+)
+
+type (
+	getParamByVersion func(version uint32) (*types.StakingParams, error)
+	ParamsGetter      interface {
+		Get(version uint32) (*types.StakingParams, error)
+	}
+	CacheVersionedParams struct {
+		sync.Mutex
+		paramsByVersion map[uint32]*types.StakingParams
+
+		getParamsByVersion getParamByVersion
+	}
+)
+
+func NewCacheVersionedParams(f getParamByVersion) ParamsGetter {
+	return &CacheVersionedParams{
+		paramsByVersion:    make(map[uint32]*types.StakingParams),
+		getParamsByVersion: f,
+	}
+}
+
+// Get returns the staking parameter from the
+func (v *CacheVersionedParams) Get(version uint32) (*types.StakingParams, error) {
+	v.Lock()
+	defer v.Unlock()
+
+	params, ok := v.paramsByVersion[version]
+	if ok {
+		return params, nil
+	}
+
+	params, err := v.getParamsByVersion(version)
+	if err != nil {
+		return nil, err
+	}
+
+	v.paramsByVersion[version] = params
+	return params, nil
+}
+
+func paramsByVersion(
+	cc clientcontroller.ClientController,
+	logger *zap.Logger,
+) getParamByVersion {
+	return func(version uint32) (*types.StakingParams, error) {
+		var (
+			err    error
+			params *types.StakingParams
+		)
+
+		if err := retry.Do(func() error {
+			params, err = cc.QueryStakingParamsByVersion(version)
+			if err != nil {
+				return err
+			}
+			return nil
+		}, RtyAtt, RtyDel, RtyErr, retry.OnRetry(func(n uint, err error) {
+			logger.Debug(
+				"failed to query the consumer chain for the staking params",
+				zap.Uint("attempt", n+1),
+				zap.Uint("max_attempts", RtyAttNum),
+				zap.Error(err),
+			)
+		})); err != nil {
+			return nil, err
+		}
+
+		return params, nil
+	}
+}

covenant/covenant.go

@@ -46,6 +46,8 @@ type CovenantEmulator struct {
 
 	config *covcfg.Config
 	logger *zap.Logger
+
+	paramCache ParamsGetter
 }
 
 func NewCovenantEmulator(
@@ -59,13 +61,15 @@ func NewCovenantEmulator(
 		return nil, fmt.Errorf("failed to get signer pub key: %w", err)
 	}
 
+	paramGetter := paramsByVersion(cc, logger)
 	return &CovenantEmulator{
-		cc:     cc,
-		signer: signer,
-		config: config,
-		logger: logger,
-		pk:     pk,
-		quit:   make(chan struct{}),
+		cc:         cc,
+		signer:     signer,
+		config:     config,
+		logger:     logger,
+		pk:         pk,
+		quit:       make(chan struct{}),
+		paramCache: NewCacheVersionedParams(paramGetter),
 	}, nil
 }
 
@@ -100,7 +104,7 @@ func (ce *CovenantEmulator) AddCovenantSignatures(btcDels []*types.Delegation) (
 		}
 
 		// 1. get the params matched to the delegation version
-		params, err := ce.getParamsByVersionWithRetry(btcDel.ParamsVersion)
+		params, err := ce.paramCache.Get(btcDel.ParamsVersion)
 		if err != nil {
 			return nil, fmt.Errorf("failed to get staking params with version %d: %w", btcDel.ParamsVersion, err)
 		}
@@ -436,30 +440,70 @@ func (ce *CovenantEmulator) delegationsToBatches(dels []*types.Delegation) [][]*
 	return batches
 }
 
-func RemoveAlreadySigned(localKey *btcec.PublicKey, dels []*types.Delegation) []*types.Delegation {
+func RemoveNotInCommittee(paramCache ParamsGetter, covenantSerializedPk []byte, dels []*types.Delegation) []*types.Delegation {
 	sanitized := make([]*types.Delegation, 0, len(dels))
-	localKeyBytes := schnorr.SerializePubKey(localKey)
 
 	for _, del := range dels {
-		delCopy := del
-		alreadySigned := false
-		for _, covSig := range delCopy.CovenantSigs {
-			remoteKey := schnorr.SerializePubKey(covSig.Pk)
-			if bytes.Equal(remoteKey, localKeyBytes) {
-				alreadySigned = true
-				break
-			}
+		if !IsKeyInCommittee(paramCache, covenantSerializedPk, del) {
+			continue
+		}
+		sanitized = append(sanitized, del)
+	}
+	return sanitized
+}
+
+// IsKeyInCommittee verifies
+func IsKeyInCommittee(paramCache ParamsGetter, covenantSerializedPk []byte, del *types.Delegation) bool {
+	stkParams, err := paramCache.Get(del.ParamsVersion)
+	if err != nil {
+		return false
+	}
+
+	for _, pk := range stkParams.CovenantPks {
+		remoteKey := schnorr.SerializePubKey(pk)
+		if !bytes.Equal(remoteKey, covenantSerializedPk) {
+			continue
 		}
-		if !alreadySigned {
-			sanitized = append(sanitized, delCopy)
+		return true
+	}
+
+	return false
+}
+
+func RemoveAlreadySigned(covenantSerializedPk []byte, dels []*types.Delegation) []*types.Delegation {
+	sanitized := make([]*types.Delegation, 0, len(dels))
+
+	for _, del := range dels {
+		if CovenantAlreadySigned(covenantSerializedPk, del) {
+			continue
 		}
+		sanitized = append(sanitized, del)
 	}
+
 	return sanitized
 }
 
-// removeAlreadySigned removes any delegations that have already been signed by the covenant
-func (ce *CovenantEmulator) removeAlreadySigned(dels []*types.Delegation) []*types.Delegation {
-	return RemoveAlreadySigned(ce.pk, dels)
+// CovenantAlreadySigned returns true if the covenant already signed the BTC Delegation
+func CovenantAlreadySigned(covenantSerializedPk []byte, del *types.Delegation) bool {
+	for _, covSig := range del.CovenantSigs {
+		remoteKey := schnorr.SerializePubKey(covSig.Pk)
+		if !bytes.Equal(remoteKey, covenantSerializedPk) {
+			continue
+		}
+		return true
+	}
+
+	return false
+}
+
+// sanitizeDelegations removes any delegations that have already been signed by the covenant and
+// remove delegations that were not constructed with this covenant public key
+func (ce *CovenantEmulator) sanitizeDelegations(dels []*types.Delegation) []*types.Delegation {
+	covenantSerializedPk := schnorr.SerializePubKey(ce.pk)
+	// 1. Remove delegations that do not need the covenant's signature
+	delsNotSigned := RemoveAlreadySigned(covenantSerializedPk, dels)
+	// 2. Remove delegations that were not constructed with this covenant public key
+	return RemoveNotInCommittee(ce.paramCache, covenantSerializedPk, delsNotSigned)
 }
 
 // covenantSigSubmissionLoop is the reactor to submit Covenant signature for BTC delegations
@@ -490,10 +534,10 @@ func (ce *CovenantEmulator) covenantSigSubmissionLoop() {
 				ce.logger.Debug("no pending delegations are found")
 			}
 			// 2. Remove delegations that do not need the covenant's signature
-			sanitizedDels := ce.removeAlreadySigned(dels)
+			sanitized := ce.sanitizeDelegations(dels)
 
 			// 3. Split delegations into batches for submission
-			batches := ce.delegationsToBatches(sanitizedDels)
+			batches := ce.delegationsToBatches(sanitized)
 			for _, delBatch := range batches {
 				_, err := ce.AddCovenantSignatures(delBatch)
 				if err != nil {
@@ -532,32 +576,6 @@ func (ce *CovenantEmulator) metricsUpdateLoop() {
 	}
 }
 
-func (ce *CovenantEmulator) getParamsByVersionWithRetry(version uint32) (*types.StakingParams, error) {
-	var (
-		params *types.StakingParams
-		err    error
-	)
-
-	if err := retry.Do(func() error {
-		params, err = ce.cc.QueryStakingParamsByVersion(version)
-		if err != nil {
-			return err
-		}
-		return nil
-	}, RtyAtt, RtyDel, RtyErr, retry.OnRetry(func(n uint, err error) {
-		ce.logger.Debug(
-			"failed to query the consumer chain for the staking params",
-			zap.Uint("attempt", n+1),
-			zap.Uint("max_attempts", RtyAttNum),
-			zap.Error(err),
-		)
-	})); err != nil {
-		return nil, err
-	}
-
-	return params, nil
-}
-
 func (ce *CovenantEmulator) recordMetricsFailedSignDelegations(n int) {
 	failedSignDelegations.WithLabelValues(ce.PublicKeyStr()).Add(float64(n))
 }

covenant/covenant_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/hex"
 	"math/rand"
 	"testing"
+	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
@@ -245,6 +246,78 @@ func TestDeduplicationWithOddKey(t *testing.T) {
 	}
 
 	// 4. After removing the already signed delegation, the list should have only one element
-	sanitized := covenant.RemoveAlreadySigned(oddKeyPub, delegations)
+	sanitized := covenant.RemoveAlreadySigned(schnorr.SerializePubKey(oddKeyPub), delegations)
 	require.Equal(t, 1, len(sanitized))
 }
+
+func TestIsKeyInCommittee(t *testing.T) {
+	r := rand.New(rand.NewSource(time.Now().Unix()))
+
+	// create a Covenant key pair in the keyring
+	covenantConfig := covcfg.DefaultConfig()
+	covenantConfig.BabylonConfig.KeyDirectory = t.TempDir()
+
+	covKeyPair, err := keyring.CreateCovenantKey(
+		covenantConfig.BabylonConfig.KeyDirectory,
+		covenantConfig.BabylonConfig.ChainID,
+		covenantConfig.BabylonConfig.Key,
+		covenantConfig.BabylonConfig.KeyringBackend,
+		passphrase,
+		hdPath,
+	)
+	require.NoError(t, err)
+	covenantSerializedPk := schnorr.SerializePubKey(covKeyPair.PublicKey)
+
+	// create params and version
+	pVersionWithoutCovenant := uint32(datagen.RandomInRange(r, 1, 10))
+	pVersionWithCovenant := pVersionWithoutCovenant + 1
+
+	paramsWithoutCovenant := testutil.GenRandomParams(r, t)
+	paramsWithCovenant := testutil.GenRandomParams(r, t)
+	paramsWithCovenant.CovenantPks = append(paramsWithCovenant.CovenantPks, covKeyPair.PublicKey)
+
+	// creates delegations to check
+	delNoCovenant := &types.Delegation{
+		ParamsVersion: pVersionWithoutCovenant,
+	}
+	delWithCovenant := &types.Delegation{
+		ParamsVersion: pVersionWithCovenant,
+	}
+
+	// simple mock with the parameter versions
+	paramsGet := NewMockParam(map[uint32]*types.StakingParams{
+		pVersionWithoutCovenant: paramsWithoutCovenant,
+		pVersionWithCovenant:    paramsWithCovenant,
+	})
+
+	// checks the case where the covenant is NOT in the committee
+	actual := covenant.IsKeyInCommittee(paramsGet, covenantSerializedPk, delNoCovenant)
+	require.False(t, actual)
+	emptyDels := covenant.RemoveNotInCommittee(paramsGet, covenantSerializedPk, []*types.Delegation{delNoCovenant, delNoCovenant})
+	require.Len(t, emptyDels, 0)
+
+	// checks the case where the covenant is in the committee
+	actual = covenant.IsKeyInCommittee(paramsGet, covenantSerializedPk, delWithCovenant)
+	require.True(t, actual)
+	dels := covenant.RemoveNotInCommittee(paramsGet, covenantSerializedPk, []*types.Delegation{delWithCovenant, delNoCovenant})
+	require.Len(t, dels, 1)
+	dels = covenant.RemoveNotInCommittee(paramsGet, covenantSerializedPk, []*types.Delegation{delWithCovenant})
+	require.Len(t, dels, 1)
+	dels = covenant.RemoveNotInCommittee(paramsGet, covenantSerializedPk, []*types.Delegation{delWithCovenant, delWithCovenant, delNoCovenant})
+	require.Len(t, dels, 2)
+}
+
+type MockParamGetter struct {
+	paramsByVersion map[uint32]*types.StakingParams
+}
+
+func NewMockParam(p map[uint32]*types.StakingParams) *MockParamGetter {
+	return &MockParamGetter{
+		paramsByVersion: p,
+	}
+}
+
+func (m *MockParamGetter) Get(version uint32) (*types.StakingParams, error) {
+	p := m.paramsByVersion[version]
+	return p, nil
+}