From 43991b99e4b2086959c8bceb4a6464808268062e Mon Sep 17 00:00:00 2001
From: Facundo Medica <14063057+facundomedica@users.noreply.github.com>
Date: Thu, 13 Jun 2024 13:27:44 +0200
Subject: [PATCH] fix(x/staking): stop validators from rotating to the same key
 on the same block (#20649)

---
 x/staking/keeper/cons_pubkey.go     | 14 +++++++++++++-
 x/staking/keeper/msg_server_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/x/staking/keeper/cons_pubkey.go b/x/staking/keeper/cons_pubkey.go
index 12a5baafe48c..966ffedeaccc 100644
--- a/x/staking/keeper/cons_pubkey.go
+++ b/x/staking/keeper/cons_pubkey.go
@@ -34,7 +34,19 @@ func (k Keeper) setConsPubKeyRotationHistory(
 		Height:          height,
 		Fee:             fee,
 	}
-	err := k.RotationHistory.Set(ctx, collections.Join(valAddr.Bytes(), height), history)
+
+	// check if there's another key rotation for this same key in the same block
+	allRotations, err := k.GetBlockConsPubKeyRotationHistory(ctx)
+	if err != nil {
+		return err
+	}
+	for _, r := range allRotations {
+		if r.NewConsPubkey.Compare(newPubKey) == 0 {
+			return types.ErrConsensusPubKeyAlreadyUsedForValidator
+		}
+	}
+
+	err = k.RotationHistory.Set(ctx, collections.Join(valAddr.Bytes(), height), history)
 	if err != nil {
 		return err
 	}
diff --git a/x/staking/keeper/msg_server_test.go b/x/staking/keeper/msg_server_test.go
index c22cc78b10a4..14163c6bcf42 100644
--- a/x/staking/keeper/msg_server_test.go
+++ b/x/staking/keeper/msg_server_test.go
@@ -1378,3 +1378,30 @@ func (s *KeeperTestSuite) TestConsKeyRotn() {
 		})
 	}
 }
+
+// TestConsKeyRotationInSameBlock tests the scenario where multiple validators try to
+// rotate to the **same** consensus key in the same block.
+func (s *KeeperTestSuite) TestConsKeyRotationInSameBlock() {
+	stakingKeeper, ctx := s.stakingKeeper, s.ctx
+
+	msgServer := stakingkeeper.NewMsgServerImpl(stakingKeeper)
+	s.setValidators(2)
+	validators, err := stakingKeeper.GetAllValidators(ctx)
+	s.Require().NoError(err)
+
+	s.Require().Len(validators, 2)
+
+	req, err := types.NewMsgRotateConsPubKey(validators[0].GetOperator(), PKs[444])
+	s.Require().NoError(err)
+
+	s.bankKeeper.EXPECT().SendCoinsFromAccountToModule(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+
+	_, err = msgServer.RotateConsPubKey(ctx, req)
+	s.Require().NoError(err)
+
+	req, err = types.NewMsgRotateConsPubKey(validators[1].GetOperator(), PKs[444])
+	s.Require().NoError(err)
+
+	_, err = msgServer.RotateConsPubKey(ctx, req)
+	s.Require().ErrorContains(err, "consensus pubkey is already used for a validator")
+}