Atomically set ERC-1967 and enforce implementation invariants

src/EIP7702Proxy.sol

@@ -4,42 +4,44 @@ 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 {Receiver} from "solady/accounts/Receiver.sol";
 
 import {NonceTracker} from "./NonceTracker.sol";
+import {IWalletValidator} from "./interfaces/IWalletValidator.sol";
 
 /// @title EIP7702Proxy
 ///
 /// @notice Proxy contract designed for EIP-7702 smart accounts
 ///
-/// @dev Implements ERC-1967 with an initial implementation address and guarded initializer function
+/// @dev Implements ERC-1967 with a guarded initializer function
 ///
 /// @author Coinbase (https://github.com/base/eip-7702-proxy)
-contract EIP7702Proxy is Proxy {
+contract EIP7702Proxy is Proxy, Receiver {
     /// @notice ERC-1271 interface constants
     bytes4 internal constant _ERC1271_MAGIC_VALUE = 0x1626ba7e;
     bytes4 internal constant _ERC1271_FAIL_VALUE = 0xffffffff;
 
     /// @notice Typehash for initialization signatures
-    bytes32 internal constant _INITIALIZATION_TYPEHASH =
-        keccak256(
-            "EIP7702ProxyInitialization(uint256 chainId,address proxy,uint256 nonce,bytes args)"
-        );
+    // bytes32 internal constant _INITIALIZATION_TYPEHASH =
+    //     keccak256(
+    //         "EIP7702ProxyInitialization(uint256 chainId,address proxy,uint256 nonce,bytes args)"
+    //     );
 
     /// @notice Typehash for resetting implementation, including chainId and current implementation
     bytes32 internal constant _IMPLEMENTATION_RESET_TYPEHASH =
         keccak256(
-            "EIP7702ProxyImplementationReset(uint256 chainId,address proxy,uint256 nonce,address currentImplementation,address newImplementation)"
+            "EIP7702ProxyImplementationReset(uint256 chainId,address proxy,uint256 nonce,address currentImplementation,address newImplementation,bytes32 initData)"
         );
 
-    /// @notice Initial implementation address set during construction
-    address public immutable INITIAL_IMPLEMENTATION;
-
     /// @notice Function selector on the implementation that is guarded from direct calls
     bytes4 public immutable GUARDED_INITIALIZER;
 
     /// @notice Address of the global nonce tracker for initialization
     NonceTracker public immutable NONCE_TRACKER;
 
+    /// @notice The validator contract for checking wallet-specific invariants
+    IWalletValidator public immutable VALIDATOR;
+
     /// @notice Address of this proxy contract delegate
     address internal immutable _PROXY;
 
@@ -52,26 +54,29 @@ contract EIP7702Proxy is Proxy {
     /// @notice Call to `GUARDED_INTIALIZER` attempted
     error InvalidInitializer();
 
-    /// @notice Initializes the proxy with an initial implementation and guarded initializer
+    /// @notice Proxy is not initialized
+    error Uninitialized();
+
+    /// @notice Initializes the proxy with a and guarded initializer
     ///
-    /// @param implementation The initial implementation address
     /// @param initializer The selector of the initializer function on `implementation` to guard
     /// @param nonceTracker The address of the nonce tracker contract
+    /// @param validator The address of the validator contract
     constructor(
-        address implementation,
         bytes4 initializer,
-        NonceTracker nonceTracker
+        NonceTracker nonceTracker,
+        IWalletValidator validator
     ) {
-        if (implementation == address(0))
-            revert ZeroValueConstructorArguments();
         if (initializer == bytes4(0)) revert ZeroValueConstructorArguments();
         if (address(nonceTracker) == address(0))
             revert ZeroValueConstructorArguments();
+        if (address(validator) == address(0))
+            revert ZeroValueConstructorArguments();
 
-        INITIAL_IMPLEMENTATION = implementation;
         GUARDED_INITIALIZER = initializer;
         NONCE_TRACKER = nonceTracker;
         _PROXY = address(this);
+        VALIDATOR = validator;
     }
 
     /// @notice Allow the account to receive ETH under any circumstances
@@ -84,43 +89,43 @@ contract EIP7702Proxy is Proxy {
     /// @param args The initialization arguments for the implementation
     /// @param signature The signature authorizing initialization
     /// @param crossChainReplayable use a chain-agnostic or chain-specific hash
-    function initialize(
-        bytes calldata args,
-        bytes calldata signature,
-        bool crossChainReplayable
-    ) external {
-        // Construct hash using typehash to prevent signature collisions
-        bytes32 initHash = keccak256(
-            abi.encode(
-                _INITIALIZATION_TYPEHASH,
-                crossChainReplayable ? 0 : block.chainid,
-                _PROXY,
-                NONCE_TRACKER.useNonce(),
-                keccak256(args)
-            )
-        );
-
-        // Verify signature is from this address (the EOA)
-        address signer = ECDSA.recover(initHash, signature);
-        if (signer != address(this)) revert InvalidSignature();
-
-        // Initialize the implementation
-        ERC1967Utils.upgradeToAndCall(
-            INITIAL_IMPLEMENTATION,
-            abi.encodePacked(GUARDED_INITIALIZER, args)
-        );
-    }
-
-    /// @notice Resets the ERC-1967 implementation slot after signature verification, allowing the account to
-    ///         correct the implementation address if it's ever changed by an unknown delegate or implementation.
-    ///
-    /// @dev Signature must be from the EOA's address that is 7702-delegating to this proxy
-    ///
+    // function initialize(
+    //     bytes calldata args,
+    //     bytes calldata signature,
+    //     bool crossChainReplayable
+    // ) external {
+    //     // Construct hash using typehash to prevent signature collisions
+    //     bytes32 initHash = keccak256(
+    //         abi.encode(
+    //             _INITIALIZATION_TYPEHASH,
+    //             crossChainReplayable ? 0 : block.chainid,
+    //             _PROXY,
+    //             NONCE_TRACKER.useNonce(),
+    //             keccak256(args)
+    //         )
+    //     );
+
+    //     // Verify signature is from this address (the EOA)
+    //     address signer = ECDSA.recover(initHash, signature);
+    //     if (signer != address(this)) revert InvalidSignature();
+
+    //     // Initialize the implementation
+    //     ERC1967Utils.upgradeToAndCall(
+    //         INITIAL_IMPLEMENTATION,
+    //         abi.encodePacked(GUARDED_INITIALIZER, args)
+    //     );
+    // }
+
+    /// @notice Resets the ERC-1967 implementation slot after signature verification and optionally executes calldata on the new implementation.
+    /// @dev Validates resulting wallet state after upgrade by calling `validateWallet` on the validator contract
+    /// @dev Signature must be from the EOA's address
     /// @param newImplementation The implementation address to set
+    /// @param initData Optional calldata to call on new implementation
     /// @param signature The EOA signature authorizing this change
     /// @param crossChainReplayable use a chain-agnostic or chain-specific hash
     function resetImplementation(
         address newImplementation,
+        bytes calldata initData,
         bytes calldata signature,
         bool crossChainReplayable
     ) external {
@@ -132,16 +137,20 @@ contract EIP7702Proxy is Proxy {
                 _PROXY,
                 NONCE_TRACKER.useNonce(),
                 ERC1967Utils.getImplementation(),
-                newImplementation
+                newImplementation,
+                keccak256(initData)
             )
         );
 
         // Verify signature is from this address (the EOA)
         address signer = ECDSA.recover(resetHash, signature);
         if (signer != address(this)) revert InvalidSignature();
 
-        // Reset the implementation slot
-        ERC1967Utils.upgradeToAndCall(newImplementation, "");
+        // Reset the implementation slot and call initialization if provided
+        ERC1967Utils.upgradeToAndCall(newImplementation, initData);
+
+        // Validate wallet state after upgrade, reverting if invalid
+        VALIDATOR.validateWallet(address(this));
     }
 
     /// @notice Handles ERC-1271 signature validation by enforcing a final `ecrecover` check if signatures fail `isValidSignature` check
@@ -186,20 +195,18 @@ contract EIP7702Proxy is Proxy {
         return _ERC1271_FAIL_VALUE;
     }
 
-    /// @notice Returns the implementation address, falling back to the initial implementation if the ERC-1967 implementation slot is not set
+    /// @notice Returns the ERC-1967 implementation address
     ///
     /// @return implementation The implementation address for this EOA
     function _implementation() internal view override returns (address) {
-        address implementation = ERC1967Utils.getImplementation();
-        return
-            implementation == address(0)
-                ? INITIAL_IMPLEMENTATION
-                : implementation;
+        return ERC1967Utils.getImplementation();
     }
 
     /// @inheritdoc Proxy
-    /// @dev Guards a specified initializer function from being called directly
+    /// @dev Guards a specified initializer function from being called directly. Reverts if the proxy is not initialized.
     function _fallback() internal override {
+        if (_implementation() == address(0)) revert Uninitialized();
+
         // block guarded initializer from being called
         if (msg.sig == GUARDED_INITIALIZER) revert InvalidInitializer();
 

src/interfaces/IWalletValidator.sol

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.23;
+
+/// @title IWalletValidator
+/// @notice Interface for wallet-specific validation logic
+interface IWalletValidator {
+    /// @notice Validates that a wallet is in a valid state
+    /// @param wallet The address of the wallet to validate
+    /// @dev Should revert if wallet state is invalid
+    function validateWallet(address wallet) external view;
+}

src/validators/CoinbaseSmartWalletValidator.sol

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.23;
+
+import {IWalletValidator} from "../interfaces/IWalletValidator.sol";
+import {MultiOwnable} from "smart-wallet/MultiOwnable.sol";
+
+/// @title CoinbaseSmartWalletValidator
+/// @notice Validates Coinbase Smart Wallet specific invariants
+contract CoinbaseSmartWalletValidator is IWalletValidator {
+    /// @notice Error thrown when a wallet has no owners
+    error Unintialized();
+
+    /// @notice Validates that a Coinbase Smart Wallet has at least one owner
+    /// @param wallet The address of the wallet to validate
+    function validateWallet(address wallet) external view override {
+        // Cast to MultiOwnable to check owner count
+        MultiOwnable walletContract = MultiOwnable(wallet);
+
+        // Ensure at least one owner exists
+        if (walletContract.ownerCount() == 0) {
+            revert Unintialized();
+        }
+    }
+}