Fix/native restaking withdraw

.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+  "[solidity]": {
+    "editor.defaultFormatter": "JuanBlanco.solidity"
+  },
+  "solidity.formatter": "forge"
+}

package.json

@@ -96,7 +96,7 @@
     "death": "^1.1.0",
     "debug": "^4.3.4",
     "decamelize": "^4.0.0",
-    "decimal.js":"10.4.3",
+    "decimal.js": "10.4.3",
     "deep-eql": "^4.1.3",
     "deep-extend": "^0.6.0",
     "deep-is": "^0.1.4",

src/.solhint.json

@@ -1,9 +1,9 @@
 {
   "extends": "solhint:recommended",
   "rules": {
-    "max-line-length": ["error", 121],
+    "max-line-length": ["error", 128],
     "compiler-version": ["error", "^0.8.0"],
-    "func-visibility": ["warn", {"ignoreConstructors": true}],
+    "func-visibility": ["warn", { "ignoreConstructors": true }],
     "no-inline-assembly": "off",
     "no-empty-blocks": "off",
     "no-unused-vars": "error",
@@ -13,6 +13,7 @@
     "max-states-count": "off",
     "reason-string": "off",
     "gas-custom-errors": "off",
-    "state-visibility": "error"
+    "state-visibility": "error",
+    "no-complex-fallback": "off"
   }
 }

src/core/ClientGatewayLzReceiver.sol

@@ -13,6 +13,9 @@ abstract contract ClientGatewayLzReceiver is PausableUpgradeable, OAppReceiverUp
     error UnexpectedResponse(uint64 nonce);
     error DepositShouldNotFailOnExocore(address token, address depositor);
 
+    // Events
+    event WithdrawFailedOnExocore(address indexed token, address indexed withdrawer);
+
     modifier onlyCalledFromThis() {
         require(
             msg.sender == address(this),
@@ -115,14 +118,24 @@ abstract contract ClientGatewayLzReceiver is PausableUpgradeable, OAppReceiverUp
 
         bool success = (uint8(bytes1(responsePayload[0])) == 1);
         uint256 lastlyUpdatedPrincipalBalance = uint256(bytes32(responsePayload[1:33]));
-        if (success) {
-            IVault vault = _getVault(token);
 
-            vault.updatePrincipalBalance(withdrawer, lastlyUpdatedPrincipalBalance);
-            vault.updateWithdrawableBalance(withdrawer, unlockPrincipalAmount, 0);
-        }
+        if (!success) {
+            emit WithdrawFailedOnExocore(token, withdrawer);
+        } else {
+            if (token == VIRTUAL_STAKED_ETH_ADDRESS) {
+                IExoCapsule capsule = _getCapsule(withdrawer);
+
+                capsule.updatePrincipalBalance(lastlyUpdatedPrincipalBalance);
+                capsule.updateWithdrawableBalance(unlockPrincipalAmount);
+            } else {
+                IVault vault = _getVault(token);
 
-        emit WithdrawPrincipalResult(success, token, withdrawer, unlockPrincipalAmount);
+                vault.updatePrincipalBalance(withdrawer, lastlyUpdatedPrincipalBalance);
+                vault.updateWithdrawableBalance(withdrawer, unlockPrincipalAmount, 0);
+            }
+
+            emit WithdrawPrincipalResult(success, token, withdrawer, unlockPrincipalAmount);
+        }
     }
 
     function afterReceiveWithdrawRewardResponse(bytes memory requestPayload, bytes calldata responsePayload)

src/core/ExoCapsule.sol

@@ -20,11 +20,27 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
     event PrincipalBalanceUpdated(address, uint256);
     event WithdrawableBalanceUpdated(address, uint256);
     event WithdrawalSuccess(address, address, uint256);
+    /// @notice Emitted when a partial withdrawal claim is successfully redeemed
+    event PartialWithdrawalRedeemed(
+        bytes32 pubkey, uint256 withdrawalTimestamp, address indexed recipient, uint64 partialWithdrawalAmountGwei
+    );
+    /// @notice Emitted when an ETH validator is prove to have fully withdrawn from the beacon chain
+    event FullWithdrawalRedeemed(
+        bytes32 pubkey, uint256 withdrawalTimestamp, address indexed recipient, uint64 withdrawalAmountGwei
+    );
+    /// @notice Emitted when capsuleOwner enables restaking
+    event RestakingActivated(address indexed capsuleOwner);
+    /// @notice Emitted when ETH is received via the `receive` fallback
+    event NonBeaconChainETHReceived(uint256 amountReceived);
+    /// @notice Emitted when ETH that was previously received via the `receive` fallback is withdrawn
+    event NonBeaconChainETHWithdrawn(address indexed recipient, uint256 amountWithdrawn);
 
     error InvalidValidatorContainer(bytes32 pubkey);
     error InvalidWithdrawalContainer(uint64 validatorIndex);
+    error InvalidHistoricalSummaries(uint64 validatorIndex);
     error DoubleDepositedValidator(bytes32 pubkey);
     error StaleValidatorContainer(bytes32 pubkey, uint256 timestamp);
+    error WithdrawalAlreadyProven(bytes32 pubkey, uint256 timestamp);
     error UnregisteredValidator(bytes32 pubkey);
     error UnregisteredOrWithdrawnValidatorContainer(bytes32 pubkey);
     error FullyWithdrawnValidatorContainer(bytes32 pubkey);
@@ -47,6 +63,11 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
         _disableInitializers();
     }
 
