have to track initialization separately

Author: amiecorso

src/EIP7702Proxy.sol

@@ -5,6 +5,7 @@ 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 +25,10 @@ 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
+    bytes32 internal constant INITIALIZED_SLOT =
+        bytes32(uint256(keccak256("EIP7702Proxy.initialized")) - 1);
+
     /// @notice Emitted when the implementation is upgraded
     event Upgraded(address indexed implementation);
 
@@ -54,7 +59,7 @@ contract EIP7702Proxy is Proxy {
 
     /// @dev Checks if proxy has been initialized by comparing implementation slot
     function _isInitialized() internal view returns (bool) {
-        return _implementation() == initialImplementation;
+        return StorageSlot.getBooleanSlot(INITIALIZED_SLOT).value;
     }
 
     /// @notice Initializes the proxy and implementation with a signed payload
@@ -72,7 +77,10 @@ contract EIP7702Proxy is Proxy {
         address recovered = ECDSA.recover(hash, signature);
         if (recovered != address(this)) revert InvalidSignature();
 
-        // Set the ERC-1967 implementation slot, emit Upgraded event, call the initializer on the initial implementation
+        // Set initialized flag before upgrading
+        StorageSlot.getBooleanSlot(INITIALIZED_SLOT).value = true;
+
+        // Set the ERC-1967 implementation slot, emit Upgraded event, call the initializer
         ERC1967Utils.upgradeToAndCall(
             initialImplementation,
             abi.encodePacked(guardedInitializer, args)
@@ -91,6 +99,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

test/EIP7702Proxy/delegate.t.sol

@@ -93,4 +93,35 @@ contract DelegateTest is EIP7702ProxyBase {
         (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
+
+        // 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 - this should work but might fail
+        vm.prank(_newOwner);
+        MockImplementation(_eoa).mockFunction();
+
+        // Verify the call succeeded
+        assertTrue(
+            MockImplementation(_eoa).mockFunctionCalled(),
+            "Should be able to call through proxy after upgrade"
+        );
+    }
 }

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);
+    }
 }
 
 /**