From 230f44930c8c1a0febc2decedc1f8e419d19c950 Mon Sep 17 00:00:00 2001
From: HuangYi <huang@crypto.com>
Date: Thu, 29 Feb 2024 14:14:35 +0800
Subject: [PATCH] Problem: hard fork logic don't check for chain-id

Solution:
- change the number to a callback
---
 simapp/app.go                        | 2 +-
 x/auth/migrations/v043/store_test.go | 2 +-
 x/gov/keeper/common_test.go          | 2 +-
 x/staking/common_test.go             | 2 +-
 x/staking/keeper/common_test.go      | 2 +-
 x/staking/keeper/grpc_query_test.go  | 2 +-
 x/staking/keeper/keeper.go           | 8 +++++---
 x/staking/keeper/slash.go            | 2 +-
 8 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/simapp/app.go b/simapp/app.go
index 7e7cf33d8efb..7237a6943f5e 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -264,7 +264,7 @@ func NewSimApp(
 		appCodec, keys[banktypes.StoreKey], app.AccountKeeper, app.GetSubspace(banktypes.ModuleName), app.ModuleAccountAddrs(),
 	)
 	stakingKeeper := stakingkeeper.NewKeeper(
-		appCodec, keys[stakingtypes.StoreKey], app.AccountKeeper, app.BankKeeper, app.GetSubspace(stakingtypes.ModuleName), 0,
+		appCodec, keys[stakingtypes.StoreKey], app.AccountKeeper, app.BankKeeper, app.GetSubspace(stakingtypes.ModuleName), nil,
 	)
 	app.MintKeeper = mintkeeper.NewKeeper(
 		appCodec, keys[minttypes.StoreKey], app.GetSubspace(minttypes.ModuleName), &stakingKeeper,
diff --git a/x/auth/migrations/v043/store_test.go b/x/auth/migrations/v043/store_test.go
index 3b29947d419a..c56389e3692d 100644
--- a/x/auth/migrations/v043/store_test.go
+++ b/x/auth/migrations/v043/store_test.go
@@ -658,7 +658,7 @@ func createValidator(t *testing.T, ctx sdk.Context, app *simapp.SimApp, powers i
 		app.AccountKeeper,
 		app.BankKeeper,
 		app.GetSubspace(stakingtypes.ModuleName),
-		0,
+		nil,
 	)
 
 	val1, err := stakingtypes.NewValidator(valAddrs[0], pks[0], stakingtypes.Description{})
diff --git a/x/gov/keeper/common_test.go b/x/gov/keeper/common_test.go
index d770c2986c8c..97fe9e536af2 100644
--- a/x/gov/keeper/common_test.go
+++ b/x/gov/keeper/common_test.go
@@ -48,7 +48,7 @@ func createValidators(t *testing.T, ctx sdk.Context, app *simapp.SimApp, powers
 		app.AccountKeeper,
 		app.BankKeeper,
 		app.GetSubspace(stakingtypes.ModuleName),
-		0,
+		nil,
 	)
 
 	val1, err := stakingtypes.NewValidator(valAddrs[0], pks[0], stakingtypes.Description{})
diff --git a/x/staking/common_test.go b/x/staking/common_test.go
index eb9f5a390426..b9c94c5bdd14 100644
--- a/x/staking/common_test.go
+++ b/x/staking/common_test.go
@@ -49,7 +49,7 @@ func getBaseSimappWithCustomKeeper(t *testing.T) (*codec.LegacyAmino, *simapp.Si
 		app.AccountKeeper,
 		app.BankKeeper,
 		app.GetSubspace(types.ModuleName),
-		0,
+		nil,
 	)
 	app.StakingKeeper.SetParams(ctx, types.DefaultParams())
 
diff --git a/x/staking/keeper/common_test.go b/x/staking/keeper/common_test.go
index c755d4d3f07f..9cff2cad93ea 100644
--- a/x/staking/keeper/common_test.go
+++ b/x/staking/keeper/common_test.go
@@ -31,7 +31,7 @@ func createTestInput(t *testing.T) (*codec.LegacyAmino, *simapp.SimApp, sdk.Cont
 		app.AccountKeeper,
 		app.BankKeeper,
 		app.GetSubspace(types.ModuleName),
-		0,
+		nil,
 	)
 	return app.LegacyAmino(), app, ctx
 }
diff --git a/x/staking/keeper/grpc_query_test.go b/x/staking/keeper/grpc_query_test.go
index fd255350cbcd..eedd95057dc1 100644
--- a/x/staking/keeper/grpc_query_test.go
+++ b/x/staking/keeper/grpc_query_test.go
@@ -796,7 +796,7 @@ func createValidators(t *testing.T, ctx sdk.Context, app *simapp.SimApp, powers
 		app.AccountKeeper,
 		app.BankKeeper,
 		app.GetSubspace(types.ModuleName),
-		0,
+		nil,
 	)
 
 	val1 := teststaking.NewValidator(t, valAddrs[0], pks[0])
diff --git a/x/staking/keeper/keeper.go b/x/staking/keeper/keeper.go
index d82c2f0b3591..ad9527fd251e 100644
--- a/x/staking/keeper/keeper.go
+++ b/x/staking/keeper/keeper.go
@@ -28,14 +28,16 @@ type Keeper struct {
 	hooks      types.StakingHooks
 	paramstore paramtypes.Subspace
 
-	hardForkHeight int64
+	// forkEnabledFunc returns true if the forked logic should be enabled.
+	// see: https://github.com/advisories/GHSA-86h5-xcpx-cfqc
+	forkEnabledFunc func(sdk.Context) bool
 }
 
 // NewKeeper creates a new staking Keeper instance
 func NewKeeper(
 	cdc codec.BinaryCodec, key storetypes.StoreKey, ak types.AccountKeeper, bk types.BankKeeper,
 	ps paramtypes.Subspace,
-	hardForkHeight int64,
+	forkEnabledFunc func(sdk.Context) bool,
 ) Keeper {
 	// set KeyTable if it has not already been set
 	if !ps.HasKeyTable() {
@@ -59,7 +61,7 @@ func NewKeeper(
 		paramstore: ps,
 		hooks:      nil,
 
-		hardForkHeight: hardForkHeight,
+		forkEnabledFunc: forkEnabledFunc,
 	}
 }
 
diff --git a/x/staking/keeper/slash.go b/x/staking/keeper/slash.go
index 4824600ac4f4..23049bdfeb49 100644
--- a/x/staking/keeper/slash.go
+++ b/x/staking/keeper/slash.go
@@ -255,7 +255,7 @@ func (k Keeper) SlashRedelegation(ctx sdk.Context, srcValidator types.Validator,
 			panic(err)
 		}
 
-		if k.hardForkHeight >= ctx.BlockHeight() {
+		if k.forkEnabledFunc != nil && k.forkEnabledFunc(ctx) {
 			// Handle undelegation after redelegation
 			// Prioritize slashing unbondingDelegation than delegation
 			unbondingDelegation, found := k.GetUnbondingDelegation(ctx, sdk.MustAccAddressFromBech32(redelegation.DelegatorAddress), valDstAddr)