feat(staking): add withdraw reward operator address store (#215)

use separate collection for withdraw, rewards, operator address store

issue: none

client/genutil/genutil.go

@@ -196,7 +196,7 @@ func addValidator(txConfig client.TxConfig, pubkey crypto.PubKey, cdc codec.Code
 		sttypes.Description{Moniker: addr.Hex()},
 		sttypes.NewCommissionRates(zero, zero, zero),
 		sdk.DefaultPowerReduction,
-		sttypes.TokenType_LOCKED,
+		sttypes.DefaultLockedTokenType,
 	)
 	if err != nil {
 		return nil, errors.Wrap(err, "create validator message")

client/x/evmengine/keeper/abci_internal_test.go

@@ -289,7 +289,7 @@ func TestKeeper_PrepareProposal(t *testing.T) {
 
 		msgDelegate := stypes.NewMsgDelegate(
 			"delAddr", "valAddr", sdk.NewInt64Coin("stake", 100),
-			stypes.FlexiblePeriodDelegationID, stypes.PeriodType_FLEXIBLE,
+			stypes.FlexiblePeriodDelegationID, stypes.DefaultFlexiblePeriodType,
 		)
 		resp.Txs[0] = appendMsgToTx(t, txConfig, resp.Txs[0], msgDelegate)
 
@@ -611,7 +611,7 @@ func (m mockVEProvider) PrepareVotes(_ context.Context, _ abci.ExtendedCommitInf
 	coin := sdk.NewInt64Coin("stake", 100)
 	msg := stypes.NewMsgDelegate(
 		"addr", "addr", coin,
-		stypes.FlexiblePeriodDelegationID, stypes.PeriodType_FLEXIBLE,
+		stypes.FlexiblePeriodDelegationID, stypes.DefaultFlexiblePeriodType,
 	)
 
 	return []sdk.Msg{msg}, nil

client/x/evmstaking/keeper/abci.go

@@ -54,7 +54,7 @@ func (k *Keeper) EndBlock(ctx context.Context) (abci.ValidatorUpdates, error) {
 
 		// This should not produce error, as all delegations are done via the evmstaking module via EL.
 		// However, we should gracefully handle in case Get fails.
-		delEvmAddr, err := k.DelegatorMap.Get(ctx, entry.DelegatorAddress)
+		delEvmAddr, err := k.DelegatorWithdrawAddress.Get(ctx, entry.DelegatorAddress)
 		if err != nil {
 			return nil, errors.Wrap(err, "map delegator pubkey to evm address")
 		}

client/x/evmstaking/keeper/delegator_address.go

@@ -0,0 +1,90 @@
+package keeper
+
+import (
+	"context"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/ethereum/go-ethereum/common"
+
+	"github.com/piplabs/story/contracts/bindings"
+	"github.com/piplabs/story/lib/errors"
+	"github.com/piplabs/story/lib/k1util"
+)
+
+func (k Keeper) ProcessSetWithdrawalAddress(ctx context.Context, ev *bindings.IPTokenStakingSetWithdrawalAddress) error {
+	delCmpPubkey, err := UncmpPubKeyToCmpPubKey(ev.DelegatorUncmpPubkey)
+	if err != nil {
+		return errors.Wrap(err, "compress depositor pubkey")
+	}
+	depositorPubkey, err := k1util.PubKeyBytesToCosmos(delCmpPubkey)
+	if err != nil {
+		return errors.Wrap(err, "depositor pubkey to cosmos")
+	}
+
+	depositorAddr := sdk.AccAddress(depositorPubkey.Address().Bytes())
+	executionAddr := common.BytesToAddress(ev.ExecutionAddress[:])
+
+	if err := k.DelegatorWithdrawAddress.Set(ctx, depositorAddr.String(), executionAddr.String()); err != nil {
+		return errors.Wrap(err, "delegator withdraw address map set")
+	}
+
+	return nil
+}
+
+func (k Keeper) ProcessSetRewardAddress(ctx context.Context, ev *bindings.IPTokenStakingSetRewardAddress) error {
+	delCmpPubkey, err := UncmpPubKeyToCmpPubKey(ev.DelegatorUncmpPubkey)
+	if err != nil {
+		return errors.Wrap(err, "compress depositor pubkey")
+	}
+	depositorPubkey, err := k1util.PubKeyBytesToCosmos(delCmpPubkey)
+	if err != nil {
+		return errors.Wrap(err, "depositor pubkey to cosmos")
+	}
+
+	depositorAddr := sdk.AccAddress(depositorPubkey.Address().Bytes())
+	executionAddr := common.BytesToAddress(ev.ExecutionAddress[:])
+
+	if err := k.DelegatorRewardAddress.Set(ctx, depositorAddr.String(), executionAddr.String()); err != nil {
+		return errors.Wrap(err, "delegator reward address map set")
+	}
+
+	return nil
+}
+
+func (k Keeper) ProcessAddOperator(ctx context.Context, ev *bindings.IPTokenStakingAddOperator) error {
+	delCmpPubkey, err := UncmpPubKeyToCmpPubKey(ev.UncmpPubkey)
+	if err != nil {
+		return errors.Wrap(err, "compress depositor pubkey")
+	}
+	depositorPubkey, err := k1util.PubKeyBytesToCosmos(delCmpPubkey)
+	if err != nil {
+		return errors.Wrap(err, "depositor pubkey to cosmos")
+	}
+
+	depositorAddr := sdk.AccAddress(depositorPubkey.Address().Bytes())
+
+	if err := k.DelegatorOperatorAddress.Set(ctx, depositorAddr.String(), ev.Operator.String()); err != nil {
+		return errors.Wrap(err, "delegator operator address map set")
+	}
+
+	return nil
+}
+
+func (k Keeper) ProcessRemoveOperator(ctx context.Context, ev *bindings.IPTokenStakingRemoveOperator) error {
+	delCmpPubkey, err := UncmpPubKeyToCmpPubKey(ev.UncmpPubkey)
+	if err != nil {
+		return errors.Wrap(err, "compress depositor pubkey")
+	}
+	depositorPubkey, err := k1util.PubKeyBytesToCosmos(delCmpPubkey)
+	if err != nil {
+		return errors.Wrap(err, "depositor pubkey to cosmos")
+	}
+
+	depositorAddr := sdk.AccAddress(depositorPubkey.Address().Bytes())
+
+	if err := k.DelegatorOperatorAddress.Remove(ctx, depositorAddr.String()); err != nil {
+		return errors.Wrap(err, "delegator operator address map remove")
+	}
+
+	return nil
+}

client/x/evmstaking/keeper/set_address_test.go → client/x/evmstaking/keeper/delegator_address_test.go

@@ -81,13 +81,13 @@ func (s *TestSuite) TestProcessSetWithdrawalAddress() {
 				require.Contains(err.Error(), tc.expectedError)
 
 				// Ensure no state change occurred
-				addr, _ := keeper.DelegatorMap.Get(cachedCtx, delAddr.String())
+				addr, _ := keeper.DelegatorWithdrawAddress.Get(cachedCtx, delAddr.String())
 				require.Empty(addr)
 			} else {
 				require.NoError(err)
 
 				// check result
-				addr, err := keeper.DelegatorMap.Get(cachedCtx, delAddr.String())
+				addr, err := keeper.DelegatorWithdrawAddress.Get(cachedCtx, delAddr.String())
 				require.NoError(err)
 				require.Equal(evmAddr.String(), addr)
 			}

client/x/evmstaking/keeper/deposit.go

@@ -77,8 +77,36 @@ func (k Keeper) ProcessDeposit(ctx context.Context, ev *bindings.IPTokenStakingD
 	// Note that, after minting, we save the mapping between delegator bech32 address and evm address, which will be used in the withdrawal queue.
 	// The saving is done regardless of any error below, as the money is already minted and sent to the delegator, who can withdraw the minted amount.
 	// TODO: Confirm that bech32 address and evm address can be used interchangeably. Must be one-to-one or many-bech32-to-one-evm.
-	if err := k.DelegatorMap.Set(ctx, depositorAddr.String(), delEvmAddr.String()); err != nil {
-		return errors.Wrap(err, "set delegator map")
+	if err := k.DelegatorWithdrawAddress.Set(ctx, depositorAddr.String(), delEvmAddr.String()); err != nil {
+		return errors.Wrap(err, "set delegator withdraw address map")
+	}
+	if err := k.DelegatorRewardAddress.Set(ctx, depositorAddr.String(), delEvmAddr.String()); err != nil {
+		return errors.Wrap(err, "set delegator reward address map")
+	}
+
+	delID := ev.DelegationId.String()
+	periodType := int32(ev.StakingPeriod.Int64())
+
+	val, err := k.stakingKeeper.GetValidator(ctx, validatorAddr)
+	if errors.Is(err, stypes.ErrNoValidatorFound) {
+		return errors.New("validator not exists")
+	} else if err != nil {
+		return errors.Wrap(err, "get validator failed")
+	}
+
+	lockedTokenType, err := k.stakingKeeper.GetLockedTokenType(ctx)
+	if err != nil {
+		return errors.Wrap(err, "get locked token type")
+	}
+
+	// locked tokens can only be staked with flexible period
+	if val.SupportTokenType == lockedTokenType {
+		flexPeriodType, err := k.stakingKeeper.GetFlexiblePeriodType(ctx)
+		if err != nil {
+			return errors.Wrap(err, "get flexible period type")
+		}
+		periodType = flexPeriodType
+		delID = stypes.FlexiblePeriodDelegationID
 	}
 
 	// TODO: Check if we can instantiate the msgServer without type assertion
@@ -87,25 +115,10 @@ func (k Keeper) ProcessDeposit(ctx context.Context, ev *bindings.IPTokenStakingD
 		return errors.New("type assertion failed")
 	}
 	skeeperMsgServer := skeeper.NewMsgServerImpl(evmstakingSKeeper)
-
-	var periodType stypes.PeriodType
-	switch ev.StakingPeriod.Int64() {
-	case int64(stypes.PeriodType_FLEXIBLE):
-		periodType = stypes.PeriodType_FLEXIBLE
-	case int64(stypes.PeriodType_THREE_MONTHS):
-		periodType = stypes.PeriodType_THREE_MONTHS
-	case int64(stypes.PeriodType_ONE_YEAR):
-		periodType = stypes.PeriodType_ONE_YEAR
-	case int64(stypes.PeriodType_EIGHTEEN_MONTHS):
-		periodType = stypes.PeriodType_EIGHTEEN_MONTHS
-	default:
-		return errors.New("invalid staking period")
-	}
-
 	// Delegation by the depositor on the validator (validator existence is checked in msgServer.Delegate)
 	msg := stypes.NewMsgDelegate(
 		depositorAddr.String(), validatorAddr.String(), amountCoin,
-		ev.DelegationId.String(), periodType,
+		delID, periodType,
 	)
 	_, err = skeeperMsgServer.Delegate(ctx, msg)
 	if err != nil {

client/x/evmstaking/keeper/genesis.go

@@ -53,9 +53,11 @@ func (k Keeper) InitGenesis(ctx context.Context, gs *types.GenesisState) error {
 			"del_addr", delAddr.String(),
 		)
 
-		err = k.DelegatorMap.Set(ctx, delAddr.String(), evmAddr.String())
-		if err != nil {
-			return errors.Wrap(err, "set delegator map")
+		if err = k.DelegatorWithdrawAddress.Set(ctx, delAddr.String(), evmAddr.String()); err != nil {
+			return errors.Wrap(err, "set delegator withdraw address map")
+		}
+		if err = k.DelegatorRewardAddress.Set(ctx, delAddr.String(), evmAddr.String()); err != nil {
+			return errors.Wrap(err, "set delegator reward address map")
 		}
 	}
 

client/x/evmstaking/keeper/keeper.go

@@ -39,8 +39,10 @@ type Keeper struct {
 
 	ipTokenStakingContract *bindings.IPTokenStaking
 
-	WithdrawalQueue addcollections.Queue[types.Withdrawal]
-	DelegatorMap    collections.Map[string, string] // bech32 to evm address (TODO: confirm that it's one-to-one or many-bech32-to-one-evm)
+	WithdrawalQueue          addcollections.Queue[types.Withdrawal]
+	DelegatorWithdrawAddress collections.Map[string, string]
+	DelegatorRewardAddress   collections.Map[string, string]
+	DelegatorOperatorAddress collections.Map[string, string]
 }
 
 // NewKeeper creates a new evmstaking Keeper instance.
@@ -74,18 +76,20 @@ func NewKeeper(
 	}
 
 	return &Keeper{
-		cdc:                    cdc,
-		storeService:           storeService,
-		authKeeper:             ak,
-		bankKeeper:             bk,
-		slashingKeeper:         slk,
-		stakingKeeper:          stk,
-		distributionKeeper:     dk,
-		authority:              authority,
-		validatorAddressCodec:  validatorAddressCodec,
-		ipTokenStakingContract: ipTokenStakingContract,
-		WithdrawalQueue:        addcollections.NewQueue(sb, types.WithdrawalQueueKey, "withdrawal_queue", codec.CollValue[types.Withdrawal](cdc)),
-		DelegatorMap:           collections.NewMap(sb, types.DelegatorMapKey, "delegator_map", collections.StringKey, collections.StringValue),
+		cdc:                      cdc,
+		storeService:             storeService,
+		authKeeper:               ak,
+		bankKeeper:               bk,
+		slashingKeeper:           slk,
+		stakingKeeper:            stk,
+		distributionKeeper:       dk,
+		authority:                authority,
+		validatorAddressCodec:    validatorAddressCodec,
+		ipTokenStakingContract:   ipTokenStakingContract,
+		WithdrawalQueue:          addcollections.NewQueue(sb, types.WithdrawalQueueKey, "withdrawal_queue", codec.CollValue[types.Withdrawal](cdc)),
+		DelegatorWithdrawAddress: collections.NewMap(sb, types.DelegatorWithdrawAddressMapKey, "delegator_withdraw_address_map", collections.StringKey, collections.StringValue),
+		DelegatorRewardAddress:   collections.NewMap(sb, types.DelegatorRewardAddressMapKey, "delegator_reward_address_map", collections.StringKey, collections.StringValue),
+		DelegatorOperatorAddress: collections.NewMap(sb, types.DelegatorOperatorAddressMapKey, "delegator_operator_address_map", collections.StringKey, collections.StringValue),
 	}
 }
 
@@ -103,7 +107,7 @@ func (k Keeper) ValidatorAddressCodec() addresscodec.Codec {
 	return k.validatorAddressCodec
 }
 
-// TODO: Return log event results to properly manage failures.
+//nolint:gocyclo // TODO
 func (k Keeper) ProcessStakingEvents(ctx context.Context, height uint64, logs []*evmenginetypes.EVMEvent) error {
 	gwei, exp := big.NewInt(10), big.NewInt(9)
 	gwei.Exp(gwei, exp, nil)
@@ -122,7 +126,7 @@ func (k Keeper) ProcessStakingEvents(ctx context.Context, height uint64, logs []
 		// Convert the amount from wei to gwei (Eth2 spec withdrawal is specified in gwei) by dividing by 10^9.
 		// TODO: consider rounding and decimal precision when dividing bigint.
 
-		switch ethlog.Topics[0] {
+		switch ethlog.Topics[0] { // TODO(rayden): update validator commission
 		case types.SetWithdrawalAddress.ID:
 			ev, err := k.ipTokenStakingContract.ParseSetWithdrawalAddress(ethlog)
 			if err != nil {
@@ -133,6 +137,36 @@ func (k Keeper) ProcessStakingEvents(ctx context.Context, height uint64, logs []
 				clog.Error(ctx, "Failed to process set withdrawal address", err)
 				continue
 			}
+		case types.SetRewardAddress.ID:
+			ev, err := k.ipTokenStakingContract.ParseSetRewardAddress(ethlog)
+			if err != nil {
+				clog.Error(ctx, "Failed to parse SetRewardAddress log", err)
+				continue
+			}
+			if err = k.ProcessSetRewardAddress(ctx, ev); err != nil {
+				clog.Error(ctx, "Failed to process set reward address", err)
+				continue
+			}
+		case types.AddOperator.ID:
+			ev, err := k.ipTokenStakingContract.ParseAddOperator(ethlog)
+			if err != nil {
+				clog.Error(ctx, "Failed to parse SetRewardAddress log", err)
+				continue
+			}
+			if err = k.ProcessAddOperator(ctx, ev); err != nil {
+				clog.Error(ctx, "Failed to process add operator", err)
+				continue
+			}
+		case types.RemoveOperator.ID:
+			ev, err := k.ipTokenStakingContract.ParseRemoveOperator(ethlog)
+			if err != nil {
+				clog.Error(ctx, "Failed to parse SetRewardAddress log", err)
+				continue
+			}
+			if err = k.ProcessRemoveOperator(ctx, ev); err != nil {
+				clog.Error(ctx, "Failed to process add operator", err)
+				continue
+			}
 		case types.CreateValidatorEvent.ID:
 			ev, err := k.ParseCreateValidatorLog(ethlog)
 			if err != nil {

client/x/evmstaking/keeper/keeper_test.go

@@ -408,6 +408,7 @@ func (s *TestSuite) TestProcessStakingEvents() {
 					cmpToUncmp(valPubKey1.Bytes()),
 					cmpToUncmp(valPubKey2.Bytes()),
 					big.NewInt(0),
+					cmpToEVM(delPubKey.Bytes()),
 					delAmtGwei,
 				)
 				require.NoError(err)
@@ -478,11 +479,11 @@ func (s *TestSuite) TestProcessStakingEvents() {
 				return evmEvents, nil
 			},
 			stateCheck: func(c context.Context) {
-				_, err := evmstakingKeeper.DelegatorMap.Get(c, delPubKey.Address().String())
+				_, err := evmstakingKeeper.DelegatorWithdrawAddress.Get(c, delPubKey.Address().String())
 				require.ErrorContains(err, "not found")
 			},
 			postStateCheck: func(c context.Context) {
-				evmDelAddr, err := evmstakingKeeper.DelegatorMap.Get(c, delAddr.String())
+				evmDelAddr, err := evmstakingKeeper.DelegatorWithdrawAddress.Get(c, delAddr.String())
 				require.NoError(err)
 				require.Equal(delEvmAddr.String(), evmDelAddr)
 			},

client/x/evmstaking/keeper/redelegation.go

@@ -3,6 +3,8 @@ package keeper
 import (
 	"context"
 
+	"cosmossdk.io/collections"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	skeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	stypes "github.com/cosmos/cosmos-sdk/x/staking/types"
@@ -59,6 +61,19 @@ func (k Keeper) ProcessRedelegate(ctx context.Context, ev *bindings.IPTokenStaki
 		return errors.Wrap(err, "dst validator pubkey to evm address")
 	}
 
+	// redelegateOnBehalf txn, need to check if it's from the operator
+	if delEvmAddr.String() != ev.OperatorAddress.String() {
+		operatorAddr, err := k.DelegatorOperatorAddress.Get(ctx, depositorAddr.String())
+		if errors.Is(err, collections.ErrNotFound) {
+			return errors.New("invalid redelegateOnBehalf txn, no operator for delegator")
+		} else if err != nil {
+			return errors.Wrap(err, "get delegator's operator address failed")
+		}
+		if operatorAddr != ev.OperatorAddress.String() {
+			return errors.New("invalid redelegateOnBehalf txn, not from operator")
+		}
+	}
+
 	amountCoin, _ := IPTokenToBondCoin(ev.Amount)
 
 	log.Info(ctx, "EVM staking relegation detected",