Add Staked Safe Factory

4337/contracts/Safe4337Module.sol

@@ -43,6 +43,8 @@ contract Safe4337Module is IAccount, HandlerContext, CompatibilityFallbackHandle
 
     address public immutable SUPPORTED_ENTRYPOINT;
 
+    mapping(address => bool) private deny;
+
     constructor(address entryPoint) {
         require(entryPoint != address(0), "Invalid entry point");
         SUPPORTED_ENTRYPOINT = entryPoint;
@@ -71,6 +73,11 @@ contract Safe4337Module is IAccount, HandlerContext, CompatibilityFallbackHandle
         // the sender is the Safe specified in the userOperation
         require(safeAddress == msg.sender, "Invalid caller");
 
+        // TODO(nlordell): Make sure to remove this! Was just using to check that factory staking was working.
+        if (deny[safeAddress]) {
+            revert("denied");
+        }
+
         // We check the execution function signature to make sure the entry point can't call any other function
         // and make sure the execution of the user operation is handled by the module
         require(

4337/contracts/StakedFactoryProxy.sol

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: LGPL-3.0-only
+pragma solidity >=0.8.0 <0.9.0;
+
+import {IEntryPoint} from "@account-abstraction/contracts/interfaces/IEntryPoint.sol";
+
+contract StakedFactoryProxy {
+    address public immutable FACTORY;
+
+    constructor(address factory) payable {
+        require(factory != address(0), "Invalid factory");
+        FACTORY = factory;
+    }
+
+    fallback() external {
+        (bool success, bytes memory result) = FACTORY.call(msg.data);
+        if (success) {
+            assembly ("memory-safe") {
+                return(add(result, 32), mload(result))
+            }
+        } else {
+            assembly ("memory-safe") {
+                revert(add(result, 32), mload(result))
+            }
+        }
+    }
+
+    function stakeEntryPoint(IEntryPoint entryPoint, uint32 unstakeDelaySecs) external payable {
+        entryPoint.addStake{value: msg.value}(unstakeDelaySecs);
+    }
+}

4337/contracts/test/TestEntryPoint.sol

@@ -2,8 +2,8 @@
 /* solhint-disable one-contract-per-file */
 pragma solidity >=0.8.0;
 
-import {INonceManager} from "@account-abstraction/contracts/interfaces/INonceManager.sol";
 import {IAccount} from "@account-abstraction/contracts/interfaces/IAccount.sol";
+import {INonceManager} from "@account-abstraction/contracts/interfaces/INonceManager.sol";
 import {UserOperation} from "@account-abstraction/contracts/interfaces/UserOperation.sol";
 
 /**

4337/src/deploy/libraries.ts

@@ -1,5 +1,7 @@
 import { DeployFunction } from 'hardhat-deploy/types'
 
+const SAFE_PROXY_FACTORY = process.env.DEPLOYMENT_SAFE_PROXY_FACTORY_ADDRESS
+
 const deploy: DeployFunction = async ({ deployments, getNamedAccounts }) => {
   const { deployer } = await getNamedAccounts()
   const { deploy } = deployments
@@ -10,8 +12,17 @@ const deploy: DeployFunction = async ({ deployments, getNamedAccounts }) => {
     log: true,
     deterministicDeployment: true,
   })
+
+  const factory = await deployments.getOrNull('SafeProxyFactory').then((deployment) => deployment?.address ?? SAFE_PROXY_FACTORY)
+  await deploy('StakedFactoryProxy', {
+    from: deployer,
+    args: [factory],
+    log: true,
+    deterministicDeployment: true,
+  })
 }
 
+deploy.dependencies = ['safe']
 deploy.tags = ['libraries']
 
 export default deploy

4337/src/deploy/safe.ts

@@ -1,5 +1,7 @@
 import { DeployFunction } from 'hardhat-deploy/types'
 
+const SAFE_PROXY_FACTORY = process.env.DEPLOYMENT_SAFE_PROXY_FACTORY_ADDRESS
+
 const deploy: DeployFunction = async ({ deployments, getNamedAccounts, network }) => {
   if (!network.tags.safe && !network.tags.test) {
     return
@@ -8,19 +10,22 @@ const deploy: DeployFunction = async ({ deployments, getNamedAccounts, network }
   const { deployer } = await getNamedAccounts()
   const { deploy } = deployments
 
-  await deploy('SafeProxyFactory', {
-    from: deployer,
-    args: [],
-    log: true,
-    deterministicDeployment: true,
-  })
-
-  await deploy('SafeL2', {
-    from: deployer,
-    args: [],
-    log: true,
-    deterministicDeployment: true,
-  })
+  if (network.tags.safe || network.tags.test) {
+    await deploy('SafeProxyFactory', {
+      from: deployer,
+      args: [],
+      log: true,
+      deterministicDeployment: true,
+    })
+    await deploy('SafeL2', {
+      from: deployer,
+      args: [],
+      log: true,
+      deterministicDeployment: true,
+    })
+  } else if (!SAFE_PROXY_FACTORY) {
+    throw new Error('DEPLOYMENT_SAFE_PROXY_FACTORY_ADDRESS must be set')
+  }
 }
 
 deploy.tags = ['safe']

4337/src/utils/safe.ts

@@ -271,10 +271,14 @@ export class Safe4337 {
     return code !== '0x'
   }
 
-  getInitCode(): string {
+  getInitCode(factoryProxy?: string): string {
     if (!this.safeConfig) throw Error('Init code not available')
-    const initParams = buildInitParamsForConfig(this.safeConfig, this.globalConfig)
-    return initParams.initCode
+    const { initCode } = buildInitParamsForConfig(this.safeConfig, this.globalConfig)
+    if (factoryProxy) {
+      return ethers.solidityPacked(['address', 'bytes'], [factoryProxy, ethers.dataSlice(initCode, 20)])
+    } else {
+      return initCode
+    }
   }
 
   async getSigners(): Promise<string[]> {

4337/test/e2e/LocalBundler.spec.ts

@@ -8,15 +8,16 @@ import { MultiProvider4337, Safe4337 } from '../../src/utils/safe'
 const BUNDLER_URL = process.env.TEST_BUNLDER_URL ?? 'http://localhost:3000/rpc'
 const BUNDLER_MNEMONIC = process.env.TEST_BUNDLER_MNEMONIC ?? 'test test test test test test test test test test test junk'
 
-describe('E2E - Local Bundler', () => {
+describe.only('E2E - Local Bundler', () => {
   before(function () {
     if (network.name !== 'localhost') {
       this.skip()
     }
   })
 
   const setupTests = async () => {
-    const { AddModulesLib, EntryPoint, HariWillibaldToken, Safe4337Module, SafeL2, SafeProxyFactory } = await deployments.run()
+    const { AddModulesLib, EntryPoint, HariWillibaldToken, Safe4337Module, SafeL2, SafeProxyFactory, StakedFactoryProxy } =
+      await deployments.run()
     const [, user] = await prepareAccounts()
     const bundler = bundlerRpc()
 
@@ -43,12 +44,23 @@ describe('E2E - Local Bundler', () => {
       },
     )
 
+    const stakedFactory = await ethers.getContractAt('StakedFactoryProxy', StakedFactoryProxy.address)
+    const stake = ethers.parseEther('1.0')
+    if ((await entryPoint.balanceOf(await stakedFactory.getAddress())) < stake) {
+      await stakedFactory
+        .stakeEntryPoint(await entryPoint.getAddress(), 0xffffffffn, {
+          value: stake,
+        })
+        .then((tx) => tx.wait())
+    }
+
     return {
       user,
       bundler,
       safe,
       validator,
       entryPoint,
+      stakedFactory,
       token,
     }
   }
@@ -86,7 +98,7 @@ describe('E2E - Local Bundler', () => {
   }
 
   it('should deploy a new Safe and execute a transaction', async () => {
-    const { user, bundler, safe, validator, entryPoint, token } = await setupTests()
+    const { user, bundler, safe, validator, entryPoint, stakedFactory, token } = await setupTests()
 
     await token.transfer(safe.address, ethers.parseUnits('4.2', 18)).then((tx) => tx.wait())
     await user.sendTransaction({ to: safe.address, value: ethers.parseEther('0.5') }).then((tx) => tx.wait())
@@ -107,7 +119,7 @@ describe('E2E - Local Bundler', () => {
       safeAddress: safe.address,
       safeOp,
       signature,
-      initCode: safe.getInitCode(),
+      initCode: safe.getInitCode(await stakedFactory.getAddress()),
     })
 
     await bundler.sendUserOperation(userOp, await entryPoint.getAddress())