Basic expander

contracts/contracts/AggregateSigValidator.sol

@@ -0,0 +1,136 @@
+//SPDX-License-Identifier: Unlicense
+pragma solidity >=0.8.4 <0.9.0;
+pragma abicoder v2;
+
+import "./lib/IBLS.sol"; // to use a deployed BLS library
+
+import "@account-abstraction/contracts/interfaces/UserOperation.sol";
+import "@account-abstraction/contracts/bls/BLSHelper.sol";
+
+import "./BLSWallet.sol";
+
+/**
+ * The 4337 Aggregator.
+ * 
+ * Validates aggregate signatures and provides functions useful for off-chain calculations.
+ */
+contract AggregateSigValidator
+{
+    /** Domain chosen arbitrarily */
+    bytes32 BLS_DOMAIN = keccak256(abi.encodePacked(uint32(0xfeedbee5)));
+    uint8 constant BLS_KEY_LEN = 4;
+
+    IBLS public immutable bls;
+
+    /**
+     * @param _bls verified bls library contract address
+     */
+    constructor(IBLS _bls) {
+        bls = _bls;
+    }
+
+    /** Throw if bundle not valid or signature verification fails */
+    function validateSignatures(
+        UserOperation[] calldata userOps,
+        bytes calldata signature
+    ) view external {
+        uint256[2][] memory messages = new uint256[2][](userOps.length);
+        uint256[BLS_KEY_LEN][] memory senderPublicKeys = new uint256[BLS_KEY_LEN][](userOps.length);
+
+        for (uint256 i = 0; i < userOps.length; i++) {
+            messages[i] = bls.hashToPoint(
+                BLS_DOMAIN,
+                abi.encodePacked(getRequestId(userOps[i]))
+            );
+
+            senderPublicKeys[i] = BLSWallet(payable(userOps[i].sender)).getBlsKey();
+        }
+
+        bool verified = bls.verifyMultiple(
+            abi.decode(signature, (uint256[2])),
+            senderPublicKeys,
+            messages
+        );
+
+        require(verified, "VG: Sig not verified");
+    }
+
+    // TODO
+    // function validateUserOpSignature(
+    //     UserOperation4337 calldata userOp,
+    //     bool offChainSigCheck
+    // ) external view returns (
+    //     bytes memory sigForUserOp,
+    //     bytes memory sigForAggregation,
+    //     bytes memory offChainSigInfo
+    // ) {}
+
+    //copied from BLS.sol
+    uint256 public constant N = 21888242871839275222246405745257275088696311157297823662689037894645226208583;
+
+    function aggregateSignatures(
+        bytes[] calldata signatures
+    ) external pure returns (bytes memory aggregateSignature) {
+        BLSHelper.XY[] memory points = new BLSHelper.XY[](signatures.length);
+        for (uint i = 0; i < points.length; i++) {
+            (uint x, uint y) = abi.decode(signatures[i], (uint, uint));
+            points[i] = BLSHelper.XY(x, y);
+        }
+        BLSHelper.XY memory sum = BLSHelper.sum(points, N);
+        return abi.encode(sum.x, sum.y);
+    }
+
+    /**
+     * get a hash of userOp
+     * NOTE: this hash is not the same as UserOperation.hash()
+     *  (slightly less efficient, since it uses memory userOp)
+     */
+    function getUserOpHash(UserOperation memory userOp) internal pure returns (bytes32) {
+        return keccak256(abi.encode(
+            userOp.sender,
+            userOp.nonce,
+            keccak256(userOp.initCode),
+            keccak256(userOp.callData),
+            userOp.callGasLimit,
+            userOp.verificationGasLimit,
+            userOp.preVerificationGas,
+            userOp.maxFeePerGas,
+            userOp.maxPriorityFeePerGas,
+            keccak256(userOp.paymasterAndData)
+        ));
+    }
+
+    /**
+     * return the BLS "message" for the given UserOp.
+     * the wallet should sign this value using its public-key
+     */
+    function userOpToMessage(UserOperation memory userOp) public view returns (uint256[2] memory) {
+        bytes32 hashPublicKey = _getUserOpPubkeyHash(userOp);
+        return _userOpToMessage(userOp, hashPublicKey);
+    }
+
+    function _userOpToMessage(UserOperation memory userOp, bytes32 publicKeyHash) internal view returns (uint256[2] memory) {
+        bytes32 requestId = _getRequestId(userOp, publicKeyHash);
+        return bls.hashToPoint(BLS_DOMAIN, abi.encodePacked(requestId));
+    }
+
+    //return the public-key hash of a userOp.
+    // if its a constructor UserOp, then return constructor hash.
+    function _getUserOpPubkeyHash(UserOperation memory userOp) internal view returns (bytes32 hashPublicKey) {
+        if (userOp.initCode.length == 0) {
+            uint256[4] memory publicKey = BLSWallet(payable(userOp.sender)).getBlsKey();
+            hashPublicKey = keccak256(abi.encode(publicKey));
+        } else {
+            hashPublicKey = keccak256(userOp.initCode);
+        }
+    }
+
+    function getRequestId(UserOperation memory userOp) public view returns (bytes32) {
+        bytes32 hashPublicKey = _getUserOpPubkeyHash(userOp);
+        return _getRequestId(userOp, hashPublicKey);
+    }
+
+    function _getRequestId(UserOperation memory userOp, bytes32 hashPublicKey) internal view returns (bytes32) {
+        return keccak256(abi.encode(getUserOpHash(userOp), hashPublicKey, address(this), block.chainid));
+    }
+}

