chore: add pagination to queryDelegationsWithStatus (#96)

* chore: add pagination to queryDelegationsWithStatus

* chore: add #96 to changelog

* chore: add test for paginated QueryPendingDelegations and add funds to covenant, also added address to covenant key chain info

* fix: lint sec

* chore: refactory to sanitize the delegations during the query of pending ones

* chore: mock gen

* fix: missing filter in QueryPendingDelegations

* chore: remove leftover comments

CHANGELOG.md

@@ -42,6 +42,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 * [#83](https://github.com/babylonlabs-io/covenant-emulator/pull/83) covenant-signer: remove go.mod
 * [#95](https://github.com/babylonlabs-io/covenant-emulator/pull/95) removed local signer option
 as the covenant emulator should only connect to a remote signer
+* [#96](https://github.com/babylonlabs-io/covenant-emulator/pull/96) add pagination to `queryDelegationsWithStatus`
 
 ## v0.11.3
 

clientcontroller/babylon.go

@@ -27,7 +27,10 @@ import (
 	"github.com/babylonlabs-io/covenant-emulator/types"
 )
 
-var _ ClientController = &BabylonController{}
+var (
+	_                  ClientController = &BabylonController{}
+	MaxPaginationLimit                  = uint64(1000)
+)
 
 type BabylonController struct {
 	bbnClient *bbnclient.Client
@@ -183,39 +186,67 @@ func (bc *BabylonController) SubmitCovenantSigs(covSigs []*types.CovenantSigs) (
 	return &types.TxResponse{TxHash: res.TxHash, Events: res.Events}, nil
 }
 
-func (bc *BabylonController) QueryPendingDelegations(limit uint64) ([]*types.Delegation, error) {
-	return bc.queryDelegationsWithStatus(btcstakingtypes.BTCDelegationStatus_PENDING, limit)
+func (bc *BabylonController) QueryPendingDelegations(limit uint64, filter FilterFn) ([]*types.Delegation, error) {
+	return bc.queryDelegationsWithStatus(btcstakingtypes.BTCDelegationStatus_PENDING, limit, filter)
 }
 
 func (bc *BabylonController) QueryActiveDelegations(limit uint64) ([]*types.Delegation, error) {
-	return bc.queryDelegationsWithStatus(btcstakingtypes.BTCDelegationStatus_ACTIVE, limit)
+	return bc.queryDelegationsWithStatus(btcstakingtypes.BTCDelegationStatus_ACTIVE, limit, nil)
 }
 
 func (bc *BabylonController) QueryVerifiedDelegations(limit uint64) ([]*types.Delegation, error) {
-	return bc.queryDelegationsWithStatus(btcstakingtypes.BTCDelegationStatus_VERIFIED, limit)
+	return bc.queryDelegationsWithStatus(btcstakingtypes.BTCDelegationStatus_VERIFIED, limit, nil)
 }
 
 // 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
-func (bc *BabylonController) queryDelegationsWithStatus(status btcstakingtypes.BTCDelegationStatus, limit uint64) ([]*types.Delegation, error) {
+func (bc *BabylonController) queryDelegationsWithStatus(status btcstakingtypes.BTCDelegationStatus, delsLimit uint64, filter FilterFn) ([]*types.Delegation, error) {
+	pgLimit := min(MaxPaginationLimit, delsLimit)
 	pagination := &sdkquery.PageRequest{
-		Limit: limit,
+		Limit: pgLimit,
 	}
 
-	res, err := bc.bbnClient.QueryClient.BTCDelegations(status, pagination)
-	if err != nil {
-		return nil, fmt.Errorf("failed to query BTC delegations: %v", err)
-	}
+	dels := make([]*types.Delegation, 0, delsLimit)
+	indexDels := uint64(0)
 
-	dels := make([]*types.Delegation, 0, len(res.BtcDelegations))
-	for _, delResp := range res.BtcDelegations {
-		del, err := DelegationRespToDelegation(delResp)
+	for indexDels < delsLimit {
+		res, err := bc.bbnClient.QueryClient.BTCDelegations(status, pagination)
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("failed to query BTC delegations: %v", err)
 		}
 
-		dels = append(dels, del)
+		for _, delResp := range res.BtcDelegations {
+			del, err := DelegationRespToDelegation(delResp)
+			if err != nil {
+				return nil, err
+			}
+
+			if filter != nil {
+				accept, err := filter(del)
+				if err != nil {
+					return nil, err
+				}
+
+				if !accept {
+					continue
+				}
+			}
+
+			dels = append(dels, del)
+			indexDels++
+
+			if indexDels == delsLimit {
+				return dels, nil
+			}
+		}
+
+		// if returned a different number of btc delegations than the pagination limit
+		// it means that there is no more delegations at the store
+		if uint64(len(res.BtcDelegations)) != pgLimit {
+			return dels, nil
+		}
+		pagination.Key = res.Pagination.NextKey
 	}
 
 	return dels, nil

clientcontroller/interface.go

@@ -14,19 +14,22 @@ const (
 	babylonConsumerChainName = "babylon"
 )
 
-type ClientController interface {
-	// SubmitCovenantSigs submits Covenant signatures to the consumer chain, each corresponding to
-	// a finality provider that the delegation is (re-)staked to
-	// it returns tx hash and error
-	SubmitCovenantSigs(covSigMsgs []*types.CovenantSigs) (*types.TxResponse, error)
+type (
+	FilterFn         func(del *types.Delegation) (accept bool, err error)
+	ClientController interface {
+		// SubmitCovenantSigs submits Covenant signatures to the consumer chain, each corresponding to
+		// a finality provider that the delegation is (re-)staked to
+		// it returns tx hash and error
+		SubmitCovenantSigs(covSigMsgs []*types.CovenantSigs) (*types.TxResponse, error)
 
-	// QueryPendingDelegations queries BTC delegations that are in status of pending
-	QueryPendingDelegations(limit uint64) ([]*types.Delegation, error)
+		// QueryPendingDelegations queries BTC delegations that are in status of pending
+		QueryPendingDelegations(limit uint64, filter FilterFn) ([]*types.Delegation, error)
 
-	QueryStakingParamsByVersion(version uint32) (*types.StakingParams, error)
+		QueryStakingParamsByVersion(version uint32) (*types.StakingParams, error)
 
-	Close() error
-}
+		Close() error
+	}
+)
 
 func NewClientController(chainName string, bbnConfig *config.BBNConfig, netParams *chaincfg.Params, logger *zap.Logger) (ClientController, error) {
 	var (

covenant-signer/keystore/cosmos/keyringcontroller.go

@@ -8,6 +8,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/crypto/keyring"
 	sdksecp256k1 "github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/go-bip39"
 )
 
@@ -18,6 +19,7 @@ const (
 
 type ChainKeyInfo struct {
 	Name       string
+	Address    sdk.AccAddress
 	Mnemonic   string
 	PublicKey  *btcec.PublicKey
 	PrivateKey *btcec.PrivateKey
@@ -98,13 +100,18 @@ func (kc *ChainKeyringController) CreateChainKey(passphrase, hdPath string) (*Ch
 		return nil, err
 	}
 
+	addr, err := record.GetAddress()
+	if err != nil {
+		return nil, err
+	}
 	privKey := record.GetLocal().PrivKey.GetCachedValue()
 
 	switch v := privKey.(type) {
 	case *sdksecp256k1.PrivKey:
 		sk, pk := btcec.PrivKeyFromBytes(v.Key)
 		return &ChainKeyInfo{
 			Name:       kc.keyName,
+			Address:    addr,
 			PublicKey:  pk,
 			PrivateKey: sk,
 			Mnemonic:   mnemonic,

covenant/covenant.go

@@ -471,40 +471,36 @@ func CovenantAlreadySigned(covenantSerializedPk []byte, del *types.Delegation) b
 	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, error) {
-	return SanitizeDelegations(ce.pk, ce.paramCache, dels)
+// acceptDelegationToSign verifies if the delegation should be accepted to sign.
+func (ce *CovenantEmulator) acceptDelegationToSign(del *types.Delegation) (accept bool, err error) {
+	return AcceptDelegationToSign(ce.pk, ce.paramCache, del)
 }
 
-// SanitizeDelegations remove the delegations in which the covenant public key already signed
-// or the delegation was not constructed with that covenant public key
-func SanitizeDelegations(
+// AcceptDelegationToSign returns true if the delegation should be accepted to be signed.
+// Returns false if the covenant public key already signed
+// or if the delegation was not constructed with that covenant public key.
+func AcceptDelegationToSign(
 	pk *btcec.PublicKey,
 	paramCache ParamsGetter,
-	dels []*types.Delegation,
-) ([]*types.Delegation, error) {
+	del *types.Delegation,
+) (accept bool, err error) {
 	covenantSerializedPk := schnorr.SerializePubKey(pk)
+	// 1. Check if the delegation does not need the covenant's signature because
+	// this covenant already signed
+	if CovenantAlreadySigned(covenantSerializedPk, del) {
+		return false, nil
+	}
 
-	sanitized := make([]*types.Delegation, 0, len(dels))
-	for _, del := range dels {
-		// 1. Remove delegations that do not need the covenant's signature because
-		// this covenant already signed
-		if CovenantAlreadySigned(covenantSerializedPk, del) {
-			continue
-		}
-		// 2. Remove delegations that were not constructed with this covenant public key
-		isInCommittee, err := IsKeyInCommittee(paramCache, covenantSerializedPk, del)
-		if err != nil {
-			return nil, fmt.Errorf("unable to verify if covenant key is in committee: %w", err)
-		}
-		if !isInCommittee {
-			continue
-		}
-		sanitized = append(sanitized, del)
+	// 2. Check if the delegation was not constructed with this covenant public key
+	isInCommittee, err := IsKeyInCommittee(paramCache, covenantSerializedPk, del)
+	if err != nil {
+		return false, fmt.Errorf("unable to verify if covenant key is in committee: %w", err)
+	}
+	if !isInCommittee {
+		return false, nil
 	}
 
-	return sanitized, nil
+	return true, nil
 }
 
 // covenantSigSubmissionLoop is the reactor to submit Covenant signature for BTC delegations
@@ -522,7 +518,7 @@ func (ce *CovenantEmulator) covenantSigSubmissionLoop() {
 		select {
 		case <-covenantSigTicker.C:
 			// 1. Get all pending delegations
-			dels, err := ce.cc.QueryPendingDelegations(limit)
+			dels, err := ce.cc.QueryPendingDelegations(limit, ce.acceptDelegationToSign)
 			if err != nil {
 				ce.logger.Debug("failed to get pending delegations", zap.Error(err))
 				continue
@@ -532,31 +528,13 @@ func (ce *CovenantEmulator) covenantSigSubmissionLoop() {
 			// record delegation metrics
 			ce.recordMetricsCurrentPendingDelegations(pendingDels)
 
-			if len(dels) == 0 {
+			if pendingDels == 0 {
 				ce.logger.Debug("no pending delegations are found")
 				continue
 			}
 
-			// 2. Remove delegations that do not need the covenant's signature
-			sanitizedDels, err := ce.sanitizeDelegations(dels)
-			if err != nil {
-				ce.logger.Error(
-					"error sanitizing delegations",
-					zap.Error(err),
-				)
-				continue
-			}
-
-			if len(sanitizedDels) == 0 {
-				ce.logger.Debug(
-					"no new delegations to sign",
-					zap.Int("pending_dels_len", pendingDels),
-				)
-				continue
-			}
-
-			// 3. Split delegations into batches for submission
-			batches := ce.delegationsToBatches(sanitizedDels)
+			// 2. Split delegations into batches for submission
+			batches := ce.delegationsToBatches(dels)
 			for _, delBatch := range batches {
 				_, err := ce.AddCovenantSignatures(delBatch)
 				if err != nil {