Nonce management + fixes

gnarvaja · gnarvaja · commit bc6c52a37332 · 2024-08-30T11:10:58.000-03:00
Nonce Management with different modes.

Fixes with respect to the gas, now baseFee is fetch from the blockchain
and maxPriorityFee is fetch from Alchemy API.

Also, there was an error when packing gasFees.
diff --git a/src/ethproto/aa_bundler.py b/src/ethproto/aa_bundler.py
@@ -1,3 +1,6 @@
+import random
+from warnings import warn
+from enum import Enum
 import requests
 from environs import Env
 from eth_abi import encode
@@ -7,11 +10,44 @@
 from web3 import Web3
 from .contracts import RevertError
 
+
 env = Env()
 
 AA_BUNDLER_SENDER = env.str("AA_BUNDLER_SENDER", None)
 AA_BUNDLER_ENTRYPOINT = env.str("AA_BUNDLER_ENTRYPOINT", "0x0000000071727De22E5E9d8BAf0edAc6f37da032")
 AA_BUNDLER_EXECUTOR_PK = env.str("AA_BUNDLER_EXECUTOR_PK", None)
+AA_BUNDLER_PROVIDER = env.str("AA_BUNDLER_PROVIDER", "alchemy")
+
+NonceMode = Enum(
+    "NonceMode",
+    [
+        "RANDOM_KEY",  # first time initializes a random key and increments nonce locally with calling the blockchain
+        "FIXED_KEY_LOCAL_NONCE",  # uses a fixed key, keeps nonce locally and fetches the nonce when receiving
+        # 'AA25 invalid account nonce'
+        "FIXED_KEY_FETCH_ALWAYS",  # uses a fixed key, always fetches unless received as parameter
+    ],
+)
+
+AA_BUNDLER_NONCE_MODE = env.enum("AA_BUNDLER_NONCE_MODE", default="FIXED_KEY_LOCAL_NONCE", type=NonceMode)
+AA_BUNDLER_NONCE_KEY = env.int("AA_BUNDLER_NONCE_KEY", 0)
+AA_BUNDLER_MAX_GETNONCE_RETRIES = env.int("AA_BUNDLER_MAX_GETNONCE_RETRIES", 3)
+
+
+GET_NONCE_ABI = [
+    {
+        "inputs": [
+            {"internalType": "address", "name": "sender", "type": "address"},
+            {"internalType": "uint192", "name": "key", "type": "uint192"},
+        ],
+        "name": "getNonce",
+        "outputs": [{"internalType": "uint256", "name": "nonce", "type": "uint256"}],
+        "stateMutability": "view",
+        "type": "function",
+    }
+]
+
+NONCE_CACHE = {}
+RANDOM_NONCE_KEY = None
 
 
 def pack_two(a, b):
@@ -37,7 +73,7 @@ def pack_user_operation(user_operation):
         "callData": user_operation["callData"],
         "accountGasLimits": pack_two(user_operation["verificationGasLimit"], user_operation["callGasLimit"]),
         "preVerificationGas": _to_uint(user_operation["preVerificationGas"]),
-        "gasFees": pack_two(user_operation["maxFeePerGas"], user_operation["maxPriorityFeePerGas"]),
+        "gasFees": pack_two(user_operation["maxPriorityFeePerGas"], user_operation["maxFeePerGas"]),
         "paymasterAndData": "0x",
         "signature": "0x",
     }
@@ -81,8 +117,64 @@ def sign_user_operation(private_key, user_operation, chain_id, entry_point):
     return signature.signature.hex()
 
 