+    receive() external payable {
+        nonBeaconChainETHBalance += msg.value;
+        emit NonBeaconChainETHReceived(msg.value);
+    }
+
     function initialize(address gateway_, address capsuleOwner_, address beaconOracle_) external initializer {
         require(gateway_ != address(0), "ExoCapsule: gateway address can not be empty");
         require(capsuleOwner_ != address(0), "ExoCapsule: capsule owner address can not be empty");
@@ -55,6 +76,8 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
         gateway = INativeRestakingController(gateway_);
         beaconOracle = IBeaconChainOracle(beaconOracle_);
         capsuleOwner = capsuleOwner_;
+
+        emit RestakingActivated(capsuleOwner);
     }
 
     function verifyDepositProof(bytes32[] calldata validatorContainer, ValidatorContainerProof calldata proof)
@@ -95,61 +118,54 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
         _capsuleValidatorsByIndex[proof.validatorIndex] = validatorPubkey;
     }
 
-    function verifyPartialWithdrawalProof(
-        bytes32[] calldata validatorContainer,
-        ValidatorContainerProof calldata validatorProof,
-        bytes32[] calldata withdrawalContainer,
-        WithdrawalContainerProof calldata withdrawalProof
-    ) external view onlyGateway {
-        bytes32 validatorPubkey = validatorContainer.getPubkey();
-        uint64 withdrawableEpoch = validatorContainer.getWithdrawableEpoch();
-
-        bool partialWithdrawal = _timestampToEpoch(validatorProof.beaconBlockTimestamp) < withdrawableEpoch;
-
-        if (!validatorContainer.verifyValidatorContainerBasic()) {
-            revert InvalidValidatorContainer(validatorPubkey);
-        }
-
-        if (!partialWithdrawal) {
-            revert NotPartialWithdrawal(validatorPubkey);
-        }
-
-        if (validatorProof.beaconBlockTimestamp != withdrawalProof.beaconBlockTimestamp) {
-            revert UnmatchedValidatorAndWithdrawal(validatorPubkey);
-        }
-
-        _verifyValidatorContainer(validatorContainer, validatorProof);
-        _verifyWithdrawalContainer(withdrawalContainer, withdrawalProof);
-    }
-
-    function verifyFullWithdrawalProof(
+    function verifyWithdrawalProof(
         bytes32[] calldata validatorContainer,
         ValidatorContainerProof calldata validatorProof,
         bytes32[] calldata withdrawalContainer,
         WithdrawalContainerProof calldata withdrawalProof
-    ) external onlyGateway {
+    ) external onlyGateway returns (bool partialWithdrawal, uint256 withdrawalAmount) {
         bytes32 validatorPubkey = validatorContainer.getPubkey();
         uint64 withdrawableEpoch = validatorContainer.getWithdrawableEpoch();
-
         Validator storage validator = _capsuleValidators[validatorPubkey];
-        bool fullyWithdrawal = _timestampToEpoch(validatorProof.beaconBlockTimestamp) > withdrawableEpoch;
+        partialWithdrawal = _timestampToEpoch(validatorProof.beaconBlockTimestamp) < withdrawableEpoch;
 
         if (!validatorContainer.verifyValidatorContainerBasic()) {
             revert InvalidValidatorContainer(validatorPubkey);
         }
-
-        if (!fullyWithdrawal) {
-            revert NotPartialWithdrawal(validatorPubkey);
+        if (validator.status == VALIDATOR_STATUS.UNREGISTERED) {
+            revert UnregisteredOrWithdrawnValidatorContainer(validatorPubkey);
         }
 
-        if (validatorProof.beaconBlockTimestamp != withdrawalProof.beaconBlockTimestamp) {
-            revert UnmatchedValidatorAndWithdrawal(validatorPubkey);
+        if (provenWithdrawal[validatorPubkey][withdrawalProof.beaconBlockTimestamp]) {
+            revert WithdrawalAlreadyProven(validatorPubkey, withdrawalProof.beaconBlockTimestamp);
         }
 
+        provenWithdrawal[validatorPubkey][withdrawalProof.beaconBlockTimestamp] = true;
+
         _verifyValidatorContainer(validatorContainer, validatorProof);
         _verifyWithdrawalContainer(withdrawalContainer, withdrawalProof);
 
-        validator.status = VALIDATOR_STATUS.WITHDRAWN;
+        uint64 withdrawalAmountGwei = withdrawalContainer.getAmount();
+
+        if (partialWithdrawal) {
+            // Immediately send ETH without sending request to Exocore side
+            emit PartialWithdrawalRedeemed(
+                validatorPubkey, withdrawalProof.beaconBlockTimestamp, capsuleOwner, withdrawalAmountGwei
+            );
+            _sendETH(capsuleOwner, withdrawalAmountGwei * GWEI_TO_WEI);
+        } else {
+            // Full withdrawal
+            validator.status = VALIDATOR_STATUS.WITHDRAWN;
+            validator.restakedBalanceGwei = 0;
+            // If over MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR = 32 * 1e9, then send remaining amount immediately
+            emit FullWithdrawalRedeemed(
+                validatorPubkey, withdrawalProof.beaconBlockTimestamp, capsuleOwner, withdrawalAmountGwei
+            );
+            if (withdrawalAmountGwei > MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR) {
+                withdrawalAmount = (withdrawalAmountGwei - MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR) * GWEI_TO_WEI;
+                _sendETH(capsuleOwner, withdrawalAmount);
+            }
+        }
     }
 
     function withdraw(uint256 amount, address payable recipient) external onlyGateway {
@@ -158,16 +174,24 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
         );
 
         withdrawableBalance -= amount;
-        (bool sent,) = recipient.call{value: amount}("");
-        if (!sent) {
-            revert WithdrawalFailure(capsuleOwner, recipient, amount);
-        }
+        _sendETH(recipient, amount);
 
         emit WithdrawalSuccess(capsuleOwner, recipient, amount);
     }
 
+    /// @notice Called by the capsule owner to withdraw the nonBeaconChainETHBalance
+    function withdrawNonBeaconChainETHBalance(address recipient, uint256 amountToWithdraw) external onlyGateway {
+        require(
+            amountToWithdraw <= nonBeaconChainETHBalance,
+            "ExoCapsule.withdrawNonBeaconChainETHBalance: amountToWithdraw is greater than nonBeaconChainETHBalance"
+        );
+        nonBeaconChainETHBalance -= amountToWithdraw;
+        _sendETH(recipient, amountToWithdraw);
+        emit NonBeaconChainETHWithdrawn(recipient, amountToWithdraw);
+    }
+
     function updatePrincipalBalance(uint256 lastlyUpdatedPrincipalBalance) external onlyGateway {
-        principalBalance = lastlyUpdatedPrincipalBalance;
+        principleBalance = lastlyUpdatedPrincipalBalance;
 
         emit PrincipalBalanceUpdated(capsuleOwner, lastlyUpdatedPrincipalBalance);
     }
@@ -215,6 +239,13 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
         return validator;
     }
 
+    function _sendETH(address recipient, uint256 amountWei) internal {
+        (bool sent,) = recipient.call{value: amountWei}("");
+        if (!sent) {
+            revert WithdrawalFailure(capsuleOwner, recipient, amountWei);
+        }
+    }
+
     function _verifyValidatorContainer(bytes32[] calldata validatorContainer, ValidatorContainerProof calldata proof)
         internal
         view
@@ -244,11 +275,19 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
             proof.withdrawalIndex,
             beaconBlockRoot,
             proof.executionPayloadRoot,
-            proof.executionPayloadRootProof
+            proof.executionPayloadRootProof,
+            proof.beaconBlockTimestamp
         );
         if (!valid) {
             revert InvalidWithdrawalContainer(withdrawalContainer.getValidatorIndex());
         }
