From 413f974d3ffb2242916ecde2f53d909257b5b722 Mon Sep 17 00:00:00 2001
From: Kartik Bhat <kartikbhatri@gmail.com>
Date: Thu, 18 Apr 2024 22:30:51 -0400
Subject: [PATCH] Update LimitSimulationGasDecorator (#45)

---
 app/ante.go                   |  4 ++--
 x/wasm/keeper/ante.go         | 17 +++++++++++++----
 x/wasm/keeper/ante_test.go    |  4 ++--
 x/wasm/keeper/recurse_test.go |  4 ++--
 4 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/app/ante.go b/app/ante.go
index bcf2608..577c1eb 100644
--- 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
index 1ffd34b..1314420 100644
--- 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
index 593613c..671045e 100644
--- 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
index cfb49e0..ff4011d 100644
--- a/x/wasm/keeper/recurse_test.go
+++ b/x/wasm/keeper/recurse_test.go
@@ -91,7 +91,7 @@ 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,
@@ -99,7 +99,7 @@ func TestGasCostOnQuery(t *testing.T) {
 				Depth: 4,
 				Work:  50,
 			},
-			expectedGas: 5*GasWork50 + 4*GasReturnHashed,
+			expectedGas: 5*GasWork50 + 4*GasReturnHashed + 1,
 		},
 	}