Skip to content

Commit

Permalink
refactor: cleanup logic
Browse files Browse the repository at this point in the history
  • Loading branch information
@failingtwice
failingtwice committed Jan 29, 2025
1 parent 938a519 commit 34d203d
Show file tree
Hide file tree
Showing 4 changed files with 111 additions and 116 deletions.
180 changes: 107 additions & 73 deletions contracts/0.8.25/vaults/PredepositDepositGuardian.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,121 +7,155 @@ pragma solidity 0.8.25;
import {StakingVault} from "./StakingVault.sol";

contract PredepositDepositGuardian {
enum ValidatorStatus {
NO_RECORD,
AWAITING_PROOF,
RESOLVED,
WITHDRAWN
}

mapping(address nodeOperator => bytes32 validatorPubkeyHash) public nodeOperatorToValidators;
mapping(bytes32 validatorPubkeyHash => ValidatorStatus validatorStatus) public validatorStatuses;
mapping(bytes32 validatorPubkeyHash => bytes32 withdrawalCredentials) public wcRecords;

function predeposit(address stakingVault, StakingVault.Deposit[] calldata deposits) external payable {
if (msg.value % 1 ether != 0) revert PredepositMustBeMultipleOfOneEther();
if (msg.value / 1 ether != deposits.length) revert PredepositMustBeOneEtherPerDeposit();
if (msg.sender != StakingVault(payable(stakingVault)).nodeOperator()) revert MustBeNodeOperatorOfStakingVault();
uint256 public constant PREDEPOSIT_AMOUNT = 1 ether;

mapping(bytes32 validatorId => bool isPreDeposited) public validatorPredeposits;
mapping(address nodeOperator => bytes32 validatorId) public nodeOperatorValidators;
mapping(bytes32 validatorId => bytes32 withdrawalCredentials) public validatorWithdrawalCredentials;

function predeposit(
StakingVault stakingVault,
address nodeOperator,
StakingVault.Deposit[] calldata deposits
) external payable {
if (deposits.length == 0) revert PredepositNoDeposits();
if (msg.value % PREDEPOSIT_AMOUNT != 0) revert PredepositValueNotMultipleOfOneEther();
if (msg.value / PREDEPOSIT_AMOUNT != deposits.length) revert PredepositValueNotMatchingNumberOfDeposits();
if (nodeOperator != stakingVault.nodeOperator()) revert PredepositNodeOperatorNotMatching();

for (uint256 i = 0; i < deposits.length; i++) {
StakingVault.Deposit calldata deposit = deposits[i];

if (validatorStatuses[keccak256(deposit.pubkey)] != ValidatorStatus.AWAITING_PROOF) {
revert MustBeNewValidatorPubkey();
bytes32 validatorId = keccak256(deposit.pubkey);

// cannot predeposit a validator that is already predeposited
if (validatorPredeposits[validatorId]) {
revert PredepositValidatorAlreadyPredeposited();
}

// cannot predeposit a validator that has withdrawal credentials already proven
if (validatorWithdrawalCredentials[validatorId] != bytes32(0)) {
revert PredepositValidatorWithdrawalCredentialsAlreadyProven();
}

nodeOperatorToValidators[msg.sender] = keccak256(deposit.pubkey);
validatorStatuses[keccak256(deposit.pubkey)] = ValidatorStatus.AWAITING_PROOF;
validatorPredeposits[validatorId] = true;
nodeOperatorValidators[nodeOperator] = validatorId;

if (deposit.amount != 1 ether) revert PredepositMustBeOneEtherPerDeposit();
if (deposit.amount != PREDEPOSIT_AMOUNT) revert PredepositDepositAmountInvalid();
}

// we don't need to pass deposit root or signature because the msg.sender is deposit guardian itself
StakingVault(payable(stakingVault)).depositToBeaconChain(deposits, bytes32(0), bytes(""));
stakingVault.depositToBeaconChain(deposits);
}

function proveWithdrawalCredentials(
bytes32[] calldata proof,
bytes calldata validatorPubkey,
bytes32 withdrawalCredentials
function proveValidatorWithdrawalCredentials(
bytes32[] calldata /* proof */,
bytes calldata _pubkey,
bytes32 _withdrawalCredentials
) external {
// TODO: proof logic
// revert if proof is invalid

bytes32 validatorPubkeyHash = keccak256(validatorPubkey);
wcRecords[validatorPubkeyHash] = withdrawalCredentials;
validatorStatuses[validatorPubkeyHash] = ValidatorStatus.RESOLVED;
validatorWithdrawalCredentials[keccak256(_pubkey)] = _withdrawalCredentials;
}

function deposit(address _stakingVault, StakingVault.Deposit[] calldata deposits) external payable {
if (msg.sender != StakingVault(payable(_stakingVault)).nodeOperator())
revert MustBeNodeOperatorOfStakingVault();
function depositToProvenValidators(
StakingVault _stakingVault,
StakingVault.Deposit[] calldata _deposits
) external payable {
if (msg.sender != _stakingVault.nodeOperator()) revert DepositSenderNotNodeOperator();

for (uint256 i = 0; i < deposits.length; i++) {
StakingVault.Deposit calldata deposit = deposits[i];
for (uint256 i = 0; i < _deposits.length; i++) {
StakingVault.Deposit calldata deposit = _deposits[i];
bytes32 validatorId = keccak256(deposit.pubkey);

if (validatorStatuses[keccak256(deposit.pubkey)] != ValidatorStatus.RESOLVED) {
revert MustBeResolvedValidatorPubkey();
if (validatorWithdrawalCredentials[validatorId] != _stakingVault.withdrawalCredentials()) {
revert DepositToUnprovenValidator();
}
}

// we don't need to pass deposit root or signature because the msg.sender is deposit guardian itself
StakingVault(payable(_stakingVault)).depositToBeaconChain(deposits, bytes32(0), bytes(""));
_stakingVault.depositToBeaconChain(_deposits);
}

function withdrawAsVaultOwner(address stakingVault, bytes[] calldata validatorPubkeys) external {
if (msg.sender != StakingVault(payable(stakingVault)).owner()) revert MustBeVaultOwner();

for (uint256 i = 0; i < validatorPubkeys.length; i++) {
bytes32 validatorPubkeyHash = keccak256(validatorPubkeys[i]);
// called by the staking vault owner if the predeposited validator has a different withdrawal credentials than the vault's withdrawal credentials,
// i.e. node operator was malicious
function withdrawDisprovenPredeposits(
StakingVault _stakingVault,
bytes32[] calldata _validatorIds,
address _recipient
) external {
if (msg.sender != _stakingVault.owner()) revert WithdrawSenderNotStakingVaultOwner();
if (_recipient == address(0)) revert WithdrawRecipientZeroAddress();

if (validatorStatuses[validatorPubkeyHash] != ValidatorStatus.RESOLVED) {
revert MustBeResolvedValidatorPubkey();
}
uint256 validatorsLength = _validatorIds.length;
for (uint256 i = 0; i < validatorsLength; i++) {
bytes32 validatorId = _validatorIds[i];

if (validatorStatuses[validatorPubkeyHash] == ValidatorStatus.WITHDRAWN) {
revert ValidatorAlreadyWithdrawn();
// cannot withdraw predeposit for a validator that is not pre-deposited
if (!validatorPredeposits[validatorId]) {
revert WithdrawValidatorNotPreDeposited();
}

if (wcRecords[validatorPubkeyHash] == StakingVault(payable(stakingVault)).withdrawalCredentials()) {
revert ValidatorWithdrawalCredentialsMatchVaultWithdrawalCredentials();
// cannot withdraw predeposit for a validator that has withdrawal credentials matching the vault's withdrawal credentials
if (validatorWithdrawalCredentials[validatorId] == _stakingVault.withdrawalCredentials()) {
revert WithdrawValidatorWithdrawalCredentialsMatchStakingVault();
}

msg.sender.call{value: 1 ether}("");
// set flag to false to prevent double withdrawal
validatorPredeposits[validatorId] = false;

validatorStatuses[validatorPubkeyHash] = ValidatorStatus.WITHDRAWN;
(bool success, ) = _recipient.call{value: 1 ether}("");
if (!success) revert WithdrawValidatorTransferFailed();
}
}

function withdrawAsNodeOperator(bytes[] calldata validatorPubkeys) external {
for (uint256 i = 0; i < validatorPubkeys.length; i++) {
bytes32 validatorPubkeyHash = keccak256(validatorPubkeys[i]);
// called by the node operator if the predeposited validator has the same withdrawal credentials as the vault's withdrawal credentials,
// i.e. node operator was honest
function withdrawProvenPredeposits(
StakingVault _stakingVault,
bytes32[] calldata _validatorIds,
address _recipient
) external {
uint256 validatorsLength = _validatorIds.length;
for (uint256 i = 0; i < validatorsLength; i++) {
bytes32 validatorId = _validatorIds[i];

if (validatorStatuses[validatorPubkeyHash] != ValidatorStatus.RESOLVED) {
revert MustBeResolvedValidatorPubkey();
if (msg.sender != _stakingVault.nodeOperator()) {
revert WithdrawSenderNotNodeOperator();
}

if (validatorStatuses[validatorPubkeyHash] == ValidatorStatus.WITHDRAWN) {
revert ValidatorAlreadyWithdrawn();
// cannot withdraw predeposit for a validator that is not pre-deposited
if (!validatorPredeposits[validatorId]) {
revert WithdrawValidatorNotPreDeposited();
}

if (nodeOperatorToValidators[msg.sender] != validatorPubkeyHash) {
revert ValidatorMustBelongToSender();
// cannot withdraw predeposit for a validator that has withdrawal credentials not matching the vault's withdrawal credentials
if (validatorWithdrawalCredentials[validatorId] != _stakingVault.withdrawalCredentials()) {
revert WithdrawValidatorWithdrawalCredentialsNotMatchingStakingVault();
}

msg.sender.call{value: 1 ether}("");
// set flag to false to prevent double withdrawal
validatorPredeposits[validatorId] = false;

validatorStatuses[validatorPubkeyHash] = ValidatorStatus.WITHDRAWN;
(bool success, ) = _recipient.call{value: 1 ether}("");
if (!success) revert WithdrawValidatorTransferFailed();
}
}

error PredepositMustBeMultipleOfOneEther();
error PredepositMustBeOneEtherPerDeposit();
error MustBeNodeOperatorOfStakingVault();
error MustBeNewValidatorPubkey();
error WithdrawalFailed();
error MustBeResolvedValidatorPubkey();
error ValidatorMustBelongToSender();
error MustBeVaultOwner();
error ValidatorWithdrawalCredentialsMatchVaultWithdrawalCredentials();
error ValidatorAlreadyWithdrawn();
error PredepositNoDeposits();
error PredepositValueNotMultipleOfOneEther();
error PredepositValueNotMatchingNumberOfDeposits();
error PredepositNodeOperatorNotMatching();
error PredepositValidatorAlreadyPredeposited();
error PredepositValidatorWithdrawalCredentialsAlreadyProven();
error PredepositDepositAmountInvalid();
error ValidatorNotPreDeposited();
error DepositSenderNotNodeOperator();
error DepositToUnprovenValidator();
error WithdrawSenderNotStakingVaultOwner();
error WithdrawRecipientZeroAddress();
error WithdrawValidatorNotPreDeposited();
error WithdrawValidatorWithdrawalCredentialsMatchStakingVault();
error WithdrawValidatorTransferFailed();
error WithdrawValidatorWithdrawalCredentialsNotMatchingStakingVault();
error WithdrawSenderNotNodeOperator();
}
35 changes: 2 additions & 33 deletions contracts/0.8.25/vaults/StakingVault.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -318,47 +318,16 @@ contract StakingVault is IStakingVault, OwnableUpgradeable {
* @param _deposits Array of deposit structs
* @dev Includes a check to ensure StakingVault is balanced before making deposits
*/
function depositToBeaconChain(
Deposit[] calldata _deposits,
bytes32 _expectedGlobalDepositRoot,
bytes calldata _signature
) external {
function depositToBeaconChain(Deposit[] calldata _deposits) external {
if (_deposits.length == 0) revert ZeroArgument("_deposits");
if (!isBalanced()) revert Unbalanced();

ERC7201Storage storage $ = _getStorage();
if ($.beaconChainDepositsPaused) revert BeaconChainDepositsArePaused();
if (msg.sender != $.depositGuardian) revert NotAuthorized("depositToBeaconChain", msg.sender);

uint256 numberOfDeposits = _deposits.length;

if (msg.sender != $.depositGuardian) {
bytes32 currentGlobalDepositRoot = BEACON_CHAIN_DEPOSIT_CONTRACT.get_deposit_root();
if (_expectedGlobalDepositRoot != currentGlobalDepositRoot)
revert GlobalDepositRootMismatch(_expectedGlobalDepositRoot, currentGlobalDepositRoot);

bytes32 depositDataBatchXorRoot;

for (uint256 i = 0; i < numberOfDeposits; i++) {
Deposit calldata deposit = _deposits[i];

depositDataBatchXorRoot ^= keccak256(abi.encodePacked(deposit.depositDataRoot));
}

if (
!SignatureChecker.isValidSignatureNow(
$.depositGuardian,
keccak256(
abi.encodePacked(
DEPOSIT_GUARDIAN_MESSAGE_PREFIX,
_expectedGlobalDepositRoot,
depositDataBatchXorRoot
)
),
_signature
)
) revert DepositGuardianSignatureInvalid();
}

uint256 totalAmount = 0;

for (uint256 i = 0; i < numberOfDeposits; i++) {
Expand Down
6 changes: 1 addition & 5 deletions contracts/0.8.25/vaults/interfaces/IStakingVault.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,11 +57,7 @@ interface IStakingVault {

function withdraw(address _recipient, uint256 _ether) external;

function depositToBeaconChain(
Deposit[] calldata _deposits,
bytes32 _expectedGlobalDepositRoot,
bytes calldata _guardianSignature
) external;
function depositToBeaconChain(Deposit[] calldata _deposits) external;

function requestValidatorExit(bytes calldata _pubkeys) external;

Expand Down
6 changes: 1 addition & 5 deletions test/0.8.25/vaults/contracts/StakingVault__HarnessForTestUpgrade.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,11 +84,7 @@ contract StakingVault__HarnessForTestUpgrade is IStakingVault, OwnableUpgradeabl
}
}

function depositToBeaconChain(
Deposit[] calldata _deposits,
bytes32 _expectedGlobalDepositRoot,
bytes calldata _guardianSignature
) external {}
function depositToBeaconChain(Deposit[] calldata _deposits) external {}

function fund() external payable {}

Expand Down

0 comments on commit 34d203d

Please sign in to comment.