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

@@ -14,6 +14,9 @@ abstract contract ClientGatewayLzReceiver is PausableUpgradeable, OAppReceiverUp
     error DepositShouldNotFailOnExocore(address token, address depositor);
     error InvalidAddWhitelistTokensRequest(uint256 expectedLength, uint256 actualLength);
 
+    // Events
+    event WithdrawFailedOnExocore(address indexed token, address indexed withdrawer);
+
     modifier onlyCalledFromThis() {
         require(
             msg.sender == address(this),
@@ -111,14 +114,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

@@ -4,27 +4,45 @@ import {IExoCapsule} from "../interfaces/IExoCapsule.sol";
 
 import {INativeRestakingController} from "../interfaces/INativeRestakingController.sol";
 import {BeaconChainProofs} from "../libraries/BeaconChainProofs.sol";
+import {Endian} from "../libraries/Endian.sol";
 import {ValidatorContainer} from "../libraries/ValidatorContainer.sol";
 import {WithdrawalContainer} from "../libraries/WithdrawalContainer.sol";
 import {ExoCapsuleStorage} from "../storage/ExoCapsuleStorage.sol";
 
 import {IBeaconChainOracle} from "@beacon-oracle/contracts/src/IBeaconChainOracle.sol";
-import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
+import {ReentrancyGuardUpgradeable} from "@openzeppelin/contracts-upgradeable/security/ReentrancyGuardUpgradeable.sol";
 
-contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
+contract ExoCapsule is ReentrancyGuardUpgradeable, ExoCapsuleStorage, IExoCapsule {
 
     using BeaconChainProofs for bytes32;
+    using Endian for bytes32;
     using ValidatorContainer for bytes32[];
     using WithdrawalContainer for bytes32[];
 
     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 withdrawalEpoch, 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, uint64 withdrawalEpoch, 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 +65,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,11 +78,14 @@ 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)
         external
         onlyGateway
+        returns (uint256 depositAmount)
     {
         bytes32 validatorPubkey = validatorContainer.getPubkey();
         bytes32 withdrawalCredentials = validatorContainer.getWithdrawalCredentials();
@@ -90,81 +116,90 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
         validator.status = VALIDATOR_STATUS.REGISTERED;
         validator.validatorIndex = proof.validatorIndex;
         validator.mostRecentBalanceUpdateTimestamp = proof.beaconBlockTimestamp;
-        validator.restakedBalanceGwei = validatorContainer.getEffectiveBalance();
-
-        _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);
+        uint64 depositAmountGwei = validatorContainer.getEffectiveBalance();
+        if (depositAmountGwei > MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR) {
+            validator.restakedBalanceGwei = MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR;
+            depositAmount = MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR * GWEI_TO_WEI;
+        } else {
+            validator.restakedBalanceGwei = depositAmountGwei;
+            depositAmount = depositAmountGwei * GWEI_TO_WEI;
         }
 
-        if (!partialWithdrawal) {
-            revert NotPartialWithdrawal(validatorPubkey);
-        }
-
-        if (validatorProof.beaconBlockTimestamp != withdrawalProof.beaconBlockTimestamp) {
-            revert UnmatchedValidatorAndWithdrawal(validatorPubkey);
-        }
-
-        _verifyValidatorContainer(validatorContainer, validatorProof);
-        _verifyWithdrawalContainer(withdrawalContainer, withdrawalProof);
+        _capsuleValidatorsByIndex[proof.validatorIndex] = validatorPubkey;
     }
 
-    function verifyFullWithdrawalProof(
+    function verifyWithdrawalProof(
         bytes32[] calldata validatorContainer,
         ValidatorContainerProof calldata validatorProof,
         bytes32[] calldata withdrawalContainer,
-        WithdrawalContainerProof calldata withdrawalProof
-    ) external onlyGateway {
+        BeaconChainProofs.WithdrawalProof calldata withdrawalProof
+    ) 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;
+        uint64 withdrawalEpoch = withdrawalProof.slotRoot.getWithdrawalEpoch();
+        partialWithdrawal = withdrawalEpoch < validatorContainer.getWithdrawableEpoch();
 
         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.withdrawalIndex]) {
+            revert WithdrawalAlreadyProven(validatorPubkey, withdrawalProof.withdrawalIndex);
         }
 
+        provenWithdrawal[validatorPubkey][withdrawalProof.withdrawalIndex] = 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, withdrawalEpoch, 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, withdrawalEpoch, capsuleOwner, withdrawalAmountGwei);
+            if (withdrawalAmountGwei > MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR) {
+                uint256 amountToSend = (withdrawalAmountGwei - MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR) * GWEI_TO_WEI;
+                _sendETH(capsuleOwner, amountToSend);
+                withdrawalAmount = MAX_RESTAKED_BALANCE_GWEI_PER_VALIDATOR * GWEI_TO_WEI;
+            } else {
+                withdrawalAmount = withdrawalAmountGwei * GWEI_TO_WEI;
+            }
+        }
     }
 
     function withdraw(uint256 amount, address payable recipient) external onlyGateway {
         require(recipient != address(0), "ExoCapsule: recipient address cannot be zero or empty");
         require(amount > 0 && amount <= withdrawableBalance, "ExoCapsule: invalid withdrawal amount");
 
         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"
+        );
+        require(recipient != address(0), "ExoCapsule: recipient address cannot be zero or empty");
+
+        nonBeaconChainETHBalance -= amountToWithdraw;
+        _sendETH(recipient, amountToWithdraw);
+        emit NonBeaconChainETHWithdrawn(recipient, amountToWithdraw);
+    }
+
     function updatePrincipalBalance(uint256 lastlyUpdatedPrincipalBalance) external onlyGateway {
         principalBalance = lastlyUpdatedPrincipalBalance;
 
@@ -214,6 +249,13 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
         return validator;
     }
 
+    function _sendETH(address recipient, uint256 amountWei) internal nonReentrant {
+        (bool sent,) = recipient.call{value: amountWei}("");
+        if (!sent) {
+            revert WithdrawalFailure(capsuleOwner, recipient, amountWei);
+        }
+    }
+
     function _verifyValidatorContainer(bytes32[] calldata validatorContainer, ValidatorContainerProof calldata proof)
         internal
         view
@@ -232,19 +274,13 @@ contract ExoCapsule is Initializable, ExoCapsuleStorage, IExoCapsule {
         }
     }
 
-    function _verifyWithdrawalContainer(bytes32[] calldata withdrawalContainer, WithdrawalContainerProof calldata proof)
-        internal
-        view
-    {
-        bytes32 beaconBlockRoot = getBeaconBlockRoot(proof.beaconBlockTimestamp);
+    function _verifyWithdrawalContainer(
+        bytes32[] calldata withdrawalContainer,
+        BeaconChainProofs.WithdrawalProof calldata proof
+    ) internal view {
+        // To-do check withdrawalContainer length is valid
         bytes32 withdrawalContainerRoot = withdrawalContainer.merklelizeWithdrawalContainer();
-        bool valid = withdrawalContainerRoot.isValidWithdrawalContainerRoot(
-            proof.withdrawalContainerRootProof,
-            proof.withdrawalIndex,
-            beaconBlockRoot,
-            proof.executionPayloadRoot,
-            proof.executionPayloadRootProof
-        );
+        bool valid = withdrawalContainerRoot.isValidWithdrawalContainerRoot(proof);
         if (!valid) {
             revert InvalidWithdrawalContainer(withdrawalContainer.getValidatorIndex());
         }
@@ -257,9 +293,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) {