-def send_transaction(w3, tx):
-    nonce = 0
+def make_nonce(nonce_key, nonce):
+    nonce_key = _to_uint(nonce_key)
+    nonce = _to_uint(nonce)
+    return (nonce_key << 64) | nonce
+
+
+def fetch_nonce(w3, account, entry_point, nonce_key):
+    ep = w3.eth.contract(abi=GET_NONCE_ABI, address=entry_point)
+    return ep.functions.getNonce(account, nonce_key).call()
+
+
+def get_random_nonce_key():
+    global RANDOM_NONCE_KEY
+    if RANDOM_NONCE_KEY is None:
+        RANDOM_NONCE_KEY = random.randint(1, 2**192 - 1)
+    return RANDOM_NONCE_KEY
+
+
+def get_nonce_and_key(w3, tx, nonce_mode, entry_point=AA_BUNDLER_ENTRYPOINT, fetch=False):
+    nonce_key = tx.get("nonceKey", None)
+    nonce = tx.get("nonce", None)
+
+    if nonce_key is None:
+        if nonce_mode == NonceMode.RANDOM_KEY:
+            nonce_key = get_random_nonce_key()
+        else:
+            nonce_key = AA_BUNDLER_NONCE_KEY
+
+    if nonce is None:
+        if fetch or nonce_mode == NonceMode.FIXED_KEY_FETCH_ALWAYS:
+            nonce = fetch_nonce(w3, tx.get("from", AA_BUNDLER_SENDER), entry_point, nonce_key)
+        elif nonce_key not in NONCE_CACHE:
+            nonce = 0
+        else:
+            nonce = NONCE_CACHE[nonce_key]
+    return nonce_key, nonce
+
+
+def handle_response_error(resp, w3, tx, retry_nonce):
+    if "AA25" in resp["error"]["message"] and AA_BUNDLER_MAX_GETNONCE_RETRIES > 0:
+        # Retry fetching the nonce
+        if retry_nonce == AA_BUNDLER_MAX_GETNONCE_RETRIES:
+            raise RevertError(resp["error"]["message"])
+        warn(f'{resp["error"]["message"]} error, I will retry fetching the nonce')
+        return send_transaction(w3, tx, retry_nonce=(retry_nonce or 0) + 1)
+    else:
+        raise RevertError(resp["error"]["message"])
+
+
+def get_base_fee(w3):
+    blk = w3.eth.get_block("latest")
+    return blk["baseFeePerGas"]
+
+
+def send_transaction(w3, tx, retry_nonce=None):
+    nonce_key, nonce = get_nonce_and_key(
+        w3, tx, AA_BUNDLER_NONCE_MODE, entry_point=AA_BUNDLER_ENTRYPOINT, fetch=retry_nonce is not None
+    )
     # "0xb61d27f6" = bytes4 hash of execute(address,uint256,bytes)
     call_data = (
         "0xb61d27f6"
@@ -94,22 +186,36 @@ def send_transaction(w3, tx):
     )
     user_operation = {
         "sender": tx.get("from", AA_BUNDLER_SENDER),
-        "nonce": hex(nonce),
+        "nonce": hex(make_nonce(nonce_key, nonce)),
         "callData": call_data,
         "signature": dummy_signature,
     }
-    resp = w3.provider.make_request("eth_estimateUserOperationGas", [user_operation, AA_BUNDLER_ENTRYPOINT])
-    if "error" in resp:
-        raise RevertError(resp["error"]["message"])
-    user_operation.update(resp["result"])
-
-    resp = w3.provider.make_request("rundler_maxPriorityFeePerGas", [])
-    if "error" in resp:
-        raise RevertError(resp["error"]["message"])
-    max_priority_fee_per_gas = resp["result"]
 
-    user_operation["maxFeePerGas"] = max_priority_fee_per_gas
-    user_operation["maxPriorityFeePerGas"] = max_priority_fee_per_gas
+    if AA_BUNDLER_PROVIDER == "alchemy":
+        resp = w3.provider.make_request(
+            "eth_estimateUserOperationGas", [user_operation, AA_BUNDLER_ENTRYPOINT]
+        )
+        if "error" in resp:
+            return handle_response_error(resp, w3, tx, retry_nonce)
+
+        user_operation.update(resp["result"])
+
+        resp = w3.provider.make_request("rundler_maxPriorityFeePerGas", [])
+        if "error" in resp:
+            raise RevertError(resp["error"]["message"])
+        max_priority_fee_per_gas = resp["result"]
+        user_operation["maxPriorityFeePerGas"] = max_priority_fee_per_gas
+        user_operation["maxFeePerGas"] = hex(_to_uint(max_priority_fee_per_gas) + _to_uint(get_base_fee(w3)))
+    elif AA_BUNDLER_PROVIDER == "gelato":
+        user_operation.update(
+            {
+                "preVerificationGas": "0x00",
+                "callGasLimit": "0x00",
+                "verificationGasLimit": "0x00",
+                "maxFeePerGas": "0x00",
+                "maxPriorityFeePerGas": "0x00",
+            }
+        )
     user_operation["signature"] = sign_user_operation(
         AA_BUNDLER_EXECUTOR_PK, user_operation, tx["chainId"], AA_BUNDLER_ENTRYPOINT
     )
