refactor ValidateUpdatePessimistically to use single state info

refactor canonical client check

x/lightclient/ante/ibc_msg_update_client.go

@@ -9,7 +9,6 @@ import (
 	ibcclienttypes "github.com/cosmos/ibc-go/v7/modules/core/02-client/types"
 	ibctm "github.com/cosmos/ibc-go/v7/modules/light-clients/07-tendermint"
 	"github.com/dymensionxyz/dymension/v3/x/lightclient/types"
-	rollapptypes "github.com/dymensionxyz/dymension/v3/x/rollapp/types"
 	sequencertypes "github.com/dymensionxyz/dymension/v3/x/sequencer/types"
 	"github.com/dymensionxyz/gerr-cosmos/gerrc"
 )
@@ -74,18 +73,17 @@ func (i IBCMessagesDecorator) HandleMsgUpdateClient(ctx sdk.Context, msg *ibccli
 	}
 
 	h := header.GetHeight().GetRevisionHeight()
-	stateInfos, err := i.getStateInfos(ctx, rollapp.RollappId, h)
-	if err != nil {
-		return errorsmod.Wrap(err, "get state infos")
-	}
+	sInfo, err := i.raK.FindStateInfoByHeight(ctx, rollapp.RollappId, h)
+	if errorsmod.IsOf(err, gerrc.ErrNotFound) {
 
-	if stateInfos.containingHPlus1 != nil {
-		// the header is pessimistic: the state update has already been received, so we check the header doesn't mismatch
-		return errorsmod.Wrap(i.validateUpdatePessimistically(ctx, stateInfos, header.ConsensusState(), h), "validate pessimistic")
+		// the header is optimistic: the state update has not yet been received, so we save optimistically
+		return errorsmod.Wrap(i.k.SaveSigner(ctx, seq.Address, msg.ClientId, h), "save updater")
+	}
+	if err != nil {
+		return errorsmod.Wrap(err, "find state info by height")
 	}
 
-	// the header is optimistic: the state update has not yet been received, so we save optimistically
-	return errorsmod.Wrap(i.k.SaveSigner(ctx, seq.Address, msg.ClientId, h), "save updater")
+	return errorsmod.Wrap(i.k.ValidateUpdatePessimistically(ctx, sInfo, header.ConsensusState(), h), "validate pessimistic")
 }
 
 var (
@@ -119,45 +117,3 @@ func getHeader(msg *ibcclienttypes.MsgUpdateClient) (*ibctm.Header, error) {
 	}
 	return header, nil
 }
-
-// if containingHPlus1 is not nil then containingH also guaranteed to not be nil
-type stateInfos struct {
-	containingH      *rollapptypes.StateInfo
-	containingHPlus1 *rollapptypes.StateInfo
-}
-
-// getStateInfos gets state infos for h and h+1
-func (i IBCMessagesDecorator) getStateInfos(ctx sdk.Context, rollapp string, h uint64) (stateInfos, error) {
-	// Check if there are existing block descriptors for the given height of client state
-	s0, err := i.raK.FindStateInfoByHeight(ctx, rollapp, h)
-	if errorsmod.IsOf(err, gerrc.ErrNotFound) {
-		return stateInfos{}, nil
-	}
-	if err != nil {
-		return stateInfos{}, err
-	}
-	s1 := s0
-	if !s1.ContainsHeight(h + 1) {
-		s1, err = i.raK.FindStateInfoByHeight(ctx, rollapp, h+1)
-		if errorsmod.IsOf(err, gerrc.ErrNotFound) {
-			return stateInfos{s0, nil}, nil
-		}
-		if err != nil {
-			return stateInfos{}, err
-		}
-	}
-	return stateInfos{s0, s1}, nil
-}
-
-func (i IBCMessagesDecorator) validateUpdatePessimistically(ctx sdk.Context, infos stateInfos, consState *ibctm.ConsensusState, h uint64) error {
-	bd, _ := infos.containingH.GetBlockDescriptor(h)
-	seq, err := i.k.SeqK.RealSequencer(ctx, infos.containingHPlus1.Sequencer)
-	if err != nil {
-		return errorsmod.Wrap(errors.Join(err, gerrc.ErrInternal), "get sequencer of state info")
-	}
-	rollappState := types.RollappState{
-		BlockDescriptor:    bd,
-		NextBlockSequencer: seq,
-	}
-	return errorsmod.Wrap(types.CheckCompatibility(*consState, rollappState), "check compatibility")
-}

