Fix gas measurement for wallets which don't yet exist

contracts/clients/src/BlsWalletWrapper.ts

@@ -32,6 +32,7 @@ export default class BlsWalletWrapper {
     public blsWalletSigner: BlsWalletSigner,
     public privateKey: string,
     public walletContract: BLSWallet,
+    public defaultGatewayAddress: string,
   ) {
     this.address = walletContract.address;
   }
@@ -149,6 +150,7 @@ export default class BlsWalletWrapper {
       blsWalletSigner,
       privateKey,
       await BlsWalletWrapper.BLSWallet(privateKey, verificationGateway),
+      verificationGateway.address,
     );
     blsWalletWrapper.blockGasLimit = (
       await provider.getBlock("latest")
@@ -235,19 +237,33 @@ export default class BlsWalletWrapper {
   }
 
   /** Add gas limit to operation using estimateGas and sign it. */
-  async signWithGasLimit(operation: Omit<Operation, "gas">) {
-    const gas = await this.estimateGas(operation);
+  async signWithGasLimit(operation: Omit<Operation, "gas">, overhead = 0) {
+    let gas = await this.measureOperationGas(operation);
+    gas = gas.add(gas.div(10000).mul(Math.ceil(overhead * 10000)));
+
     return this.sign({ ...operation, gas });
   }
 
-  /** Estimate the gas needed for an operation. */
-  async estimateGas(operation: Omit<Operation, "gas">) {
-    const bundle = this.sign({ ...operation, gas: this.blockGasLimit });
-    const gatewayAddress = await this.walletContract.trustedBLSGateway();
+  /** Measure the gas needed for an operation. */
+  async measureOperationGas(operation: Omit<Operation, "gas">) {
+    const exists =
+      (await this.walletContract.provider.getCode(this.address)) !== "0x";
 
-    const gas = await this.walletContract
-      .connect(gatewayAddress)
-      .estimateGas.performOperation(bundle.operations[0]);
+    const gatewayAddress = exists
+      ? await this.walletContract.trustedBLSGateway()
+      : this.defaultGatewayAddress;
+
+    const gateway = VerificationGateway__factory.connect(
+      gatewayAddress,
+      this.walletContract.provider,
+    );
+
+    const gas = await gateway
+      .connect(ethers.constants.AddressZero)
+      .callStatic.measureOperationGas(this.PublicKey(), {
+        ...operation,
+        gas: this.blockGasLimit,
+      });
 
     return gas;
   }

contracts/contracts/VerificationGateway.sol

@@ -151,7 +151,7 @@ contract VerificationGateway
     function setPendingBLSKeyForWallet() public {
         IWallet wallet = IWallet(msg.sender);
         bytes32 existingHash = hashFromWallet[wallet];
-        require(existingHash != bytes32(0), "VG: hash does not exist for caller");
+        require(existingHash != bytes32(0), "VG: hash not found");
         if (
             (pendingBLSPublicKeyTimeFromHash[existingHash] != 0) &&
             (block.timestamp > pendingBLSPublicKeyTimeFromHash[existingHash])
@@ -196,7 +196,7 @@ contract VerificationGateway
             for (uint256 i=0; i<32; i++) {
                 require(
                     (encodedFunction[selectorOffset+i] == encodedAddress[i]),
-                    "VG: first param to proxy admin is not calling wallet"
+                    "VG: first param is not wallet"
                 );
             }
         }
@@ -255,7 +255,7 @@ contract VerificationGateway
         assembly { size := extcodesize(blsGateway) }
         require(
             (blsGateway != address(0)) && (size > 0),
-            "BLSWallet: gateway address param not valid"
+            "VG: invalid gateway"
         );
 
         IWallet wallet = walletFromHash[hash];
@@ -340,6 +340,30 @@ contract VerificationGateway
         return IWallet(blsWallet);
     }
 
+    /**
+     * Utility for measuring the gas used by an operation, suitable for use in
+     * the operation's required gas parameter.
+     *
+     * This has two important differences over standard estimateGas methods:
+     *   1. It does not include gas for calldata, which has already been paid
+     *   2. It works for wallets which don't yet exist
+     */
+    function measureOperationGas(
+        uint256[BLS_KEY_LEN] memory publicKey,
+        IWallet.Operation calldata op
+    ) external returns (uint) {
+        // Don't allow this to actually be executed on chain. Static calls only.
+        require(msg.sender == address(0));
+
+        IWallet wallet = getOrCreateWallet(publicKey);
+
+        uint gasBefore = gasleft();
+        wallet.performOperation(op);
+        uint gasUsed = gasBefore - gasleft();
+
+        return gasUsed;
+    }
+
     /**
     @dev safely sets/overwrites the wallet for the given public key, ensuring it is properly signed
     @param wallletAddressSignature signature of message containing only the wallet address
@@ -358,7 +382,7 @@ contract VerificationGateway
         );
         require(
             blsLib.verifySingle(wallletAddressSignature, publicKey, addressMsg),
-            "VG: Signature not verified for wallet address."
+            "VG: Sig not verified"
         );
         bytes32 publicKeyHash = keccak256(abi.encodePacked(
             publicKey

contracts/shared/helpers/Fixture.ts

@@ -160,6 +160,14 @@ export default class Fixture {
     );
   }
 
+  async createBLSWallet(): Promise<BlsWalletWrapper> {
+    return BlsWalletWrapper.connect(
+      `0x${Math.floor(Math.random() * 0xffffffff).toString(16)}`,
+      this.verificationGateway.address,
+      this.provider,
+    );
+  }
+
   bundleFrom(
     wallet: BlsWalletWrapper,
     contract: Contract,

contracts/shared/helpers/callProxyAdmin.ts

@@ -41,9 +41,12 @@ export async function proxyAdminBundle(
     contractAddress: fx.verificationGateway.address,
     encodedFunction: encodedWalletAdminCall,
   };
-  const bundle: Bundle = await wallet.signWithGasLimit({
-    nonce: await wallet.Nonce(),
-    actions: [action],
-  });
+  const bundle: Bundle = await wallet.signWithGasLimit(
+    {
+      nonce: await wallet.Nonce(),
+      actions: [action],
+    },
+    0.6, // Add 60% (nested call issue)
+  );
   return bundle;
 }

contracts/test/recovery-test.ts

@@ -348,7 +348,7 @@ describe("Recovery", async function () {
       fx.verificationGateway
         .connect(recoverySigner)
         .recoverWallet(addressSignature, hashAttacker, salt, wallet1Key),
-    ).to.be.rejectedWith("VG: Signature not verified for wallet address");
+    ).to.be.rejectedWith("VG: Sig not verified");
   });
 
   it("should NOT allow a bundle to be executed on a wallet with the same BLS pubkey but different address (replay attack)", async function () {

contracts/test/upgrade-test.ts

@@ -391,10 +391,7 @@ describe("Upgrade", async function () {
       wallet2.address,
       wallet2.address,
     ]);
-    expectOperationFailure(
-      txnReceipt,
-      "VG: first param to proxy admin is not calling wallet",
-    );
+    expectOperationFailure(txnReceipt, "VG: first param is not wallet");
   });
 
   it("should NOT allow walletAdminCall to ProxyAdmin.transferOwnership", async function () {

contracts/test/walletAction-test.ts

@@ -99,8 +99,8 @@ describe("WalletActions", async function () {
 
   it("should send ETH (empty call)", async function () {
     // send money to sender bls wallet
-    const sendWallet = await fx.lazyBlsWallets[0]();
-    const recvWallet = await fx.lazyBlsWallets[1]();
+    const sendWallet = await fx.createBLSWallet();
+    const recvWallet = await fx.createBLSWallet();
     const ethToTransfer = parseEther("0.0001");
     await fx.signers[0].sendTransaction({
       to: sendWallet.address,