@@ -121,6 +227,8 @@ def send_transaction(w3, tx):
 
     resp = w3.provider.make_request("eth_sendUserOperation", [user_operation, AA_BUNDLER_ENTRYPOINT])
     if "error" in resp:
-        raise RevertError(resp["error"]["message"])
+        return handle_response_error(resp, w3, tx, retry_nonce)
 
+    # Store nonce in the cache, so next time uses a new nonce
+    NONCE_CACHE[nonce_key] = nonce + 1
     return resp["result"]
diff --git a/src/ethproto/w3wrappers.py b/src/ethproto/w3wrappers.py
@@ -108,13 +108,15 @@ def transact(provider, function, tx_kwargs):
     elif W3_TRANSACT_MODE == "defender-async":
         from .defender_relay import send_transaction
 
-        tx_kwargs = {**provider.tx_kwargs, **tx_kwargs}
+        tx_kwargs |= provider.tx_kwargs
         tx = function.build_transaction(tx_kwargs)
         return send_transaction(tx)
     elif W3_TRANSACT_MODE == "aa-bundler-async":
         from .aa_bundler import send_transaction
 
-        tx_kwargs = {**provider.tx_kwargs, **tx_kwargs}
+        tx_kwargs |= provider.tx_kwargs
+        # To avoid fetching gas and gasPrice in a standard way, when it's not relevant for user ops
+        tx_kwargs.update(dict(gas=0, gasPrice=0))
         tx = function.build_transaction(tx_kwargs)
         return send_transaction(provider.w3, tx)
     else:
diff --git a/tests/test_aa_bundler.py b/tests/test_aa_bundler.py
@@ -1,7 +1,7 @@
 from hexbytes import HexBytes
 from ethproto import aa_bundler
 from web3.constants import HASH_ZERO
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 
 def test_pack_two():
@@ -21,9 +21,10 @@ def test_pack_two():
 CHAIN_ID = 31337
 ENTRYPOINT = "0x0000000071727De22E5E9d8BAf0edAc6f37da032"
 