+        // Verify historical summaries
+        bool validHistoricalSummaries = proof.stateRoot.isValidHistoricalSummaryRoot(
+            proof.historicalSummaryBlockRootProof, proof.historicalSummaryIndex, beaconBlockRoot, proof.blockRootIndex
+        );
+        if (!validHistoricalSummaries) {
+            revert InvalidHistoricalSummaries(withdrawalContainer.getValidatorIndex());
+        }
     }
 
     function _isActivatedAtEpoch(bytes32[] calldata validatorContainer, uint256 atTimestamp)
@@ -258,9 +297,8 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
     {
         uint64 atEpoch = _timestampToEpoch(atTimestamp);
         uint64 activationEpoch = validatorContainer.getActivationEpoch();
-        uint64 exitEpoch = validatorContainer.getExitEpoch();
 
-        return (atEpoch >= activationEpoch && atEpoch < exitEpoch);
+        return atEpoch >= activationEpoch;
     }
 
     function _isStaleProof(Validator storage validator, uint256 proofTimestamp) internal view returns (bool) {

src/core/NativeRestakingController.sol

@@ -5,7 +5,6 @@ import {INativeRestakingController} from "../interfaces/INativeRestakingControll
 
 import {ValidatorContainer} from "../libraries/ValidatorContainer.sol";
 import {BaseRestakingController} from "./BaseRestakingController.sol";
-import {ExoCapsule} from "./ExoCapsule.sol";
 
 import {PausableUpgradeable} from "@openzeppelin-upgradeable/contracts/utils/PausableUpgradeable.sol";
 import {Create2} from "@openzeppelin/contracts/utils/Create2.sol";
@@ -39,7 +38,7 @@ abstract contract NativeRestakingController is
             address(ownerToCapsule[msg.sender]) == address(0),
             "NativeRestakingController: message sender has already created the capsule"
         );
