fix!: constant jail patch

go.mod

@@ -8,6 +8,10 @@ replace cosmossdk.io/core => cosmossdk.io/core v0.11.0
 // breaks SDK app.toml parsing in ictest.
 replace github.com/spf13/viper => github.com/spf13/viper v1.17.0
 
+// TODO: not for production, seeing if removing .Jailed check panic from `ApplyAndReturnValidatorSetUpdates` fixes jailing issue.
+// https://github.com/Reecepbcups/cosmos-sdk/commit/26ff44b40ec78ca3e003aa49d9db1ce8199ae7dd
+replace github.com/cosmos/cosmos-sdk => github.com/Reecepbcups/cosmos-sdk v0.50.3-0.20240416135850-26ff44b40ec7
+
 require (
 	cosmossdk.io/api v0.7.3
 	cosmossdk.io/collections v0.4.0

keeper/expected_keepers.go

@@ -17,38 +17,6 @@ type BankKeeper interface {
 type SlashingKeeper interface {
 	DeleteMissedBlockBitmap(ctx context.Context, addr sdk.ConsAddress) error
 	SetValidatorSigningInfo(ctx context.Context, address sdk.ConsAddress, info slashingtypes.ValidatorSigningInfo) error
-}
-
-// We use almost all the methods from StakingKeeper. Just using it directly
-/*
-type StakingKeeper interface {
-	Hooks() stakingtypes.StakingHooks
-
-	SetParams(ctx context.Context, params stakingtypes.Params) error
-	SetLastValidatorPower(ctx context.Context, operator sdk.ValAddress, power int64) error
-	SetDelegation(ctx context.Context, delegation stakingtypes.Delegation) error
-	SetValidator(ctx context.Context, validator stakingtypes.Validator) error
-	SetValidatorByConsAddr(ctx context.Context, validator stakingtypes.Validator) error
-	SetNewValidatorByPowerIndex(ctx context.Context, validator stakingtypes.Validator) error
-	SetLastTotalPower(ctx context.Context, power math.Int) error
-	DeleteLastValidatorPower(ctx context.Context, operator sdk.ValAddress) error
-	GetLastTotalPower(ctx context.Context) (math.Int, error)
-
-	Slash(ctx context.Context, consAddr sdk.ConsAddress, infractionHeight, power int64, slashFactor math.LegacyDec) (math.Int, error)
-
-	MinCommissionRate(ctx context.Context) (math.LegacyDec, error)
-	BondDenom(ctx context.Context) (string, error)
-
-	GetValidator(ctx context.Context, addr sdk.ValAddress) (validator stakingtypes.Validator, err error)
-	GetAllValidators(ctx context.Context) (validators []stakingtypes.Validator, err error)
-	GetValidatorByConsAddr(ctx context.Context, consAddr sdk.ConsAddress) (validator stakingtypes.Validator, err error)
-	GetLastValidatorPower(ctx context.Context, operator sdk.ValAddress) (power int64, err error)
-
-	GetAllDelegations(ctx context.Context) (delegations []stakingtypes.Delegation, err error)
-
-	TokensFromConsensusPower(ctx context.Context, power int64) math.Int
-	TokensToConsensusPower(ctx context.Context, tokens math.Int) int64
 
-	PowerReduction(ctx context.Context) math.Int
+	GetValidatorSigningInfo(ctx context.Context, address sdk.ConsAddress) (slashingtypes.ValidatorSigningInfo, error)
 }
-*/

keeper/keeper.go

@@ -34,7 +34,10 @@ type Keeper struct {
 	Params                 collections.Item[poa.Params]
 	PendingValidators      collections.Item[poa.Validators]
 	UpdatedValidatorsCache collections.KeySet[string]
-	BeforeJailedValidators collections.KeySet[string]
+
+	// CheckForJailedValidators is a set of validators which were modified in the current block & captured with staking hooks (BeforeValidatorModified).
+	// where valoper -> height set at (this way after so many blocks, if they are not jailed, then remove from the map)
+	CheckForJailedValidators collections.Map[string, int64]
 
 	CachedBlockPower            collections.Item[poa.PowerCache]
 	AbsoluteChangedInBlockPower collections.Item[poa.PowerCache]
@@ -61,10 +64,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)),
-		UpdatedValidatorsCache: collections.NewKeySet(sb, poa.UpdatedValidatorsCacheKey, "updated_validators", collections.StringKey),
-		BeforeJailedValidators: collections.NewKeySet(sb, poa.BeforeJailedValidatorsKey, "before_jailed", collections.StringKey),
+		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),
+		CheckForJailedValidators: collections.NewMap(sb, poa.BeforeJailedValidatorsKey, "check_for_jailed", collections.StringKey, collections.Int64Value),
 
 		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)),