x/lightclient/keeper/canonical_client.go

@@ -5,29 +5,34 @@ import (
 
 	errorsmod "cosmossdk.io/errors"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/types/query"
-	ibcclienttypes "github.com/cosmos/ibc-go/v7/modules/core/02-client/types"
 	"github.com/cosmos/ibc-go/v7/modules/core/exported"
 	ibctm "github.com/cosmos/ibc-go/v7/modules/light-clients/07-tendermint"
 	"github.com/dymensionxyz/dymension/v3/x/lightclient/types"
+	rollapptypes "github.com/dymensionxyz/dymension/v3/x/rollapp/types"
+	"github.com/dymensionxyz/gerr-cosmos/gerrc"
 )
 
-// GetProspectiveCanonicalClient returns the client id of the first IBC client which can be set as the canonical client for the given rollapp.
+// FindMatchingClient returns the client id of the first IBC client which can be set as the canonical client for the given rollapp.
 // The canonical client criteria are:
 // 1. The client must be a tendermint client.
 // 2. The client state must match the expected client params as configured by the module
 // 3. All the existing consensus states much match the corresponding height rollapp block descriptors
-func (k Keeper) GetProspectiveCanonicalClient(ctx sdk.Context, rollappId string, maxHeight uint64) (clientID string, stateCompatible bool) {
+func (k Keeper) FindMatchingClient(ctx sdk.Context, sInfo *rollapptypes.StateInfo) (clientID string, stateCompatible bool) {
 	k.ibcClientKeeper.IterateClientStates(ctx, nil, func(client string, cs exported.ClientState) bool {
-		err := k.validClient(ctx, client, cs, rollappId, maxHeight)
-		if err != nil && !errorsmod.IsOf(err, errChainIDMismatch) {
-			ctx.Logger().Debug("tried to validate rollapp against light client for same chain id: rollapp: %s: client: %s", rollappId, client, "err", err)
-		}
+		err := k.validClient(ctx, client, cs, sInfo)
 		if err == nil {
 			clientID = client
 			stateCompatible = true
 			return true
 		}
+		if !errorsmod.IsOf(err, errChainIDMismatch) {
+			// Log the error with key-value pairs
+			ctx.Logger().Debug("tried to validate rollapp against light client for same chain id",
+				"rollapp", sInfo.GetRollappId(),
+				"client", client,
+				"err", err,
+			)
+		}
 		return false
 	})
 	return
@@ -68,7 +73,11 @@ func (k Keeper) expectedClient(ctx sdk.Context) ibctm.ClientState {
 
 var errChainIDMismatch = errors.New("chain id mismatch")
 
-func (k Keeper) validClient(ctx sdk.Context, clientID string, cs exported.ClientState, rollappId string, maxHeight uint64) error {
+func (k Keeper) validClient(ctx sdk.Context, clientID string, cs exported.ClientState, sInfo *rollapptypes.StateInfo) error {
+	maxHeight := sInfo.GetLatestHeight()
+	minHeight := sInfo.StartHeight
+	rollappId := sInfo.GetRollappId()
+
 	tmClientState, ok := cs.(*ibctm.ClientState)
 	if !ok {
 		return errors.New("not tm client")
@@ -78,56 +87,71 @@ func (k Keeper) validClient(ctx sdk.Context, clientID string, cs exported.Client
 	}
 
 	expClient := k.expectedClient(ctx)
-
 	if err := types.IsCanonicalClientParamsValid(tmClientState, &expClient); err != nil {
 		return errorsmod.Wrap(err, "params")
 	}
 
-	// FIXME: No need to get all consensus states. should iterate over the consensus states
-	res, err := k.ibcClientKeeper.ConsensusStateHeights(ctx, &ibcclienttypes.QueryConsensusStateHeightsRequest{
-		ClientId:   clientID,
-		Pagination: &query.PageRequest{Limit: maxHeight},
-	})
-	if err != nil {
-		return errorsmod.Wrap(err, "cons state heights")
-	}
 	atLeastOneMatch := false
-	for _, consensusHeight := range res.ConsensusStateHeights {
-		h := consensusHeight.GetRevisionHeight()
-		if maxHeight < h {
-			break
+	csStore := k.ibcClientKeeper.ClientStore(ctx, clientID)
+	var err error
+	IterateConsensusStateDescending(csStore, func(h exported.Height) bool {
+		// skip future heights
+		if h.GetRevisionHeight() >= maxHeight {
+			return false
 		}
-		consensusState, _ := k.ibcClientKeeper.GetClientConsensusState(ctx, clientID, consensusHeight)
-		tmConsensusState, _ := consensusState.(*ibctm.ConsensusState)
-		stateInfoH, err := k.rollappKeeper.FindStateInfoByHeight(ctx, rollappId, h)
-		if err != nil {
-			return errorsmod.Wrapf(err, "find state info by height h: %d", h)
+
+		// iterate until we pass the fraud height
+		if h.GetRevisionHeight() < minHeight {
+			return true
 		}
-		stateInfoHplus1, err := k.rollappKeeper.FindStateInfoByHeight(ctx, rollappId, h+1)
-		if err != nil {
-			return errorsmod.Wrapf(err, "find state info by height h+1: %d", h+1)
+
+		consensusState, ok := k.ibcClientKeeper.GetClientConsensusState(ctx, clientID, h)
+		if !ok {
+			return false
+		}
+		tmConsensusState, ok := consensusState.(*ibctm.ConsensusState)
+		if !ok {
+			return false
 		}
-		bd, _ := stateInfoH.GetBlockDescriptor(h)
 
-		nextSeq, err := k.SeqK.RealSequencer(ctx, stateInfoHplus1.Sequencer)
+		var stateInfo *rollapptypes.StateInfo
+		stateInfo, err = k.rollappKeeper.FindStateInfoByHeight(ctx, rollappId, h.GetRevisionHeight())
 		if err != nil {
-			return errorsmod.Wrap(err, "get sequencer")
-		}
-		rollappState := types.RollappState{
-			BlockDescriptor:    bd,
-			NextBlockSequencer: nextSeq,
+			err = errorsmod.Wrapf(err, "find state info by height h: %d", h.GetRevisionHeight())
+			return true
 		}
-		err = types.CheckCompatibility(*tmConsensusState, rollappState)
+
+		err = k.ValidateUpdatePessimistically(ctx, stateInfo, tmConsensusState, h.GetRevisionHeight())
 		if err != nil {
-			return errorsmod.Wrapf(err, "check compatibility: height: %d", h)
+			err = errorsmod.Wrapf(err, "validate pessimistic h: %d", h.GetRevisionHeight())
+			return true
 		}
+
 		atLeastOneMatch = true
-	}
+		return false
+	})
 	// Need to be sure that at least one consensus state agrees with a state update
 	// (There are also no disagreeing consensus states. There may be some consensus states
 	// for future state updates, which will incur a fraud if they disagree.)
 	if !atLeastOneMatch {
-		return errors.New("no matching consensus state found")
+		err = errors.Join(errors.New("no consensus state matches"), err)
+	}
+
+	if err != nil {
+		return errorsmod.Wrapf(err, "testing client %s for rollapp %s", clientID, rollappId)
 	}
 	return nil
 }
+
+func (k Keeper) ValidateUpdatePessimistically(ctx sdk.Context, sInfo *rollapptypes.StateInfo, consState *ibctm.ConsensusState, h uint64) error {
+	bd, _ := sInfo.GetBlockDescriptor(h)
+	nextSeq, err := k.SeqK.RealSequencer(ctx, sInfo.NextSequencerForHeight(h))
+	if err != nil {
+		return errorsmod.Wrap(errors.Join(err, gerrc.ErrInternal), "get sequencer of state info")
+	}
+	rollappState := types.RollappState{
+		BlockDescriptor:    bd,
+		NextBlockSequencer: nextSeq,
+	}
+	return errorsmod.Wrap(types.CheckCompatibility(*consState, rollappState), "check compatibility")
+}

x/lightclient/keeper/client_store.go

@@ -9,6 +9,21 @@ import (
 	ibctm "github.com/cosmos/ibc-go/v7/modules/light-clients/07-tendermint"
 )
 
+// IterateConsensusStateDescending iterates through all consensus states in descending order
+// until cb returns true.
+func IterateConsensusStateDescending(clientStore sdk.KVStore, cb func(height exported.Height) (stop bool)) {
+	iterator := sdk.KVStoreReversePrefixIterator(clientStore, []byte(ibctm.KeyIterateConsensusStatePrefix))
+	defer iterator.Close() // nolint: errcheck
+
+	for ; iterator.Valid(); iterator.Next() {
+		iterKey := iterator.Key()
+		height := ibctm.GetHeightFromIterationKey(iterKey)
+		if cb(height) {
+			break
+		}
+	}
+}
+
 // functions here copied from ibc-go/modules/core/02-client/keeper/
 // as we need direct access to the client store
 

x/lightclient/keeper/hook_listener.go

@@ -41,7 +41,7 @@ func (hook rollappHook) AfterUpdateState(
 
 	client, ok := hook.k.GetCanonicalClient(ctx, rollappId)
 	if !ok {
-		client, ok = hook.k.GetProspectiveCanonicalClient(ctx, rollappId, stateInfo.GetLatestHeight()-1)
+		client, ok = hook.k.FindMatchingClient(ctx, stateInfo)
 		if ok {
 			hook.k.SetCanonicalClient(ctx, rollappId, client)
 		}

x/lightclient/keeper/rollback.go

@@ -108,18 +108,3 @@ func (k Keeper) resetClientToValidState(clientStore sdk.KVStore, height uint64)
 
 	setClientState(clientStore, k.cdc, tmClientState)
 }
-
-// IterateConsensusStateDescending iterates through all consensus states in descending order
-// until cb returns true.
-func IterateConsensusStateDescending(clientStore sdk.KVStore, cb func(height exported.Height) (stop bool)) {
-	iterator := sdk.KVStoreReversePrefixIterator(clientStore, []byte(ibctm.KeyIterateConsensusStatePrefix))
-	defer iterator.Close() // nolint: errcheck
-
-	for ; iterator.Valid(); iterator.Next() {
-		iterKey := iterator.Key()
-		height := ibctm.GetHeightFromIterationKey(iterKey)
-		if cb(height) {
-			break
-		}
-	}
-}

x/rollapp/keeper/hard_fork.go

@@ -124,9 +124,9 @@ func (k Keeper) UpdateLastStateInfo(ctx sdk.Context, stateInfo *types.StateInfo,
 	if stateInfo.StartHeight == fraudHeight {
 		// If fraud height is at the beginning of the state info, return the previous index to keep
 		var ok bool
-		*stateInfo, ok = k.GetStateInfo(ctx, stateInfo.StateInfoIndex.RollappId, stateInfo.StateInfoIndex.Index-1)
+		*stateInfo, ok = k.GetStateInfo(ctx, stateInfo.GetRollappId(), stateInfo.StateInfoIndex.Index-1)
 		if !ok {
-			return nil, errorsmod.Wrapf(gerrc.ErrFailedPrecondition, "no state info found for rollapp: %s", stateInfo.StateInfoIndex.RollappId)
+			return nil, errorsmod.Wrapf(gerrc.ErrFailedPrecondition, "no state info found for rollapp: %s", stateInfo.GetRollappId())
 		}
 	} else if stateInfo.GetLatestHeight() >= fraudHeight {
 		// Remove block descriptors until the one we need to rollback to

x/rollapp/types/state_info.go

@@ -45,6 +45,10 @@ func (s *StateInfo) GetIndex() StateInfoIndex {
 	return s.StateInfoIndex
 }
 
+func (s *StateInfo) GetRollappId() string {
+	return s.StateInfoIndex.RollappId
+}
+
 func (s *StateInfo) GetLatestHeight() uint64 {
 	if s.StartHeight+s.NumBlocks > 0 {
 		return s.StartHeight + s.NumBlocks - 1
@@ -68,9 +72,16 @@ func (s *StateInfo) GetLatestBlockDescriptor() BlockDescriptor {
 	return s.BDs.BD[len(s.BDs.BD)-1]
 }
 
+func (s *StateInfo) NextSequencerForHeight(height uint64) string {
+	if height != s.GetLatestHeight() {
+		return s.Sequencer
+	}
+	return s.NextProposer
+}
+
 func (s *StateInfo) GetEvents() []sdk.Attribute {
 	eventAttributes := []sdk.Attribute{
-		sdk.NewAttribute(AttributeKeyRollappId, s.StateInfoIndex.RollappId),
+		sdk.NewAttribute(AttributeKeyRollappId, s.GetRollappId()),
 		sdk.NewAttribute(AttributeKeyStateInfoIndex, strconv.FormatUint(s.StateInfoIndex.Index, 10)),
 		sdk.NewAttribute(AttributeKeyStartHeight, strconv.FormatUint(s.StartHeight, 10)),
 		sdk.NewAttribute(AttributeKeyNumBlocks, strconv.FormatUint(s.NumBlocks, 10)),