⚡️ Optimize EIP7702Proxy for zero address default admin (#1368)

Author: Vectorized

src/accounts/EIP7702Proxy.sol

@@ -19,6 +19,10 @@ contract EIP7702Proxy {
     /// @dev For allowing the differentiation of the EOA and the proxy itself.
     uint256 internal immutable __self = uint256(uint160(address(this)));
 
+    /// @dev The default implementation. Provided for optimization.
+    /// Set if the `initialAdmin == address(0) && initialImplementation != address(0)`.
+    uint256 internal immutable _defaultImplementation;
+
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
     /*                          STORAGE                           */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
@@ -45,11 +49,18 @@ contract EIP7702Proxy {
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
 
     constructor(address initialImplementation, address initialAdmin) payable {
+        uint256 defaultImplementation;
         /// @solidity memory-safe-assembly
         assembly {
-            sstore(_ERC1967_IMPLEMENTATION_SLOT, shr(96, shl(96, initialImplementation)))
-            sstore(_ERC1967_ADMIN_SLOT, shr(96, shl(96, initialAdmin)))
+            let implementation := shr(96, shl(96, initialImplementation))
+            let admin := shr(96, shl(96, initialAdmin))
+            // We will store the implementation in the storage regardless,
+            // to aid proxy detection on block explorers.
+            sstore(_ERC1967_IMPLEMENTATION_SLOT, implementation)
+            sstore(_ERC1967_ADMIN_SLOT, admin)
+            defaultImplementation := mul(lt(admin, iszero(iszero(implementation))), implementation)
         }
+        _defaultImplementation = defaultImplementation;
     }
 
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
@@ -58,6 +69,7 @@ contract EIP7702Proxy {
 
     fallback() external payable virtual {
         uint256 s = __self;
+        uint256 defaultImplementation = _defaultImplementation;
         /// @solidity memory-safe-assembly
         assembly {
             mstore(0x40, returndatasize()) // Optimization trick to change `6040608052` into `3d604052`.
@@ -106,35 +118,41 @@ contract EIP7702Proxy {
             // Workflow for the EIP7702 authority (i.e. the EOA).
             let impl := sload(_ERC1967_IMPLEMENTATION_SLOT) // The preferred implementation on the EOA.
             calldatacopy(0x00, 0x00, calldatasize()) // Copy the calldata for the delegatecall.
-            // If the preferred implementation is `address(0)`, perform the initialization workflow.
+            // If the preferred implementation is `address(0)`.
             if iszero(shl(96, impl)) {
-                if iszero(
-                    and( // The arguments of `and` are evaluated from right to left.
-                        delegatecall(
-                            gas(), mload(calldatasize()), 0x00, calldatasize(), calldatasize(), 0x00
-                        ),
-                        // Fetch the implementation from the proxy.
-                        staticcall(gas(), s, calldatasize(), 0x00, calldatasize(), 0x20)
-                    )
-                ) {
+                // If `defaultImplementation` is `address(0)`, perform the initialization workflow.
+                if iszero(defaultImplementation) {
+                    if iszero(
+                        and( // The arguments of `and` are evaluated from right to left.
+                            delegatecall(
+                                gas(), mload(calldatasize()), 0x00, calldatasize(), 0x00, 0x00
+                            ),
+                            // Fetch the implementation from the proxy.
+                            staticcall(gas(), s, calldatasize(), 0x00, calldatasize(), 0x20)
+                        )
+                    ) {
+                        returndatacopy(0x00, 0x00, returndatasize())
+                        revert(0x00, returndatasize())
+                    }
+                    // Because we cannot reliably and efficiently tell if the call is made
+                    // via staticcall or call, we shall ask the delegation to make a proxy delegation
+                    // initialization request to signal that we should initialize the storage slot with
+                    // the actual implementation. This also gives flexibility on whether to let the
+                    // proxy auto-upgrade, or let the authority manually upgrade (via 7702 or passkey).
+                    // A non-zero value in the transient storage denotes a initialization request.
+                    if tload(_EIP7702_PROXY_DELEGATION_INITIALIZATION_REQUEST_SLOT) {
+                        let implSlot := _ERC1967_IMPLEMENTATION_SLOT
+                        // The `implementation` is still at `calldatasize()` in memory.
+                        // Preserve the upper 96 bits when updating in case they are used for some stuff.
+                        sstore(
+                            implSlot, or(shl(160, shr(160, sload(implSlot))), mload(calldatasize()))
+                        )
+                        tstore(_EIP7702_PROXY_DELEGATION_INITIALIZATION_REQUEST_SLOT, 0) // Clear.
+                    }
                     returndatacopy(0x00, 0x00, returndatasize())
-                    revert(0x00, returndatasize())
+                    return(0x00, returndatasize())
                 }
-                // Because we cannot reliably and efficiently tell if the call is made
-                // via staticcall or call, we shall ask the delegation to make a proxy delegation
-                // initialization request to signal that we should initialize the storage slot with
-                // the actual implementation. This also gives flexibility on whether to let the
-                // proxy auto-upgrade, or let the authority manually upgrade (via 7702 or passkey).
-                // A non-zero value in the transient storage denotes a initialization request.
-                if tload(_EIP7702_PROXY_DELEGATION_INITIALIZATION_REQUEST_SLOT) {
-                    let implSlot := _ERC1967_IMPLEMENTATION_SLOT
-                    // The `implementation` is still at `calldatasize()` in memory.
-                    // Preserve the upper 96 bits when updating in case they are used for some stuff.
-                    sstore(implSlot, or(shl(160, shr(160, sload(implSlot))), mload(calldatasize())))
-                    tstore(_EIP7702_PROXY_DELEGATION_INITIALIZATION_REQUEST_SLOT, 0) // Clear.
-                }
-                returndatacopy(0x00, 0x00, returndatasize())
-                return(0x00, returndatasize())
+                impl := defaultImplementation
             }
             // Otherwise, just delegatecall and bubble up the results without initialization.
             if iszero(delegatecall(gas(), impl, 0x00, calldatasize(), calldatasize(), 0x00)) {

test/EIP7702Proxy.t.sol

@@ -24,6 +24,10 @@ contract Implementation2 {
         value = value_;
         LibEIP7702.requestProxyDelegationInitialization();
     }
+
+    function upgradeProxyDelegation(address newImplementation) public {
+        LibEIP7702.upgradeProxyDelegation(newImplementation);
+    }
 }
 
 contract EIP7702ProxyTest is SoladyTest {
@@ -59,7 +63,7 @@ contract EIP7702ProxyTest is SoladyTest {
 
         assertEq(EIP7702ProxyTest(instance).version(), 1);
 
-        uint256 v = _random();
+        uint256 v = _randomUniform();
         uint256 thisValue = this.value();
         if (thisValue == v) {
             v ^= 1;
@@ -72,6 +76,10 @@ contract EIP7702ProxyTest is SoladyTest {
         EIP7702ProxyTest(instance).revertWithError();
     }
 
+    function upgradeProxyDelegation(address newImplementation) public {
+        LibEIP7702.upgradeProxyDelegation(newImplementation);
+    }
+
     function testEIP7702Proxy(bytes32, bool f) public {
         vm.pauseGasMetering();
 
@@ -121,6 +129,8 @@ contract EIP7702ProxyTest is SoladyTest {
             vm.stopPrank();
         }
 
+        // Generate some random value that has the lower 160 bits zeroized,
+        // to test if the proxy can handle dirty bits.
         uint256 r = (_random() >> 160) << 160;
         vm.store(address(this), _ERC1967_IMPLEMENTATION_SLOT, bytes32(r));
 
@@ -142,20 +152,20 @@ contract EIP7702ProxyTest is SoladyTest {
         emit LogAddress("proxy", address(eip7702Proxy));
         emit LogAddress("address(this)", address(this));
 
-        vm.resumeGasMetering();
-
         // Runtime REVM detection.
         // If this check fails, then we are not ready to test it in CI.
         // The exact length is 23 at the time of writing as of the EIP7702 spec,
         // but we give our heuristic some leeway.
         if (authority.code.length > 0x20) return;
 
-        if (!f) assertEq(LibEIP7702.delegation(authority), address(eip7702Proxy));
+        vm.resumeGasMetering();
 
         _checkBehavesLikeProxy(authority);
 
         vm.pauseGasMetering();
 
+        if (!f) assertEq(LibEIP7702.delegation(authority), address(eip7702Proxy));
+
         // Check that upgrading the proxy won't cause the authority's implementation to change.
         if (!f && _randomChance(2)) {
             vm.startPrank(admin);
@@ -164,6 +174,12 @@ contract EIP7702ProxyTest is SoladyTest {
 
         _checkBehavesLikeProxy(authority);
 
+        if (!f && _randomChance(2)) {
+            EIP7702ProxyTest(authority).upgradeProxyDelegation(address(new Implementation2()));
+            assertEq(Implementation2(authority).version(), 2);
+            Implementation2(authority).upgradeProxyDelegation(address(this));
+        }
+
         if (!f && _randomChance(2) && (r >> 160) > 0) {
             vm.startPrank(admin);
             eip7702Proxy.upgrade(address(new Implementation2()));
@@ -187,6 +203,47 @@ contract EIP7702ProxyTest is SoladyTest {
     }
 
     function testEIP7702Proxy() public {
-        this.testEIP7702Proxy(0, true);
+        testEIP7702Proxy(0, true);
+    }
+
+    function testEIP7702ProxyWithDefaultImplementation(bytes32, bool f) public {
+        vm.pauseGasMetering();
+
+        IEIP7702ProxyWithAdminABI eip7702Proxy =
+            IEIP7702ProxyWithAdminABI(address(new EIP7702Proxy(address(this), address(0))));
+
+        assertEq(eip7702Proxy.admin(), address(0));
+        assertEq(LibEIP7702.proxyAdmin(address(eip7702Proxy)), address(0));
+        assertEq(eip7702Proxy.implementation(), address(this));
+        assertEq(LibEIP7702.proxyImplementation(address(eip7702Proxy)), address(this));
+
+        address authority = _randomUniqueHashedAddress();
+        assertEq(LibEIP7702.delegation(authority), address(0));
+        vm.etch(authority, abi.encodePacked(hex"ef0100", address(eip7702Proxy)));
+
+        // Generate some random value that has the lower 160 bits zeroized,
+        // to test if the proxy can handle dirty bits.
+        uint256 r = (_random() >> 160) << 160;
+        vm.store(authority, _ERC1967_IMPLEMENTATION_SLOT, bytes32(r));
+
+        if (authority.code.length > 0x20) return;
+
+        vm.resumeGasMetering();
+
+        _checkBehavesLikeProxy(authority);
+
+        vm.pauseGasMetering();
+
+        if (!f && _randomChance(2)) {
+            EIP7702ProxyTest(authority).upgradeProxyDelegation(address(new Implementation2()));
+            assertEq(Implementation2(authority).version(), 2);
+            Implementation2(authority).upgradeProxyDelegation(address(this));
+        }
+
+        vm.resumeGasMetering();
+    }
+
+    function testEIP7702ProxyWithDefaultImplementation() public {
+        testEIP7702ProxyWithDefaultImplementation(0, true);
     }
 }