From 0ad4cb205d90c5334fdd76bcce74e9957da175a5 Mon Sep 17 00:00:00 2001
From: KonradStaniec <konrad.staniec@gmail.com>
Date: Mon, 6 Jan 2025 10:18:39 +0100
Subject: [PATCH] backport/fix epoching ante handler (#385) (#386)

- fixes epoching ante handler that returned early from ante handler even
in case of success validation of staking message
---
 CHANGELOG.md                                  |  2 +
 .../keeper/drop_validator_msg_decorator.go    |  5 ++-
 .../drop_validator_msg_decorator_test.go      | 45 +++++++++++++++++++
 3 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c7832d75a..b769ad203 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -43,6 +43,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 of the block in `TallyBlocks` function
 - [#378](https://github.com/babylonlabs-io/babylon/pull/378) Fix give out rewards
 with gaps of unfinalized blocks
+- [#385](https://github.com/babylonlabs-io/babylon/pull/385) Fix epoching module
+ante handler to return from antehandler chain only in case of error
 
 ## v1.0.0-rc2
 
diff --git a/x/epoching/keeper/drop_validator_msg_decorator.go b/x/epoching/keeper/drop_validator_msg_decorator.go
index 36479145d..07c9359ea 100644
--- a/x/epoching/keeper/drop_validator_msg_decorator.go
+++ b/x/epoching/keeper/drop_validator_msg_decorator.go
@@ -33,8 +33,9 @@ func (qmd DropValidatorMsgDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simu
 	}
 	// after genesis, if validator-related message, reject msg
 	for _, msg := range tx.GetMsgs() {
-		err := qmd.ValidateMsg(msg)
-		return ctx, err
+		if err := qmd.ValidateMsg(msg); err != nil {
+			return sdk.Context{}, err
+		}
 	}
 
 	return next(ctx, tx, simulate)
diff --git a/x/epoching/keeper/drop_validator_msg_decorator_test.go b/x/epoching/keeper/drop_validator_msg_decorator_test.go
index b1094e353..9b16e7c3b 100644
--- a/x/epoching/keeper/drop_validator_msg_decorator_test.go
+++ b/x/epoching/keeper/drop_validator_msg_decorator_test.go
@@ -3,6 +3,7 @@ package keeper
 import (
 	"testing"
 
+	appparams "github.com/babylonlabs-io/babylon/app/params"
 	"github.com/babylonlabs-io/babylon/x/epoching/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	authz "github.com/cosmos/cosmos-sdk/x/authz"
@@ -55,3 +56,47 @@ func TestDropValidatorMsgDecorator(t *testing.T) {
 		}
 	}
 }
+
+func TestReturnNilErrorNotCallNext(t *testing.T) {
+	decorator := NewDropValidatorMsgDecorator(&Keeper{})
+	ctx := sdk.Context{}.WithBlockHeight(1)
+	createValidatorMsg := &stakingtypes.MsgCreateValidator{}
+	encCfg := appparams.DefaultEncodingConfig()
+
+	builder := encCfg.TxConfig.NewTxBuilder()
+	builder.SetMsgs(createValidatorMsg)
+	tx := builder.GetTx()
+
+	nextAnteHandlerCalled := false
+	next := func(ctx sdk.Context, tx sdk.Tx, simulate bool) (sdk.Context, error) {
+		nextAnteHandlerCalled = true
+		return ctx, nil
+	}
+
+	ctx, err := decorator.AnteHandle(ctx, tx, false, next)
+	require.Error(t, err)
+	require.Equal(t, sdk.Context{}, ctx)
+	require.False(t, nextAnteHandlerCalled)
+}
+
+func TestReturnSuccessCallNext(t *testing.T) {
+	decorator := NewDropValidatorMsgDecorator(&Keeper{})
+	ctx := sdk.Context{}.WithBlockHeight(1)
+	createValidatorMsg := &stakingtypes.MsgEditValidator{}
+	encCfg := appparams.DefaultEncodingConfig()
+
+	builder := encCfg.TxConfig.NewTxBuilder()
+	builder.SetMsgs(createValidatorMsg)
+	tx := builder.GetTx()
+	nextAnteHandlerCalled := false
+
+	next := func(ctx sdk.Context, tx sdk.Tx, simulate bool) (sdk.Context, error) {
+		nextAnteHandlerCalled = true
+		return ctx, nil
+	}
+
+	ctx, err := decorator.AnteHandle(ctx, tx, false, next)
+	require.NoError(t, err)
+	require.Equal(t, ctx, ctx)
+	require.True(t, nextAnteHandlerCalled)
+}