From e8ac4dbc88b92bc1cbc0b867b541d3648b4a4eb7 Mon Sep 17 00:00:00 2001
From: Michael Tsitrin <michael@dymension.xyz>
Date: Fri, 17 Jan 2025 09:58:18 +0200
Subject: [PATCH] fixed consensus states iterator

---
 ibctesting/light_client_test.go          | 53 ++++++++++++++++++++++++
 x/lightclient/keeper/canonical_client.go |  6 +--
 x/lightclient/keeper/client_store.go     |  5 ++-
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/ibctesting/light_client_test.go b/ibctesting/light_client_test.go
index 6a53f91b7..81d8b7455 100644
--- a/ibctesting/light_client_test.go
+++ b/ibctesting/light_client_test.go
@@ -227,6 +227,59 @@ func (s *lightClientSuite) TestSetCanonicalClient_Succeeds() {
 	s.Equal(endpointA.ClientID, canonClientID)
 }
 
+func (s *lightClientSuite) TestSetCanonicalClient_MultipleClients_Succeeds() {
+	s.createRollapp(false, nil)
+	s.registerSequencer()
+
+	// Create multiple IBC endpoints
+	endpointA := ibctesting.NewEndpoint(s.hubChain(), &canonicalClientConfig, ibctesting.NewConnectionConfig(), ibctesting.NewChannelConfig())
+	endpointB := ibctesting.NewEndpoint(s.rollappChain(), ibctesting.NewTendermintConfig(), ibctesting.NewConnectionConfig(), ibctesting.NewChannelConfig())
+	endpointC := ibctesting.NewEndpoint(s.rollappChain(), ibctesting.NewTendermintConfig(), ibctesting.NewConnectionConfig(), ibctesting.NewChannelConfig()) // Additional client
+	endpointD := ibctesting.NewEndpoint(s.hubChain(), &canonicalClientConfig, ibctesting.NewConnectionConfig(), ibctesting.NewChannelConfig())               // Additional client
+
+	endpointA.Counterparty = endpointB
+	endpointB.Counterparty = endpointA
+	s.path = &ibctesting.Path{EndpointA: endpointA, EndpointB: endpointB}
+
+	// create dummy channels
+	endpointC.Counterparty = endpointD
+	endpointD.Counterparty = endpointC
+	s.NoError(endpointC.CreateClient())
+	s.NoError(endpointD.CreateClient())
+
+	currentHeader := s.rollappChain().CurrentHeader
+	startHeight := uint64(currentHeader.Height)
+	bd := rollapptypes.BlockDescriptor{Height: startHeight, StateRoot: currentHeader.AppHash, Timestamp: currentHeader.Time}
+
+	// Create the IBC clients
+	s.NoError(s.path.EndpointA.CreateClient())
+
+	currentHeader = s.rollappChain().CurrentHeader
+	bdNext := rollapptypes.BlockDescriptor{Height: uint64(currentHeader.Height), StateRoot: currentHeader.AppHash, Timestamp: currentHeader.Time}
+
+	// Update the rollapp state
+	msgUpdateState := rollapptypes.NewMsgUpdateState(
+		s.hubChain().SenderAccount.GetAddress().String(),
+		rollappChainID(),
+		"mock-da-path",
+		startHeight,
+		2,
+		&rollapptypes.BlockDescriptors{BD: []rollapptypes.BlockDescriptor{bd, bdNext}},
+	)
+	_, err := s.rollappMsgServer().UpdateState(s.hubCtx(), msgUpdateState)
+	s.Require().NoError(err)
+
+	setCanonMsg := &types.MsgSetCanonicalClient{
+		Signer: s.hubChain().SenderAccount.GetAddress().String(), ClientId: s.path.EndpointA.ClientID,
+	}
+	_, err = s.lightclientMsgServer().SetCanonicalClient(s.hubCtx(), setCanonMsg)
+	s.Require().NoError(err)
+
+	canonClientID, found := s.hubApp().LightClientKeeper.GetCanonicalClient(s.hubCtx(), s.rollappChain().ChainID)
+	s.Require().True(found)
+	s.Equal(endpointA.ClientID, canonClientID)
+}
+
 func (s *lightClientSuite) TestMsgUpdateClient_StateUpdateDoesntExist() {
 	s.createRollapp(false, nil)
 	s.registerSequencer()
diff --git a/x/lightclient/keeper/canonical_client.go b/x/lightclient/keeper/canonical_client.go
index a97e73f53..2e8adcd91 100644
--- a/x/lightclient/keeper/canonical_client.go
+++ b/x/lightclient/keeper/canonical_client.go
@@ -46,7 +46,7 @@ func (k *Keeper) TrySetCanonicalClient(ctx sdk.Context, clientID string) error {
 
 	err := k.validClient(ctx, clientID, clientState, rollappID)
 	if err != nil {
-		return errorsmod.Wrap(err, "unsafe to mark client canonical")
+		return errorsmod.Wrapf(err, "set client canonical")
 	}
 
 	k.SetCanonicalClient(ctx, rollappID, clientID)
@@ -124,8 +124,8 @@ func (k Keeper) validClient(ctx sdk.Context, clientID string, cs *ibctm.ClientSt
 			atLeastOneMatch = true
 		}
 
-		// break point with the lowest height of the consensus states
-		if sInfo.StartHeight > baseHeight {
+		// break point when we validate the state info for the first height of the client
+		if sInfo.StartHeight < baseHeight {
 			break
 		}
 	}
diff --git a/x/lightclient/keeper/client_store.go b/x/lightclient/keeper/client_store.go
index fa3f58ffd..9a6208cc6 100644
--- a/x/lightclient/keeper/client_store.go
+++ b/x/lightclient/keeper/client_store.go
@@ -110,7 +110,10 @@ func deleteIterationKey(clientStore sdk.KVStore, height exported.Height) {
 // GetFirstHeight returns the lowest height available for a client.
 func (k Keeper) GetFirstConsensusStateHeight(ctx sdk.Context, clientID string) uint64 {
 	height := clienttypes.Height{}
-	k.ibcClientKeeper.IterateConsensusStates(ctx, func(clientID string, cs clienttypes.ConsensusStateWithHeight) bool {
+	k.ibcClientKeeper.IterateConsensusStates(ctx, func(id string, cs clienttypes.ConsensusStateWithHeight) bool {
+		if id != clientID {
+			return false
+		}
 		height = cs.Height
 		return true
 	})