Update LimitSimulationGasDecorator (#45)

sei-protocol · Apr 19, 2024 · 413f974 · 413f974
1 parent 63a4fb1
commit 413f974
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 10 deletions.
diff --git a/app/ante.go b/app/ante.go
@@ -48,8 +48,8 @@ func NewAnteHandler(options HandlerOptions) (sdk.AnteHandler, sdk.AnteDepGenerat
 	}
 
 	anteDecorators := []sdk.AnteFullDecorator{
-		sdk.DefaultWrappedAnteDecorator(ante.NewDefaultSetUpContextDecorator()),                                           // outermost AnteDecorator. SetUpContext must be called first
-		sdk.DefaultWrappedAnteDecorator(wasmkeeper.NewLimitSimulationGasDecorator(options.WasmConfig.SimulationGasLimit)), // after setup context to enforce limits early
+		sdk.DefaultWrappedAnteDecorator(ante.NewDefaultSetUpContextDecorator()),                                                                               // outermost AnteDecorator. SetUpContext must be called first
+		sdk.DefaultWrappedAnteDecorator(wasmkeeper.NewLimitSimulationGasDecorator(options.WasmConfig.SimulationGasLimit, wasmkeeper.DefaultGasMeterSetter())), // after setup context to enforce limits early
 		sdk.DefaultWrappedAnteDecorator(wasmkeeper.NewCountTXDecorator(options.TXCounterStoreKey)),
 		sdk.DefaultWrappedAnteDecorator(ante.NewRejectExtensionOptionsDecorator()),
 		sdk.DefaultWrappedAnteDecorator(ante.NewValidateBasicDecorator()),

diff --git a/x/wasm/keeper/ante.go b/x/wasm/keeper/ante.go
@@ -57,14 +57,22 @@ func decodeHeightCounter(bz []byte) (int64, uint32) {
 // LimitSimulationGasDecorator ante decorator to limit gas in simulation calls
 type LimitSimulationGasDecorator struct {
 	gasLimit *sdk.Gas
+	// inputs: simulate, ctx, gas limit, tx
+	gasMeterSetter func(bool, sdk.Context, uint64, sdk.Tx) sdk.Context
 }
 
 // NewLimitSimulationGasDecorator constructor accepts nil value to fallback to block gas limit.
-func NewLimitSimulationGasDecorator(gasLimit *sdk.Gas) *LimitSimulationGasDecorator {
+func NewLimitSimulationGasDecorator(gasLimit *sdk.Gas, gasMeterSetter func(bool, sdk.Context, uint64, sdk.Tx) sdk.Context) *LimitSimulationGasDecorator {
 	if gasLimit != nil && *gasLimit == 0 {
 		panic("gas limit must not be zero")
 	}
-	return &LimitSimulationGasDecorator{gasLimit: gasLimit}
+	return &LimitSimulationGasDecorator{gasLimit, gasMeterSetter}
+}
+
+func DefaultGasMeterSetter() func(bool, sdk.Context, uint64, sdk.Tx) sdk.Context {
+	return func(simulate bool, ctx sdk.Context, gasLimit uint64, tx sdk.Tx) sdk.Context {
+		return ctx.WithGasMeter(sdk.NewGasMeter(gasLimit))
+	}
 }
 
 // AnteHandle that limits the maximum gas available in simulations only.
@@ -85,12 +93,13 @@ func (d LimitSimulationGasDecorator) AnteHandle(ctx sdk.Context, tx sdk.Tx, simu
 
 	// apply custom node gas limit
 	if d.gasLimit != nil {
-		return next(ctx.WithGasMeter(sdk.NewGasMeter(*d.gasLimit)), tx, simulate)
+		return next(d.gasMeterSetter(simulate, ctx, *d.gasLimit, tx), tx, simulate)
 	}
 
 	// default to max block gas when set, to be on the safe side
 	if maxGas := ctx.ConsensusParams().GetBlock().MaxGas; maxGas > 0 {
-		return next(ctx.WithGasMeter(sdk.NewGasMeter(sdk.Gas(maxGas))), tx, simulate)
+		return next(d.gasMeterSetter(simulate, ctx, sdk.Gas(maxGas), tx), tx, simulate)
 	}
+
 	return next(ctx, tx, simulate)
 }
diff --git a/x/wasm/keeper/ante_test.go b/x/wasm/keeper/ante_test.go
@@ -168,12 +168,12 @@ func TestLimitSimulationGasDecorator(t *testing.T) {
 			// when
 			if spec.expErr != nil {
 				require.PanicsWithValue(t, spec.expErr, func() {
-					ante := keeper.NewLimitSimulationGasDecorator(spec.customLimit)
+					ante := keeper.NewLimitSimulationGasDecorator(spec.customLimit, keeper.DefaultGasMeterSetter())
 					ante.AnteHandle(ctx, nil, spec.simulation, nextAnte)
 				})
 				return
 			}
-			ante := keeper.NewLimitSimulationGasDecorator(spec.customLimit)
+			ante := keeper.NewLimitSimulationGasDecorator(spec.customLimit, keeper.DefaultGasMeterSetter())
 			ante.AnteHandle(ctx, nil, spec.simulation, nextAnte)
 		})
 	}

diff --git a/x/wasm/keeper/recurse_test.go b/x/wasm/keeper/recurse_test.go
@@ -91,15 +91,15 @@ func TestGasCostOnQuery(t *testing.T) {
 				Depth: 1,
 				Work:  50,
 			},
-			expectedGas: 2*GasWork50 + GasReturnHashed,
+			expectedGas: 2*GasWork50 + GasReturnHashed + 1,
 		},
 		"recursion 4, some work": {
 			gasLimit: 400_000,
 			msg: Recurse{
 				Depth: 4,
 				Work:  50,
 			},
-			expectedGas: 5*GasWork50 + 4*GasReturnHashed,
+			expectedGas: 5*GasWork50 + 4*GasReturnHashed + 1,
 		},
 	}