Wallet key recovery

contracts/contracts/BLSWallet.sol

@@ -14,31 +14,75 @@ import "./interfaces/IWallet.sol";
 contract BLSWallet is Initializable, IBLSWallet
 {
     uint256 public nonce;
+    bytes32 public recoveryHash;
+    bytes32 pendingRecoveryHash;
+    uint256 pendingRecoveryHashTime;
+    bytes public approvedProxyAdminFunction;
+    bytes pendingPAFunction;
+    uint256 pendingPAFunctionTime;
 
     // BLS variables
     uint256[4] public blsPublicKey;
+    uint256[4] pendingBLSPublicKey;
+    uint256 pendingBLSPublicKeyTime;
     address public trustedBLSGateway;
+    address pendingBLSGateway;
+    uint256 pendingGatewayTime;
+
+    event PendingRecoveryHashSet(
+        bytes32 pendingRecoveryHash
+    );
+    event PendingBLSKeySet(
+        uint256[4] pendingBLSKey
+    );
+    event PendingGatewaySet(
+        address pendingGateway
+    );
+    event PendingProxyAdminFunctionSet(
+        bytes pendingProxyAdminFunction
+    );
+
+    event RecoveryHashUpdated(
+        bytes32 oldHash,
+        bytes32 newHash
+    );
+    event BLSKeySet(
+        uint256[4] oldBLSKey,
+        uint256[4] newBLSKey
+    );
+    event GatewayUpdated(
+        address oldGateway,
+        address newGateway
+    );
+    event ProxyAdminFunctionApproved(
+        bytes approvedProxyAdmin
+    );
 
     function initialize(
         address blsGateway
     ) external initializer {
         nonce = 0;
         trustedBLSGateway = blsGateway;
+        pendingGatewayTime = type(uint256).max;
+        pendingPAFunctionTime = type(uint256).max;
     }
 
     /** */
     function latchBLSPublicKey(
         uint256[4] memory blsKey
     ) public onlyTrustedGateway {
-        for (uint256 i=0; i<4; i++) {
-            require(
-                blsPublicKey[i] == 0,
-                "BLSWallet: public key already set"
-            );
-        }
+        require(isZeroBLSKey(blsPublicKey), "BLSWallet: public key already set");
         blsPublicKey = blsKey;
     }
 
+    function isZeroBLSKey(uint256[4] memory blsKey) public pure returns (bool) {
+        bool isZero = true;
+        for (uint256 i=0; isZero && i<4; i++) {
+            isZero = (blsKey[i] == 0);
+        }
+        return isZero;
+    }
+
     receive() external payable {}
     fallback() external payable {}
 
@@ -49,18 +93,100 @@ contract BLSWallet is Initializable, IBLSWallet
         return blsPublicKey;
     }
 
+    /**
+    Wallet can update its recovery hash
+     */
+    function setRecoveryHash(bytes32 hash) public onlyThis {
+        if (recoveryHash == bytes32(0)) {
+            recoveryHash = hash;
+            clearPendingRecoveryHash();
+            emit RecoveryHashUpdated(bytes32(0), recoveryHash);
+        }
+        else {
+            pendingRecoveryHash = hash;
+            pendingRecoveryHashTime = block.timestamp + 604800; // 1 week from now
+            emit PendingRecoveryHashSet(pendingRecoveryHash);
+        }
+    }
+
+    /**
+    Wallet can update its BLS key
+     */
+    function setBLSPublicKey(uint256[4] memory blsKey) public onlyThis {
+        require(isZeroBLSKey(blsKey) == false, "BLSWallet: blsKey must be non-zero");
+        pendingBLSPublicKey = blsKey;
+        pendingBLSPublicKeyTime = block.timestamp + 604800; // 1 week from now
+        emit PendingBLSKeySet(pendingBLSPublicKey);
+    }
+
     /**
     Wallet can migrate to a new gateway, eg additional signature support
      */
