From dc4ad457f5d3fc8ead91eb5ec46091d0935b489c Mon Sep 17 00:00:00 2001
From: Reece Williams <31943163+Reecepbcups@users.noreply.github.com>
Date: Sun, 14 Apr 2024 18:13:54 -0500
Subject: [PATCH] fix(abci v2): Fix  (#161)

* `UpdatedValidatorsCacheKey`

* ABCI: `DeleteValidatorByPowerIndex` logic

* validate ABCI events are not stuck

* wait for blocks jail test

* BeforeJailedValidatorsKey

* lint + fix remove val test (check bonded)

* lint

* logs

* touchups
---
 INTEGRATION.md          |  11 ++++
 e2e/poa_add_test.go     |   7 +++
 e2e/poa_jail_test.go    |  11 +++-
 e2e/poa_remove_test.go  |  18 +++++-
 keeper/keeper.go        |  14 +++--
 keeper/poa.go           |  47 +++++++++++---
 keeper/staking_hooks.go |  84 +++++++++++++++++++++++++
 keys.go                 |   5 ++
 module/abci.go          | 132 +++++++++++++++++++++++++++++++++-------
 module/module.go        |   5 ++
 simapp/app.go           |  10 ++-
 simapp/test_node.sh     |   7 ++-
 12 files changed, 310 insertions(+), 41 deletions(-)
 create mode 100644 keeper/staking_hooks.go

diff --git a/INTEGRATION.md b/INTEGRATION.md
index 1a855ca..dd93574 100644
--- a/INTEGRATION.md
+++ b/INTEGRATION.md
@@ -71,6 +71,16 @@ app.ModuleManager = module.NewManager(
 
 ...
 
+// Register POA Staking Hooks
+app.StakingKeeper.SetHooks(
+    stakingtypes.NewMultiStakingHooks(
+        ...
+        app.POAKeeper.Hooks(),
+    ),
+)
+
+...
+
 // Add PoA to BeginBlock logic
 // NOTE: This must be before the staking module begin blocker
 app.ModuleManager.SetOrderBeginBlockers(
@@ -80,6 +90,7 @@ app.ModuleManager.SetOrderBeginBlockers(
 )
 
 // Add PoA to end blocker logic
+// NOTE: This must be before the staking module end blocker
 app.ModuleManager.SetOrderEndBlockers(
     ...
     poa.ModuleName,
diff --git a/e2e/poa_add_test.go b/e2e/poa_add_test.go
index 4df990d..c465385 100644
--- a/e2e/poa_add_test.go
+++ b/e2e/poa_add_test.go
@@ -11,6 +11,7 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/strangelove-ventures/interchaintest/v8"
 	"github.com/strangelove-ventures/interchaintest/v8/chain/cosmos"
+	"github.com/strangelove-ventures/interchaintest/v8/testutil"
 	"github.com/strangelove-ventures/poa/e2e/helpers"
 	"github.com/stretchr/testify/require"
 )
@@ -95,4 +96,10 @@ func TestPOAAddValidator(t *testing.T) {
 	assertSignatures(t, ctx, chain, 1) // They are not signing, should be 0
 	assertConsensus(t, ctx, chain, 2)  // ensure they were added to CometBFT
 
+	// validate that the ABCI events are not "stuck" (where the event is not cleared by POA or x/staking)
+	_, err = helpers.POASetPower(t, ctx, chain, admin, pending[0].OperatorAddress, 2_000_000)
+	require.NoError(t, err)
+
+	require.NoError(t, testutil.WaitForBlocks(ctx, 4, chain))
+
 }
diff --git a/e2e/poa_jail_test.go b/e2e/poa_jail_test.go
index 0a72d81..2c4d811 100644
--- a/e2e/poa_jail_test.go
+++ b/e2e/poa_jail_test.go
@@ -1,6 +1,7 @@
 package e2e
 
 import (
+	"fmt"
 	"testing"
 	"time"
 
@@ -28,7 +29,7 @@ func TestPOAJailing(t *testing.T) {
 	updatedSlashingCfg.ModifyGenesis = cosmos.ModifyGenesis(append(defaultGenesis, []cosmos.GenesisKV{
 		{
 			Key:   "app_state.slashing.params.signed_blocks_window",
-			Value: "3",
+			Value: "10",
 		},
 		{
 			Key:   "app_state.slashing.params.min_signed_per_window",
@@ -67,7 +68,8 @@ func TestPOAJailing(t *testing.T) {
 	}
 
 	// Wait for the stopped node to be jailed & persist
-	require.NoError(t, testutil.WaitForBlocks(ctx, 5, chain.Validators[0]))
+	t.Log("Waiting for validator to become jailed")
+	require.NoError(t, testutil.WaitForBlocks(ctx, 15, chain.Validators[0]))
 
 	// Validate 1 validator is jailed (status 1)
 	vals := helpers.GetValidators(t, ctx, chain)
@@ -75,6 +77,7 @@ func TestPOAJailing(t *testing.T) {
 	require.True(t, func() bool {
 		for _, v := range vals.Validators {
 			if v.Status == int(stakingtypes.Unbonded) || v.Status == int(stakingtypes.Unbonding) {
+				fmt.Println("Validator", v.OperatorAddress, "is jailed", v.Status)
 				jailedValAddr = v.OperatorAddress
 				return true
 			}
@@ -94,4 +97,8 @@ func TestPOAJailing(t *testing.T) {
 			require.True(t, unjailTime.After(now))
 		}
 	}
+
+	// wait to ensure the chain is not halted
+	t.Log("Waiting for chain to progress")
+	require.NoError(t, testutil.WaitForBlocks(ctx, 5, chain))
 }
diff --git a/e2e/poa_remove_test.go b/e2e/poa_remove_test.go
index f253daf..956d7fc 100644
--- a/e2e/poa_remove_test.go
+++ b/e2e/poa_remove_test.go
@@ -1,6 +1,7 @@
 package e2e
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -8,6 +9,8 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/stretchr/testify/require"
 
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+
 	"github.com/strangelove-ventures/interchaintest/v8"
 	"github.com/strangelove-ventures/interchaintest/v8/chain/cosmos"
 	"github.com/strangelove-ventures/interchaintest/v8/testutil"
@@ -75,7 +78,7 @@ func TestPOARemoval(t *testing.T) {
 
 	assertSignatures(t, ctx, chain, initialValidators-1)
 	require.Equal(t, initialValidators-1, len(consensus.Validators), "BFT consensus should have one less validator")
-	require.Equal(t, initialValidators-1, len(vals), "Validators should have one less validator")
+	require.Equal(t, initialValidators-1, getBondedValidators(t, ctx, chain), "Bonded validators should have one less active validator")
 }
 
 func getValToRemove(t *testing.T, vals helpers.Validators, delegationAmt int64) string {
@@ -90,3 +93,16 @@ func getValToRemove(t *testing.T, vals helpers.Validators, delegationAmt int64)
 
 	return valToRemove
 }
+
+func getBondedValidators(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain) int {
+	vals := helpers.GetValidators(t, ctx, chain).Validators
+
+	bonded := 0
+	for _, v := range vals {
+		if v.Status == int(stakingtypes.Bonded) {
+			bonded++
+		}
+	}
+
+	return bonded
+}
diff --git a/keeper/keeper.go b/keeper/keeper.go
index 4df7495..d96be96 100644
--- a/keeper/keeper.go
+++ b/keeper/keeper.go
@@ -30,9 +30,11 @@ type Keeper struct {
 	logger log.Logger
 
 	// state management
-	Schema            collections.Schema
-	Params            collections.Item[poa.Params]
-	PendingValidators collections.Item[poa.Validators]
+	Schema                 collections.Schema
+	Params                 collections.Item[poa.Params]
+	PendingValidators      collections.Item[poa.Validators]
+	UpdatedValidatorsCache collections.KeySet[string]
+	BeforeJailedValidators collections.KeySet[string]
 
 	CachedBlockPower            collections.Item[poa.PowerCache]
 	AbsoluteChangedInBlockPower collections.Item[poa.PowerCache]
@@ -59,8 +61,10 @@ func NewKeeper(
 		logger:        logger,
 
 		// Stores
-		Params:            collections.NewItem(sb, poa.ParamsKey, "params", codec.CollValue[poa.Params](cdc)),
-		PendingValidators: collections.NewItem(sb, poa.PendingValidatorsKey, "pending", codec.CollValue[poa.Validators](cdc)),
+		Params:                 collections.NewItem(sb, poa.ParamsKey, "params", codec.CollValue[poa.Params](cdc)),
+		PendingValidators:      collections.NewItem(sb, poa.PendingValidatorsKey, "pending", codec.CollValue[poa.Validators](cdc)),
+		UpdatedValidatorsCache: collections.NewKeySet(sb, poa.UpdatedValidatorsCacheKey, "updated_validators", collections.StringKey),
+		BeforeJailedValidators: collections.NewKeySet(sb, poa.BeforeJailedValidatorsKey, "before_jailed", collections.StringKey),
 
 		CachedBlockPower:            collections.NewItem(sb, poa.CachedPreviousBlockPowerKey, "cached_block", codec.CollValue[poa.PowerCache](cdc)),
 		AbsoluteChangedInBlockPower: collections.NewItem(sb, poa.AbsoluteChangedInBlockPowerKey, "absolute_changed_power", codec.CollValue[poa.PowerCache](cdc)),
diff --git a/keeper/poa.go b/keeper/poa.go
index d4495b4..ca43c9e 100644
--- a/keeper/poa.go
+++ b/keeper/poa.go
@@ -72,6 +72,15 @@ func (k Keeper) SetPOAPower(ctx context.Context, valOpBech32 string, newShares i
 		return stakingtypes.Validator{}, err
 	}
 
+	// no need to process anything, same values
+	if newBFTConsensusPower == currentPower {
+		return val, fmt.Errorf("current power (%d) is the same as the new power (%d) for %s", currentPower, newBFTConsensusPower, valOpBech32)
+	}
+
+	// When we SetValidatorByPowerIndex, the Tokens are used to get the shares of power for CometBFT consensus (voting_power).
+	// We don't `k.stakingKeeper.SetValidator` since we only use this for CometBFT consensus power.
+	val.Tokens = sdkmath.NewIntFromUint64(uint64(newShares))
+
 	// slash all the validator's tokens (100%)
 	if newShares == 0 && currentPower > 0 {
 		pk, ok := val.ConsensusPubkey.GetCachedValue().(cryptotypes.PubKey)
@@ -79,6 +88,8 @@ func (k Keeper) SetPOAPower(ctx context.Context, valOpBech32 string, newShares i
 			return stakingtypes.Validator{}, fmt.Errorf("issue getting consensus pubkey for %s", valOpBech32)
 		}
 
+		consAddr := sdk.GetConsAddress(pk)
+
 		height := sdk.UnwrapSDKContext(ctx).BlockHeight()
 
 		normalizedToken := k.stakingKeeper.TokensFromConsensusPower(ctx, currentPower)
@@ -86,11 +97,31 @@ func (k Keeper) SetPOAPower(ctx context.Context, valOpBech32 string, newShares i
 		if _, err := k.stakingKeeper.Slash(ctx, sdk.GetConsAddress(pk), height, normalizedToken.Int64(), sdkmath.LegacyOneDec()); err != nil {
 			return stakingtypes.Validator{}, err
 		}
-	}
+		// TODO:
+		if err := k.stakingKeeper.DeleteLastValidatorPower(ctx, valAddr); err != nil {
+			return stakingtypes.Validator{}, err
+		}
+		if err := k.stakingKeeper.DeleteValidatorByPowerIndex(ctx, val); err != nil {
+			return stakingtypes.Validator{}, err
+		}
+		if err := k.slashKeeper.DeleteMissedBlockBitmap(ctx, consAddr); err != nil {
+			return stakingtypes.Validator{}, err
+		}
+	} else {
+		// Sets the new consensus power for the validator (this is executed in the x/staking ApplyAndReturnValidatorUpdates method)
+		if err := k.GetStakingKeeper().SetLastValidatorPower(ctx, valAddr, newBFTConsensusPower); err != nil {
+			return stakingtypes.Validator{}, err
+		}
+		if err := k.GetStakingKeeper().SetValidatorByPowerIndex(ctx, val); err != nil {
+			return stakingtypes.Validator{}, err
+		}
 
-	// Sets the new consensus power for the validator (this is executed in the x/staking ApplyAndReturnValidatorUpdates method)
-	if err := k.stakingKeeper.SetLastValidatorPower(ctx, valAddr, newBFTConsensusPower); err != nil {
-		return stakingtypes.Validator{}, err
+		// A cache to handle updated validators power. Once set, the next begin block will remove it from the cache.
+		// This allows us to Delete the validator index on the staking side, and ensures power updates do not persist over many blocks.
+		// This multi block persistence would break ABCI updates via x/staking if the validator's power is updated again, or removed.
+		if err := k.UpdatedValidatorsCache.Set(ctx, val.OperatorAddress); err != nil {
+			return stakingtypes.Validator{}, err
+		}
 	}
 
 	absPowerDiff := uint64(math.Abs(float64(newBFTConsensusPower - currentPower)))
@@ -129,6 +160,10 @@ func (k Keeper) AcceptNewValidator(ctx context.Context, operatingAddress string,
 		return err
 	}
 
+	if err := k.stakingKeeper.SetNewValidatorByPowerIndex(ctx, val); err != nil {
+		return err
+	}
+
 	// since the validator is set, remove it from the pending set
 	if err := k.RemovePendingValidator(ctx, val.OperatorAddress); err != nil {
 		return err
@@ -172,10 +207,6 @@ func (k Keeper) setValidatorInternals(ctx context.Context, val stakingtypes.Vali
 		return err
 	}
 
-	if err := k.stakingKeeper.SetNewValidatorByPowerIndex(ctx, val); err != nil {
-		return err
-	}
-
 	return k.stakingKeeper.Hooks().AfterValidatorCreated(ctx, valAddr)
 }
 
diff --git a/keeper/staking_hooks.go b/keeper/staking_hooks.go
new file mode 100644
index 0000000..60c354e
--- /dev/null
+++ b/keeper/staking_hooks.go
@@ -0,0 +1,84 @@
+package keeper
+
+import (
+	"context"
+
+	"github.com/cosmos/cosmos-sdk/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+
+	"cosmossdk.io/math"
+)
+
+// Before a validator is jailed, we must delete it from the power index. else:
+// - ERR CONSENSUS FAILURE!!! err="should never retrieve a jailed validator from the power store"
+
+var _ stakingtypes.StakingHooks = Hooks{}
+
+// Hooks wrapper struct for staking keeper
+type Hooks struct {
+	k Keeper
+}
+
+// Hooks return the mesh-security hooks
+func (k Keeper) Hooks() Hooks {
+	return Hooks{k}
+}
+
+// BeforeValidatorModified implements types.StakingHooks.
+func (h Hooks) BeforeValidatorModified(ctx context.Context, valAddr types.ValAddress) error {
+	h.k.logger.Info("BeforeValidatorModified", "valAddr", valAddr.String())
+	return h.k.BeforeJailedValidators.Set(ctx, valAddr.String())
+}
+
+// BeforeValidatorSlashed implements types.StakingHooks.
+func (h Hooks) BeforeValidatorSlashed(ctx context.Context, valAddr types.ValAddress, fraction math.LegacyDec) error {
+	h.k.logger.Info("BeforeValidatorSlashed", valAddr.String(), fraction.String())
+	return nil
+}
+
+// ----------------------------
+
+// AfterDelegationModified implements types.StakingHooks.
+func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr types.AccAddress, valAddr types.ValAddress) error {
+	return nil
+}
+
+// AfterUnbondingInitiated implements types.StakingHooks.
+func (h Hooks) AfterUnbondingInitiated(ctx context.Context, id uint64) error {
+	return nil
+}
+
+// AfterValidatorBeginUnbonding implements types.StakingHooks.
+func (h Hooks) AfterValidatorBeginUnbonding(ctx context.Context, consAddr types.ConsAddress, valAddr types.ValAddress) error {
+	return nil
+}
+
+// AfterValidatorBonded implements types.StakingHooks.
+func (h Hooks) AfterValidatorBonded(ctx context.Context, consAddr types.ConsAddress, valAddr types.ValAddress) error {
+	return nil
+}
+
+// AfterValidatorCreated implements types.StakingHooks.
+func (h Hooks) AfterValidatorCreated(ctx context.Context, valAddr types.ValAddress) error {
+	return nil
+}
+
+// AfterValidatorRemoved implements types.StakingHooks.
+func (h Hooks) AfterValidatorRemoved(ctx context.Context, consAddr types.ConsAddress, valAddr types.ValAddress) error {
+	return nil
+}
+
+// BeforeDelegationCreated implements types.StakingHooks.
+func (h Hooks) BeforeDelegationCreated(ctx context.Context, delAddr types.AccAddress, valAddr types.ValAddress) error {
+	return nil
+}
+
+// BeforeDelegationRemoved implements types.StakingHooks.
+func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr types.AccAddress, valAddr types.ValAddress) error {
+	return nil
+}
+
+// BeforeDelegationSharesModified implements types.StakingHooks.
+func (h Hooks) BeforeDelegationSharesModified(ctx context.Context, delAddr types.AccAddress, valAddr types.ValAddress) error {
+	return nil
+}
diff --git a/keys.go b/keys.go
index 84a602c..dcd8a1c 100644
--- a/keys.go
+++ b/keys.go
@@ -17,6 +17,11 @@ var (
 	// AbsoluteChangedInBlockPowerKey tracks the current blocks total power amount.
 	// If this becomes >30% of CachedPreviousBlockPowerKey, messages will fail to limit IBC issues.
 	AbsoluteChangedInBlockPowerKey = collections.NewPrefix(3)
+
+	// UpdatedValidatorsCacheKey tracks recently updated validators from SetPower.
+	UpdatedValidatorsCacheKey = collections.NewPrefix(4)
+
+	BeforeJailedValidatorsKey = collections.NewPrefix(5)
 )
 
 const (
diff --git a/module/abci.go b/module/abci.go
index c98890f..961770b 100644
--- a/module/abci.go
+++ b/module/abci.go
@@ -2,50 +2,71 @@ package module
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/cosmos/cosmos-sdk/telemetry"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	"github.com/strangelove-ventures/poa"
 )
 
+func (am AppModule) EndBlocker(ctx context.Context) error {
+	sk := am.keeper.GetStakingKeeper()
+
+	// Front running x/staking maturity ?
+	if err := sk.UnbondAllMatureValidators(ctx); err != nil {
+		return err
+	}
+
+	if err := am.handleBeforeJailedValidators(ctx); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // BeginBlocker updates the validator set without applying updates.
 // Since this module depends on staking, that module will `ApplyAndReturnValidatorSetUpdates` from x/staking.
 func (am AppModule) BeginBlocker(ctx context.Context) error {
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	defer telemetry.ModuleMeasureSince(poa.ModuleName, sdkCtx.BlockTime(), telemetry.MetricKeyBeginBlocker)
 
-	vals, err := am.keeper.GetStakingKeeper().GetAllValidators(ctx)
+	iterator, err := am.keeper.UpdatedValidatorsCache.Iterate(ctx, nil)
 	if err != nil {
 		return err
 	}
+	defer iterator.Close()
+
+	sk := am.keeper.GetStakingKeeper()
+
+	for ; iterator.Valid(); iterator.Next() {
+		valOperAddr, err := iterator.Key()
+		if err != nil {
+			return err
+		}
+		am.keeper.Logger().Info("UpdatedValidatorsCache: %s\n", valOperAddr)
 
-	for _, v := range vals {
-		switch v.GetStatus() {
-		case stakingtypes.Unbonding:
-			// if the validator is unbonding, force it to be unbonded. (H+1)
-			v.Status = stakingtypes.Unbonded
-			if err := am.keeper.GetStakingKeeper().SetValidator(ctx, v); err != nil {
-				return err
-			}
+		valAddr, err := sk.ValidatorAddressCodec().StringToBytes(valOperAddr)
+		if err != nil {
+			return err
+		}
 
-		case stakingtypes.Unbonded:
-			// if the validator is unbonded (above case), delete the last validator power. (H+2)
-			valAddr, err := sdk.ValAddressFromBech32(v.OperatorAddress)
-			if err != nil {
-				return err
-			}
+		val, err := sk.GetValidator(ctx, valAddr)
+		if err != nil {
+			return err
+		}
 
-			if err := am.keeper.GetStakingKeeper().DeleteLastValidatorPower(ctx, valAddr); err != nil {
-				return err
-			}
+		// Remove it from persisting across many blocks
+		if err := sk.DeleteValidatorByPowerIndex(ctx, val); err != nil {
+			return err
+		}
 
-		case stakingtypes.Unspecified, stakingtypes.Bonded:
-			continue
+		if err := am.keeper.UpdatedValidatorsCache.Remove(ctx, valOperAddr); err != nil {
+			return err
 		}
 	}
 
+	// reset caches
 	if sdkCtx.BlockHeight() > 1 {
 		// non gentx messages reset the cached block powers for IBC validations.
 		if err := am.keeper.ResetCachedTotalPower(ctx); err != nil {
@@ -57,5 +78,74 @@ func (am AppModule) BeginBlocker(ctx context.Context) error {
 		}
 	}
 
+	// Event Debugging
+	events, err := am.keeper.GetStakingKeeper().GetValidatorUpdates(ctx)
+	if err != nil {
+		return err
+	}
+	if len(events) == 0 {
+		return nil
+	}
+
+	am.keeper.Logger().Info("BeginBlocker events:\n")
+	for _, e := range events {
+		e := e
+		am.keeper.Logger().Info(fmt.Sprintf("PubKey: %s, Power: %d", &e.PubKey, e.Power))
+	}
+	am.keeper.Logger().Info("\n")
+
+	return nil
+}
+
+func (am AppModule) handleBeforeJailedValidators(ctx context.Context) error {
+	sk := am.keeper.GetStakingKeeper()
+
+	iterator, err := am.keeper.BeforeJailedValidators.Iterate(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer iterator.Close()
+
+	// Why? we don't want it in the store w/ the val state change in x/staking
+	for ; iterator.Valid(); iterator.Next() {
+		valOperAddr, err := iterator.Key()
+		if err != nil {
+			return err
+		}
+		am.keeper.Logger().Info("EndBlocker BeforeJailedValidators", valOperAddr, "\n")
+
+		valAddr, err := sk.ValidatorAddressCodec().StringToBytes(valOperAddr)
+		if err != nil {
+			return err
+		}
+
+		val, err := sk.GetValidator(ctx, valAddr)
+		if err != nil {
+			return err
+		}
+
+		if err := sk.DeleteValidatorByPowerIndex(ctx, val); err != nil {
+			return err
+		}
+
+		// TODO: If this is used here, it persist ABCI Updates. When removes, it looks like the validator gets slashed every block in x/staking? (when we do the hack and force set jailed = false)
+		// if err := sk.DeleteLastValidatorPower(ctx, valAddr); err != nil {
+		// 	return err
+		// }
+
+		// !IMPORTANT HACK: Set validator from jailed to not jailed to see what happens
+		// Okay so this like kind of worked for a split second
+		// Issue: the validator keeps trying to be converted to a jailed validator every single block when x/staking is calling it
+		val.Jailed = false
+		if err := sk.SetValidator(ctx, val); err != nil {
+			return err
+		}
+
+		// remove it from persisting
+		if err := am.keeper.BeforeJailedValidators.Remove(ctx, valOperAddr); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
diff --git a/module/module.go b/module/module.go
index 496b667..d5cd1c2 100644
--- a/module/module.go
+++ b/module/module.go
@@ -27,6 +27,7 @@ var (
 	_ module.AppModule          = AppModule{}
 	_ module.AppModuleGenesis   = AppModule{}
 	_ appmodule.HasBeginBlocker = AppModule{}
+	_ appmodule.HasEndBlocker   = AppModule{}
 )
 
 // ConsensusVersion defines the current module consensus version.
@@ -132,3 +133,7 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 func (am AppModule) BeginBlock(ctx context.Context) error {
 	return am.BeginBlocker(ctx)
 }
+
+func (am AppModule) EndBlock(ctx context.Context) error {
+	return am.EndBlocker(ctx)
+}
diff --git a/simapp/app.go b/simapp/app.go
index f1a78d2..42b2fa5 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -326,7 +326,11 @@ func NewSimApp(
 	// register the staking hooks
 	// NOTE: stakingKeeper above is passed by reference, so that it will contain these hooks
 	app.StakingKeeper.SetHooks(
-		stakingtypes.NewMultiStakingHooks(app.DistrKeeper.Hooks(), app.SlashingKeeper.Hooks()),
+		stakingtypes.NewMultiStakingHooks(
+			app.DistrKeeper.Hooks(),
+			app.SlashingKeeper.Hooks(),
+			app.POAKeeper.Hooks(),
+		),
 	)
 
 	app.CircuitKeeper = circuitkeeper.NewKeeper(appCodec, runtime.NewKVStoreService(keys[circuittypes.StoreKey]), govModAddress, app.AccountKeeper.AddressCodec())
@@ -450,16 +454,16 @@ func NewSimApp(
 		distrtypes.ModuleName,
 		slashingtypes.ModuleName,
 		evidencetypes.ModuleName,
+		poa.ModuleName, // TODO: has to be first
 		stakingtypes.ModuleName,
-		poa.ModuleName,
 		genutiltypes.ModuleName,
 		authz.ModuleName,
 	)
 	app.ModuleManager.SetOrderEndBlockers(
 		crisistypes.ModuleName,
 		govtypes.ModuleName,
+		poa.ModuleName, // TODO: pretty sure this has to go first so we can remove jailed before BlockValidatorUpdates
 		stakingtypes.ModuleName,
-		poa.ModuleName,
 		genutiltypes.ModuleName,
 		feegrant.ModuleName,
 		group.ModuleName,
diff --git a/simapp/test_node.sh b/simapp/test_node.sh
index 5f70608..0e4e66f 100644
--- a/simapp/test_node.sh
+++ b/simapp/test_node.sh
@@ -3,7 +3,7 @@
 # Example:
 : '
 cd simapp
-BINARY="poad" CHAIN_ID="poa-1" HOME_DIR="$HOME/.poad" TIMEOUT_COMMIT="2500ms" CLEAN=true sh test_node.sh
+BINARY="poad" CHAIN_ID="poa-1" HOME_DIR="$HOME/.poad" TIMEOUT_COMMIT="1200ms" CLEAN=true sh test_node.sh
 
 FLAGS="--keyring-backend=test --chain-id=poa-1 --home="$HOME/.poad" --yes"
 
@@ -118,6 +118,11 @@ from_scratch () {
   # staking
   update_test_genesis '.app_state["staking"]["params"]["bond_denom"]="stake"'
   update_test_genesis '.app_state["staking"]["params"]["min_commission_rate"]="0.050000000000000000"'
+  # slashing
+  update_test_genesis '.app_state["slashing"]["params"]["signed_blocks_window"]="10"'
+  update_test_genesis '.app_state["slashing"]["params"]["min_signed_per_window"]="1.000000000000000000"'
+  update_test_genesis '.app_state["slashing"]["params"]["slash_fraction_double_sign"]="0.000000000000000000"' # no need to slash, no stake
+  update_test_genesis '.app_state["slashing"]["params"]["slash_fraction_downtime"]="0.000000000000000000"'
   # mint
   update_test_genesis '.app_state["mint"]["params"]["mint_denom"]="stake"'
   # crisis