prevent all delegate calls unless implementation has been set to proxy's expected

Author: amiecorso

src/EIP7702Proxy.sol

@@ -33,8 +33,8 @@ contract EIP7702Proxy is Proxy {
     /// @notice Emitted when the `guardedInitializer` is called
     error InvalidInitializer();
 
-    /// @notice Emitted when initialization is attempted on a non-initial implementation
-    error InvalidImplementation();
+    /// @notice Emitted when trying to delegate before initialization
+    error ProxyNotInitialized();
 
     /// @notice Emitted when constructor arguments are zero
     error ZeroValueConstructorArguments();
@@ -52,6 +52,11 @@ contract EIP7702Proxy is Proxy {
         guardedInitializer = initializer;
     }
 
+    /// @dev Checks if proxy has been initialized by comparing implementation slot
+    function _isInitialized() internal view returns (bool) {
+        return _implementation() == initialImplementation;
+    }
+
     /// @notice Initializes the proxy and implementation with a signed payload
     ///
     /// @dev Signature must be from this contract's address
@@ -67,11 +72,6 @@ contract EIP7702Proxy is Proxy {
         address recovered = ECDSA.recover(hash, signature);
         if (recovered != address(this)) revert InvalidSignature();
 
-        // enforce initialization only on initial implementation
-        address implementation = _implementation();
-        if (implementation != initialImplementation)
-            revert InvalidImplementation();
-
         // Set the ERC-1967 implementation slot, emit Upgraded event, call the initializer on the initial implementation
         ERC1967Utils.upgradeToAndCall(
             initialImplementation,
@@ -117,22 +117,19 @@ contract EIP7702Proxy is Proxy {
         return ERC1271_FAIL_VALUE;
     }
 
-    /// @inheritdoc Proxy
-    function _implementation() internal view override returns (address) {
-        address implementation = ERC1967Utils.getImplementation();
-        return
-            implementation != address(0)
-                ? implementation
-                : initialImplementation;
-    }
-
     /// @inheritdoc Proxy
     /// @dev Handles ERC-1271 signature validation by enforcing an ecrecover check if signatures fail `isValidSignature` check
     /// @dev Guards a specified initializer function from being called directly
     function _fallback() internal override {
+        if (!_isInitialized()) revert ProxyNotInitialized();
+
         // block guarded initializer from being called
         if (msg.sig == guardedInitializer) revert InvalidInitializer();
 
         _delegate(_implementation());
     }
+
+    function _implementation() internal view override returns (address) {
+        return ERC1967Utils.getImplementation();
+    }
 }

test/EIP7702Proxy/coinbaseImplementation.t.sol

@@ -167,6 +167,7 @@ contract CoinbaseImplementationTest is Test {
         uint256 amount
     ) public {
         vm.assume(recipient != address(0));
+        vm.assume(recipient != address(_eoa));
         assumeNotPrecompile(recipient);
         assumePayable(recipient);
         vm.assume(amount > 0 && amount <= 100 ether);
@@ -193,6 +194,7 @@ contract CoinbaseImplementationTest is Test {
     ) public {
         vm.assume(nonOwner != address(0));
         vm.assume(nonOwner != _newOwner); // Ensure caller isn't the actual owner
+        vm.assume(nonOwner != _eoa); // Ensure caller isn't the EOA address
 
         address newImpl = address(new CoinbaseSmartWallet());
 

test/EIP7702Proxy/delegate.t.sol

@@ -68,4 +68,29 @@ contract DelegateTest is EIP7702ProxyBase {
             "State should not change when write fails"
         );
     }
+
+    function test_reverts_whenCallingBeforeInitialization() public {
+        // Deploy a fresh proxy without initializing it
+        address payable uninitProxy = payable(makeAddr("uninitProxy"));
+        _deployProxy(uninitProxy);
+
+        vm.expectRevert(EIP7702Proxy.ProxyNotInitialized.selector);
+        MockImplementation(payable(uninitProxy)).owner();
+    }
+
+    function test_reverts_whenCallingWithArbitraryDataBeforeInitialization(
+        bytes memory arbitraryCalldata
+    ) public {
+        // Deploy a fresh proxy without initializing it
+        address payable uninitProxy = payable(makeAddr("uninitProxy"));
+        _deployProxy(uninitProxy);
+
+        // Test that it reverts with the correct error
+        vm.expectRevert(EIP7702Proxy.ProxyNotInitialized.selector);
+        address(uninitProxy).call(arbitraryCalldata);
+
+        // Also verify the low-level call fails
+        (bool success, ) = address(uninitProxy).call(arbitraryCalldata);
+        assertFalse(success, "Low-level call should fail");
+    }
 }

test/EIP7702Proxy/initialize.t.sol

@@ -5,6 +5,7 @@ import {EIP7702ProxyBase} from "../base/EIP7702ProxyBase.sol";
 import {EIP7702Proxy} from "../../src/EIP7702Proxy.sol";
 import {MockImplementation, RevertingInitializerMockImplementation} from "../mocks/MockImplementation.sol";
 import {ECDSA} from "openzeppelin-contracts/contracts/utils/cryptography/ECDSA.sol";
+import {ERC1967Utils} from "openzeppelin-contracts/contracts/proxy/ERC1967/ERC1967Utils.sol";
 
 contract InitializeTest is EIP7702ProxyBase {
     function test_succeeds_withValidSignatureAndArgs(address newOwner) public {
@@ -157,4 +158,70 @@ contract InitializeTest is EIP7702ProxyBase {
         vm.expectRevert(EIP7702Proxy.ZeroValueConstructorArguments.selector);
         new EIP7702Proxy(address(_implementation), bytes4(0));
     }
+
+    function test_succeeds_whenImplementationSlotAlreadySetToDifferentAddress(
+        address mockPreviousImpl,
+        address newOwner,
+        uint128 uninitProxyPk
+    ) public {
+        vm.assume(mockPreviousImpl != address(0));
+        vm.assume(mockPreviousImpl != address(_implementation));
+        vm.assume(mockPreviousImpl != address(_eoa));
+        vm.assume(newOwner != address(0));
+        vm.assume(newOwner != mockPreviousImpl);
+        vm.assume(newOwner != _eoa);
+        assumeNotPrecompile(mockPreviousImpl);
+        assumeNotPrecompile(newOwner);
+        vm.assume(uninitProxyPk != 0);
+        vm.assume(uninitProxyPk != _EOA_PRIVATE_KEY);
+
+        // Derive address for uninitProxy from private key
+        address payable uninitProxy = payable(vm.addr(uninitProxyPk));
+
+        // Deploy proxy template and etch its code at the target address
+        EIP7702Proxy proxyTemplate = new EIP7702Proxy(
+            address(_implementation),
+            _initSelector
+        );
+        bytes memory proxyCode = address(proxyTemplate).code;
+        vm.etch(uninitProxy, proxyCode);
+
+        // Set the implementation slot to some other address, simulating a previous implementation
+        vm.store(
+            uninitProxy,
+            ERC1967Utils.IMPLEMENTATION_SLOT,
+            bytes32(uint256(uint160(mockPreviousImpl)))
+        );
+
+        // Verify implementation slot is set to the previous implementation
+        assertEq(
+            _getERC1967Implementation(uninitProxy),
+            mockPreviousImpl,
+            "Implementation slot should be set to previous implementation"
+        );
+
+        // Initialize the proxy
+        bytes memory initArgs = _createInitArgs(_newOwner);
+        bytes32 initHash = keccak256(
+            abi.encode(address(proxyTemplate), initArgs)
+        );
+        (uint8 v, bytes32 r, bytes32 s) = vm.sign(uninitProxyPk, initHash);
+        bytes memory signature = abi.encodePacked(r, s, v);
+
+        EIP7702Proxy(uninitProxy).initialize(initArgs, signature);
+
+        // Verify implementation slot was changed to the correct implementation
+        assertEq(
+            _getERC1967Implementation(uninitProxy),
+            address(_implementation),
+            "Implementation slot should be set to correct implementation"
+        );
+
+        // Verify we can make calls through the proxy now
+        assertEq(
+            MockImplementation(payable(uninitProxy)).owner(),
+            _newOwner,
+            "Should be able to call through proxy after initialization"
+        );
+    }
 }