zoneconcierge: API for querying headers in a given epoch (#261)

proto/babylon/zoneconcierge/query.proto

@@ -38,6 +38,10 @@ service Query {
   rpc ListHeaders(QueryListHeadersRequest) returns (QueryListHeadersResponse) {
     option (google.api.http).get = "/babylon/zoneconcierge/v1/headers/{chain_id}";
   }
+  // ListEpochHeaders queries the headers of a chain timestamped in a given epoch of Babylon, with pagination support
+  rpc ListEpochHeaders(QueryListEpochHeadersRequest) returns (QueryListEpochHeadersResponse) {
+    option (google.api.http).get = "/babylon/zoneconcierge/v1/headers/{chain_id}/epochs/{epoch_num}";
+  }
   // FinalizedChainInfo queries the BTC-finalised info of a chain, with proofs
   rpc FinalizedChainInfo(QueryFinalizedChainInfoRequest) returns (QueryFinalizedChainInfoResponse) {
     option (google.api.http).get = "/babylon/zoneconcierge/v1/finalized_chain_info/{chain_id}";
@@ -99,6 +103,18 @@ message QueryListHeadersResponse {
   cosmos.base.query.v1beta1.PageResponse pagination = 2;
 }
 
+// QueryListEpochHeadersRequest is request type for the Query/ListEpochHeaders RPC method.
+message QueryListEpochHeadersRequest {
+  uint64 epoch_num = 1;
+  string chain_id = 2;
+}
+
+// QueryListEpochHeadersResponse is response type for the Query/ListEpochHeaders RPC method.
+message QueryListEpochHeadersResponse {
+  // headers is the list of headers
+  repeated babylon.zoneconcierge.v1.IndexedHeader headers = 1;
+}
+
 // QueryFinalizedChainInfoRequest is request type for the Query/FinalizedChainInfo RPC method.
 message QueryFinalizedChainInfoRequest {
   // chain_id is the ID of the CZ

x/zoneconcierge/keeper/canonical_chain_indexer_test.go

@@ -14,7 +14,8 @@ func FuzzCanonicalChainIndexer(f *testing.F) {
 	f.Fuzz(func(t *testing.T, seed int64) {
 		rand.Seed(seed)
 
-		_, babylonChain, czChain, zcKeeper := SetupTest(t)
+		_, babylonChain, czChain, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
 
 		ctx := babylonChain.GetContext()
 		hooks := zcKeeper.Hooks()

x/zoneconcierge/keeper/epoch_chain_info_indexer.go

@@ -1,6 +1,7 @@
 package keeper
 
 import (
+	bbn "github.com/babylonchain/babylon/types"
 	"github.com/babylonchain/babylon/x/zoneconcierge/types"
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -19,6 +20,45 @@ func (k Keeper) GetEpochChainInfo(ctx sdk.Context, chainID string, epochNumber u
 	return &chainInfo, nil
 }
 
+// GetEpochHeaders gets the headers timestamped in a given epoch, in the ascending order
+func (k Keeper) GetEpochHeaders(ctx sdk.Context, chainID string, epochNumber uint64) ([]*types.IndexedHeader, error) {
+	headers := []*types.IndexedHeader{}
+
+	// find the last timestamped header of this chain in the epoch
+	epochChainInfo, err := k.GetEpochChainInfo(ctx, chainID, epochNumber)
+	if err != nil {
+		return nil, err
+	}
+	// it's possible that this epoch's snapshot is not updated for many epochs
+	// this implies that this epoch does not timestamp any header for this chain at all
+	if epochChainInfo.LatestHeader.BabylonEpoch < epochNumber {
+		return nil, types.ErrEpochHeadersNotFound
+	}
+	// now we have the last header in this epoch
+	headers = append(headers, epochChainInfo.LatestHeader)
+
+	// append all previous headers until reaching the previous epoch
+	canonicalChainStore := k.canonicalChainStore(ctx, chainID)
+	lastHeaderKey := sdk.Uint64ToBigEndian(epochChainInfo.LatestHeader.Height)
+	// NOTE: even in ReverseIterator, start and end should still be specified in ascending order
+	canonicalChainIter := canonicalChainStore.ReverseIterator(nil, lastHeaderKey)
+	defer canonicalChainIter.Close()
+	for ; canonicalChainIter.Valid(); canonicalChainIter.Next() {
+		var prevHeader types.IndexedHeader
+		k.cdc.MustUnmarshal(canonicalChainIter.Value(), &prevHeader)
+		if prevHeader.BabylonEpoch < epochNumber {
+			// we have reached the previous epoch, break the loop
+			break
+		}
+		headers = append(headers, &prevHeader)
+	}
+
+	// reverse the list so that it remains ascending order
+	bbn.Reverse(headers)
+
+	return headers, nil
+}
+
 // recordEpochChainInfo records the chain info for a given epoch number of given chain ID
 // where the latest chain info is retrieved from the chain info indexer
 func (k Keeper) recordEpochChainInfo(ctx sdk.Context, chainID string, epochNumber uint64) {

x/zoneconcierge/keeper/epoch_chain_info_indexer_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/babylonchain/babylon/testutil/datagen"
+	ibctmtypes "github.com/cosmos/ibc-go/v5/modules/light-clients/07-tendermint/types"
 	"github.com/stretchr/testify/require"
 )
 
@@ -14,7 +15,8 @@ func FuzzEpochChainInfoIndexer(f *testing.F) {
 	f.Fuzz(func(t *testing.T, seed int64) {
 		rand.Seed(seed)
 
-		_, babylonChain, czChain, zcKeeper := SetupTest(t)
+		_, babylonChain, czChain, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
 
 		ctx := babylonChain.GetContext()
 		hooks := zcKeeper.Hooks()
@@ -35,3 +37,68 @@ func FuzzEpochChainInfoIndexer(f *testing.F) {
 		require.Equal(t, numForkHeaders, uint64(len(chainInfo.LatestForks.Headers)))
 	})
 }
+
+func FuzzGetEpochHeaders(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+
+	f.Fuzz(func(t *testing.T, seed int64) {
+		rand.Seed(seed)
+
+		_, babylonChain, czChain, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
+		epochingKeeper := babylonApp.EpochingKeeper
+
+		ctx := babylonChain.GetContext()
+		hooks := zcKeeper.Hooks()
+
+		numReqs := datagen.RandomInt(5) + 1
+
+		epochNumList := []uint64{datagen.RandomInt(10) + 1}
+		nextHeightList := []uint64{0}
+		numHeadersList := []uint64{}
+		expectedHeadersMap := map[uint64][]*ibctmtypes.Header{}
+		numForkHeadersList := []uint64{}
+
+		// we test the scenario of ending an epoch for multiple times, in order to ensure that
+		// consecutive epoch infos do not affect each other.
+		for i := uint64(0); i < numReqs; i++ {
+			epochNum := epochNumList[i]
+			// enter a random epoch
+			if i == 0 {
+				for j := uint64(0); j < epochNum; j++ {
+					epochingKeeper.IncEpoch(ctx)
+				}
+			} else {
+				for j := uint64(0); j < epochNum-epochNumList[i-1]; j++ {
+					epochingKeeper.IncEpoch(ctx)
+				}
+			}
+
+			// generate a random number of headers and fork headers
+			numHeadersList = append(numHeadersList, datagen.RandomInt(100)+1)
+			numForkHeadersList = append(numForkHeadersList, datagen.RandomInt(10)+1)
+			// trigger hooks to append these headers and fork headers
+			expectedHeaders, _ := SimulateHeadersAndForksViaHook(ctx, hooks, czChain.ChainID, nextHeightList[i], numHeadersList[i], numForkHeadersList[i])
+			expectedHeadersMap[epochNum] = expectedHeaders
+			// prepare nextHeight for the next request
+			nextHeightList = append(nextHeightList, nextHeightList[i]+numHeadersList[i])
+
+			// simulate the scenario that a random epoch has ended
+			hooks.AfterEpochEnds(ctx, epochNum)
+			// prepare epochNum for the next request
+			epochNumList = append(epochNumList, epochNum+datagen.RandomInt(10)+1)
+		}
+
+		// attest the correctness of epoch info for each tested epoch
+		for i := uint64(0); i < numReqs; i++ {
+			epochNum := epochNumList[i]
+			// check if the headers are same as expected
+			headers, err := zcKeeper.GetEpochHeaders(ctx, czChain.ChainID, epochNum)
+			require.NoError(t, err)
+			require.Equal(t, len(expectedHeadersMap[epochNum]), len(headers))
+			for j := 0; j < len(expectedHeadersMap[epochNum]); j++ {
+				require.Equal(t, expectedHeadersMap[epochNum][j].Header.LastCommitHash, headers[j].Hash)
+			}
+		}
+	})
+}

x/zoneconcierge/keeper/fork_indexer_test.go

@@ -14,7 +14,8 @@ func FuzzForkIndexer(f *testing.F) {
 	f.Fuzz(func(t *testing.T, seed int64) {
 		rand.Seed(seed)
 
-		_, babylonChain, czChain, zcKeeper := SetupTest(t)
+		_, babylonChain, czChain, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
 
 		ctx := babylonChain.GetContext()
 		hooks := zcKeeper.Hooks()

x/zoneconcierge/keeper/grpc_query.go

@@ -97,6 +97,30 @@ func (k Keeper) ListHeaders(c context.Context, req *types.QueryListHeadersReques
 	return resp, nil
 }
 
+// ListEpochHeaders returns all headers of a chain with given ID
+// TODO: support pagination in this RPC
+func (k Keeper) ListEpochHeaders(c context.Context, req *types.QueryListEpochHeadersRequest) (*types.QueryListEpochHeadersResponse, error) {
+	if req == nil {
+		return nil, status.Error(codes.InvalidArgument, "invalid request")
+	}
+
+	if len(req.ChainId) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "chain ID cannot be empty")
+	}
+
+	ctx := sdk.UnwrapSDKContext(c)
+
+	headers, err := k.GetEpochHeaders(ctx, req.ChainId, req.EpochNum)
+	if err != nil {
+		return nil, err
+	}
+
+	resp := &types.QueryListEpochHeadersResponse{
+		Headers: headers,
+	}
+	return resp, nil
+}
+
 func (k Keeper) FinalizedChainInfo(c context.Context, req *types.QueryFinalizedChainInfoRequest) (*types.QueryFinalizedChainInfoResponse, error) {
 	if req == nil {
 		return nil, status.Error(codes.InvalidArgument, "invalid request")
@@ -122,7 +146,7 @@ func (k Keeper) FinalizedChainInfo(c context.Context, req *types.QueryFinalizedC
 
 	// It's possible that the chain info's epoch is way before the last finalised epoch
 	// e.g., when there is no relayer for many epochs
-	// NOTE: if an epoch is finalisedm then all of its previous epochs are also finalised
+	// NOTE: if an epoch is finalised then all of its previous epochs are also finalised
 	if chainInfo.LatestHeader.BabylonEpoch < finalizedEpoch {
 		finalizedEpoch = chainInfo.LatestHeader.BabylonEpoch
 	}

x/zoneconcierge/keeper/grpc_query_test.go

@@ -11,6 +11,7 @@ import (
 	checkpointingtypes "github.com/babylonchain/babylon/x/checkpointing/types"
 	zctypes "github.com/babylonchain/babylon/x/zoneconcierge/types"
 	"github.com/cosmos/cosmos-sdk/types/query"
+	ibctmtypes "github.com/cosmos/ibc-go/v5/modules/light-clients/07-tendermint/types"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 	tmcrypto "github.com/tendermint/tendermint/proto/tendermint/crypto"
@@ -24,7 +25,8 @@ func FuzzChainList(f *testing.F) {
 	f.Fuzz(func(t *testing.T, seed int64) {
 		rand.Seed(seed)
 
-		_, babylonChain, _, zcKeeper := SetupTest(t)
+		_, babylonChain, _, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
 
 		ctx := babylonChain.GetContext()
 		hooks := zcKeeper.Hooks()
@@ -66,7 +68,8 @@ func FuzzChainInfo(f *testing.F) {
 	f.Fuzz(func(t *testing.T, seed int64) {
 		rand.Seed(seed)
 
-		_, babylonChain, czChain, zcKeeper := SetupTest(t)
+		_, babylonChain, czChain, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
 
 		ctx := babylonChain.GetContext()
 		hooks := zcKeeper.Hooks()
@@ -91,7 +94,8 @@ func FuzzEpochChainInfo(f *testing.F) {
 	f.Fuzz(func(t *testing.T, seed int64) {
 		rand.Seed(seed)
 
-		_, babylonChain, czChain, zcKeeper := SetupTest(t)
+		_, babylonChain, czChain, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
 
 		ctx := babylonChain.GetContext()
 		hooks := zcKeeper.Hooks()
@@ -137,7 +141,8 @@ func FuzzListHeaders(f *testing.F) {
 	f.Fuzz(func(t *testing.T, seed int64) {
 		rand.Seed(seed)
 
-		_, babylonChain, czChain, zcKeeper := SetupTest(t)
+		_, babylonChain, czChain, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
 
 		ctx := babylonChain.GetContext()
 		hooks := zcKeeper.Hooks()
@@ -164,6 +169,78 @@ func FuzzListHeaders(f *testing.F) {
 	})
 }
 
+func FuzzListEpochHeaders(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+
+	f.Fuzz(func(t *testing.T, seed int64) {
+		rand.Seed(seed)
+
+		_, babylonChain, czChain, babylonApp := SetupTest(t)
+		zcKeeper := babylonApp.ZoneConciergeKeeper
+		epochingKeeper := babylonApp.EpochingKeeper
+
+		ctx := babylonChain.GetContext()
+		hooks := zcKeeper.Hooks()
+
+		numReqs := datagen.RandomInt(5) + 1
+
+		epochNumList := []uint64{datagen.RandomInt(10) + 1}
+		nextHeightList := []uint64{0}
+		numHeadersList := []uint64{}
+		expectedHeadersMap := map[uint64][]*ibctmtypes.Header{}
+		numForkHeadersList := []uint64{}
+
+		// we test the scenario of ending an epoch for multiple times, in order to ensure that
+		// consecutive epoch infos do not affect each other.
+		for i := uint64(0); i < numReqs; i++ {
+			epochNum := epochNumList[i]
+			// enter a random epoch
+			if i == 0 {
+				for j := uint64(0); j < epochNum; j++ {
+					epochingKeeper.IncEpoch(ctx)
+				}
+			} else {
+				for j := uint64(0); j < epochNum-epochNumList[i-1]; j++ {
+					epochingKeeper.IncEpoch(ctx)
+				}
+			}
+
+			// generate a random number of headers and fork headers
+			numHeadersList = append(numHeadersList, datagen.RandomInt(100)+1)
+			numForkHeadersList = append(numForkHeadersList, datagen.RandomInt(10)+1)
+			// trigger hooks to append these headers and fork headers
+			expectedHeaders, _ := SimulateHeadersAndForksViaHook(ctx, hooks, czChain.ChainID, nextHeightList[i], numHeadersList[i], numForkHeadersList[i])
+			expectedHeadersMap[epochNum] = expectedHeaders
+			// prepare nextHeight for the next request
+			nextHeightList = append(nextHeightList, nextHeightList[i]+numHeadersList[i])
+
+			// simulate the scenario that a random epoch has ended
+			hooks.AfterEpochEnds(ctx, epochNum)
+			// prepare epochNum for the next request
+			epochNumList = append(epochNumList, epochNum+datagen.RandomInt(10)+1)
+		}
+
+		// attest the correctness of epoch info for each tested epoch
+		for i := uint64(0); i < numReqs; i++ {
+			epochNum := epochNumList[i]
+			// make request
+			req := &zctypes.QueryListEpochHeadersRequest{
+				ChainId:  czChain.ChainID,
+				EpochNum: epochNum,
+			}
+			resp, err := zcKeeper.ListEpochHeaders(ctx, req)
+			require.NoError(t, err)
+
+			// check if the headers are same as expected
+			headers := resp.Headers
+			require.Equal(t, len(expectedHeadersMap[epochNum]), len(headers))
+			for j := 0; j < len(expectedHeadersMap[epochNum]); j++ {
+				require.Equal(t, expectedHeadersMap[epochNum][j].Header.LastCommitHash, headers[j].Hash)
+			}
+		}
+	})
+}
+
 func FuzzFinalizedChainInfo(f *testing.F) {
 	datagen.AddRandomSeedsToFuzzer(f, 10)
 

x/zoneconcierge/keeper/keeper_test.go

@@ -13,13 +13,13 @@ import (
 )
 
 // SetupTest creates a coordinator with 2 test chains, and a ZoneConcierge keeper.
-func SetupTest(t *testing.T) (*ibctesting.Coordinator, *ibctesting.TestChain, *ibctesting.TestChain, zckeeper.Keeper) {
-	var zcKeeper zckeeper.Keeper
+func SetupTest(t *testing.T) (*ibctesting.Coordinator, *ibctesting.TestChain, *ibctesting.TestChain, *app.BabylonApp) {
+	var bbnApp *app.BabylonApp
 	coordinator := ibctesting.NewCoordinator(t, 2)
 	// replace the first test chain with a Babylon chain
 	ibctesting.DefaultTestingAppInit = func() (ibctesting.TestingApp, map[string]json.RawMessage) {
 		babylonApp := app.Setup(t, false)
-		zcKeeper = babylonApp.ZoneConciergeKeeper
+		bbnApp = babylonApp
 		encCdc := app.MakeTestEncodingConfig()
 		genesis := app.NewDefaultGenesisState(encCdc.Marshaler)
 		return babylonApp, genesis
@@ -30,7 +30,7 @@ func SetupTest(t *testing.T) (*ibctesting.Coordinator, *ibctesting.TestChain, *i
 	babylonChain := coordinator.GetChain(ibctesting.GetChainID(1))
 	czChain := coordinator.GetChain(ibctesting.GetChainID(2))
 
-	return coordinator, babylonChain, czChain, zcKeeper
+	return coordinator, babylonChain, czChain, bbnApp
 }
 
 // SimulateHeadersViaHook generates a non-zero number of canonical headers via the hook

x/zoneconcierge/keeper/proof_tx_in_block_test.go

@@ -31,7 +31,9 @@ func TestProveTxInBlock(t *testing.T) {
 	err = testNetwork.WaitForNextBlock()
 	require.NoError(t, err)
 
-	_, babylonChain, _, zcKeeper := SetupTest(t)
+	_, babylonChain, _, babylonApp := SetupTest(t)
+	zcKeeper := babylonApp.ZoneConciergeKeeper
+
 	ctx := babylonChain.GetContext()
 
 	// construct client context