-    function setTrustedBLSGateway(address blsGateway) public onlyThis {
-        uint256 size;
-        // solhint-disable-next-line no-inline-assembly
-        assembly { size := extcodesize(blsGateway) }
-        require(
-            (blsGateway != address(0)) && (size > 0),
-            "BLSWallet: gateway address param not valid"
-        );
-        trustedBLSGateway = blsGateway;
+    function setTrustedGateway(address blsGateway) public onlyTrustedGateway {
+        pendingBLSGateway = blsGateway;
+        pendingGatewayTime = block.timestamp + 604800; // 1 week from now
+        emit PendingGatewaySet(pendingBLSGateway);
+    }
+
+    /**
+    Prepare wallet with desired implementation contract to upgrade to.
+    */
+    function setProxyAdminFunction(bytes calldata encodedFunction) public onlyTrustedGateway {
+        pendingPAFunction = encodedFunction;
+        pendingPAFunctionTime = block.timestamp + 604800; // 1 week from now
+        emit PendingProxyAdminFunctionSet(pendingPAFunction);
+    }
+
+    /**
+    Set results of any pending set operation if their respective timestamp has elapsed.
+     */
+    function setAnyPending() public {
+        if (block.timestamp > pendingRecoveryHashTime) {
+            bytes32 previousRecoveryHash = recoveryHash;
+            recoveryHash = pendingRecoveryHash;
+            clearPendingRecoveryHash();
+            emit RecoveryHashUpdated(previousRecoveryHash, recoveryHash);
+        }
+        if (block.timestamp > pendingBLSPublicKeyTime) {
+            uint256[4] memory previousBLSPublicKey = blsPublicKey;
+            blsPublicKey = pendingBLSPublicKey;
+            pendingBLSPublicKeyTime = type(uint256).max;
+            pendingBLSPublicKey = [0,0,0,0];
+            emit BLSKeySet(previousBLSPublicKey, blsPublicKey);
+        }
+        if (block.timestamp > pendingGatewayTime) {
+            address previousGateway = trustedBLSGateway;
+            trustedBLSGateway = pendingBLSGateway;
+            pendingGatewayTime = type(uint256).max;
+            pendingBLSGateway = address(0);
+            emit GatewayUpdated(previousGateway, trustedBLSGateway);
+        }
+        if (block.timestamp > pendingPAFunctionTime) {
+            approvedProxyAdminFunction = pendingPAFunction;
+            pendingPAFunctionTime = type(uint256).max;
+            pendingPAFunction = new bytes(0);
+            emit ProxyAdminFunctionApproved(approvedProxyAdminFunction);
+        }
+    }
+
+    function clearPendingRecoveryHash() internal {
+        pendingRecoveryHashTime = type(uint256).max;
+        pendingRecoveryHash = bytes32(0);
+    }
+
+    function recover(
+        uint256[4] calldata newBLSKey
+    ) public onlyTrustedGateway {
+        // set new bls key
+        blsPublicKey = newBLSKey;
+        // clear any pending operations
+        clearPendingRecoveryHash();
+        pendingBLSPublicKeyTime = type(uint256).max;
+        pendingBLSPublicKey = [0,0,0,0];
+        pendingGatewayTime = type(uint256).max;
+        pendingBLSGateway = address(0);
+        pendingPAFunctionTime = type(uint256).max;
+        pendingPAFunction = new bytes(0);
     }
 
     /**
@@ -113,6 +239,10 @@ contract BLSWallet is Initializable, IBLSWallet
         }
     }
 
+    function clearApprovedProxyAdminFunction() public onlyTrustedGateway {
+        approvedProxyAdminFunction = new bytes(0);
+    }
+
     /**
     Consecutive nonce increment, contract can be upgraded for other types
      */

contracts/contracts/VerificationGateway.sol

@@ -141,28 +141,81 @@ contract VerificationGateway
     }
 
     /**
-    Calls to proxy admin, exclusively from a wallet.
+    Calls to proxy admin, exclusively from a wallet. Must be called twice.
+    Once to set the function in the wallet as pending, then again after the recovery time.
     @param hash calling wallet's bls public key hash
     @param encodedFunction the selector and params to call (first encoded param must be calling wallet)
      */
     function walletAdminCall(
         bytes32 hash,
         bytes calldata encodedFunction
-    ) public onlyWallet(hash) returns (
-        bytes memory
-    ) {
-        // ensure first parameter is the calling wallet
-        bytes memory encodedAddress = abi.encode(address(walletFromHash(hash)));
+    ) public onlyWallet(hash) {
+        IWallet wallet = walletFromHash(hash);
+
+        // ensure first parameter is the calling wallet address
+        bytes memory encodedAddress = abi.encode(address(wallet));
         uint8 selectorOffset = 4;
         for (uint256 i=0; i<32; i++) {
             require(
                 (encodedFunction[selectorOffset+i] == encodedAddress[i]),
                 "VG: first param to proxy admin is not calling wallet"
             );
         }
-        (bool success, bytes memory result) = address(walletProxyAdmin).call(encodedFunction);
-        require(success, "VG: call to proxy admin failed");
-        return result;
+
+        wallet.setAnyPending();
+
+        // ensure wallet has pre-approved encodedFunction
+        bytes memory approvedFunction = wallet.approvedProxyAdminFunction();
+        bool matchesApproved = (encodedFunction.length == approvedFunction.length);
+        for (uint i=0; matchesApproved && i<approvedFunction.length; i++) {
+            matchesApproved = (encodedFunction[i] == approvedFunction[i]);
+        }
+
+        if (matchesApproved == false) {
+            // prepare for a future call
+            wallet.setProxyAdminFunction(encodedFunction);
+        }
+        else {
+            // call approved function
+            (bool success, ) = address(walletProxyAdmin).call(encodedFunction);
+            require(success, "VG: call to proxy admin failed");
+            wallet.clearApprovedProxyAdminFunction();
+        }
+    }
+
+    function recoverWallet(
+        bytes32 blsKeyHash,
+        bytes32 salt,
+        uint256[4] memory newBLSKey
+    ) public {
+        IWallet wallet = walletFromHash(blsKeyHash);
+        bytes32 recoveryHash = keccak256(
+            abi.encodePacked(msg.sender, blsKeyHash, salt)
+        );
+        if (recoveryHash == wallet.recoveryHash()) {
+            // override mapping of old key hash (takes precedence over create2 address)
+            externalWalletsFromHash[blsKeyHash] = IWallet(0xFFfFfFffFFfffFFfFFfFFFFFffFFFffffFfFFFfF);
+            bytes32 newKeyHash = keccak256(abi.encodePacked(newBLSKey));
+            externalWalletsFromHash[newKeyHash] = wallet;
+            wallet.recover(newBLSKey);
+        }
+    }
+
+    /**
+    Wallet can migrate to a new gateway, eg additional signature support
+     */
+    function setTrustedBLSGateway(
+        bytes32 hash,
+        address blsGateway
+    ) public onlyWallet(hash) {
+        uint256 size;
+        // solhint-disable-next-line no-inline-assembly
+        assembly { size := extcodesize(blsGateway) }
+        require(
+            (blsGateway != address(0)) && (size > 0),
+            "BLSWallet: gateway address param not valid"
+        );
+        walletFromHash(hash).setTrustedGateway(blsGateway);
     }
 
     /** 
@@ -189,12 +242,11 @@ contract VerificationGateway
             // create wallet if not found
             createNewWallet(bundle.senderPublicKeys[i]);
 
-            // construct params for signature verification
+            // calculate public key hash
             publicKeyHash = keccak256(abi.encodePacked(
                 bundle.senderPublicKeys[i]
             ));
             wallet = walletFromHash(publicKeyHash);
-
             // check nonce then perform action
             if (bundle.operations[i].nonce == wallet.nonce()) {
                 // request wallet perform operation
@@ -221,7 +273,6 @@ contract VerificationGateway
     ) private {
         bytes32 publicKeyHash = keccak256(abi.encodePacked(publicKey));
         address blsWallet = address(walletFromHash(publicKeyHash));
-
         // wallet with publicKeyHash doesn't exist at expected create2 address
         if (blsWallet == address(0)) {
             blsWallet = address(new TransparentUpgradeableProxy{salt: publicKeyHash}(

contracts/contracts/interfaces/IWallet.sol

@@ -26,6 +26,18 @@ interface IWallet {
         bool success,
         bytes[] memory results
     );
+
+    function recoveryHash() external returns (bytes32);
+    function recover(uint256[4] calldata newBLSKey) external;
+
+    // prepares gateway to be set (after pending timestamp)
+    function setTrustedGateway(address gateway) external;
+    // checks any pending variables and sets them if past their timestamp
+    function setAnyPending() external;
+
+    function setProxyAdminFunction(bytes memory) external;
+    function approvedProxyAdminFunction() external view returns (bytes memory);
+    function clearApprovedProxyAdminFunction() external;
 }
 
 /** Interface for bls-specific functions
@@ -38,7 +50,5 @@ interface IBLSWallet is IWallet {
     ) external;
 
     function getBLSPublicKey() external view returns (uint256[4] memory);
-    function setTrustedBLSGateway(address blsGateway) external;
-
  }