save nonce to the safe storage

goodmai · Jun 15, 2022 · 76383a6 · 76383a6
1 parent aaf978c
commit 76383a6
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 10 deletions.
diff --git a/contracts/SafeEIP4337Diatomic.sol b/contracts/SafeEIP4337Diatomic.sol
@@ -43,8 +43,7 @@ contract SafeEIP4337Diatomic is HandlerContext {
         );
 
     bytes32 private constant TRANSACTION_TO_EXECUTE_SLOT = keccak256("eip4337diatomic.transaction_hash_to_execute");
-
-    mapping(address => uint256) public safeNonces;
+    bytes32 private constant SAFE_EIP4337_NONCE_SLOT = keccak256("eip4337diatomic.nonce");
 
     constructor() {
         diatomicAddress = address(this);
@@ -65,8 +64,9 @@ contract SafeEIP4337Diatomic is HandlerContext {
         if (safeAddress != msg.sender) revert InvalidCaller();
 
         // We need to increase the nonce to make it impossible to drain the safe by making it send prefunds for the same transaction
-        uint256 safeNonce = safeNonces[safeAddress]++;
+        uint256 safeNonce = getSafeEip4337Nonce(safeAddress);
         if (safeNonce != userOp.nonce) revert InvalidNonce(userOp.nonce, safeNonce);
+        Safe(safeAddress).execTransactionFromModule(address(this), 0, abi.encodeWithSelector(this.setNonce.selector, safeNonce + 1), 1);
 
         // We need to make sure that the entryPoint's requested prefund is in bounds
         if (requiredPrefund > userOp.requiredPreFund()) revert InvalidPrefund();
@@ -75,7 +75,6 @@ contract SafeEIP4337Diatomic is HandlerContext {
         _validateSignatures(entryPoint, userOp);
 
         bytes32 intermediateTxHash = getIntermediateTransactionHash(userOp.callData, safeNonce, entryPoint, block.chainid);
-
         Safe(safeAddress).execTransactionFromModule(
             address(this),
             0,
@@ -119,7 +118,7 @@ contract SafeEIP4337Diatomic is HandlerContext {
         address payable safeAddress = payable(msg.sender);
         address entryPoint = _msgSender();
         // `validateUserOp` increased the nonce, so we need to use nonce - 1 for hash calculation
-        uint256 safeNonce = safeNonces[safeAddress] - 1;
+        uint256 safeNonce = getSafeEip4337Nonce(safeAddress) - 1;
 
         Safe safe = Safe(safeAddress);
         if (
@@ -251,6 +250,14 @@ contract SafeEIP4337Diatomic is HandlerContext {
         Safe(payable(userOp.sender)).checkSignatures(operationHash, operationData, userOp.signature);
     }
 
+    function getSafeEip4337Nonce(address safe) internal view returns (uint256 nonce) {
+        bytes memory nonceBytes = Safe(safe).getStorageAt(uint256(SAFE_EIP4337_NONCE_SLOT), 32);
+
+        assembly {
+            nonce := mload(add(nonceBytes, 32))
+        }
+    }
+
     function setTransactionToExecute(bytes32 txHash) public {
         if (address(this) == diatomicAddress) {
             revert InvalidCallOpcode();
@@ -262,4 +269,16 @@ contract SafeEIP4337Diatomic is HandlerContext {
             sstore(slot, txHash)
         }
     }
+
+    function setNonce(uint256 nonce) public {
+        if (address(this) == diatomicAddress) {
+            revert InvalidCallOpcode();
+        }
+
+        bytes32 slot = SAFE_EIP4337_NONCE_SLOT;
+        // solhint-disable-next-line no-inline-assembly
+        assembly {
+            sstore(slot, nonce)
+        }
+    }
 }
diff --git a/scripts/runOpGoerli.ts b/scripts/runOpGoerli.ts
@@ -43,7 +43,7 @@ const runOp = async () => {
   const safeOp = buildSafeUserOpContractCall(
     storageSetter,
     'setStorage',
-    ['123456789'],
+    [safe.address],
     eip4337Safe.address,
     '0',
     '0',
@@ -52,7 +52,9 @@ const runOp = async () => {
     { maxFeePerGas: '10', maxPriorityFeePerGas: '5' },
   )
   const safeOpHash = calculateSafeOperationHash(eip4337Diatomic.address, safeOp, await chainId())
-  const signature = buildSignatureBytes([await signHash(user1, safeOpHash)])
+  let signature = buildSignatureBytes([await signHash(user1, safeOpHash)])
+  signature = `${signature.slice(0, -2)}20`
+
   const userOp = buildUserOperationFromSafeUserOperation({
     safeAddress: eip4337Safe.address,
     safeOp,

diff --git a/src/utils/execution.ts b/src/utils/execution.ts
@@ -3,6 +3,7 @@ import { TypedDataSigner } from '@ethersproject/abstract-signer'
 import { AddressZero } from '@ethersproject/constants'
 
 export const TRANSACTION_TO_EXECUTE_SLOT = utils.keccak256(utils.toUtf8Bytes('eip4337diatomic.transaction_hash_to_execute'))
+export const SAFE_EIP4337_NONCE_SLOT = utils.keccak256(utils.toUtf8Bytes('eip4337diatomic.nonce'))
 
 export const EIP_DOMAIN = {
   EIP712Domain: [

diff --git a/test/eip4337/SafeEIP4337Diatomic.spec.ts b/test/eip4337/SafeEIP4337Diatomic.spec.ts
@@ -5,7 +5,7 @@ import { expect } from 'chai'
 import hre, { deployments, waffle } from 'hardhat'
 import '@nomiclabs/hardhat-ethers'
 import { deployContract, getTestSafe, getEip4337Diatomic, getEip4337DiatomicExposed } from '../utils/setup'
-import { buildSignatureBytes, signHash, TRANSACTION_TO_EXECUTE_SLOT } from '../../src/utils/execution'
+import { buildSignatureBytes, signHash, TRANSACTION_TO_EXECUTE_SLOT, SAFE_EIP4337_NONCE_SLOT } from '../../src/utils/execution'
 import {
   buildSafeUserOp,
   buildSafeUserOpContractCall,
@@ -222,8 +222,9 @@ describe('SafeEIP4337Diatomic', async () => {
 
     it('should increase the nonce', async () => {
       const { eip4337Safe, eip4337Diatomic } = await setupTests()
+      const provider = hre.ethers.provider
 
-      expect(await eip4337Safe.safeNonces(eip4337Safe.address)).to.equal(0)
+      expect(await provider.getStorageAt(eip4337Safe.address, SAFE_EIP4337_NONCE_SLOT)).to.equal(`0x${'0'.repeat(64)}`)
       const safeOp = buildSafeUserOp({ safe: eip4337Safe.address, nonce: '0', entryPoint: user1.address })
       const safeOpHash = calculateSafeOperationHash(eip4337Diatomic.address, safeOp, await chainId())
       const signature = buildSignatureBytes([await signHash(user1, safeOpHash)])
@@ -235,7 +236,7 @@ describe('SafeEIP4337Diatomic', async () => {
       const requiredPrefund = getRequiredPrefund(userOp)
 
       await eip4337Safe.validateUserOp(userOp, `0x${'0'.repeat(64)}`, requiredPrefund)
-      expect(await eip4337Safe.safeNonces(eip4337Safe.address)).to.equal(1)
+      expect(await provider.getStorageAt(eip4337Safe.address, SAFE_EIP4337_NONCE_SLOT)).to.equal(`0x${'1'.padStart(64, '0')}`)
     })
 
     it('should mark the transaction as ready to be executed', async () => {
@@ -251,6 +252,7 @@ describe('SafeEIP4337Diatomic', async () => {
         signature,
       })
       const requiredPrefund = getRequiredPrefund(userOp)
+
       const expectedIntermediateTxHash = calculateIntermediateTxHash(userOp.callData, userOp.nonce, user1.address, await chainId())
 
       expect(await provider.getStorageAt(eip4337Safe.address, TRANSACTION_TO_EXECUTE_SLOT)).to.equal(`0x${'0'.repeat(64)}`)