+TEST_SENDER = "0x8961423b54f06bf6D57F8dD3dD1184FA6F3aac3f"
 
 user_operation = {
-    "sender": "0x515f3Db6c4249919B74eA55915969944fEA4B311",
+    "sender": TEST_SENDER,
     "nonce": 0,
     "initCode": "0x",
     "callData": TEST_CALL_DATA,
@@ -41,7 +42,7 @@ def test_pack_two():
 
 def test_pack_user_operation():
     expected = {
-        "sender": "0x515f3Db6c4249919B74eA55915969944fEA4B311",
+        "sender": TEST_SENDER,
         "nonce": 0,
         "initCode": "0x",
         "callData": TEST_CALL_DATA,
@@ -57,20 +58,126 @@ def test_pack_user_operation():
 def test_hash_packed_user_operation():
     packed = aa_bundler.pack_user_operation(user_operation)
     hash = aa_bundler.hash_packed_user_operation_only(packed)
-    assert hash == "0xb3c6cda6d25de5a793bc280200673119f76f92017c97dacd26bc1329771b96a4"
+    assert hash == "0xa2c19765d18b0d690c05b20061bd23d066201aff1833a51bd28af115fbd4bcd9"
     hash = aa_bundler.hash_packed_user_operation(packed, CHAIN_ID, ENTRYPOINT)
-    assert hash == "0x213b6b5f785983fa3310d6ae06e63ff883915ad5454dd422e15d9778a9e1da48"
+    assert hash == "0xb365ad4d366e9081718e926912da7a78a2faae592286fda0cc11923bd141b7cf"
 
 
 def test_sign_user_operation():
     signature = aa_bundler.sign_user_operation(TEST_PRIVATE_KEY, user_operation, CHAIN_ID, ENTRYPOINT)
     assert (
         signature
-        == "0x9a2e58cbe1d7c79b933c115e6d041fca080c5a1f572b78116c36b956faf9bf660b4fc10f339fd608d11b56072407bb29d311edb3a79f312f6f8375a97692870d1b"  # noqa
+        == "0xb9b872bfe4e90f4628e8ec24879a5b01045f91da8457f3ce2b417d2e5774b508261ec1147a820e75a141cb61b884a78d7e88996ceddafb9a7016cfe7a48a1f4f1b"  # noqa
     )
 
 
-def test_send_transaction():
+def test_sign_user_operation_gas_diff():
+    user_operation_2 = dict(user_operation)
+    user_operation_2["maxPriorityFeePerGas"] -= 1
+    signature = aa_bundler.sign_user_operation(TEST_PRIVATE_KEY, user_operation_2, CHAIN_ID, ENTRYPOINT)
+    assert (
+        signature
+        == "0x8162479d2dbd18d7fe93a2f51e283021d6e4eae4f57d20cdd553042723a0b0ea690ab3903d45126b0047da08ab53dfdf86656e4f258ac4936ba96a759ccb77f61b"  # noqa
+    )
+
+
+def test_make_nonce():
+    assert aa_bundler.make_nonce(0, 0) == 0
+    assert aa_bundler.make_nonce(0, 1) == 1
+    assert aa_bundler.make_nonce(1, 1) == (1 << 64) + 1
+
+
+FAIL_IF_USED = object()
+
+
+@patch.object(aa_bundler.random, "randint")
+@patch.object(aa_bundler, "fetch_nonce")
+def test_get_nonce_force_fetch(fetch_nonce_mock, randint_mock):
+    # Test fetch=True
+    fetch_nonce_mock.return_value = 123
+    assert aa_bundler.get_nonce_and_key(
+        FAIL_IF_USED,
+        {"nonceKey": 12, "from": TEST_SENDER},
+        nonce_mode=aa_bundler.NonceMode.FIXED_KEY_LOCAL_NONCE,
+        fetch=True,
+    ) == (12, 123)
+    fetch_nonce_mock.assert_called_once_with(FAIL_IF_USED, TEST_SENDER, ENTRYPOINT, 12)
+    randint_mock.assert_not_called()
+
+
+@patch.object(aa_bundler.random, "randint")
+@patch.object(aa_bundler, "fetch_nonce")
+def test_get_nonce_fetch_always_mode(fetch_nonce_mock, randint_mock):
+    # Test nonce_mode=NonceMode.FIXED_KEY_FETCH_ALWAYS
+    fetch_nonce_mock.return_value = 111
+    assert aa_bundler.get_nonce_and_key(
+        FAIL_IF_USED,
+        {"nonceKey": 22, "from": TEST_SENDER},
+        nonce_mode=aa_bundler.NonceMode.FIXED_KEY_FETCH_ALWAYS,
+    ) == (22, 111)
+    fetch_nonce_mock.assert_called_once_with(FAIL_IF_USED, TEST_SENDER, ENTRYPOINT, 22)
+    randint_mock.assert_not_called()
+    fetch_nonce_mock.reset_mock()
+
+
+@patch.object(aa_bundler.random, "randint")
+@patch.object(aa_bundler, "fetch_nonce")
+def test_get_nonce_nonce_key_in_tx(fetch_nonce_mock, randint_mock):
+    # Test nonce_mode=NonceMode.FIXED_KEY_LOCAL_NONCE
+    assert aa_bundler.get_nonce_and_key(
+        FAIL_IF_USED,
+        {"nonceKey": 22, "from": TEST_SENDER},
+        nonce_mode=aa_bundler.NonceMode.FIXED_KEY_LOCAL_NONCE,
+    ) == (22, 0)
+    randint_mock.assert_not_called()
+    fetch_nonce_mock.assert_not_called()
+
+    # Same if nonce_mode=NonceMode.RANDOM_KEY but nonceKey in the tx
+    assert aa_bundler.get_nonce_and_key(
+        FAIL_IF_USED,
+        {"nonceKey": 22, "from": TEST_SENDER},
+        nonce_mode=aa_bundler.NonceMode.RANDOM_KEY,
+    ) == (22, 0)
+    randint_mock.assert_not_called()
+    fetch_nonce_mock.assert_not_called()
+
+
+@patch.object(aa_bundler.random, "randint")
+@patch.object(aa_bundler, "fetch_nonce")
+def test_get_nonce_random_key_mode(fetch_nonce_mock, randint_mock):
+    # If nonce_mode=NonceMode.RANDOM_KEY creates a random key and stores it
+    randint_mock.return_value = 444
+    assert aa_bundler.get_nonce_and_key(
+        FAIL_IF_USED,
+        {"from": TEST_SENDER},
+        nonce_mode=aa_bundler.NonceMode.RANDOM_KEY,
+    ) == (444, 0)
+    fetch_nonce_mock.assert_not_called()
+    randint_mock.assert_called_with(1, 2**192 - 1)
+    randint_mock.reset_mock()
+    assert aa_bundler.RANDOM_NONCE_KEY == 444
+    aa_bundler.RANDOM_NONCE_KEY = None  # cleanup
+
+
+@patch.object(aa_bundler.random, "randint")
+@patch.object(aa_bundler, "fetch_nonce")
+def test_get_nonce_with_local_cache(fetch_nonce_mock, randint_mock):
+    with patch.object(aa_bundler, "AA_BUNDLER_NONCE_KEY", new=55), patch.object(
+        aa_bundler, "NONCE_CACHE", new={55: 33}
+    ):
+        # Test nonce_mode=NonceMode.FIXED_KEY_LOCAL_NONCE
+        assert aa_bundler.get_nonce_and_key(
+            FAIL_IF_USED,
+            {"from": TEST_SENDER},
+            nonce_mode=aa_bundler.NonceMode.FIXED_KEY_LOCAL_NONCE,
+        ) == (55, 33)
+        randint_mock.assert_not_called()
+        fetch_nonce_mock.assert_not_called()
+
+
+@patch.object(aa_bundler, "get_base_fee")
+def test_send_transaction(get_base_fee_mock):
+    get_base_fee_mock.return_value = 0
     w3 = MagicMock()
     w3.eth.chain_id = CHAIN_ID
 
@@ -91,8 +198,8 @@ def make_request(method, params):
             assert params[0] == {
                 "sender": "0xE8B412158c205B0F605e0FC09dCdA27d3F140FE9",
                 "nonce": "0x0",
-                "callData": "0xb61d27f60000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa84174000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000600000000000000000000000000000000000000000000000000000000000000044095ea7b30000000000000000000000007ace242f32208d836a2245df957c08547059bf45ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff00000000000000000000000000000000000000000000000000000000",
-                "signature": "0xfffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c",
+                "callData": "0xb61d27f60000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa84174000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000600000000000000000000000000000000000000000000000000000000000000044095ea7b30000000000000000000000007ace242f32208d836a2245df957c08547059bf45ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff00000000000000000000000000000000000000000000000000000000",  # noqa
+                "signature": "0xfffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c",  # noqa
             }
             return {
                 "jsonrpc": "2.0",
@@ -113,13 +220,13 @@ def make_request(method, params):
             assert params[0] == {
                 "sender": "0xE8B412158c205B0F605e0FC09dCdA27d3F140FE9",
                 "nonce": "0x0",
-                "callData": "0xb61d27f60000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa84174000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000600000000000000000000000000000000000000000000000000000000000000044095ea7b30000000000000000000000007ace242f32208d836a2245df957c08547059bf45ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff00000000000000000000000000000000000000000000000000000000",
+                "callData": "0xb61d27f60000000000000000000000002791bca1f2de4661ed88a30c99a7a9449aa84174000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000600000000000000000000000000000000000000000000000000000000000000044095ea7b30000000000000000000000007ace242f32208d836a2245df957c08547059bf45ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff00000000000000000000000000000000000000000000000000000000",  # noqa
                 "callGasLimit": "0xcbb8",
                 "verificationGasLimit": "0x13664",
                 "preVerificationGas": "0xb430",
                 "maxFeePerGas": "0x7ffffffff",
                 "maxPriorityFeePerGas": "0x7ffffffff",
-                "signature": "0x7980544d044bc1202fed7edec96f2fa795ab8670b439935e6bbb5104e95d84ea32af8bff187913ff7eb2b442baab06d0c300273942e312332659ab0a194bbbe81c",
+                "signature": "0x7980544d044bc1202fed7edec96f2fa795ab8670b439935e6bbb5104e95d84ea32af8bff187913ff7eb2b442baab06d0c300273942e312332659ab0a194bbbe81c",  # noqa
             }
             return {
                 "jsonrpc": "2.0",
@@ -130,4 +237,6 @@ def make_request(method, params):
     w3.provider.make_request.side_effect = make_request
 
     ret = aa_bundler.send_transaction(w3, tx)
+    get_base_fee_mock.assert_called_once_with(w3)
+    assert aa_bundler.NONCE_CACHE[0] == 1
     assert ret == "0xa950a17ca1ed83e974fb1aa227360a007cb65f566518af117ffdbb04d8d2d524"
