From 558ab4e311f4d757f05066ffdca30a394cfb59d1 Mon Sep 17 00:00:00 2001 From: Hanjun Kim Date: Tue, 14 Jan 2025 12:25:44 +0900 Subject: [PATCH 1/3] test: add gas consumption test --- x/restaking/keeper/end_blocker_test.go | 96 ++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) diff --git a/x/restaking/keeper/end_blocker_test.go b/x/restaking/keeper/end_blocker_test.go index 1fafa116..7ab85903 100644 --- a/x/restaking/keeper/end_blocker_test.go +++ b/x/restaking/keeper/end_blocker_test.go @@ -1,6 +1,7 @@ package keeper_test import ( + "fmt" "time" sdkmath "cosmossdk.io/math" @@ -8,6 +9,7 @@ import ( operatorstypes "github.com/milkyway-labs/milkyway/v7/x/operators/types" poolstypes "github.com/milkyway-labs/milkyway/v7/x/pools/types" + "github.com/milkyway-labs/milkyway/v7/x/restaking/keeper" "github.com/milkyway-labs/milkyway/v7/x/restaking/types" servicestypes "github.com/milkyway-labs/milkyway/v7/x/services/types" ) @@ -314,3 +316,97 @@ func (suite *KeeperTestSuite) TestKeeper_CompleteMatureUnbondingDelegations() { }) } } + +func (suite *KeeperTestSuite) TestGasConsumption_UndelegateVsEndBlockProcessing() { + for _, numDenomsPerDelegation := range []int{1, 10, 30, 100, 500} { + suite.Run(fmt.Sprintf("with%dDenoms", numDenomsPerDelegation), func() { + ctx, _ := suite.ctx.CacheContext() + + ctx = ctx. + WithBlockHeight(10). + WithBlockTime(time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)) + + // Set the unbonding time to 7 days + params, err := suite.k.GetParams(ctx) + suite.Require().NoError(err) + params.UnbondingTime = 7 * 24 * time.Hour + err = suite.k.SetParams(ctx, params) + suite.Require().NoError(err) + + // Create a delegator address + delegator := "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4" + + // Generate denoms to delegate + denoms := make([]string, numDenomsPerDelegation) + for i := range denoms { + denoms[i] = fmt.Sprintf("denom%d", i) + } + + // Create a service to delegate to + const serviceID = 1 + serviceAddress := servicestypes.GetServiceAddress(serviceID).String() + err = suite.sk.SaveService(ctx, servicestypes.Service{ + ID: serviceID, + Status: servicestypes.SERVICE_STATUS_ACTIVE, + Address: serviceAddress, + Tokens: sdk.NewCoins(), + DelegatorShares: sdk.NewDecCoins(), + }) + suite.Require().NoError(err) + + // Fund the delegator account with sufficient balance + initialBalance := sdk.NewCoins() + for _, denom := range denoms { + initialBalance = initialBalance.Add(sdk.NewInt64Coin(denom, 1000_000000)) + } + suite.fundAccount(ctx, delegator, initialBalance) + + // Prepare the total amount to delegate + delAmt := sdk.NewCoins() + for _, denom := range denoms { + delAmt = delAmt.Add(sdk.NewInt64Coin(denom, 100_000000)) + } + + // Delegate multiple denominations to the service using MsgServer + // + // NOTE: We don't include this as part of the gas comparison because delegations can + // be cheaply accumulated over arbitrary amounts of time and undelegated in a batch in a single block. + // The attack vector hinges on this burst of undelegations to pack >1 block worth of unmetered gas + // into a single block of gas consumption. + msgServer := keeper.NewMsgServer(suite.k) + _, err = msgServer.DelegateService(ctx, types.NewMsgDelegateService(serviceID, delAmt, delegator)) + suite.Require().NoError(err) + + // --- Undelegation Gas Tracking --- + + // Measure gas consumption during initial undelegation + undelegationGasStart := ctx.GasMeter().GasConsumed() + + // Undelegate the denominations using MsgServer + _, err = msgServer.UndelegateService(ctx, types.NewMsgUndelegateService(serviceID, delAmt, delegator)) + suite.Require().NoError(err) + + // Calculate gas used during undelegations + gasUsedForUndelegation := ctx.GasMeter().GasConsumed() - undelegationGasStart + + // --- EndBlock Unbond Completion Gas Tracking --- + + // Advance context time to when undelegations mature + ctx = ctx.WithBlockTime(ctx.BlockTime().Add(7 * 24 * time.Hour)) + + endBlockGasStart := ctx.GasMeter().GasConsumed() + + // Measure gas consumption during end block processing + // NOTE: we isolate the component of the EndBlock call we want to test + err = suite.k.CompleteMatureUnbondingDelegations(ctx) + suite.Require().NoError(err) + + // Calculate gas used during end block processing + gasUsedForEndBlock := ctx.GasMeter().GasConsumed() - endBlockGasStart + + // --- Gas Consumption Comparison --- + + suite.Require().GreaterOrEqual(gasUsedForUndelegation, gasUsedForEndBlock) + }) + } +} From ec489876f44543fc14a68682a6cf39c5b8694565 Mon Sep 17 00:00:00 2001 From: Hanjun Kim Date: Wed, 15 Jan 2025 14:37:23 +0900 Subject: [PATCH 2/3] refactor!: remove scaling gas cost entirely for delegations since delegations don't have an impact on chain's performance anymore, we can safely remove it now --- x/restaking/keeper/alias_functions.go | 44 ----------------- x/restaking/keeper/msg_server_test.go | 49 ------------------- x/restaking/keeper/operator_restaking_test.go | 16 ------ x/restaking/keeper/pool_restaking_test.go | 7 --- x/restaking/keeper/service_restaking_test.go | 16 ------ x/restaking/types/const.go | 4 -- 6 files changed, 136 deletions(-) diff --git a/x/restaking/keeper/alias_functions.go b/x/restaking/keeper/alias_functions.go index 19033a87..91927720 100644 --- a/x/restaking/keeper/alias_functions.go +++ b/x/restaking/keeper/alias_functions.go @@ -188,35 +188,6 @@ func (k *Keeper) SetDelegation(ctx context.Context, delegation types.Delegation) return nil } -// GetDelegationsForTarget returns all the delegations for the given target -func (k *Keeper) GetDelegationsForTarget(ctx context.Context, target types.DelegationTarget) ([]types.Delegation, error) { - store := k.storeService.OpenKVStore(ctx) - - // Get the function used to build the store prefix to get the delegations - var buildStorePrefix func(targetID uint32) []byte - switch target.(type) { - case poolstypes.Pool: - buildStorePrefix = types.DelegationsByPoolIDStorePrefix - case operatorstypes.Operator: - buildStorePrefix = types.DelegationsByOperatorIDStorePrefix - case servicestypes.Service: - buildStorePrefix = types.DelegationsByServiceIDStorePrefix - default: - return nil, fmt.Errorf("invalid target type %T", target) - } - - iterator := storetypes.KVStorePrefixIterator(runtime.KVStoreAdapter(store), buildStorePrefix(target.GetID())) - defer iterator.Close() - - var delegations []types.Delegation - for ; iterator.Valid(); iterator.Next() { - delegation := types.MustUnmarshalDelegation(k.cdc, iterator.Value()) - delegations = append(delegations, delegation) - } - - return delegations, nil -} - // GetDelegationForTarget returns the delegation for the given delegator and target. func (k *Keeper) GetDelegationForTarget(ctx context.Context, target types.DelegationTarget, delegator string) (types.Delegation, bool, error) { switch target.(type) { @@ -598,21 +569,6 @@ func (k *Keeper) PerformDelegation(ctx context.Context, data types.DelegationDat } } - // ----------------------------------------- - // --- Scaling gas costs - // ----------------------------------------- - sdkCtx := sdk.UnwrapSDKContext(ctx) - - // Charge gas based on the number of delegations that this target has - delegations, err := k.GetDelegationsForTarget(ctx, data.Target) - if err != nil { - return nil, err - } - sdkCtx.GasMeter().ConsumeGas(types.BaseDelegationGasCost*uint64(len(delegations)), "delegation update gas cost") - - // Charge gas based on the number of denoms that are being delegated - sdkCtx.GasMeter().ConsumeGas(types.BaseDelegationDenomCost*uint64(len(data.Amount)), "multi-denom delegation gas cost") - // Convert the addresses to sdk.AccAddress delegatorAddress, err := k.accountKeeper.AddressCodec().StringToBytes(delegator) if err != nil { diff --git a/x/restaking/keeper/msg_server_test.go b/x/restaking/keeper/msg_server_test.go index ee114e61..acd3c48b 100644 --- a/x/restaking/keeper/msg_server_test.go +++ b/x/restaking/keeper/msg_server_test.go @@ -1392,11 +1392,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegatePool() { sdk.NewAttribute(types.AttributeKeyNewShares, "500.000000000000000000pool/1/umilk"), ), }, - check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationDenomCost - // The 34450 is obtained by running this test with BaseDelegationDenomCost set to 0 - suite.Require().GreaterOrEqual(ctx.GasMeter().GasConsumed(), 34450+types.BaseDelegationDenomCost) - }, }, { name: "allowed denom is delegated properly", @@ -1435,11 +1430,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegatePool() { sdk.NewAttribute(types.AttributeKeyNewShares, "500.000000000000000000pool/1/umilk"), ), }, - check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationDenomCost - // The 34565 value is obtained by running this test with BaseDelegationDenomCost set to 0 - suite.Require().GreaterOrEqual(ctx.GasMeter().GasConsumed(), 34565+types.BaseDelegationDenomCost) - }, }, } @@ -1458,9 +1448,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegatePool() { tc.store(ctx) } - // Reset the gas meter - ctx = ctx.WithGasMeter(storetypes.NewInfiniteGasMeter()) - msgServer := keeper.NewMsgServer(suite.k) _, err := msgServer.DelegatePool(ctx, tc.msg) if tc.shouldErr { @@ -1715,11 +1702,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegateOperator() { sdk.NewAttribute(types.AttributeKeyNewShares, "500.000000000000000000operator/1/umilk"), ), }, - check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationDenomCost - // The 49110 is obtained by running this test with BaseDelegationDenomCost set to 0 - suite.Require().GreaterOrEqual(ctx.GasMeter().GasConsumed(), 49110+types.BaseDelegationDenomCost) - }, }, { name: "allowed denom is delegated properly", @@ -1764,11 +1746,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegateOperator() { sdk.NewAttribute(types.AttributeKeyNewShares, "500.000000000000000000operator/1/umilk"), ), }, - check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationDenomCost - // The 49210 is obtained by running this test with BaseDelegationDenomCost set to 0 - suite.Require().GreaterOrEqual(ctx.GasMeter().GasConsumed(), 49210+types.BaseDelegationDenomCost) - }, }, } @@ -1787,9 +1764,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegateOperator() { tc.store(ctx) } - // Reset the gas meter - ctx = ctx.WithGasMeter(storetypes.NewInfiniteGasMeter()) - msgServer := keeper.NewMsgServer(suite.k) _, err := msgServer.DelegateOperator(ctx, tc.msg) if tc.shouldErr { @@ -2122,11 +2096,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegateService() { sdk.NewAttribute(types.AttributeKeyNewShares, "500.000000000000000000service/1/umilk"), ), }, - check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationDenomCost - // The 50060 is obtained by running this test with BaseDelegationDenomCost set to 0 - suite.Require().GreaterOrEqual(ctx.GasMeter().GasConsumed(), 50060+types.BaseDelegationDenomCost) - }, }, { name: "allowed denom is delegated properly", @@ -2171,11 +2140,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegateService() { sdk.NewAttribute(types.AttributeKeyNewShares, "500.000000000000000000service/1/umilk"), ), }, - check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationDenomCost - // The 50170 is obtained by running this test with BaseDelegationDenomCost set to 0 - suite.Require().GreaterOrEqual(ctx.GasMeter().GasConsumed(), 50170+types.BaseDelegationDenomCost) - }, }, { name: "allowed service denom is delegated properly", @@ -2220,11 +2184,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegateService() { sdk.NewAttribute(types.AttributeKeyNewShares, "500.000000000000000000service/1/umilk"), ), }, - check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationDenomCost - // The 50080 is obtained by running this test with BaseDelegationDenomCost set to 0 - suite.Require().GreaterOrEqual(ctx.GasMeter().GasConsumed(), 50080+types.BaseDelegationDenomCost) - }, }, { name: "allowed denom after intersecting allowed denoms is delegated properly", @@ -2273,11 +2232,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegateService() { sdk.NewAttribute(types.AttributeKeyNewShares, "500.000000000000000000service/1/umilk"), ), }, - check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationDenomCost - // The 50230 is obtained by running this test with BaseDelegationDenomCost set to 0 - suite.Require().GreaterOrEqual(ctx.GasMeter().GasConsumed(), 50230+types.BaseDelegationDenomCost) - }, }, } @@ -2296,9 +2250,6 @@ func (suite *KeeperTestSuite) TestMsgServer_DelegateService() { tc.store(ctx) } - // Reset the gas meter - ctx = ctx.WithGasMeter(storetypes.NewInfiniteGasMeter()) - msgServer := keeper.NewMsgServer(suite.k) _, err := msgServer.DelegateService(ctx, tc.msg) if tc.shouldErr { diff --git a/x/restaking/keeper/operator_restaking_test.go b/x/restaking/keeper/operator_restaking_test.go index f6f389ca..3c03ade7 100644 --- a/x/restaking/keeper/operator_restaking_test.go +++ b/x/restaking/keeper/operator_restaking_test.go @@ -460,14 +460,6 @@ func (suite *KeeperTestSuite) TestKeeper_DelegateToOperator() { sdk.NewDecCoinFromDec("operator/1/uinit", sdkmath.LegacyNewDec(100)), ), check: func(ctx sdk.Context) { - // Make sure the gas charged is at least - // BaseDelegationGasCost + BaseDelegationDenomCost - // since a delegation already exists and we are delegating two denoms - suite.Require().GreaterOrEqual( - suite.ctx.GasMeter().GasConsumed(), - types.BaseDelegationGasCost+types.BaseDelegationDenomCost, - ) - // Make sure the operator now exists operator, err := suite.ok.GetOperator(ctx, 1) suite.Require().NoError(err) @@ -574,14 +566,6 @@ func (suite *KeeperTestSuite) TestKeeper_DelegateToOperator() { sdk.NewDecCoinFromDec("operator/1/uinit", sdkmath.LegacyNewDec(600)), ), check: func(ctx sdk.Context) { - // Make sure the gas charged is at least - // BaseDelegationGasCost + (2 * BaseDelegationDenomCost) - // since a delegation already exists and we are delegating two denoms - suite.Require().GreaterOrEqual( - suite.ctx.GasMeter().GasConsumed(), - types.BaseDelegationGasCost+2*types.BaseDelegationDenomCost, - ) - // Make sure the operator now exists operator, err := suite.ok.GetOperator(ctx, 1) suite.Require().NoError(err) diff --git a/x/restaking/keeper/pool_restaking_test.go b/x/restaking/keeper/pool_restaking_test.go index 7673ba69..bc6cb2a1 100644 --- a/x/restaking/keeper/pool_restaking_test.go +++ b/x/restaking/keeper/pool_restaking_test.go @@ -457,13 +457,6 @@ func (suite *KeeperTestSuite) TestKeeper_DelegateToPool() { shouldErr: false, expShares: sdk.NewDecCoins(sdk.NewDecCoinFromDec("pool/1/umilk", sdkmath.LegacyNewDecWithPrec(15625, 2))), check: func(ctx sdk.Context) { - // Make sure the gas charged is at least BaseDelegationGasCost + BaseDelegationDenomCost - // since it's the first delegation and we are delegating one denom - suite.Require().GreaterOrEqual( - suite.ctx.GasMeter().GasConsumed(), - types.BaseDelegationGasCost+types.BaseDelegationDenomCost, - ) - // Make sure the pool now exists pool, err := suite.pk.GetPool(ctx, 1) suite.Require().NoError(err) diff --git a/x/restaking/keeper/service_restaking_test.go b/x/restaking/keeper/service_restaking_test.go index 93153ba6..1df13b90 100644 --- a/x/restaking/keeper/service_restaking_test.go +++ b/x/restaking/keeper/service_restaking_test.go @@ -757,14 +757,6 @@ func (suite *KeeperTestSuite) TestKeeper_DelegateToService() { { name: "delegating another token denom works properly", store: func(ctx sdk.Context) { - // Make sure the gas charged is at least - // BaseDelegationGasCost + BaseDelegationDenomCost - // since a delegation already exists and we are delegating two denoms - suite.Require().GreaterOrEqual( - suite.ctx.GasMeter().GasConsumed(), - types.BaseDelegationGasCost+types.BaseDelegationDenomCost, - ) - // Create the service err := suite.sk.SaveService(ctx, servicestypes.Service{ ID: 1, @@ -919,14 +911,6 @@ func (suite *KeeperTestSuite) TestKeeper_DelegateToService() { sdk.NewDecCoinFromDec("service/1/uinit", sdkmath.LegacyNewDec(600)), ), check: func(ctx sdk.Context) { - // Make sure the gas charged is at least - // BaseDelegationGasCost + (2 * BaseDelegationDenomCost) - // since a delegation already exists and we are delegating two denoms - suite.Require().GreaterOrEqual( - suite.ctx.GasMeter().GasConsumed(), - types.BaseDelegationGasCost+2*types.BaseDelegationDenomCost, - ) - // Make sure the service now exists service, err := suite.sk.GetService(ctx, 1) suite.Require().NoError(err) diff --git a/x/restaking/types/const.go b/x/restaking/types/const.go index fc82aa42..230635a0 100644 --- a/x/restaking/types/const.go +++ b/x/restaking/types/const.go @@ -1,10 +1,6 @@ package types const ( - // BaseDelegationGasCost is the gas cost for a delegation transaction that will be paid per each delegation - // already existing for the same target. This is used to implement delegations count-based scaling costs. - BaseDelegationGasCost uint64 = 20_000 - // BaseDelegationDenomCost is the gas cost for storing or deleting a coin denom for each delegation. // Examples: // * if a user wants to create a new delegation with 3 denoms, they will be charged 3 * BaseDelegationDenomCost From 475ac9e92162dba19021f52a44290c0e6e42b5b5 Mon Sep 17 00:00:00 2001 From: Hanjun Kim Date: Thu, 16 Jan 2025 16:23:14 +0900 Subject: [PATCH 3/3] test: apply code review suggestions --- x/restaking/keeper/end_blocker_test.go | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/x/restaking/keeper/end_blocker_test.go b/x/restaking/keeper/end_blocker_test.go index 7ab85903..4e445981 100644 --- a/x/restaking/keeper/end_blocker_test.go +++ b/x/restaking/keeper/end_blocker_test.go @@ -5,6 +5,7 @@ import ( "time" sdkmath "cosmossdk.io/math" + storetypes "cosmossdk.io/store/types" sdk "github.com/cosmos/cosmos-sdk/types" operatorstypes "github.com/milkyway-labs/milkyway/v7/x/operators/types" @@ -380,21 +381,22 @@ func (suite *KeeperTestSuite) TestGasConsumption_UndelegateVsEndBlockProcessing( // --- Undelegation Gas Tracking --- // Measure gas consumption during initial undelegation - undelegationGasStart := ctx.GasMeter().GasConsumed() + ctx = ctx.WithGasMeter(storetypes.NewInfiniteGasMeter()) // Undelegate the denominations using MsgServer _, err = msgServer.UndelegateService(ctx, types.NewMsgUndelegateService(serviceID, delAmt, delegator)) suite.Require().NoError(err) // Calculate gas used during undelegations - gasUsedForUndelegation := ctx.GasMeter().GasConsumed() - undelegationGasStart + gasUsedForUndelegation := ctx.GasMeter().GasConsumed() + fmt.Println("Gas used for undelegation:", gasUsedForUndelegation) // --- EndBlock Unbond Completion Gas Tracking --- // Advance context time to when undelegations mature ctx = ctx.WithBlockTime(ctx.BlockTime().Add(7 * 24 * time.Hour)) - endBlockGasStart := ctx.GasMeter().GasConsumed() + ctx = ctx.WithGasMeter(storetypes.NewInfiniteGasMeter()) // Measure gas consumption during end block processing // NOTE: we isolate the component of the EndBlock call we want to test @@ -402,7 +404,8 @@ func (suite *KeeperTestSuite) TestGasConsumption_UndelegateVsEndBlockProcessing( suite.Require().NoError(err) // Calculate gas used during end block processing - gasUsedForEndBlock := ctx.GasMeter().GasConsumed() - endBlockGasStart + gasUsedForEndBlock := ctx.GasMeter().GasConsumed() + fmt.Println("Gas used for end block:", gasUsedForEndBlock) // --- Gas Consumption Comparison ---