From c82c524e139d7960794e08bca838a1a1b78a3532 Mon Sep 17 00:00:00 2001
From: joaquim-verges <3353417+joaquim-verges@users.noreply.github.com>
Date: Wed, 24 Jul 2024 22:37:44 +0000
Subject: [PATCH] feat(smart): allow overriding nonce for smart accounts
 (#3782)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

<!-- start pr-codex -->

## PR-Codex overview
The focus of this PR is to allow overriding the nonce for smart accounts in the `thirdweb` package.

### Detailed summary
- Added the ability to override nonce for smart accounts
- Exported `generateRandomUint192` function
- Updated imports and added new functions in `userop.ts` file
- Modified `createUnsignedUserOp` function to use custom nonce logic

> ✨ Ask PR-Codex anything about this PR by commenting with `/codex {your question}`

<!-- end pr-codex -->
---
 .changeset/nervous-fans-clap.md               |  5 ++
 .../thirdweb/src/wallets/smart/lib/userop.ts  | 46 +++++++++++++++++--
 .../thirdweb/src/wallets/smart/lib/utils.ts   |  2 +-
 packages/thirdweb/src/wallets/smart/types.ts  |  1 +
 4 files changed, 49 insertions(+), 5 deletions(-)
 create mode 100644 .changeset/nervous-fans-clap.md

diff --git a/.changeset/nervous-fans-clap.md b/.changeset/nervous-fans-clap.md
new file mode 100644
index 00000000000..cf375d65edf
--- /dev/null
+++ b/.changeset/nervous-fans-clap.md
@@ -0,0 +1,5 @@
+---
+"thirdweb": patch
+---
+
+Allow overriding nonce for smart accounts
diff --git a/packages/thirdweb/src/wallets/smart/lib/userop.ts b/packages/thirdweb/src/wallets/smart/lib/userop.ts
index 98730d092c5..8757734e211 100644
--- a/packages/thirdweb/src/wallets/smart/lib/userop.ts
+++ b/packages/thirdweb/src/wallets/smart/lib/userop.ts
@@ -1,6 +1,11 @@
 import { concat } from "viem";
 import type { Chain } from "../../../chains/types.js";
-import type { ThirdwebContract } from "../../../contract/contract.js";
+import type { ThirdwebClient } from "../../../client/client.js";
+import {
+  type ThirdwebContract,
+  getContract,
+} from "../../../contract/contract.js";
+import { getNonce } from "../../../extensions/erc4337/__generated__/IEntryPoint/read/getNonce.js";
 import { getDefaultGasOverrides } from "../../../gas/fee-data.js";
 import { encode } from "../../../transaction/actions/encode.js";
 import type { PreparedTransaction } from "../../../transaction/prepare-transaction.js";
@@ -30,7 +35,7 @@ import {
   getDefaultBundlerUrl,
 } from "./constants.js";
 import { getPaymasterAndData } from "./paymaster.js";
-import { randomNonce } from "./utils.js";
+import { generateRandomUint192 } from "./utils.js";
 
 /**
  * Wait for the user operation to be mined.
@@ -156,8 +161,13 @@ export async function createUnsignedUserOp(args: {
     }
   }
 
-  // const nonce = BigInt(transaction.nonce || randomNonce());
-  const nonce = randomNonce(); // FIXME getNonce should be overrideable by the wallet
+  const nonce = await getAccountNonce({
+    accountContract,
+    chain,
+    client,
+    entrypointAddress: overrides?.entrypointAddress,
+    getNonceOverride: overrides?.getAccountNonce,
+  });
 
   const partialOp: UserOperation = {
     sender: accountContract.address,
@@ -299,6 +309,34 @@ async function getAccountInitCode(options: {
   return concat([factoryContract.address as Hex, await encode(deployTx)]);
 }
 
+async function getAccountNonce(options: {
+  accountContract: ThirdwebContract;
+  chain: Chain;
+  client: ThirdwebClient;
+  entrypointAddress?: string;
+  getNonceOverride?: (accountContract: ThirdwebContract) => Promise<bigint>;
+}) {
+  const {
+    accountContract,
+    chain,
+    client,
+    entrypointAddress,
+    getNonceOverride,
+  } = options;
+  if (getNonceOverride) {
+    return getNonceOverride(accountContract);
+  }
+  return getNonce({
+    contract: getContract({
+      address: entrypointAddress || ENTRYPOINT_ADDRESS_v0_6,
+      chain,
+      client,
+    }),
+    key: generateRandomUint192(),
+    sender: accountContract.address,
+  });
+}
+
 /**
  * Get the hash of a user operation.
  * @param args - The user operation, entrypoint address, and chain ID
diff --git a/packages/thirdweb/src/wallets/smart/lib/utils.ts b/packages/thirdweb/src/wallets/smart/lib/utils.ts
index 3027f099cc1..d367a518f49 100644
--- a/packages/thirdweb/src/wallets/smart/lib/utils.ts
+++ b/packages/thirdweb/src/wallets/smart/lib/utils.ts
@@ -3,7 +3,7 @@ import type { Chain } from "../../../chains/types.js";
 import { isHex, numberToHex, toHex } from "../../../utils/encoding/hex.js";
 import type { UserOperation, UserOperationHexed } from "../types.js";
 
-const generateRandomUint192 = (): bigint => {
+export const generateRandomUint192 = (): bigint => {
   const rand1 = BigInt(Math.floor(Math.random() * 0x100000000));
   const rand2 = BigInt(Math.floor(Math.random() * 0x100000000));
   const rand3 = BigInt(Math.floor(Math.random() * 0x100000000));
diff --git a/packages/thirdweb/src/wallets/smart/types.ts b/packages/thirdweb/src/wallets/smart/types.ts
index 5875bb9f14a..96a7aefc0ab 100644
--- a/packages/thirdweb/src/wallets/smart/types.ts
+++ b/packages/thirdweb/src/wallets/smart/types.ts
@@ -34,6 +34,7 @@ export type SmartWalletOptions = Prettify<
         accountContract: ThirdwebContract,
         transactions: SendTransactionOption[],
       ) => PreparedTransaction;
+      getAccountNonce?: (accountContract: ThirdwebContract) => Promise<bigint>;
     };
   } & (
     | {