-        ExoCapsule capsule = ExoCapsule(
+        IExoCapsule capsule = IExoCapsule(
             Create2.deploy(
                 0,
                 bytes32(uint256(uint160(msg.sender))),
@@ -70,18 +69,24 @@ abstract contract NativeRestakingController is
         _processRequest(Action.REQUEST_DEPOSIT, actionArgs, encodedRequest);
     }
 
-    function processBeaconChainPartialWithdrawal(
+    function processBeaconChainWithdrawal(
         bytes32[] calldata validatorContainer,
         IExoCapsule.ValidatorContainerProof calldata validatorProof,
         bytes32[] calldata withdrawalContainer,
         IExoCapsule.WithdrawalContainerProof calldata withdrawalProof
-    ) external payable whenNotPaused {}
-
-    function processBeaconChainFullWithdrawal(
-        bytes32[] calldata validatorContainer,
-        IExoCapsule.ValidatorContainerProof calldata validatorProof,
-        bytes32[] calldata withdrawalContainer,
-        IExoCapsule.WithdrawalContainerProof calldata withdrawalProof
-    ) external payable whenNotPaused {}
+    ) external payable whenNotPaused {
+        IExoCapsule capsule = _getCapsule(msg.sender);
+        (bool partialWithdrawal, uint256 withdrawalAmount) =
+            capsule.verifyWithdrawalProof(validatorContainer, validatorProof, withdrawalContainer, withdrawalProof);
+        if (!partialWithdrawal) {
+            // request full withdraw
+            bytes memory actionArgs = abi.encodePacked(
+                bytes32(bytes20(VIRTUAL_STAKED_ETH_ADDRESS)), bytes32(bytes20(msg.sender)), withdrawalAmount
+            );
+            bytes memory encodedRequest = abi.encode(VIRTUAL_STAKED_ETH_ADDRESS, msg.sender, withdrawalAmount);
+
+            _processRequest(Action.REQUEST_WITHDRAW_PRINCIPAL_FROM_EXOCORE, actionArgs, encodedRequest);
+        }
+    }
 
 }

src/interfaces/IExoCapsule.sol

@@ -13,28 +13,39 @@ interface IExoCapsule {
 
     struct WithdrawalContainerProof {
         uint256 beaconBlockTimestamp;
+        bytes32 stateRoot;
         bytes32 executionPayloadRoot;
         bytes32[] executionPayloadRootProof;
         bytes32[] withdrawalContainerRootProof;
+        bytes32[] historicalSummaryBlockRootProof;
+        uint256 historicalSummaryIndex;
+        uint256 blockRootIndex;
         uint256 withdrawalIndex;
     }
 
+    struct HistoricalBlockRootProof {
+        uint256 beaconBlockTimestamp;
+        bytes32 stateRoot;
+        bytes32 executionPayloadRoot;
+        bytes32[] executionPayloadRootProof;
+        bytes32[] withdrawalContainerRootProof;
+        bytes32[] historicalSummaryBlockRootProof;
+        uint256 historicalSummaryIndex;
+        uint256 blockRootIndex;
+        uint256 withdrawalIndex;
+    }
+
+    function initialize(address gateway, address capsuleOwner, address beaconOracle) external;
+
     function verifyDepositProof(bytes32[] calldata validatorContainer, ValidatorContainerProof calldata proof)
         external;
 
-    function verifyPartialWithdrawalProof(
-        bytes32[] calldata validatorContainer,
-        ValidatorContainerProof calldata validatorProof,
-        bytes32[] calldata withdrawalContainer,
-        WithdrawalContainerProof calldata withdrawalProof
-    ) external;
-
-    function verifyFullWithdrawalProof(
+    function verifyWithdrawalProof(
         bytes32[] calldata validatorContainer,
         ValidatorContainerProof calldata validatorProof,
         bytes32[] calldata withdrawalContainer,
         WithdrawalContainerProof calldata withdrawalProof
-    ) external;
+    ) external returns (bool partialWithdrawal, uint256 withdrawalAmount);
 
     function withdraw(uint256 amount, address payable recipient) external;