From f611c494d5e98a17e1fbdd179a42240d5a0c0cd4 Mon Sep 17 00:00:00 2001
From: Michael Tsitrin <114929630+mtsitrin@users.noreply.github.com>
Date: Sun, 17 Nov 2024 20:37:19 +0200
Subject: [PATCH] fix(rotation): complete rotation on `afterStateUpdate` hook
 (#1493)

---
 x/rollapp/keeper/expected_keepers.go        |  1 +
 x/rollapp/keeper/liveness_test.go           |  5 ++++
 x/rollapp/keeper/msg_server_update_state.go |  8 +++---
 x/sequencer/keeper/hook_listener.go         | 28 +++++++++++++++++----
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/x/rollapp/keeper/expected_keepers.go b/x/rollapp/keeper/expected_keepers.go
index 46c5dbabf..b09d8bc90 100644
--- a/x/rollapp/keeper/expected_keepers.go
+++ b/x/rollapp/keeper/expected_keepers.go
@@ -17,6 +17,7 @@ type SequencerKeeper interface {
 	SlashLiveness(ctx sdk.Context, rollappID string) error
 	PunishSequencer(ctx sdk.Context, seqAddr string, rewardee *sdk.AccAddress) error
 	GetProposer(ctx sdk.Context, rollappId string) types.Sequencer
+	GetSuccessor(ctx sdk.Context, rollapp string) types.Sequencer
 }
 
 // BankKeeper defines the expected interface needed to retrieve account balances.
diff --git a/x/rollapp/keeper/liveness_test.go b/x/rollapp/keeper/liveness_test.go
index 272acebfc..6f4001c47 100644
--- a/x/rollapp/keeper/liveness_test.go
+++ b/x/rollapp/keeper/liveness_test.go
@@ -180,6 +180,11 @@ type livenessMockSequencerKeeper struct {
 	slashes map[string]int
 }
 
+// GetSuccessor implements keeper.SequencerKeeper.
+func (l livenessMockSequencerKeeper) GetSuccessor(ctx sdk.Context, rollapp string) seqtypes.Sequencer {
+	panic("unimplemented")
+}
+
 func (l livenessMockSequencerKeeper) PunishSequencer(ctx sdk.Context, seqAddr string, rewardee *sdk.AccAddress) error {
 	// TODO implement me
 	panic("implement me")
diff --git a/x/rollapp/keeper/msg_server_update_state.go b/x/rollapp/keeper/msg_server_update_state.go
index d199e9e38..1ab268cea 100644
--- a/x/rollapp/keeper/msg_server_update_state.go
+++ b/x/rollapp/keeper/msg_server_update_state.go
@@ -21,9 +21,7 @@ func (k msgServer) UpdateState(goCtx context.Context, msg *types.MsgUpdateState)
 	}
 
 	// call the before-update-state hook
-	// currently used by `x/sequencer` to:
-	// 1. validate the state update submitter
-	// 2. complete the rotation of the proposer if needed
+	// currently used by `x/sequencer` to validate the proposer
 	err := k.hooks.BeforeUpdateState(ctx, msg.Creator, msg.RollappId, msg.Last)
 	if err != nil {
 		return nil, errorsmod.Wrap(err, "before update state")
@@ -82,6 +80,9 @@ func (k msgServer) UpdateState(goCtx context.Context, msg *types.MsgUpdateState)
 	// by the sequencer rotation in k.hooks.BeforeUpdateState
 	// the proposer we get is the one that will propose the next block.
 	val := k.sequencerKeeper.GetProposer(ctx, msg.RollappId)
+	if msg.Last {
+		val = k.sequencerKeeper.GetSuccessor(ctx, msg.RollappId)
+	}
 
 	creationHeight := uint64(ctx.BlockHeight())
 	blockTime := ctx.BlockTime()
@@ -115,6 +116,7 @@ func (k msgServer) UpdateState(goCtx context.Context, msg *types.MsgUpdateState)
 
 	// call the after-update-state hook
 	// currently used by `x/lightclient` to validate the state update in regards to the light client
+	// x/sequencer will complete the rotation if needed
 	err = k.hooks.AfterUpdateState(ctx, msg.RollappId, stateInfo)
 	if err != nil {
 		return nil, errorsmod.Wrap(err, "hook: after update state")
diff --git a/x/sequencer/keeper/hook_listener.go b/x/sequencer/keeper/hook_listener.go
index 4dc208ddb..3c0e65f4c 100644
--- a/x/sequencer/keeper/hook_listener.go
+++ b/x/sequencer/keeper/hook_listener.go
@@ -5,6 +5,7 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	rollapptypes "github.com/dymensionxyz/dymension/v3/x/rollapp/types"
 	"github.com/dymensionxyz/dymension/v3/x/sequencer/types"
+	"github.com/dymensionxyz/gerr-cosmos/gerrc"
 )
 
 var _ rollapptypes.RollappHooks = rollappHook{}
@@ -18,17 +19,34 @@ func (k Keeper) RollappHooks() rollapptypes.RollappHooks {
 	return rollappHook{k: k}
 }
 
-// BeforeUpdateState will reject if the caller is not proposer, or if they are proposer but haven't
-// finished their rotation notice period.
-// If valid, it will set the successor as proposer
+// BeforeUpdateState will reject if the caller is not proposer
+// if lastStateUpdateBySequencer is true, validate that the sequencer is in the middle of a rotation
 func (hook rollappHook) BeforeUpdateState(ctx sdk.Context, seqAddr, rollappId string, lastStateUpdateBySequencer bool) error {
 	proposer := hook.k.GetProposer(ctx, rollappId)
 	if seqAddr != proposer.Address {
 		return types.ErrNotProposer
 	}
 
-	if lastStateUpdateBySequencer {
-		return errorsmod.Wrap(hook.k.OnProposerLastBlock(ctx, proposer), "on proposer last block")
+	// if lastStateUpdateBySequencer is true, validate that the sequencer is in the middle of a rotation
+	if lastStateUpdateBySequencer && !hook.k.AwaitingLastProposerBlock(ctx, rollappId) {
+		return errorsmod.Wrap(gerrc.ErrInvalidArgument, "sequencer is not in the middle of a rotation")
+	}
+
+	return nil
+}
+
+// AfterUpdateState checks if rotation is completed and the nextProposer is changed
+func (hook rollappHook) AfterUpdateState(ctx sdk.Context, rollappID string, stateInfo *rollapptypes.StateInfo) error {
+	// no proposer changed - no op
+	if stateInfo.Sequencer == stateInfo.NextProposer {
+		return nil
+	}
+
+	// handle proposer rotation completion
+	proposer := hook.k.GetProposer(ctx, rollappID)
+	err := hook.k.OnProposerLastBlock(ctx, proposer)
+	if err != nil {
+		return errorsmod.Wrap(err, "on proposer last block")
 	}
 
 	return nil