feat: add proper SMA support to factory

Zer0dot · Zer0dot · commit e23a8563aedd · 2024-08-15T15:27:46.000-04:00
diff --git a/script/Deploy.s.sol b/script/Deploy.s.sol
@@ -7,6 +7,8 @@ import {Script, console} from "forge-std/Script.sol";
 import {Create2} from "@openzeppelin/contracts/utils/Create2.sol";
 
 import {AccountFactory} from "../src/account/AccountFactory.sol";
+
+import {SemiModularAccount} from "../src/account/SemiModularAccount.sol";
 import {UpgradeableModularAccount} from "../src/account/UpgradeableModularAccount.sol";
 import {SingleSignerValidation} from "../src/modules/validation/SingleSignerValidation.sol";
 
@@ -16,10 +18,12 @@ contract DeployScript is Script {
     address public owner = vm.envAddress("OWNER");
 
     address public accountImpl = vm.envOr("ACCOUNT_IMPL", address(0));
+    address public semiModularAccountImpl = vm.envOr("SMA_IMPL", address(0));
     address public factory = vm.envOr("FACTORY", address(0));
     address public singleSignerValidation = vm.envOr("SINGLE_SIGNER_VALIDATION", address(0));
 
     bytes32 public accountImplSalt = bytes32(vm.envOr("ACCOUNT_IMPL_SALT", uint256(0)));
+    bytes32 public semiModularAccountImplSalt = bytes32(vm.envOr("SMA_IMPL_SALT", uint256(0)));
     bytes32 public factorySalt = bytes32(vm.envOr("FACTORY_SALT", uint256(0)));
     bytes32 public singleSignerValidationSalt = bytes32(vm.envOr("SINGLE_SIGNER_VALIDATION_SALT", uint256(0)));
 
@@ -34,6 +38,7 @@ contract DeployScript is Script {
 
         vm.startBroadcast();
         _deployAccountImpl(accountImplSalt, accountImpl);
+        _deploySemiModularAccountImpl(semiModularAccountImplSalt, semiModularAccountImpl);
         _deploySingleSignerValidation(singleSignerValidationSalt, singleSignerValidation);
         _deployAccountFactory(factorySalt, factory);
         _addStakeForFactory(uint32(requiredUnstakeDelay), requiredStakeAmount);
@@ -72,6 +77,38 @@ contract DeployScript is Script {
         }
     }
 
+    function _deploySemiModularAccountImpl(bytes32 salt, address expected) internal {
+        console.log(string.concat("Deploying SemiModularAccountImpl with salt: ", vm.toString(salt)));
+
+        address addr = Create2.computeAddress(
+            salt,
+            keccak256(abi.encodePacked(type(SemiModularAccount).creationCode, abi.encode(entryPoint))),
+            CREATE2_FACTORY
+        );
+        if (addr != expected) {
+            console.log("Expected address mismatch");
+            console.log("Expected: ", expected);
+            console.log("Actual: ", addr);
+            revert();
+        }
+
+        if (addr.code.length == 0) {
+            console.log("No code found at expected address, deploying...");
+            SemiModularAccount deployed = new SemiModularAccount{salt: salt}(entryPoint);
+
+            if (address(deployed) != expected) {
+                console.log("Deployed address mismatch");
+                console.log("Expected: ", expected);
+                console.log("Deployed: ", address(deployed));
+                revert();
+            }
+
+            console.log("Deployed SemiModularAccount at: ", address(deployed));
+        } else {
+            console.log("Code found at expected address, skipping deployment");
+        }
+    }
+
     function _deploySingleSignerValidation(bytes32 salt, address expected) internal {
         console.log(string.concat("Deploying SingleSignerValidation with salt: ", vm.toString(salt)));
 
@@ -110,7 +147,7 @@ contract DeployScript is Script {
             keccak256(
                 abi.encodePacked(
                     type(AccountFactory).creationCode,
-                    abi.encode(entryPoint, accountImpl, singleSignerValidation, owner)
+                    abi.encode(entryPoint, accountImpl, semiModularAccountImpl, singleSignerValidation, owner)
                 )
             ),
             CREATE2_FACTORY
@@ -124,8 +161,13 @@ contract DeployScript is Script {
 
         if (addr.code.length == 0) {
             console.log("No code found at expected address, deploying...");
+            // Patched
             AccountFactory deployed = new AccountFactory{salt: salt}(
-                entryPoint, UpgradeableModularAccount(payable(accountImpl)), singleSignerValidation, owner
+                entryPoint,
+                UpgradeableModularAccount(payable(accountImpl)),
+                SemiModularAccount(payable(semiModularAccountImpl)),
+                singleSignerValidation,
+                owner
             );
 
             if (address(deployed) != expected) {
diff --git a/src/account/AccountFactory.sol b/src/account/AccountFactory.sol
@@ -8,26 +8,30 @@ import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.s
 import {Create2} from "@openzeppelin/contracts/utils/Create2.sol";
 
 import {UpgradeableModularAccount} from "../account/UpgradeableModularAccount.sol";
+import {SemiModularAccount} from "../account/SemiModularAccount.sol";
 import {ValidationConfigLib} from "../helpers/ValidationConfigLib.sol";
 
 import {LibClone} from "solady/utils/LibClone.sol";
 
 contract AccountFactory is Ownable {
     UpgradeableModularAccount public immutable ACCOUNT_IMPL;
+    SemiModularAccount public immutable SEMI_MODULAR_ACCOUNT_IMPL;
     bytes32 private immutable _PROXY_BYTECODE_HASH;
     IEntryPoint public immutable ENTRY_POINT;
     address public immutable SINGLE_SIGNER_VALIDATION;
 
     constructor(
         IEntryPoint _entryPoint,
         UpgradeableModularAccount _accountImpl,
+        SemiModularAccount _semiModularImpl,
         address _singleSignerValidation,
         address owner
     ) Ownable(owner) {
         ENTRY_POINT = _entryPoint;
         _PROXY_BYTECODE_HASH =
             keccak256(abi.encodePacked(type(ERC1967Proxy).creationCode, abi.encode(address(_accountImpl), "")));
         ACCOUNT_IMPL = _accountImpl;
+        SEMI_MODULAR_ACCOUNT_IMPL = _semiModularImpl;
         SINGLE_SIGNER_VALIDATION = _singleSignerValidation;
     }
 
@@ -62,7 +66,7 @@ contract AccountFactory is Ownable {
         return UpgradeableModularAccount(payable(addr));
     }
 
-    function createAccountWithFallbackValidation(address owner, uint256 salt)
+    function createSemiModularAccount(address owner, uint256 salt)
         external
         returns (UpgradeableModularAccount)
     {
@@ -76,7 +80,7 @@ contract AccountFactory is Ownable {
         // short circuit if exists
         if (addr.code.length == 0) {
             // not necessary to check return addr since next call will fail if so
-            LibClone.createDeterministicERC1967(address(ACCOUNT_IMPL), immutables, fullSalt);
+            LibClone.createDeterministicERC1967(address(SEMI_MODULAR_ACCOUNT_IMPL), immutables, fullSalt);
         }
 
         return UpgradeableModularAccount(payable(addr));
diff --git a/test/account/AccountFactory.t.sol b/test/account/AccountFactory.t.sol
@@ -2,16 +2,23 @@
 pragma solidity ^0.8.19;
 
 import {AccountFactory} from "../../src/account/AccountFactory.sol";
+
+import {SemiModularAccount} from "../../src/account/SemiModularAccount.sol";
 import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";
 import {AccountTestBase} from "../utils/AccountTestBase.sol";
 
 contract AccountFactoryTest is AccountTestBase {
     AccountFactory internal _factory;
     UpgradeableModularAccount internal _account;
+    SemiModularAccount internal _semiModularAccount;
 
     function setUp() public {
         _account = new UpgradeableModularAccount(entryPoint);
-        _factory = new AccountFactory(entryPoint, _account, address(singleSignerValidation), address(this));
+        _semiModularAccount = new SemiModularAccount(entryPoint);
+
+        _factory = new AccountFactory(
+            entryPoint, _account, _semiModularAccount, address(singleSignerValidation), address(this)
+        );
     }
 
     function test_createAccount() public {
diff --git a/test/mocks/SingleSignerFactoryFixture.sol b/test/mocks/SingleSignerFactoryFixture.sol
@@ -49,7 +49,7 @@ contract SingleSignerFactoryFixture is OptimizedTest {
      */
     function createAccount(address owner, uint256 salt) public returns (UpgradeableModularAccount) {
         if (vm.envBool("SMA_TEST")) {
-            return createAccountWithFallbackValidation(owner, salt);
+            return createSemiModularAccount(owner, salt);
         }
 
         address addr = Create2.computeAddress(getSalt(owner, salt), _PROXY_BYTECODE_HASH);
@@ -74,7 +74,7 @@ contract SingleSignerFactoryFixture is OptimizedTest {
         return UpgradeableModularAccount(payable(addr));
     }
 
-    function createAccountWithFallbackValidation(address owner, uint256 salt)
+    function createSemiModularAccount(address owner, uint256 salt)
         public
         returns (UpgradeableModularAccount)
     {
diff --git a/test/script/Deploy.s.t.sol b/test/script/Deploy.s.t.sol
@@ -10,6 +10,8 @@ import {Create2} from "@openzeppelin/contracts/utils/Create2.sol";
 import {DeployScript} from "../../script/Deploy.s.sol";
 
 import {AccountFactory} from "../../src/account/AccountFactory.sol";
+
+import {SemiModularAccount} from "../../src/account/SemiModularAccount.sol";
 import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";
 import {SingleSignerValidation} from "../../src/modules/validation/SingleSignerValidation.sol";
 
@@ -21,6 +23,7 @@ contract DeployTest is Test {
     address internal _owner;
 
     address internal _accountImpl;
+    address internal _smaImpl;
     address internal _singleSignerValidation;
     address internal _factory;
 
@@ -42,6 +45,12 @@ contract DeployTest is Test {
             CREATE2_FACTORY
         );
 
+        _smaImpl = Create2.computeAddress(
+            bytes32(0),
+            keccak256(abi.encodePacked(type(SemiModularAccount).creationCode, abi.encode(address(_entryPoint)))),
+            CREATE2_FACTORY
+        );
+
         _singleSignerValidation = Create2.computeAddress(
             bytes32(0), keccak256(abi.encodePacked(type(SingleSignerValidation).creationCode)), CREATE2_FACTORY
         );
@@ -51,17 +60,25 @@ contract DeployTest is Test {
             keccak256(
                 abi.encodePacked(
                     type(AccountFactory).creationCode,
-                    abi.encode(address(_entryPoint), _accountImpl, _singleSignerValidation, _owner)
+                    abi.encode(
+                        address(_entryPoint),
+                        _accountImpl,
+                        _smaImpl,
+                        _singleSignerValidation,
+                        _owner
+                    )
                 )
             ),
             CREATE2_FACTORY
         );
 
         vm.setEnv("ACCOUNT_IMPL", vm.toString(address(_accountImpl)));
+        vm.setEnv("SMA_IMPL", vm.toString(address(_smaImpl)));
         vm.setEnv("FACTORY", vm.toString(address(_factory)));
         vm.setEnv("SINGLE_SIGNER_VALIDATION", vm.toString(address(_singleSignerValidation)));
 
         vm.setEnv("ACCOUNT_IMPL_SALT", vm.toString(uint256(0)));
+        vm.setEnv("SMA_IMPL_SALT", vm.toString(uint256(0)));
         vm.setEnv("FACTORY_SALT", vm.toString(uint256(0)));
         vm.setEnv("SINGLE_SIGNER_VALIDATION_SALT", vm.toString(uint256(0)));
 

Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,7 @@ contract SingleSignerFactoryFixture is OptimizedTest {`
`49`	`49`	`*/`
`50`	`50`	`function createAccount(address owner, uint256 salt) public returns (UpgradeableModularAccount) {`
`51`	`51`	`if (vm.envBool("SMA_TEST")) {`
`52`		`- return createAccountWithFallbackValidation(owner, salt);`
	`52`	`+ return createSemiModularAccount(owner, salt);`
`53`	`53`	`}`
`54`	`54`
`55`	`55`	`address addr = Create2.computeAddress(getSalt(owner, salt), _PROXY_BYTECODE_HASH);`
`@@ -74,7 +74,7 @@ contract SingleSignerFactoryFixture is OptimizedTest {`
`74`	`74`	`return UpgradeableModularAccount(payable(addr));`
`75`	`75`	`}`
`76`	`76`
`77`		`- function createAccountWithFallbackValidation(address owner, uint256 salt)`
	`77`	`+ function createSemiModularAccount(address owner, uint256 salt)`
`78`	`78`	`public`
`79`	`79`	`returns (UpgradeableModularAccount)`
`80`	`80`	`{`