contracts/contracts/BLSExpander.sol

@@ -16,23 +16,23 @@ contract BLSExpander {
         verificationGateway = VerificationGateway(gateway);
     }
 
-    // eg approve and transfers of a token contract
-    function blsCallMultiCheckRewardIncrease(
-        IERC20 tokenRewardAddress,
-        uint256 tokenRewardAmount,
-        VerificationGateway.Bundle calldata bundle
-        // uint256[4][] calldata publicKeys,
-        // uint256[2] memory signature,
-        // VerificationGateway.TxSet[] calldata txs
-    ) external returns (uint256 balanceIncrease) {
-        uint256 balanceBefore = tokenRewardAddress.balanceOf(tx.origin);
+    // // eg approve and transfers of a token contract
+    // function blsCallMultiCheckRewardIncrease(
+    //     IERC20 tokenRewardAddress,
+    //     uint256 tokenRewardAmount,
+    //     VerificationGateway.Bundle calldata bundle
+    //     // uint256[4][] calldata publicKeys,
+    //     // uint256[2] memory signature,
+    //     // VerificationGateway.TxSet[] calldata txs
+    // ) external returns (uint256 balanceIncrease) {
+    //     uint256 balanceBefore = tokenRewardAddress.balanceOf(tx.origin);
 
-        verificationGateway.processBundle(bundle);
+    //     verificationGateway.processBundle(bundle);
 
-        uint256 balanceAfter = tokenRewardAddress.balanceOf(tx.origin);
-        balanceIncrease = balanceAfter - balanceBefore;
-        require(balanceIncrease >= tokenRewardAmount, "BLSExpander: Insufficient reward");
-    }
+    //     uint256 balanceAfter = tokenRewardAddress.balanceOf(tx.origin);
+    //     balanceIncrease = balanceAfter - balanceBefore;
+    //     require(balanceIncrease >= tokenRewardAmount, "BLSExpander: Insufficient reward");
+    // }
 
 
     // eg approve and transfers of a token contract
