Merge pull request #15 from base-org/amie/block-delegatecalls-until-initialized

amiecorso · web-flow · commit dd72de1e1d67 · 2025-02-02T17:12:21.000-08:00
Prevent all delegated calls unless proxy has been initialized
diff --git a/src/EIP7702Proxy.sol b/src/EIP7702Proxy.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.23;
 import {Proxy} from "openzeppelin-contracts/contracts/proxy/Proxy.sol";
 import {ERC1967Utils} from "openzeppelin-contracts/contracts/proxy/ERC1967/ERC1967Utils.sol";
 import {ECDSA} from "openzeppelin-contracts/contracts/utils/cryptography/ECDSA.sol";
-import {Address} from "openzeppelin-contracts/contracts/utils/Address.sol";
+import {StorageSlot} from "openzeppelin-contracts/contracts/utils/StorageSlot.sol";
 
 /// @title EIP7702Proxy
 /// @notice Proxy contract designed for EIP-7702 smart accounts
@@ -24,6 +24,12 @@ contract EIP7702Proxy is Proxy {
     /// @notice Function selector on the implementation that is guarded from direct calls
     bytes4 immutable guardedInitializer;
 
+    /// @dev Storage slot with the initialized flag, conforms to ERC-7201
+    bytes32 internal constant INITIALIZED_SLOT =
+        keccak256(
+            abi.encode(uint256(keccak256("EIP7702Proxy.initialized")) - 1)
+        ) & ~bytes32(uint256(0xff));
+
     /// @notice Emitted when the implementation is upgraded
     event Upgraded(address indexed implementation);
 
@@ -33,8 +39,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 +58,11 @@ contract EIP7702Proxy is Proxy {
         guardedInitializer = initializer;
     }
 
+    /// @dev Checks if proxy has been initialized by checking the initialized flag
+    function _isInitialized() internal view returns (bool) {
+        return StorageSlot.getBooleanSlot(INITIALIZED_SLOT).value;
+    }
+
     /// @notice Initializes the proxy and implementation with a signed payload
     ///
     /// @dev Signature must be from this contract's address
@@ -62,17 +73,20 @@ contract EIP7702Proxy is Proxy {
         bytes calldata args,
         bytes calldata signature
     ) external {
-        // construct hash incompatible with wallet RPCs to avoid phishing
+        // Construct hash without Ethereum signed message prefix to prevent phishing via standard wallet signing.
+        // Since this proxy is designed for EIP-7702 (where the proxy address is an EOA),
+        // using a raw hash ensures that initialization signatures cannot be obtained through normal
+        // wallet "Sign Message" prompts. This prevents malicious dapps from tricking users into
+        // initializing their account via standard wallet signing flows.
+        // Wallets must implement custom signing logic at a lower level to support initialization.
         bytes32 hash = keccak256(abi.encode(proxy, args));
         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 initialized flag before upgrading
+        StorageSlot.getBooleanSlot(INITIALIZED_SLOT).value = true;
 
-        // Set the ERC-1967 implementation slot, emit Upgraded event, call the initializer on the initial implementation
+        // Set the ERC-1967 implementation slot, emit Upgraded event, call the initializer
         ERC1967Utils.upgradeToAndCall(
             initialImplementation,
             abi.encodePacked(guardedInitializer, args)
@@ -91,6 +105,9 @@ contract EIP7702Proxy is Proxy {
         bytes32 hash,
         bytes calldata signature
     ) external returns (bytes4) {
+        // Check initialization status first
+        if (!_isInitialized()) revert ProxyNotInitialized();
+
         // First try delegatecall to implementation
         (bool success, bytes memory result) = _implementation().delegatecall(
             msg.data
@@ -117,22 +134,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();
+    }
 }
diff --git a/test/EIP7702Proxy/coinbaseImplementation.t.sol b/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());
 
diff --git a/test/EIP7702Proxy/delegate.t.sol b/test/EIP7702Proxy/delegate.t.sol
@@ -68,4 +68,60 @@ 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");
+    }
+
+    function test_continues_delegating_afterUpgrade() public {
+        // Setup will have already initialized the proxy with initial implementation and an owner
+
+        // Deploy a new implementation
+        MockImplementation newImplementation = new MockImplementation();
+
+        // Upgrade to the new implementation
+        vm.prank(_newOwner);
+        MockImplementation(_eoa).upgradeToAndCall(
+            address(newImplementation),
+            ""
+        );
+
+        // Verify the implementation was changed
+        assertEq(
+            _getERC1967Implementation(_eoa),
+            address(newImplementation),
+            "Implementation should be updated"
+        );
+
+        // Try to make a call through the proxy
+        vm.prank(_newOwner);
+        MockImplementation(_eoa).mockFunction();
+
+        // Verify the call succeeded
+        assertTrue(
+            MockImplementation(_eoa).mockFunctionCalled(),
+            "Should be able to call through proxy after upgrade"
+        );
+    }
 }
diff --git a/test/EIP7702Proxy/initialize.t.sol b/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"
+        );
+    }
 }
diff --git a/test/EIP7702Proxy/isValidSignature.t.sol b/test/EIP7702Proxy/isValidSignature.t.sol
@@ -114,6 +114,19 @@ contract FailingImplementationTest is IsValidSignatureTestBase {
         );
         assertEq(result, ERC1271_FAIL_VALUE, "Should reject empty signature");
     }
+
+    function test_reverts_whenCalledBeforeInitialization() public {
+        // Deploy a fresh proxy without initializing
+        address payable uninitProxy = payable(makeAddr("uninitProxy"));
+        _deployProxy(uninitProxy);
+
+        // Try to call isValidSignature
+        bytes32 hash = keccak256("test message");
+        bytes memory signature = new bytes(65);
+
+        vm.expectRevert(EIP7702Proxy.ProxyNotInitialized.selector);
+        EIP7702Proxy(uninitProxy).isValidSignature(hash, signature);
+    }
 }
 
 /**