keeper/staking_hooks.go

@@ -3,10 +3,9 @@ package keeper
 import (
 	"context"
 
-	"github.com/cosmos/cosmos-sdk/types"
-	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
-
 	"cosmossdk.io/math"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 )
 
 // Before a validator is jailed, we must delete it from the power index. else:
@@ -24,61 +23,72 @@ func (k Keeper) Hooks() Hooks {
 	return Hooks{k}
 }
 
-// BeforeValidatorModified implements types.StakingHooks.
-func (h Hooks) BeforeValidatorModified(ctx context.Context, valAddr types.ValAddress) error {
+// BeforeValidatorModified implements sdk.StakingHooks.
+func (h Hooks) BeforeValidatorModified(ctx context.Context, valAddr sdk.ValAddress) error {
 	h.k.logger.Info("BeforeValidatorModified", "valAddr", valAddr.String())
-	return h.k.BeforeJailedValidators.Set(ctx, valAddr.String())
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+
+	// ignore setting a new value if we already have in the cache
+	// Could also do a jail check here but that is redundant imo
+	if ok, err := h.k.CheckForJailedValidators.Has(ctx, valAddr.String()); err != nil {
+		return err
+	} else if ok {
+		return nil
+	}
+
+	// we increase by 5 for when we actually check this in the end blocker
+	return h.k.CheckForJailedValidators.Set(ctx, valAddr.String(), sdkCtx.BlockHeight())
 }
 
-// BeforeValidatorSlashed implements types.StakingHooks.
-func (h Hooks) BeforeValidatorSlashed(ctx context.Context, valAddr types.ValAddress, fraction math.LegacyDec) error {
+// BeforeValidatorSlashed implements sdk.StakingHooks.
+func (h Hooks) BeforeValidatorSlashed(ctx context.Context, valAddr sdk.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 {
+// AfterDelegationModified implements sdk.StakingHooks.
+func (h Hooks) AfterDelegationModified(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error {
 	return nil
 }
 
-// AfterUnbondingInitiated implements types.StakingHooks.
+// AfterUnbondingInitiated implements sdk.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 {
+// AfterValidatorBeginUnbonding implements sdk.StakingHooks.
+func (h Hooks) AfterValidatorBeginUnbonding(ctx context.Context, consAddr sdk.ConsAddress, valAddr sdk.ValAddress) error {
 	return nil
 }
 
-// AfterValidatorBonded implements types.StakingHooks.
-func (h Hooks) AfterValidatorBonded(ctx context.Context, consAddr types.ConsAddress, valAddr types.ValAddress) error {
+// AfterValidatorBonded implements sdk.StakingHooks.
+func (h Hooks) AfterValidatorBonded(ctx context.Context, consAddr sdk.ConsAddress, valAddr sdk.ValAddress) error {
 	return nil
 }
 
-// AfterValidatorCreated implements types.StakingHooks.
-func (h Hooks) AfterValidatorCreated(ctx context.Context, valAddr types.ValAddress) error {
+// AfterValidatorCreated implements sdk.StakingHooks.
+func (h Hooks) AfterValidatorCreated(ctx context.Context, valAddr sdk.ValAddress) error {
 	return nil
 }
 
-// AfterValidatorRemoved implements types.StakingHooks.
-func (h Hooks) AfterValidatorRemoved(ctx context.Context, consAddr types.ConsAddress, valAddr types.ValAddress) error {
+// AfterValidatorRemoved implements sdk.StakingHooks.
+func (h Hooks) AfterValidatorRemoved(ctx context.Context, consAddr sdk.ConsAddress, valAddr sdk.ValAddress) error {
 	return nil
 }
 
-// BeforeDelegationCreated implements types.StakingHooks.
-func (h Hooks) BeforeDelegationCreated(ctx context.Context, delAddr types.AccAddress, valAddr types.ValAddress) error {
+// BeforeDelegationCreated implements sdk.StakingHooks.
+func (h Hooks) BeforeDelegationCreated(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error {
 	return nil
 }
 
-// BeforeDelegationRemoved implements types.StakingHooks.
-func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr types.AccAddress, valAddr types.ValAddress) error {
+// BeforeDelegationRemoved implements sdk.StakingHooks.
+func (h Hooks) BeforeDelegationRemoved(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error {
 	return nil
 }
 
-// BeforeDelegationSharesModified implements types.StakingHooks.
-func (h Hooks) BeforeDelegationSharesModified(ctx context.Context, delAddr types.AccAddress, valAddr types.ValAddress) error {
+// BeforeDelegationSharesModified implements sdk.StakingHooks.
+func (h Hooks) BeforeDelegationSharesModified(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) error {
 	return nil
 }

keys.go

@@ -21,6 +21,7 @@ var (
 	// UpdatedValidatorsCacheKey tracks recently updated validators from SetPower.
 	UpdatedValidatorsCacheKey = collections.NewPrefix(4)
 
+	// TODO: change name
 	// BeforeJailedValidatorsKey tracks validators that are about to be jailed (from staking hooks).
 	BeforeJailedValidatorsKey = collections.NewPrefix(5)
 )

module/abci.go

@@ -92,8 +92,12 @@ func (am AppModule) BeginBlocker(ctx context.Context) error {
 
 func (am AppModule) handleBeforeJailedValidators(ctx context.Context) error {
 	sk := am.keeper.GetStakingKeeper()
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	curHeight := sdkCtx.BlockHeight()
+	bt := sdkCtx.BlockTime()
+	logger := am.keeper.Logger()
 
-	iterator, err := am.keeper.BeforeJailedValidators.Iterate(ctx, nil)
+	iterator, err := am.keeper.CheckForJailedValidators.Iterate(ctx, nil)
 	if err != nil {
 		return err
 	}
@@ -105,7 +109,14 @@ func (am AppModule) handleBeforeJailedValidators(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
-		am.keeper.Logger().Info("EndBlocker BeforeJailedValidators", valOperAddr, "\n")
+
+		height, err := iterator.Value()
+		if err != nil {
+			return err
+		}
+
+		// .Error for viewability
+		logger.Error("EndBlocker BeforeJailedValidators", "operator", valOperAddr, "height", height)
 
 		valAddr, err := sk.ValidatorAddressCodec().StringToBytes(valOperAddr)
 		if err != nil {
@@ -117,27 +128,57 @@ func (am AppModule) handleBeforeJailedValidators(ctx context.Context) error {
 			return err
 		}
 
-		if err := sk.DeleteValidatorByPowerIndex(ctx, val); err != nil {
+		consBz, err := val.GetConsAddr()
+		if 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 {
+		si, err := am.keeper.GetSlashingKeeper().GetValidatorSigningInfo(ctx, consBz)
+		if err != nil {
 			return err
 		}
 
-		// remove it from persisting
-		if err := am.keeper.BeforeJailedValidators.Remove(ctx, valOperAddr); err != nil {
-			return err
+		if height == curHeight {
+			// if si.JailedUntil.After(bt) is false, then we remove from the keeper store (false positive for the val modification wrt jailing from staking hook)
+			if !si.JailedUntil.After(bt) {
+				logger.Error("EndBlocker BeforeJailedValidators validator was not jailed, removing from cache",
+					"height", height, "blockHeight", curHeight, "si.JailedUntil", si.JailedUntil, "block_time", bt, "operator", valOperAddr,
+				)
+				if err := am.keeper.CheckForJailedValidators.Remove(ctx, valOperAddr); err != nil {
+					return err
+				}
+
+				continue
+			}
+
+			logger.Error("handleBeforeJailedValidators deleting jailed validator", "height", height, "blockHeight", curHeight)
+
+			if err := sk.DeleteValidatorByPowerIndex(ctx, val); err != nil {
+				return err
+			}
+			if err := sk.DeleteLastValidatorPower(ctx, valAddr); err != nil {
+				return err
+			}
+
+			val.Jailed = false
+			if err := sk.SetValidator(ctx, val); err != nil {
+				return err
+			}
+
+		} else if height+2 == curHeight {
+			// Why is staking / slashing not handling this for us anyways?
+			logger.Error("handleBeforeJailedValidators setting val to jailed", "height", height, "blockHeight", curHeight)
+
+			val.Jailed = true
+			if err := sk.SetValidator(ctx, val); err != nil {
+				return err
+			}
+
+			if err := am.keeper.CheckForJailedValidators.Remove(ctx, valOperAddr); err != nil {
+				return err
+			}
 		}
+
 	}
 
 	return nil