@@ -89,34 +89,34 @@ contract BLSExpander {
     //     );
     // }
 
-    // eg airdrop
-    function blsCallMultiSameCallerContractFunction(
-        uint256[4] calldata publicKey,
-        uint256 nonce,
-        uint256[2] calldata signature,
-        address contractAddress,
-        bytes4 methodId,
-        bytes[] calldata encodedParamSets
-    ) external {
-        uint256 length = encodedParamSets.length;
+    // // eg airdrop
+    // function blsCallMultiSameCallerContractFunction(
+    //     uint256[4] calldata publicKey,
+    //     uint256 nonce,
+    //     uint256[2] calldata signature,
+    //     address contractAddress,
+    //     bytes4 methodId,
+    //     bytes[] calldata encodedParamSets
+    // ) external {
+    //     uint256 length = encodedParamSets.length;
 
-        VerificationGateway.Bundle memory bundle;
-        bundle.signature = signature;
+    //     VerificationGateway.Bundle memory bundle;
+    //     bundle.signature = signature;
 
-        bundle.senderPublicKeys = new uint256[4][](1);
-        bundle.senderPublicKeys[0] = publicKey;
+    //     bundle.senderPublicKeys = new uint256[4][](1);
+    //     bundle.senderPublicKeys[0] = publicKey;
 
-        bundle.operations = new IWallet.Operation[](1);
-        bundle.operations[0].nonce = nonce;
-        bundle.operations[0].actions = new IWallet.ActionData[](length);
-        for (uint256 i=0; i<length; i++) {
-            bundle.operations[0].actions[i].ethValue = 0;
-            bundle.operations[0].actions[i].contractAddress = contractAddress;
-            bundle.operations[0].actions[i].encodedFunction = abi.encodePacked(methodId, encodedParamSets[i]);
-        }
+    //     bundle.operations = new IWallet.Operation[](1);
+    //     bundle.operations[0].nonce = nonce;
+    //     bundle.operations[0].actions = new IWallet.ActionData[](length);
+    //     for (uint256 i=0; i<length; i++) {
+    //         bundle.operations[0].actions[i].ethValue = 0;
+    //         bundle.operations[0].actions[i].contractAddress = contractAddress;
+    //         bundle.operations[0].actions[i].encodedFunction = abi.encodePacked(methodId, encodedParamSets[i]);
+    //     }
 
-        verificationGateway.processBundle(bundle);
-    }
+    //     verificationGateway.processBundle(bundle);
+    // }
 
     // eg identical txs from multiple accounts
     // function blsCallMultiSameContractFunctionParams(

contracts/contracts/BLSWallet.sol

@@ -2,11 +2,12 @@
 pragma solidity >=0.8.4 <0.9.0;
 pragma abicoder v2;
 
-
 //To avoid constructor params having forbidden evm bytecodes on Optimism
 import "@openzeppelin/contracts/proxy/utils/Initializable.sol";
-import "./interfaces/IWallet.sol";
 
+import "@account-abstraction/contracts/interfaces/UserOperation.sol";
+
+import "./interfaces/IWallet.sol";
 
 /** Minimal upgradable smart contract wallet.
     Generic calls can only be requested by its trusted gateway.
@@ -14,6 +15,7 @@ import "./interfaces/IWallet.sol";
 contract BLSWallet is Initializable, IWallet
 {
     uint256 public nonce;
+    uint256[4] public blsKey;
     bytes32 public recoveryHash;
     bytes32 pendingRecoveryHash;
     uint256 pendingRecoveryHashTime;
@@ -22,7 +24,9 @@ contract BLSWallet is Initializable, IWallet
     uint256 pendingPAFunctionTime;
 
     // BLS variables
-    address public trustedBLSGateway;
+    address public blsGateway;
+    address public entryPoint;
+    address public aggregator;
     address pendingBLSGateway;
     uint256 pendingGatewayTime;
 
@@ -49,15 +53,29 @@ contract BLSWallet is Initializable, IWallet
     );
 
     function initialize(
-        address blsGateway
+        uint256[4] memory _blsKey,
+        address _blsGateway,
+        address _entryPoint,
+        address _aggregator
     ) external initializer {
         nonce = 0;
-        trustedBLSGateway = blsGateway;
+        blsKey = _blsKey;
+        blsGateway = _blsGateway;
+        entryPoint = _entryPoint;
+        aggregator = _aggregator;
     }
 
     receive() external payable {}
     fallback() external payable {}
 
+    function getBlsKey() external view returns (uint256[4] memory) {
+        return blsKey;
+    }
+
+    function setBlsKey(uint256[4] memory newBlsKey) external onlyGateway {
+        blsKey = newBlsKey;
+    }
+
     /**
     Wallet can update its recovery hash
      */
@@ -77,8 +95,8 @@ contract BLSWallet is Initializable, IWallet
     /**
     Wallet can migrate to a new gateway, eg additional signature support
      */
-    function setTrustedGateway(address blsGateway) public onlyTrustedGateway {
-        pendingBLSGateway = blsGateway;
+    function setTrustedGateway(address _blsGateway) public onlyGateway {
+        pendingBLSGateway = _blsGateway;
         pendingGatewayTime = block.timestamp + 604800; // 1 week from now
         emit PendingGatewaySet(pendingBLSGateway);
     }
@@ -86,7 +104,7 @@ contract BLSWallet is Initializable, IWallet
     /**
     Prepare wallet with desired implementation contract to upgrade to.
     */
-    function setProxyAdminFunctionHash(bytes32 encodedFunctionHash) public onlyTrustedGateway {
+    function setProxyAdminFunctionHash(bytes32 encodedFunctionHash) public onlyGateway {
         pendingPAFunctionHash = encodedFunctionHash;
         pendingPAFunctionTime = block.timestamp + 604800; // 1 week from now
         emit PendingProxyAdminFunctionHashSet(encodedFunctionHash);
@@ -107,11 +125,11 @@ contract BLSWallet is Initializable, IWallet
         if (pendingGatewayTime != 0 &&
             block.timestamp > pendingGatewayTime
         ) {
-            address previousGateway = trustedBLSGateway;
-            trustedBLSGateway = pendingBLSGateway;
+            address previousGateway = blsGateway;
+            blsGateway = pendingBLSGateway;
             pendingGatewayTime = 0;
             pendingBLSGateway = address(0);
-            emit GatewayUpdated(previousGateway, trustedBLSGateway);
+            emit GatewayUpdated(previousGateway, blsGateway);
         }
         if (
             pendingPAFunctionTime != 0 &&
@@ -129,7 +147,7 @@ contract BLSWallet is Initializable, IWallet
         pendingRecoveryHash = bytes32(0);
     }
 
-    function recover() public onlyTrustedGateway {
+    function recover() public onlyGateway {
         // clear any pending operations
         clearPendingRecoveryHash();
         pendingGatewayTime = 0;
@@ -144,7 +162,7 @@ contract BLSWallet is Initializable, IWallet
      */
     function performOperation(
         IWallet.Operation calldata op
-    ) public payable onlyTrustedGateway thisNonce(op.nonce) returns (
+    ) public payable onlyEntryPoint thisNonce(op.nonce) returns (
         bool success,
         bytes[] memory results
     ) {
@@ -210,7 +228,7 @@ contract BLSWallet is Initializable, IWallet
         return params;
     }
 
-    function clearApprovedProxyAdminFunctionHash() public onlyTrustedGateway {
+    function clearApprovedProxyAdminFunctionHash() public onlyGateway {
         approvedProxyAdminFunctionHash = 0;
     }
 
@@ -221,22 +239,44 @@ contract BLSWallet is Initializable, IWallet
         nonce++;
     }
 
+    function validateUserOp(
+        UserOperation calldata userOp,
+        bytes32 requestId,
+        address _aggregator,
+        uint256 missingWalletFunds
+    ) external view {
+        require(_aggregator == aggregator);
+        require(userOp.nonce == nonce);
+        require(missingWalletFunds == 0);
+    }
+
+    function getAggregator() public view returns (address) {
+        return aggregator;
+    }
+
     modifier onlyThis() {
         require(msg.sender == address(this), "BLSWallet: only callable from this");
          _;
     }
 
-    modifier onlyTrustedGateway() {
-        bool isTrustedGateway =
-            (msg.sender == trustedBLSGateway)
+    modifier onlyGateway() {
+        bool isGateway =
+            (msg.sender == blsGateway)
         ;
-        require(isTrustedGateway, "BLSWallet: only callable from trusted gateway");
+        require(isGateway, "BLSWallet: only callable from gateway");
          _;
     }
 
+    modifier onlyEntryPoint() {
+        require(
+            msg.sender == entryPoint,
+            "BLSWallet: only callable from 4337 entry point"
+        );
+        _;
+    }
+
     modifier thisNonce(uint256 opNonce) {
         require(opNonce == nonce, "BLSWallet: only callable with current nonce");
         _;
     }
-
 }

contracts/contracts/Expander.sol

@@ -0,0 +1,64 @@
+//SPDX-License-Identifier: Unlicense
+pragma solidity >=0.8.4 <0.9.0;
+
+import "@account-abstraction/contracts/interfaces/IAggregator.sol";
+import "@account-abstraction/contracts/interfaces/IEntryPoint.sol";
+import "@account-abstraction/contracts/interfaces/UserOperation.sol";
+// (See https://github.com/eth-infinitism/account-abstraction/tree/develop/contracts/interfaces)
+
+contract Expander {
+    uint constant ARBITRARY_GAS_LIMIT = 1000000000;
+
+    // (Needs real deployed address)
+    IEntryPoint constant ENTRY_POINT = IEntryPoint(
+        0x0001020304050607080910111213141516171819
+    );
+
+    // (Needs real deployed address)
+    IAggregator constant AGGREGATE_SIG_VALIDATOR = IAggregator(
+        0x0001020304050607080910111213141516171819
+    );
+
+    struct SmallUserOperation {
+        address sender;
+        uint256 nonce;
+        bytes callData;
+    }
+
+    function handleAggregatedOps(
+        SmallUserOperation[] calldata smallUserOps,
+        bytes calldata signature
+    ) public {
+        IEntryPoint.UserOpsPerAggregator[] memory opsPerAggregator =
+            new IEntryPoint.UserOpsPerAggregator[](1);
+
+        uint len = smallUserOps.length;
+
+        UserOperation[] memory userOps = new UserOperation[](len);
+
+        for (uint i = 0; i < len; i++) {
+            userOps[i].sender = smallUserOps[i].sender;
+            userOps[i].nonce = smallUserOps[i].nonce;
+            // Leaving .initCode as empty
+            userOps[i].callData = smallUserOps[i].callData;
+            userOps[i].callGasLimit = ARBITRARY_GAS_LIMIT;
+            userOps[i].verificationGasLimit = ARBITRARY_GAS_LIMIT;
+            // Leaving .preVerificationGas as zero
+            // Leaving .maxFeePerGas as zero
+            // Leaving .maxPriorityFeePerGas as zero
+            // Leaving .paymasterAndData as empty
+            // Leaving .signature as empty
+        }
+
+        opsPerAggregator[0] = IEntryPoint.UserOpsPerAggregator(
+            userOps,
+            AGGREGATE_SIG_VALIDATOR,
+            signature
+        );
+
+        ENTRY_POINT.handleAggregatedOps(
+            opsPerAggregator,
+            payable(msg.sender)
+        );
+    }
+}

contracts/contracts/VerificationGateway.sol

@@ -8,6 +8,7 @@ import "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol";
 import "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
 
 import "./interfaces/IWallet.sol";
+import "./BLSWallet.sol";
 
 /**
 A non-upgradable gateway used to create BLSWallets and call them with
@@ -25,21 +26,15 @@ contract VerificationGateway
     IBLS public immutable blsLib;
     ProxyAdmin public immutable walletProxyAdmin;
     address public immutable blsWalletLogic;
-    mapping(bytes32 => IWallet) public walletFromHash;
-    mapping(IWallet => bytes32) public hashFromWallet;
+    address public immutable entryPoint;
+    address public immutable aggregateSigValidator;
+    mapping(bytes32 => BLSWallet) public walletFromHash;
 
     //mapping from an existing wallet's bls key hash to pending variables when setting a new BLS key
     mapping(bytes32 => uint256[BLS_KEY_LEN]) public pendingBLSPublicKeyFromHash;
     mapping(bytes32 => uint256[2]) public pendingMessageSenderSignatureFromHash;
     mapping(bytes32 => uint256) public pendingBLSPublicKeyTimeFromHash;
 
-    /** Aggregated signature with corresponding senders + operations */
-    struct Bundle {
-        uint256[2] signature;
-        uint256[BLS_KEY_LEN][] senderPublicKeys;
-        IWallet.Operation[] operations;
-    }
-
     event WalletCreated(
         address indexed wallet,
         uint256[BLS_KEY_LEN] publicKey
@@ -68,36 +63,25 @@ contract VerificationGateway
     constructor(
         IBLS bls,
         address blsWalletImpl,
-        address proxyAdmin
+        address proxyAdmin,
+        address _entryPoint,
+        address _aggregateSigValidator
     ) {
         blsLib = bls;
         blsWalletLogic = blsWalletImpl;
         walletProxyAdmin = ProxyAdmin(proxyAdmin);
+        entryPoint = _entryPoint;
+        aggregateSigValidator = _aggregateSigValidator;
     }
 
-    /** Throw if bundle not valid or signature verification fails */
-    function verify(
-        Bundle memory bundle
-    ) public view {
-        uint256 opLength = bundle.operations.length;
-        require(
-            opLength == bundle.senderPublicKeys.length,
-            "VG: Sender/op length mismatch"
-        );
-        uint256[2][] memory messages = new uint256[2][](opLength);
+    function hashFromWallet(BLSWallet wallet) public view returns (bytes32) {
+        uint256[BLS_KEY_LEN] memory blsKey = wallet.getBlsKey();
 
-        for (uint256 i = 0; i<opLength; i++) {
-            // construct params for signature verification
-            messages[i] = messagePoint(bundle.operations[i]);
+        if (blsLib.isZeroBLSKey(blsKey)) {
+            return bytes32(0);
         }
 
-        bool verified = blsLib.verifyMultiple(
-            bundle.signature,
-            bundle.senderPublicKeys,
-            messages
-        );
-
-        require(verified, "VG: Sig not verified");
+        return keccak256(abi.encodePacked(blsKey));
     }
 
     /**
@@ -114,8 +98,8 @@ contract VerificationGateway
         uint256[BLS_KEY_LEN] memory publicKey
     ) public {
         require(blsLib.isZeroBLSKey(publicKey) == false, "VG: publicKey must be non-zero");
-        IWallet wallet = IWallet(msg.sender);
-        bytes32 existingHash = hashFromWallet[wallet];
+        BLSWallet wallet = BLSWallet(payable(msg.sender));
+        bytes32 existingHash = hashFromWallet(wallet);
         if (existingHash == bytes32(0)) { // wallet does not yet have a bls key registered with this gateway
             // set it instantly
             safeSetWallet(messageSenderSignature, publicKey, wallet);
@@ -129,8 +113,8 @@ contract VerificationGateway
     }
 
     function setPendingBLSKeyForWallet() public {
-        IWallet wallet = IWallet(msg.sender);
-        bytes32 existingHash = hashFromWallet[wallet];
+        BLSWallet wallet = BLSWallet(payable(msg.sender));
+        bytes32 existingHash = hashFromWallet(wallet);
         require(existingHash != bytes32(0), "VG: hash does not exist for caller");
         if (
             (pendingBLSPublicKeyTimeFromHash[existingHash] != 0) &&
@@ -201,7 +185,7 @@ contract VerificationGateway
         bytes32 salt,
         uint256[BLS_KEY_LEN] memory newBLSKey
     ) public {
-        IWallet wallet = walletFromHash[blsKeyHash];
+        BLSWallet wallet = walletFromHash[blsKeyHash];
         bytes32 recoveryHash = keccak256(
             abi.encodePacked(msg.sender, blsKeyHash, salt)
         );
@@ -243,46 +227,6 @@ contract VerificationGateway
         wallet.setTrustedGateway(blsGateway);
     }
 
-    /**
-    Base function for verifying and processing BLS-signed transactions.
-    Creates a new contract wallet per bls key if existing wallet not found.
-    Can be called with a single operation with no actions.
-    */
-    function processBundle(
-        Bundle memory bundle
-    ) external returns (
-        bool[] memory successes,
-        bytes[][] memory results
-    ) {
-        // revert if signature not verified
-        verify(bundle);
-
-        uint256 opLength = bundle.operations.length;
-        successes = new bool[](opLength);
-        results = new bytes[][](opLength);
-        for (uint256 i = 0; i<opLength; i++) {
-            IWallet wallet = getOrCreateWallet(bundle.senderPublicKeys[i]);
-
-            // check nonce then perform action
-            if (bundle.operations[i].nonce == wallet.nonce()) {
-                // request wallet perform operation
-                (
-                    bool success,
-                    bytes[] memory resultSet
-                ) = wallet.performOperation(bundle.operations[i]);
-                successes[i] = success;
-                results[i] = resultSet;
-                emit WalletOperationProcessed(
-                    address(wallet),
-                    bundle.operations[i].nonce,
-                    bundle.operations[i].actions,
-                    successes[i],
-                    results[i]
-                );
-            }
-        }
-    }
-
     /**
     Gets the wallet contract associated with the public key, creating it if
     needed.
@@ -291,15 +235,15 @@ contract VerificationGateway
         uint256[BLS_KEY_LEN] memory publicKey
     ) private returns (IWallet) {
         bytes32 publicKeyHash = keccak256(abi.encodePacked(publicKey));
-        IWallet blsWallet = walletFromHash[publicKeyHash];
+        BLSWallet blsWallet = walletFromHash[publicKeyHash];
         // publicKeyHash does not yet refer to a wallet, create one then update mappings.
         if (address(blsWallet) == address(0)) {
-            blsWallet = IWallet(address(new TransparentUpgradeableProxy{salt: publicKeyHash}(
+            blsWallet = BLSWallet(payable(new TransparentUpgradeableProxy{salt: publicKeyHash}(
                 address(blsWalletLogic),
                 address(walletProxyAdmin),
-                getInitializeData()
+                getInitializeData(publicKey)
             )));
-            updateWalletHashMappings(publicKeyHash, blsWallet);
+            updateWalletHashMappings(publicKey, blsWallet);
             emit WalletCreated(
                 address(blsWallet),
                 publicKey
@@ -317,7 +261,7 @@ contract VerificationGateway
     function safeSetWallet(
         uint256[2] memory wallletAddressSignature,
         uint256[BLS_KEY_LEN] memory publicKey,
-        IWallet wallet
+        BLSWallet wallet
     ) private {
         // verify the given wallet was signed for by the bls key
         uint256[2] memory addressMsg = blsLib.hashToPoint(
@@ -328,29 +272,32 @@ contract VerificationGateway
             blsLib.verifySingle(wallletAddressSignature, publicKey, addressMsg),
             "VG: Signature not verified for wallet address."
         );
-        bytes32 publicKeyHash = keccak256(abi.encodePacked(
-            publicKey
-        ));
         emit BLSKeySetForWallet(publicKey, wallet);
-        updateWalletHashMappings(publicKeyHash, wallet);
+        updateWalletHashMappings(publicKey, wallet);
     }
 
     /** @dev Only to be called on wallet creation, and in `safeSetWallet` */
     function updateWalletHashMappings(
-        bytes32 publicKeyHash,
-        IWallet wallet
+        uint256[BLS_KEY_LEN] memory blsKey,
+        BLSWallet wallet
     ) private {
         // remove reference from old hash
-        bytes32 oldHash = hashFromWallet[wallet];
-        walletFromHash[oldHash] = IWallet(address(0));
+        bytes32 oldHash = hashFromWallet(wallet);
+        delete walletFromHash[oldHash];
 
         // update new hash / wallet mappings
-        walletFromHash[publicKeyHash] = wallet;
-        hashFromWallet[wallet] = publicKeyHash;
+        walletFromHash[keccak256(abi.encodePacked(blsKey))] = wallet;
+        wallet.setBlsKey(blsKey);
     }
 
-    function getInitializeData() private view returns (bytes memory) {
-        return abi.encodeWithSignature("initialize(address)", address(this));
+    function getInitializeData(uint256[BLS_KEY_LEN] memory blsKey) private view returns (bytes memory) {
+        return abi.encodeWithSelector(
+            BLSWallet.initialize.selector,
+            blsKey,
+            address(this),
+            entryPoint,
+            aggregateSigValidator
+        );
     }
 
     modifier onlyWallet(bytes32 hash) {
@@ -360,31 +307,4 @@ contract VerificationGateway
         );
         _;
     }
-
-    function messagePoint(
-        IWallet.Operation memory op
-    ) internal view returns (
-        uint256[2] memory
-    ) {
-        bytes memory encodedActionData;
-        IWallet.ActionData memory a;
-        for (uint256 i=0; i<op.actions.length; i++) {
-            a = op.actions[i];
-            encodedActionData = abi.encodePacked(
-                encodedActionData,
-                a.ethValue,
-                a.contractAddress,
-                keccak256(a.encodedFunction)
-            );
-        }
-        return blsLib.hashToPoint(
-            BLS_DOMAIN,
-            abi.encodePacked(
-                block.chainid,
-                op.nonce,
-                keccak256(encodedActionData)
-            )
-        );
-    }
-
 }

contracts/contracts/interfaces/IWallet.sol

@@ -5,7 +5,6 @@ pragma abicoder v2;
 /** Interface for a contract wallet that can perform Operations
  */
 interface IWallet {
-
     struct Operation {
         uint256 nonce;
         IWallet.ActionData[] actions;
@@ -17,7 +16,15 @@ interface IWallet {
         bytes encodedFunction;
     }
 
-    function initialize(address gateway) external;
+    function initialize(
+        uint256[4] memory _blsKey,
+        address _blsGateway,
+        address _entryPoint,
+        address _aggregator
+    ) external;
+
+    function getBlsKey() external returns (uint256[4] memory);
+    function setBlsKey(uint256[4] memory newBlsKey) external;
     function nonce() external returns (uint256);
 
     function performOperation(

contracts/package.json

@@ -17,6 +17,7 @@
   "author": "James Zaki",
   "license": "MIT",
   "dependencies": {
+    "@account-abstraction/contracts": "^0.2.0",
     "@openzeppelin/contracts": "^4.3.2"
   },
   "devDependencies": {

contracts/shared/helpers/Fixture.ts

@@ -20,6 +20,8 @@ import Range from "./Range";
 import assert from "./assert";
 import Create2Fixture from "./Create2Fixture";
 import { VerificationGateway, BLSOpen, ProxyAdmin } from "../../typechain";
+import { EntryPoint } from "@account-abstraction/contracts";
+import { AggregateSigValidator } from "../../clients/typechain";
 
 export default class Fixture {
   static readonly ECDSA_ACCOUNTS_LENGTH = 5;
@@ -34,6 +36,8 @@ export default class Fixture {
 
     public lazyBlsWallets: (() => Promise<BlsWalletWrapper>)[],
 
+    public entryPoint: EntryPoint,
+    public aggregateSigValidator: AggregateSigValidator,
     public verificationGateway: VerificationGateway,
 
     public blsLibrary: BLSOpen,
@@ -61,9 +65,16 @@ export default class Fixture {
 
     // deploy wallet implementation contract
     const blsWalletImpl = await create2Fixture.create2Contract("BLSWallet");
+
+    // TODO: Why do we ignore errors here?
     try {
       await (
-        await blsWalletImpl.initialize(ethers.constants.AddressZero)
+        await blsWalletImpl.initialize(
+          [0, 0, 0, 0],
+          ethers.constants.AddressZero,
+          ethers.constants.AddressZero,
+          ethers.constants.AddressZero,
+        )
       ).wait();
     } catch (e) {}
 
@@ -71,12 +82,32 @@ export default class Fixture {
     const ProxyAdmin = await ethers.getContractFactory("ProxyAdmin");
     const proxyAdmin = (await ProxyAdmin.deploy()) as ProxyAdmin;
     await proxyAdmin.deployed();
+
+    const entryPoint = (await create2Fixture.create2Contract(
+      "EntryPoint",
+      ethers.utils.defaultAbiCoder.encode(
+        ["uint256", "uint32"],
+        [ethers.utils.parseEther("2"), 2],
+      ),
+    )) as EntryPoint;
+
+    const aggregateSigValidator = (await create2Fixture.create2Contract(
+      "AggregateSigValidator",
+      ethers.utils.defaultAbiCoder.encode(["address"], [bls.address]),
+    )) as EntryPoint;
+
     // deploy Verification Gateway
     const verificationGateway = (await create2Fixture.create2Contract(
       "VerificationGateway",
       ethers.utils.defaultAbiCoder.encode(
         ["address", "address", "address"],
-        [bls.address, blsWalletImpl.address, proxyAdmin.address],
+        [
+          bls.address,
+          blsWalletImpl.address,
+          proxyAdmin.address,
+          entryPoint.address,
+          aggregateSigValidator.address,
+        ],
       ),
     )) as VerificationGateway;
     await (
@@ -133,6 +164,8 @@ export default class Fixture {
       signers,
       addresses,
       lazyBlsWallets,
+      entryPoint,
+      aggregateSigValidator,
       verificationGateway,
       bls,
       blsExpander,

contracts/yarn.lock

@@ -2,6 +2,11 @@
 # yarn lockfile v1
 
 
+"@account-abstraction/contracts@^0.2.0":
+  version "0.2.0"
+  resolved "https://registry.yarnpkg.com/@account-abstraction/contracts/-/contracts-0.2.0.tgz#d353fd4b26d8867a3c0067fc0e20da03c99ca2dd"
+  integrity sha512-CbPYA7o55u1S5Fee5qIWffD/4dxJ7v+86Si8TOUOE/tAMUWP+nzU1kHKUwFrM/HT8OxAQpeXccCD4UtROLtfwA==
+
 "@babel/code-frame@7.12.11":
   version "7.12.11"
   resolved "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.12.11.tgz"