fix: fix scaling gas problem

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 {

x/restaking/keeper/end_blocker_test.go

@@ -1,13 +1,16 @@
 package keeper_test
 
 import (
+	"fmt"
 	"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"
 	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 +317,99 @@ 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
+			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()
+			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))
+
+			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
+			err = suite.k.CompleteMatureUnbondingDelegations(ctx)
+			suite.Require().NoError(err)
+
+			// Calculate gas used during end block processing
+			gasUsedForEndBlock := ctx.GasMeter().GasConsumed()
+			fmt.Println("Gas used for end block:", gasUsedForEndBlock)
+
+			// --- Gas Consumption Comparison ---
+
+			suite.Require().GreaterOrEqual(gasUsedForUndelegation, gasUsedForEndBlock)
+		})
+	}
+}

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 {

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)

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)

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)

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