zoneconcierge/epoching: proof that a header is in an epoch (#248)

proto/babylon/epoching/v1/epoching.proto

@@ -16,9 +16,12 @@ message Epoch {
   // Babylon needs to remember the last header of each epoch to complete unbonding validators/delegations when a previous epoch's checkpoint is finalised.
   // The last_block_header field is nil in the epoch's beginning, and is set upon the end of this epoch.
   tendermint.types.Header last_block_header = 4;
+  // app_hash_root is the Merkle root of all AppHashs in this epoch
+  // It will be used for proving a block is in an epoch
+  bytes app_hash_root = 5;
   // sealer_header is the 2nd header of the next epoch
   // This validator set has generated a BLS multisig on `last_commit_hash` of the sealer header
-  tendermint.types.Header sealer_header = 5;
+  tendermint.types.Header sealer_header = 6;
 }
 
 // QueuedMessage is a message that can change the validator set and is delayed to the epoch boundary

proto/babylon/zoneconcierge/query.proto

@@ -92,8 +92,8 @@ message QueryFinalizedChainInfoResponse {
   */
   // proof_tx_in_block is the proof that tx that carries the header is included in a certain Babylon block
   tendermint.types.TxProof proof_tx_in_block = 5;
-  // proof_block_in_epoch is the proof that the Babylon block is in a certain epoch
-  tendermint.crypto.ProofOps proof_block_in_epoch = 6;
+  // proof_header_in_epoch is the proof that the Babylon header is in a certain epoch
+  tendermint.crypto.Proof proof_header_in_epoch = 6;
   // proof_epoch_sealed is the proof that the epoch is sealed
   babylon.zoneconcierge.v1.ProofEpochSealed proof_epoch_sealed = 7;
   // proof_epoch_submitted is the proof that the epoch's checkpoint is included in BTC ledger

proto/babylon/zoneconcierge/zoneconcierge.proto

@@ -17,8 +17,8 @@ message IndexedHeader {
   // height is the height of this header on CZ ledger
   // (hash, height) jointly provides the position of the header on CZ ledger
   uint64 height = 3;
-  // babylon_block_height is the height of the Babylon block that includes this header
-  uint64 babylon_block_height = 4;
+  // babylon_header is the header of the babylon block that includes this CZ header
+  tendermint.types.Header babylon_header = 4;
   // epoch is the epoch number of this header on Babylon ledger
   uint64 babylon_epoch = 5;
   // babylon_tx_hash is the hash of the tx that includes this header

x/epoching/abci.go

@@ -14,6 +14,7 @@ import (
 
 // BeginBlocker is called at the beginning of every block.
 // Upon each BeginBlock,
+// - record the current AppHash
 // - if reaching the epoch beginning, then
 //    - increment epoch number
 //    - trigger AfterEpochBegins hook
@@ -24,6 +25,9 @@ import (
 func BeginBlocker(ctx sdk.Context, k keeper.Keeper, req abci.RequestBeginBlock) {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyBeginBlocker)
 
+	// record the current AppHash
+	k.RecordAppHash(ctx)
+
 	// if this block is the first block of the next epoch
 	// note that we haven't incremented the epoch number yet
 	epoch := k.GetEpoch(ctx)
@@ -68,8 +72,10 @@ func EndBlocker(ctx sdk.Context, k keeper.Keeper) []abci.ValidatorUpdate {
 	// if reaching an epoch boundary, then
 	epoch := k.GetEpoch(ctx)
 	if epoch.IsLastBlock(ctx) {
-		// finalise this epoch, i.e., record the current header
-		k.RecordLastBlockHeader(ctx)
+		// finalise this epoch, i.e., record the current header and the Merkle root of all AppHashs in this epoch
+		if err := k.RecordLastHeaderAndAppHashRoot(ctx); err != nil {
+			panic(err)
+		}
 		// get all msgs in the msg queue
 		queuedMsgs := k.GetCurrentEpochMsgs(ctx)
 		// forward each msg in the msg queue to the right keeper

x/epoching/keeper/apphash_chain.go

@@ -0,0 +1,109 @@
+package keeper
+
+import (
+	"crypto/sha256"
+	"fmt"
+
+	"github.com/babylonchain/babylon/x/epoching/types"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/tendermint/tendermint/crypto/merkle"
+	tmcrypto "github.com/tendermint/tendermint/proto/tendermint/crypto"
+)
+
+func (k Keeper) setAppHash(ctx sdk.Context, height uint64, appHash []byte) {
+	store := k.appHashStore(ctx)
+	heightBytes := sdk.Uint64ToBigEndian(height)
+	store.Set(heightBytes, appHash)
+}
+
+// GetAppHash gets the AppHash of the header at the given height
+func (k Keeper) GetAppHash(ctx sdk.Context, height uint64) ([]byte, error) {
+	store := k.appHashStore(ctx)
+	heightBytes := sdk.Uint64ToBigEndian(height)
+	appHash := store.Get(heightBytes)
+	if appHash == nil {
+		return nil, sdkerrors.Wrapf(types.ErrInvalidHeight, "height %d is now known in DB yet", height)
+	}
+	return appHash, nil
+}
+
+// RecordAppHash stores the AppHash of the current header to KVStore
+func (k Keeper) RecordAppHash(ctx sdk.Context) {
+	header := ctx.BlockHeader()
+	height := uint64(header.Height)
+	k.setAppHash(ctx, height, header.AppHash)
+}
+
+// GetAllAppHashsForEpoch fetches all AppHashs in the given epoch
+func (k Keeper) GetAllAppHashsForEpoch(ctx sdk.Context, epoch *types.Epoch) ([][]byte, error) {
+	// if this epoch is the most recent AND has not ended, then we cannot get all AppHashs for this epoch
+	if k.GetEpoch(ctx).EpochNumber == epoch.EpochNumber && !epoch.IsLastBlock(ctx) {
+		return nil, sdkerrors.Wrapf(types.ErrInvalidHeight, "GetAllAppHashsForEpoch can only be invoked when this epoch has ended")
+	}
+
+	// fetch each AppHash in this epoch
+	appHashs := [][]byte{}
+	for i := epoch.FirstBlockHeight; i <= uint64(epoch.LastBlockHeader.Height); i++ {
+		appHash, err := k.GetAppHash(ctx, i)
+		if err != nil {
+			return nil, err
+		}
+		appHashs = append(appHashs, appHash)
+	}
+
+	return appHashs, nil
+}
+
+// ProveAppHashInEpoch generates a proof that the given appHash is in a given epoch
+func (k Keeper) ProveAppHashInEpoch(ctx sdk.Context, height uint64, epochNumber uint64) (*tmcrypto.Proof, error) {
+	// ensure height is inside this epoch
+	epoch, err := k.GetHistoricalEpoch(ctx, epochNumber)
+	if err != nil {
+		return nil, err
+	}
+	if !epoch.WithinBoundary(height) {
+		return nil, sdkerrors.Wrapf(types.ErrInvalidHeight, "the given height %d is not in epoch %d (interval [%d, %d])", height, epoch.EpochNumber, epoch.FirstBlockHeight, uint64(epoch.LastBlockHeader.Height))
+	}
+
+	// calculate index of this height in this epoch
+	idx := height - epoch.FirstBlockHeight
+
+	// fetch all AppHashs, calculate Merkle tree and proof
+	appHashs, err := k.GetAllAppHashsForEpoch(ctx, epoch)
+	if err != nil {
+		return nil, err
+	}
+	_, proofs := merkle.ProofsFromByteSlices(appHashs)
+
+	return proofs[idx].ToProto(), nil
+}
+
+// VerifyAppHashInclusion verifies whether the given appHash is in the Merkle tree w.r.t. the appHashRoot
+func VerifyAppHashInclusion(appHash []byte, appHashRoot []byte, proof *tmcrypto.Proof) error {
+	if len(appHash) != sha256.Size {
+		return fmt.Errorf("appHash with length %d is not a Sha256 hash", len(appHash))
+	}
+	if len(appHashRoot) != sha256.Size {
+		return fmt.Errorf("appHash with length %d is not a Sha256 hash", len(appHashRoot))
+	}
+	if proof == nil {
+		return fmt.Errorf("proof is nil")
+	}
+
+	unwrappedProof, err := merkle.ProofFromProto(proof)
+	if err != nil {
+		return fmt.Errorf("failed to unwrap proof: %w", err)
+	}
+	return unwrappedProof.Verify(appHashRoot, appHash)
+}
+
+// appHashStore returns the KVStore for the AppHash of each header
+// prefix: AppHashKey
+// key: height
+// value: AppHash in bytes
+func (k Keeper) appHashStore(ctx sdk.Context) prefix.Store {
+	store := ctx.KVStore(k.storeKey)
+	return prefix.NewStore(store, types.AppHashKey)
+}

x/epoching/keeper/apphash_chain_test.go

@@ -0,0 +1,58 @@
+package keeper_test
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/babylonchain/babylon/testutil/datagen"
+	"github.com/babylonchain/babylon/x/epoching/keeper"
+	"github.com/babylonchain/babylon/x/epoching/testepoching"
+	"github.com/babylonchain/babylon/x/epoching/types"
+	"github.com/stretchr/testify/require"
+)
+
+func FuzzAppHashChain(f *testing.F) {
+	datagen.AddRandomSeedsToFuzzer(f, 10)
+
+	f.Fuzz(func(t *testing.T, seed int64) {
+		rand.Seed(seed)
+
+		helper := testepoching.NewHelper(t)
+		ctx, k := helper.Ctx, helper.EpochingKeeper
+		// ensure that the epoch info is correct at the genesis
+		epoch := k.GetEpoch(ctx)
+		require.Equal(t, epoch.EpochNumber, uint64(0))
+		require.Equal(t, epoch.FirstBlockHeight, uint64(0))
+
+		// set a random epoch interval
+		epochInterval := rand.Uint64()%100 + 2 // the epoch interval should at at least 2
+		k.SetParams(ctx, types.Params{
+			EpochInterval: epochInterval,
+		})
+
+		// reach the end of the 1st epoch
+		expectedHeight := epochInterval
+		expectedAppHashs := [][]byte{}
+		for i := uint64(0); i < expectedHeight; i++ {
+			ctx = helper.GenAndApplyEmptyBlock()
+			expectedAppHashs = append(expectedAppHashs, ctx.BlockHeader().AppHash)
+		}
+		// ensure epoch number is 1
+		epoch = k.GetEpoch(ctx)
+		require.Equal(t, uint64(1), epoch.EpochNumber)
+
+		// ensure appHashs are same as expectedAppHashs
+		appHashs, err := k.GetAllAppHashsForEpoch(ctx, epoch)
+		require.NoError(t, err)
+		require.Equal(t, expectedAppHashs, appHashs)
+
+		// ensure prover and verifier are correct
+		randomHeightInEpoch := uint64(rand.Intn(int(expectedHeight)) + 1)
+		randomAppHash, err := k.GetAppHash(ctx, randomHeightInEpoch)
+		require.NoError(t, err)
+		proof, err := k.ProveAppHashInEpoch(ctx, randomHeightInEpoch, epoch.EpochNumber)
+		require.NoError(t, err)
+		err = keeper.VerifyAppHashInclusion(randomAppHash, epoch.AppHashRoot, proof)
+		require.NoError(t, err)
+	})
+}

x/epoching/keeper/epochs.go

@@ -4,6 +4,8 @@ import (
 	"github.com/babylonchain/babylon/x/epoching/types"
 	"github.com/cosmos/cosmos-sdk/store/prefix"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/tendermint/tendermint/crypto/merkle"
 )
 
 const (
@@ -75,15 +77,26 @@ func (k Keeper) GetHistoricalEpoch(ctx sdk.Context, epochNumber uint64) (*types.
 	return epoch, err
 }
 
-func (k Keeper) RecordLastBlockHeader(ctx sdk.Context) *types.Epoch {
+// RecordLastHeaderAndAppHashRoot records the last header and Merkle root of all AppHashs
+// for the current epoch, and stores the epoch metadata to KVStore
+func (k Keeper) RecordLastHeaderAndAppHashRoot(ctx sdk.Context) error {
 	epoch := k.GetEpoch(ctx)
 	if !epoch.IsLastBlock(ctx) {
-		panic("RecordLastBlockHeader can only be invoked at the last block of an epoch")
+		return sdkerrors.Wrapf(types.ErrInvalidHeight, "RecordLastBlockHeader can only be invoked at the last block of an epoch")
 	}
+	// record last block header
 	header := ctx.BlockHeader()
 	epoch.LastBlockHeader = &header
+	// calculate and record the Merkle root
+	appHashs, err := k.GetAllAppHashsForEpoch(ctx, epoch)
+	if err != nil {
+		return err
+	}
+	appHashRoot := merkle.HashFromByteSlices(appHashs)
+	epoch.AppHashRoot = appHashRoot
+	// save back to KVStore
 	k.setEpochInfo(ctx, epoch.EpochNumber, epoch)
-	return epoch
+	return nil
 }
 
 // RecordSealerHeaderForPrevEpoch records the sealer header for the previous epoch,

x/epoching/testepoching/helper.go

@@ -1,9 +1,11 @@
 package testepoching
 
 import (
-	"github.com/babylonchain/babylon/crypto/bls12381"
 	"testing"
 
+	"github.com/babylonchain/babylon/crypto/bls12381"
+	"github.com/babylonchain/babylon/testutil/datagen"
+
 	"cosmossdk.io/math"
 	appparams "github.com/babylonchain/babylon/app/params"
 
@@ -110,7 +112,7 @@ func (h *Helper) GenAndApplyEmptyBlock() sdk.Context {
 	valhash := CalculateValHash(valSet)
 	newHeader := tmproto.Header{
 		Height:             newHeight,
-		AppHash:            h.App.LastCommitID().Hash,
+		AppHash:            datagen.GenRandomByteArray(32),
 		ValidatorsHash:     valhash,
 		NextValidatorsHash: valhash,
 	}
@@ -129,7 +131,7 @@ func (h *Helper) BeginBlock() sdk.Context {
 	valhash := CalculateValHash(valSet)
 	newHeader := tmproto.Header{
 		Height:             newHeight,
-		AppHash:            h.App.LastCommitID().Hash,
+		AppHash:            datagen.GenRandomByteArray(32),
 		ValidatorsHash:     valhash,
 		NextValidatorsHash: valhash,
 	}

x/epoching/types/epoching.go

@@ -70,6 +70,15 @@ func (e Epoch) IsFirstBlockOfNextEpoch(ctx sdk.Context) bool {
 	}
 }
 
+// WithinBoundary checks whether the given height is within this epoch or not
+func (e Epoch) WithinBoundary(height uint64) bool {
+	if height < e.FirstBlockHeight || height > uint64(e.LastBlockHeader.Height) {
+		return false
+	} else {
+		return true
+	}
+}
+
 // ValidateBasic does sanity checks on Epoch
 func (e Epoch) ValidateBasic() error {
 	if e.CurrentEpochInterval < 2 {