Add erc20 override utils (#308)

* Add erc20 override utils * Add balanceOverride in prepareUserOperationForErc20Paymaster * Add changeset * chore: format * Change approvalSlot to allowanceSlot * make slots optional * allow overrides by user * Change approval to allowance * update bun * chore: format --------- Co-authored-by: plusminushalf <[email protected]>
pimlicolabs · Oct 23, 2024 · bfc278b · bfc278b
1 parent bc110b0
commit bfc278b
Show file tree

Hide file tree

Showing 14 changed files with 421 additions and 4 deletions.
diff --git a/.changeset/long-walls-hang.md b/.changeset/long-walls-hang.md
@@ -0,0 +1,5 @@
+---
+"permissionless": patch
+---
+
+Added utils to create erc20 state overrides
diff --git a/.changeset/modern-lies-count.md b/.changeset/modern-lies-count.md
@@ -0,0 +1,5 @@
+---
+"permissionless": patch
+---
+
+Added balanceOverride to prepareUserOperationForErc20Paymaster
diff --git a/bun.lockb b/bun.lockb
diff --git a/package.json b/package.json
@@ -29,7 +29,7 @@
         "get-port": "^7.0.0",
         "tsc-alias": "^1.8.8",
         "vitest": "^1.2.0",
-        "viem": "^2.21.2",
+        "viem": "2.21.22",
         "wagmi": "^2.12.8",
         "@permissionless/wagmi": "workspace:packages/wagmi",
         "@types/react": "^18.3.1",

diff --git a/packages/permissionless/actions/pimlico/getTokenQuotes.ts b/packages/permissionless/actions/pimlico/getTokenQuotes.ts
@@ -25,6 +25,8 @@ export type GetTokenQuotesReturnType = {
     postOpGas: bigint
     exchangeRate: bigint
     exchangeRateNativeToUsd: bigint
+    balanceSlot?: bigint
+    allowanceSlot?: bigint
 }[]
 
 /**
@@ -62,6 +64,12 @@ export const getTokenQuotes = async <
 
     return res.quotes.map((quote) => ({
         ...quote,
+        balanceSlot: quote.balanceSlot
+            ? hexToBigInt(quote.balanceSlot)
+            : undefined,
+        allowanceSlot: quote.allowanceSlot
+            ? hexToBigInt(quote.allowanceSlot)
+            : undefined,
         postOpGas: hexToBigInt(quote.postOpGas),
         exchangeRate: hexToBigInt(quote.exchangeRate),
         exchangeRateNativeToUsd: hexToBigInt(quote.exchangeRateNativeToUsd)

diff --git a/...s/permissionless/experimental/pimlico/utils/prepareUserOperationForErc20Paymaster.test.ts b/...s/permissionless/experimental/pimlico/utils/prepareUserOperationForErc20Paymaster.test.ts
@@ -191,5 +191,93 @@ describe.each(getCoreSmartAccounts())(
                 expect(FINAL_ETH_BALANCE).toEqual(INTIAL_ETH_BALANCE) // There should be no ETH balance change
             }
         )
+
+        testWithRpc.skipIf(!supportsEntryPointV07)(
+            "prepareUserOperationForErc20Paymaster_v07",
+            async ({ rpc }) => {
+                const { anvilRpc } = rpc
+
+                const account = (
+                    await getSmartAccountClient({
+                        entryPoint: {
+                            version: "0.7"
+                        },
+                        ...rpc
+                    })
+                ).account
+
+                const publicClient = getPublicClient(anvilRpc)
+
+                const pimlicoClient = createPimlicoClient({
+                    transport: http(rpc.paymasterRpc),
+                    entryPoint: {
+                        address: entryPoint07Address,
+                        version: "0.7"
+                    }
+                })
+
+                const smartAccountClient = createSmartAccountClient({
+                    // @ts-ignore
+                    client: getPublicClient(anvilRpc),
+                    account,
+                    paymaster: pimlicoClient,
+                    chain: foundry,
+                    userOperation: {
+                        prepareUserOperation:
+                            prepareUserOperationForErc20Paymaster(
+                                pimlicoClient,
+                                {
+                                    balanceOverride: true
+                                }
+                            )
+                    },
+                    bundlerTransport: http(rpc.altoRpc)
+                })
+
+                const INITIAL_TOKEN_BALANCE = parseEther("100")
+                const INTIAL_ETH_BALANCE = await publicClient.getBalance({
+                    address: smartAccountClient.account.address
+                })
+
+                sudoMintTokens({
+                    amount: INITIAL_TOKEN_BALANCE,
+                    to: smartAccountClient.account.address,
+                    anvilRpc
+                })
+
+                const opHash = await smartAccountClient.sendUserOperation({
+                    calls: [
+                        {
+                            to: zeroAddress,
+                            data: "0x",
+                            value: 0n
+                        }
+                    ],
+                    paymasterContext: {
+                        token: ERC20_ADDRESS
+                    }
+                })
+
+                const receipt =
+                    await smartAccountClient.waitForUserOperationReceipt({
+                        hash: opHash
+                    })
+
+                expect(receipt).toBeTruthy()
+                expect(receipt).toBeTruthy()
+                expect(receipt.success).toBeTruthy()
+
+                const FINAL_TOKEN_BALANCE = await tokenBalanceOf(
+                    smartAccountClient.account.address,
+                    rpc.anvilRpc
+                )
+                const FINAL_ETH_BALANCE = await publicClient.getBalance({
+                    address: smartAccountClient.account.address
+                })
+
+                expect(FINAL_TOKEN_BALANCE).toBeLessThan(INITIAL_TOKEN_BALANCE) // Token balance should be deducted
+                expect(FINAL_ETH_BALANCE).toEqual(INTIAL_ETH_BALANCE) // There should be no ETH balance change
+            }
+        )
     }
 )
diff --git a/packages/permissionless/experimental/pimlico/utils/prepareUserOperationForErc20Paymaster.ts b/packages/permissionless/experimental/pimlico/utils/prepareUserOperationForErc20Paymaster.ts
@@ -3,6 +3,7 @@ import {
     type Chain,
     type Client,
     type ContractFunctionParameters,
+    RpcError,
     type Transport,
     encodeFunctionData,
     erc20Abi,
@@ -24,9 +25,21 @@ import { getChainId as getChainId_ } from "viem/actions"
 import { readContract } from "viem/actions"
 import { getAction, parseAccount } from "viem/utils"
 import { getTokenQuotes } from "../../../actions/pimlico"
+import { erc20AllowanceOverride, erc20BalanceOverride } from "../../../utils"
 
 export const prepareUserOperationForErc20Paymaster =
-    (pimlicoClient: Client) =>
+    (
+        pimlicoClient: Client,
+        {
+            balanceOverride = false,
+            balanceSlot: _balanceSlot,
+            allowanceSlot: _allowanceSlot
+        }: {
+            balanceOverride?: boolean
+            balanceSlot?: bigint
+            allowanceSlot?: bigint
+        } = {}
+    ) =>
     async <
         account extends SmartAccount | undefined,
         const calls extends readonly unknown[],
@@ -95,6 +108,13 @@ export const prepareUserOperationForErc20Paymaster =
                 entryPointAddress: account.entryPoint.address
             })
 
+            if (quotes.length === 0) {
+                throw new RpcError(new Error("Quotes not found"), {
+                    shortMessage:
+                        "client didn't return token quotes, check if the token is supported"
+                })
+            }
+
             const {
                 postOpGas,
                 exchangeRate,
@@ -118,9 +138,55 @@ export const prepareUserOperationForErc20Paymaster =
             }
 
             ////////////////////////////////////////////////////////////////////////////////
+
             // Call prepareUserOperation
             ////////////////////////////////////////////////////////////////////////////////
 
+            const allowanceSlot = _allowanceSlot ?? quotes[0].allowanceSlot
+            const balanceSlot = _balanceSlot ?? quotes[0].balanceSlot
+
+            const hasSlot = allowanceSlot && balanceSlot
+
+            if (!hasSlot && balanceOverride) {
+                throw new Error(
+                    `balanceOverride is not supported for token ${token}, provide custom slot for balance & allowance overrides`
+                )
+            }
+
+            const balanceStateOverride =
+                balanceOverride && balanceSlot
+                    ? erc20BalanceOverride({
+                          token,
+                          owner: account.address,
+                          slot: balanceSlot
+                      })[0]
+                    : undefined
+
+            const allowanceStateOverride =
+                balanceOverride && allowanceSlot
+                    ? erc20AllowanceOverride({
+                          token,
+                          owner: account.address,
+                          spender: paymasterERC20Address,
+                          slot: allowanceSlot
+                      })[0]
+                    : undefined
+
+            parameters.stateOverride =
+                balanceOverride &&
+                balanceStateOverride &&
+                allowanceStateOverride // allowanceSlot && balanceSlot is cuz of TypeScript :/
+                    ? (parameters.stateOverride ?? []).concat([
+                          {
+                              address: token,
+                              stateDiff: [
+                                  ...(allowanceStateOverride.stateDiff ?? []),
+                                  ...(balanceStateOverride.stateDiff ?? [])
+                              ]
+                          }
+                      ])
+                    : parameters.stateOverride
+
             const userOperation = await getAction(
                 client,
                 prepareUserOperation,

diff --git a/packages/permissionless/package.json b/packages/permissionless/package.json
@@ -71,6 +71,6 @@
         }
     },
     "peerDependencies": {
-        "viem": "^2.21.2"
+        "viem": "^2.21.22"
     }
 }
diff --git a/packages/permissionless/types/pimlico.ts b/packages/permissionless/types/pimlico.ts
@@ -35,6 +35,8 @@ type GetTokenQuotesWithBigIntAsHex = {
         postOpGas: Hex
         exchangeRate: Hex
         exchangeRateNativeToUsd: Hex
+        balanceSlot?: Hex
+        allowanceSlot?: Hex
     }[]
 }
 

diff --git a/packages/permissionless/utils/erc20AllowanceOverride.test.ts b/packages/permissionless/utils/erc20AllowanceOverride.test.ts
@@ -0,0 +1,59 @@
+import { toHex } from "viem"
+import { describe, expect, test } from "vitest"
+import {
+    type Erc20AllowanceOverrideParameters,
+    erc20AllowanceOverride
+} from "./erc20AllowanceOverride"
+
+describe("erc20AllowanceOverride", () => {
+    test("should return the correct structure for valid inputs", () => {
+        const params = {
+            token: "0xTokenAddress",
+            owner: "0xOwnerAddress",
+            spender: "0xSpenderAddress",
+            slot: BigInt(1),
+            amount: BigInt(100)
+        } as const
+
+        const result = erc20AllowanceOverride(params)
+
+        expect(result).toEqual([
+            {
+                address: params.token,
+                stateDiff: [
+                    {
+                        slot: expect.any(String), // Slot will be a keccak256 hash
+                        value: toHex(params.amount)
+                    }
+                ]
+            }
+        ])
+    })
+
+    test("should use the default amount when none is provided", () => {
+        const params: Erc20AllowanceOverrideParameters = {
+            token: "0xTokenAddress",
+            owner: "0xOwnerAddress",
+            spender: "0xSpenderAddress",
+            slot: BigInt(1)
+        }
+
+        const result = erc20AllowanceOverride(params)
+
+        const expectedDefaultAmount = BigInt(
+            "0x7FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF"
+        )
+
+        expect(result).toEqual([
+            {
+                address: params.token,
+                stateDiff: [
+                    {
+                        slot: expect.any(String), // Slot will be a keccak256 hash
+                        value: toHex(expectedDefaultAmount)
+                    }
+                ]
+            }
+        ])
+    })
+})
diff --git a/packages/permissionless/utils/erc20AllowanceOverride.ts b/packages/permissionless/utils/erc20AllowanceOverride.ts
@@ -0,0 +1,66 @@
+import {
+    type Address,
+    type StateOverride,
+    encodeAbiParameters,
+    keccak256,
+    toHex
+} from "viem"
+
+export type Erc20AllowanceOverrideParameters = {
+    token: Address
+    owner: Address
+    spender: Address
+    slot: bigint
+    amount?: bigint
+}
+
+export function erc20AllowanceOverride({
+    token,
+    owner,
+    spender,
+    slot,
+    amount = BigInt(
+        "0x7FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF"
+    )
+}: Erc20AllowanceOverrideParameters): StateOverride {
+    const smartAccountErc20AllowanceSlot = keccak256(
+        encodeAbiParameters(
+            [
+                {
+                    type: "address"
+                },
+                {
+                    type: "bytes32"
+                }
+            ],
+            [
+                spender,
+                keccak256(
+                    encodeAbiParameters(
+                        [
+                            {
+                                type: "address"
+                            },
+                            {
+                                type: "uint256"
+                            }
+                        ],
+                        [owner, BigInt(slot)]
+                    )
+                )
+            ]
+        )
+    )
+
+    return [
+        {
+            address: token,
+            stateDiff: [
+                {
+                    slot: smartAccountErc20AllowanceSlot,
+                    value: toHex(amount)
+                }
+            ]
+        }
+    ]
+}