override paymaster deposit with max value for gas estimation

Author: sherifahmed990

voltaire_bundler/gas/gas_manager.py

@@ -3,7 +3,10 @@
 import logging
 import math
 from typing import Generic, Any
+from functools import cache
+
 from eth_abi import encode, decode
+from eth_utils import keccak
 
 from voltaire_bundler.user_operation.models import UserOperationType
 from voltaire_bundler.bundle.exceptions import \
@@ -299,3 +302,15 @@ async def calc_l1_gas_estimate_arbitrum(
     @abstractmethod
     def calc_base_preverification_gas(self, user_operation: UserOperationType) -> int:
         pass
+
+
+@cache
+def calculate_deposit_slot_index(address: str) -> str:
+    # same deposit value slot for all entrypoints(for ep 0.6 this slot also has
+    # the staked and stake values)
+    return "0x" + keccak(
+        encode(
+            ["uint256", "uint256"],
+            [int(address, 16), 0]  # slot = 0
+        )
+    ).hex()

voltaire_bundler/gas/gas_manager_v6.py

@@ -6,7 +6,7 @@
 
 from voltaire_bundler.bundle.exceptions import ExecutionException, \
         ExecutionExceptionCode, ValidationException, ValidationExceptionCode
-from voltaire_bundler.gas.gas_manager import GasManager
+from voltaire_bundler.gas.gas_manager import GasManager, calculate_deposit_slot_index
 from voltaire_bundler.user_operation.models import FailedOp
 from voltaire_bundler.user_operation.user_operation_handler import \
         decode_failed_op_event
@@ -197,11 +197,21 @@ async def simulate_handle_op_mod(
         call_data = function_selector + call_data_params.hex()
         # if there is no paymaster, override the sender's balance for gas estimation
         if len(user_operation.paymaster_and_data) == 0:
-            # if the target is zero, simulate_handle_op is called to estimate
-            # gas limits override the sender balance with the high value of 10^15 eth
             default_state_overrides[user_operation.sender_address] = {
+                # override the sender balance with the high value of 10^15 eth
+                # not needed as gas prices can be set to zero, but keeping as
+                # is for ep 0.6 only as changing it might be a breaking change
                 "balance": "0x314dc6448d9338c15b0a00000000"
             }
+        else:
+            slot_index = calculate_deposit_slot_index(
+                    user_operation.paymaster_address_lowercase)
+            default_state_overrides[entrypoint]["stateDiff"] = {
+                # override paymaster deposit with max value as verificationgas
+                # is set to max value for estimation, this will override staked
+                # and stake values as well
+                slot_index: "0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
+            }
 
         params: list[Any] = [
             {

voltaire_bundler/gas/gas_manager_v7v8.py

@@ -8,7 +8,7 @@
 from voltaire_bundler.bundle.exceptions import \
     (ExecutionException, ExecutionExceptionCode,
      ValidationException, ValidationExceptionCode)
-from voltaire_bundler.gas.gas_manager import GasManager
+from voltaire_bundler.gas.gas_manager import GasManager, calculate_deposit_slot_index
 from voltaire_bundler.user_operation.models import FailedOp, FailedOpWithRevert
 from voltaire_bundler.user_operation.user_operation_handler import \
     decode_failed_op_event, decode_failed_op_with_revert_event
@@ -201,16 +201,18 @@ async def simulate_handle_op_mod(
         }
         eip7702_auth = user_operation.eip7702_auth
 
-        if eip7702_auth is not None or user_operation.paymaster is not None:
-            default_state_overrides[user_operation.sender_address] = {}
-            if eip7702_auth is not None:
-                default_state_overrides[user_operation.sender_address][
-                    "code"] = "0xef0100" + eip7702_auth["address"][2:]
-
-            if user_operation.paymaster is not None:
-                default_state_overrides[user_operation.sender_address][
-                    "balance"] = "0x314dc6448d9338c15b0a00000000"
-
+        if eip7702_auth is not None:
+            default_state_overrides[user_operation.sender_address][
+                "code"] = "0xef0100" + eip7702_auth["address"][2:]
+
+        if user_operation.paymaster is not None:
+            slot_index = calculate_deposit_slot_index(
+                user_operation.paymaster_address_lowercase)
+            default_state_overrides[entrypoint]["stateDiff"] = {
+                # override paymaster deposit with max value as verificationgas
+                # is set to max value for estimation
+                slot_index: "0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"
+            }
         call_data = function_selector + call_data_params.hex()
         params: list[Any] = [
             {