diff --git a/contracts/0.8.25/lib/SSZ.sol b/contracts/0.8.25/lib/SSZ.sol index 67f057ccd..59cf7f68d 100644 --- a/contracts/0.8.25/lib/SSZ.sol +++ b/contracts/0.8.25/lib/SSZ.sol @@ -1,124 +1,20 @@ // SPDX-FileCopyrightText: 2024 Lido // SPDX-License-Identifier: GPL-3.0 -/* - Cut version of SSZ library from CSM, only supports Validator container - original: https://github.com/lidofinance/community-staking-module/blob/7071c2096983a7780a5f147963aaa5405c0badb1/src/lib/SSZ.sol -*/ - // See contracts/COMPILERS.md pragma solidity 0.8.25; import {GIndex} from "./GIndex.sol"; -// As defined in phase0/beacon-chain.md:356 -struct Validator { - bytes pubkey; - bytes32 withdrawalCredentials; - uint64 effectiveBalance; - bool slashed; - uint64 activationEligibilityEpoch; - uint64 activationEpoch; - uint64 exitEpoch; - uint64 withdrawableEpoch; -} - +/* + Cut and modified version of SSZ library from CSM only has methods for merkilized SSZ proof validation + original: https://github.com/lidofinance/community-staking-module/blob/7071c2096983a7780a5f147963aaa5405c0badb1/src/lib/SSZ.sol +*/ library SSZ { error BranchHasMissingItem(); error BranchHasExtraItem(); error InvalidProof(); - - function hashTreeRoot(Validator calldata validator) internal view returns (bytes32 root) { - bytes32 pubkeyRoot; - - assembly { - // In calldata, a dynamic field is encoded as an offset (relative to the start - // of the struct’s calldata) followed by its contents. The first 32 bytes of - // `validator` is the offset for `pubkey`. (Remember that `pubkey` is expected - // to be exactly 48 bytes long.) - let pubkeyOffset := calldataload(validator) - // The pubkey’s actual data is encoded at: - // validator + pubkeyOffset + 32 - // because the first word at that location is the length. - // Copy 48 bytes of pubkey data into memory at 0x00. - calldatacopy(0x00, add(validator, add(pubkeyOffset, 32)), 48) - // Zero the remaining 16 bytes to form a 64‐byte block. - // (0x30 = 48, so mstore at 0x30 will zero 32 bytes covering addresses 48–79; - // only bytes 48–63 matter for our 64-byte input.) - mstore(0x30, 0) - // Call the SHA‑256 precompile (at address 0x02) with the 64-byte block. - if iszero(staticcall(gas(), 0x02, 0x00, 0x40, 0x00, 0x20)) { - revert(0, 0) - } - pubkeyRoot := mload(0x00) - } - - bytes32[8] memory nodes = [ - pubkeyRoot, - validator.withdrawalCredentials, - toLittleEndian(validator.effectiveBalance), - toLittleEndian(validator.slashed), - toLittleEndian(validator.activationEligibilityEpoch), - toLittleEndian(validator.activationEpoch), - toLittleEndian(validator.exitEpoch), - toLittleEndian(validator.withdrawableEpoch) - ]; - - /// @solidity memory-safe-assembly - assembly { - // Count of nodes to hash - let count := 8 - - // Loop over levels - // prettier-ignore - for { } 1 { } { - // Loop over nodes at the given depth - - // Initialize `offset` to the offset of `proof` elements in memory. - let target := nodes - let source := nodes - let end := add(source, shl(5, count)) - - // prettier-ignore - for { } 1 { } { - // Read next two hashes to hash - mcopy(0x00, source, 0x40) - - // Call sha256 precompile - let result := staticcall( - gas(), - 0x02, - 0x00, - 0x40, - 0x00, - 0x20 - ) - - if iszero(result) { - // Precompiles returns no data on OutOfGas error. - revert(0, 0) - } - - // Store the resulting hash at the target location - mstore(target, mload(0x00)) - - // Advance the pointers - target := add(target, 0x20) - source := add(source, 0x40) - - if iszero(lt(source, end)) { - break - } - } - - count := shr(1, count) - if eq(count, 1) { - root := mload(0x00) - break - } - } - } - } + error InvalidPubkeyLength(); /// @notice Modified version of `verify` from Solady `MerkleProofLib` to support generalized indices and sha256 precompile. /// @dev Reverts if `leaf` doesn't exist in the Merkle tree with `root`, given `proof`. @@ -190,25 +86,47 @@ library SSZ { } } - // See https://github.com/succinctlabs/telepathy-contracts/blob/5aa4bb7/src/libraries/SimpleSerialize.sol#L17-L28 - function toLittleEndian(uint256 v) internal pure returns (bytes32) { - v = - ((v & 0xFF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00) >> 8) | - ((v & 0x00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF) << 8); - v = - ((v & 0xFFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000) >> 16) | - ((v & 0x0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF) << 16); - v = - ((v & 0xFFFFFFFF00000000FFFFFFFF00000000FFFFFFFF00000000FFFFFFFF00000000) >> 32) | - ((v & 0x00000000FFFFFFFF00000000FFFFFFFF00000000FFFFFFFF00000000FFFFFFFF) << 32); - v = - ((v & 0xFFFFFFFFFFFFFFFF0000000000000000FFFFFFFFFFFFFFFF0000000000000000) >> 64) | - ((v & 0x0000000000000000FFFFFFFFFFFFFFFF0000000000000000FFFFFFFFFFFFFFFF) << 64); - v = (v >> 128) | (v << 128); - return bytes32(v); + /// @notice Extracted part from `verifyProof` for hashing two leaves + /// @dev Combines 2 bytes32 in 64 bytes input for sha256 precompile + function sha256Pair(bytes32 left, bytes32 right) internal view returns (bytes32 result) { + /// @solidity memory-safe-assembly + assembly { + // Store `left` at memory position 0x00 + mstore(0x00, left) + // Store `right` at memory position 0x20 + mstore(0x20, right) + + // Call SHA-256 precompile (0x02) with 64-byte input at memory 0x00 + let success := staticcall(gas(), 0x02, 0x00, 0x40, 0x00, 0x20) + if iszero(success) { + revert(0, 0) + } + + // Load the resulting hash from memory + result := mload(0x00) + } } - function toLittleEndian(bool v) internal pure returns (bytes32) { - return bytes32(v ? 1 << 248 : 0); + /// @notice Extracted and modified part from `hashTreeRoot` for hashing validator pubkey from calldata + /// @dev Reverts if `pubkey` length is not 48 + function pubkeyRoot(bytes calldata pubkey) internal view returns (bytes32 _pubkeyRoot) { + if (pubkey.length != 48) revert InvalidPubkeyLength(); + + /// @solidity memory-safe-assembly + assembly { + // Copy 48 bytes of `pubkey` to memory at 0x00 + calldatacopy(0x00, pubkey.offset, 48) + + // Zero the remaining 16 bytes to form a 64-byte input block + mstore(0x30, 0) + + // Call the SHA-256 precompile (0x02) with the 64-byte input + if iszero(staticcall(gas(), 0x02, 0x00, 0x40, 0x00, 0x20)) { + revert(0, 0) + } + + // Load the resulting SHA-256 hash + _pubkeyRoot := mload(0x00) + } } } diff --git a/contracts/0.8.25/vaults/predeposit_guarantee/CLProofVerifier.sol b/contracts/0.8.25/vaults/predeposit_guarantee/CLProofVerifier.sol index 01a6e2041..4f0022632 100644 --- a/contracts/0.8.25/vaults/predeposit_guarantee/CLProofVerifier.sol +++ b/contracts/0.8.25/vaults/predeposit_guarantee/CLProofVerifier.sol @@ -4,38 +4,79 @@ // See contracts/COMPILERS.md pragma solidity 0.8.25; -import {Validator, SSZ, GIndex} from "contracts/0.8.25/lib/SSZ.sol"; - -struct ValidatorWitness { - Validator validator; - bytes32[] proof; - uint256 validatorIndex; - uint64 beaconBlockTimestamp; -} +import {GIndex, pack, concat} from "contracts/0.8.25/lib/GIndex.sol"; +import {SSZ} from "contracts/0.8.25/lib/SSZ.sol"; +/** + * @title CLProofVerifier + * @author Lido + * @notice + * + * CLProofVerifier is base abstract contract that provides internal method to verify + * merkle proofs of validator entry in CL. It uses concatenated proofs that prove + * validator existence in CL just from pubkey and withdrawalCredentials againts Beacon block root + * stored in BeaconRoots contract. + * + * + * NB!: GI_FIRST_VALIDATOR must be updated if Ethereum hardfork changes order of CL state tree + * (e.g. Pectra, Altair, etc.) + * + */ abstract contract CLProofVerifier { - using SSZ for Validator; + struct ValidatorWitness { + bytes32[] proof; + bytes pubkey; + uint256 validatorIndex; + uint64 childBlockTimestamp; + } + // See `BEACON_ROOTS_ADDRESS` constant in the EIP-4788. - address public immutable BEACON_ROOTS; + address public immutable BEACON_ROOTS = 0x000F3df6D732807Ef1319fB7B8bB8522d0Beac02; + + // Index of parent node for (Pubkey,WC) in validator container + GIndex public immutable GI_PUBKEY_WC_PARENT = pack(1 << 2, 2); + // Index of stateRoot in Beacon Block state + GIndex public immutable GI_STATE_VIEW = pack((1 << 3) + 3, 3); + // Index of first validator in CL state GIndex public immutable GI_FIRST_VALIDATOR; constructor(GIndex _gIFirstValidator) { - BEACON_ROOTS = 0x000F3df6D732807Ef1319fB7B8bB8522d0Beac02; GI_FIRST_VALIDATOR = _gIFirstValidator; } - function _validateWCProof(ValidatorWitness calldata _witness) internal view { + /** + * @notice validates proof of validator in CL with withdrawalCredentials and pubkey against Beacon block root + * @param _witness object containing user input passed as calldata + * `proof` - array of hashes for concatenated merkle proof from parent(pubkey,wc) node to the Beacon block root + * `pubkey` - pubkey of the validator + * `validatorIndex` - numerical index of validator in CL + * `childBlockTimestamp` - timestamp of EL block that has Beacon root corresponding to proof + * @param _withdrawalCredentials to verify proof with + * @dev reverts with `InvalidProof` when provided input cannot be proven to Beacon block root + */ + function _validatePubKeyWCProof(ValidatorWitness calldata _witness, bytes32 _withdrawalCredentials) internal view { + // parent node for first two leaves in validator container tree + // pubkey + wc + bytes32 _leaf = SSZ.sha256Pair(SSZ.pubkeyRoot(_witness.pubkey), _withdrawalCredentials); + // concatenated index for parent(pubkey + wc) -> Validator Index in state tree -> stateView Index in Beacon block Tree + GIndex _gIndex = concat(GI_STATE_VIEW, concat(_getValidatorGI(_witness.validatorIndex), GI_PUBKEY_WC_PARENT)); + SSZ.verifyProof({ proof: _witness.proof, - root: _getParentBlockRoot(_witness.beaconBlockTimestamp), - leaf: _witness.validator.hashTreeRoot(), - gIndex: _getValidatorGI(_witness.validatorIndex) + root: _getParentBlockRoot(_witness.childBlockTimestamp), + leaf: _leaf, + gIndex: _gIndex }); } - // virtual for testing - function _getParentBlockRoot(uint64 blockTimestamp) internal view virtual returns (bytes32) { - (bool success, bytes memory data) = BEACON_ROOTS.staticcall(abi.encode(blockTimestamp)); + /** + * @notice returns parent CL block root for given child block timestamp + * @param _childBlockTimestamp timestamp of child block + * @return parent block root + * @dev reverts with `RootNotFound` if timestamp is not found in Beacon Block roots + */ + function _getParentBlockRoot(uint64 _childBlockTimestamp) internal view returns (bytes32) { + (bool success, bytes memory data) = BEACON_ROOTS.staticcall(abi.encode(_childBlockTimestamp)); if (!success || data.length == 0) { revert RootNotFound(); @@ -44,11 +85,14 @@ abstract contract CLProofVerifier { return abi.decode(data, (bytes32)); } - function _getValidatorGI(uint256 offset) internal view returns (GIndex) { - return GI_FIRST_VALIDATOR.shr(offset); + /** + * @notice calculates general validator index in CL state tree by provided offset + * @param _offset from first validator (Validator Index) + * @return gIndex of container in CL state tree + */ + function _getValidatorGI(uint256 _offset) internal view returns (GIndex) { + return GI_FIRST_VALIDATOR.shr(_offset); } - // proving errors - error InvalidGeneralIndex(uint256); error RootNotFound(); } diff --git a/contracts/0.8.25/vaults/predeposit_guarantee/PredepositGuarantee.sol b/contracts/0.8.25/vaults/predeposit_guarantee/PredepositGuarantee.sol index e41f60765..26f5a8cad 100644 --- a/contracts/0.8.25/vaults/predeposit_guarantee/PredepositGuarantee.sol +++ b/contracts/0.8.25/vaults/predeposit_guarantee/PredepositGuarantee.sol @@ -4,7 +4,7 @@ // See contracts/COMPILERS.md pragma solidity 0.8.25; -import {CLProofVerifier, ValidatorWitness, GIndex} from "./CLProofVerifier.sol"; +import {CLProofVerifier, GIndex} from "./CLProofVerifier.sol"; import {IStakingVaultOwnable} from "../interfaces/IStakingVault.sol"; @@ -92,9 +92,9 @@ contract PredepositGuarantee is CLProofVerifier { if (bond.locked != 0) revert BondMustBeFullyUnlocked(); if (bond.total > 0 && nodeOperatorVoucher[msg.sender] != address(0)) { - uint256 bondAmount = bond.total; + uint256 ejected = nodeOperatorBonds[msg.sender].total; nodeOperatorBonds[msg.sender].total = 0; - (bool success, ) = nodeOperatorVoucher[msg.sender].call{value: bondAmount}(""); + (bool success, ) = nodeOperatorVoucher[msg.sender].call{value: ejected}(""); // voucher can block change? if (!success) revert WithdrawalFailed(); @@ -121,6 +121,11 @@ contract PredepositGuarantee is CLProofVerifier { _topUpNodeOperatorCollateral(_nodeOperator, msg.value); } + // ensures vault fair play + if (address(_stakingVault) != _wcToAddress(_stakingVault.withdrawalCredentials())) { + revert stakingVaultWithdrawalCredentialsMismatch(); + } + uint128 totalDepositAmount = PREDEPOSIT_AMOUNT * uint128(_deposits.length); uint256 unlockedCollateral = nodeOperatorBonds[_nodeOperator].total - nodeOperatorBonds[_nodeOperator].locked; @@ -150,8 +155,14 @@ contract PredepositGuarantee is CLProofVerifier { emit ValidatorPreDeposited(_nodeOperator, address(_stakingVault), _deposits.length, totalDepositAmount); } + /* + * + * POSITIVE PROOF METHODS + * + */ + function proveValidatorWC(ValidatorWitness calldata _witness) external { - _processWCProof(_witness); + _processWitnessProof(_witness); } function depositToProvenValidators( @@ -189,15 +200,43 @@ contract PredepositGuarantee is CLProofVerifier { IStakingVaultOwnable _stakingVault ) external payable { for (uint256 i = 0; i < _witnesses.length; i++) { - _processWCProof(_witnesses[i]); + _processWitnessProof(_witnesses[i]); } depositToProvenValidators(_stakingVault, _deposits); } + /* + * + * NEGATIVE PROOF METHODS + * + */ + + function proveInvalidValidatorWC(ValidatorWitness calldata _witness, bytes32 _invalidWithdrawalCredentials) public { + ValidatorStatus storage validatorStatus = validatorStatuses[_witness.pubkey]; + + if (validatorStatus.bondStatus != BondStatus.AWAITING_PROOF) { + revert ValidatorNotPreDeposited(_witness.pubkey, validatorStatus.bondStatus); + } + + if (address(validatorStatus.stakingVault) == _wcToAddress(_invalidWithdrawalCredentials)) { + revert WithdrawalCredentialsAreValid(); + } + + _validatePubKeyWCProof(_witness, _invalidWithdrawalCredentials); + + // reduces total&locked NO deposit + nodeOperatorBonds[validatorStatus.nodeOperator].total -= PREDEPOSIT_AMOUNT; + nodeOperatorBonds[validatorStatus.nodeOperator].locked -= PREDEPOSIT_AMOUNT; + // freed ether only will returned to owner of the vault with this validator + validatorStatus.bondStatus = BondStatus.PROVED_INVALID; + + emit ValidatorDisproven(validatorStatus.nodeOperator, _witness.pubkey, address(validatorStatus.stakingVault), _invalidWithdrawalCredentials); + } + // called by the staking vault owner if the predeposited validator was proven invalid // i.e. node operator was malicious and has stolen vault ether - function withdrawDisprovenCollateral(bytes calldata validatorPubkey, address _recipient) external { + function withdrawDisprovenPredeposit(bytes calldata validatorPubkey, address _recipient) public { ValidatorStatus storage validatorStatus = validatorStatuses[validatorPubkey]; if (_recipient == address(0)) revert ZeroArgument("_recipient"); @@ -217,6 +256,15 @@ contract PredepositGuarantee is CLProofVerifier { emit ValidatorDisprovenWithdrawn(validatorStatus.nodeOperator, validatorPubkey, address(validatorStatus.stakingVault), _recipient); } + function disproveAndWithdraw( + ValidatorWitness calldata _witness, + bytes32 _invalidWithdrawalCredentials, + address _recipient + ) external { + proveInvalidValidatorWC(_witness, _invalidWithdrawalCredentials); + withdrawDisprovenPredeposit(_witness.pubkey, _recipient); + } + /// Internal functions function _validateNodeOperatorCaller(address _nodeOperator) internal view { @@ -236,46 +284,36 @@ contract PredepositGuarantee is CLProofVerifier { emit NodeOperatorBondToppedUp(_nodeOperator, _amount); } - function _wcToAddress(bytes32 _withdrawalCredentials) internal pure returns (address) { - return address(uint160(uint256(_withdrawalCredentials))); - } + function _wcToAddress(bytes32 _withdrawalCredentials) internal pure returns (address _wcAddress) { + uint64 _wcVersion = uint8(_withdrawalCredentials[0]); + + if (_wcVersion < 1) { + revert WithdrawalCredentialsAreInvalid(); + } - function _deconstructWC(bytes32 _withdrawalCredentials) internal pure returns (uint64, address) { - return (uint8(_withdrawalCredentials[0]), address(uint160(uint256(_withdrawalCredentials)))); + _wcAddress = address(uint160(uint256(_withdrawalCredentials))); } - function _processWCProof(ValidatorWitness calldata _witness) internal { - ValidatorStatus storage validatorStatus = validatorStatuses[_witness.validator.pubkey]; + function _processWitnessProof(ValidatorWitness calldata _witness) internal { + ValidatorStatus storage validatorStatus = validatorStatuses[_witness.pubkey]; if (validatorStatus.bondStatus != BondStatus.AWAITING_PROOF) { - revert ValidatorNotPreDeposited(_witness.validator.pubkey, validatorStatus.bondStatus); + revert ValidatorNotPreDeposited(_witness.pubkey, validatorStatus.bondStatus); } - (uint64 _wcVersion, address _wcAddress) = _deconstructWC(_witness.validator.withdrawalCredentials); + bytes32 _withdrawalCredentials = validatorStatus.stakingVault.withdrawalCredentials(); - if (_wcVersion < 1) { - revert WithdrawalCredentialsAreInvalid(_witness.validator.pubkey); + // ensures vault fair play + if (address(validatorStatus.stakingVault) != _wcToAddress(_withdrawalCredentials)) { + revert WithdrawalCredentialsAreInvalid(); } - _validateWCProof(_witness); - - // determine proof direction - if (address(validatorStatus.stakingVault) == _wcAddress) { - // stricter WC check to ensure WC version matches - if (validatorStatus.stakingVault.withdrawalCredentials() != _witness.validator.withdrawalCredentials) { - revert WithdrawalCredentialsAreInvalid(_witness.validator.pubkey); - } - - validatorStatus.bondStatus = BondStatus.PROVED; + _validatePubKeyWCProof(_witness, _withdrawalCredentials); - emit ValidatorProven(validatorStatus.nodeOperator, _witness.validator.pubkey, address(validatorStatus.stakingVault), _witness.validator.withdrawalCredentials); - } else { - validatorStatus.bondStatus = BondStatus.PROVED_INVALID; - nodeOperatorBonds[validatorStatus.nodeOperator].total -= PREDEPOSIT_AMOUNT; - - emit ValidatorDisproven(validatorStatus.nodeOperator, _witness.validator.pubkey, address(validatorStatus.stakingVault), _witness.validator.withdrawalCredentials); - } + validatorStatus.bondStatus = BondStatus.PROVED; nodeOperatorBonds[validatorStatus.nodeOperator].locked -= PREDEPOSIT_AMOUNT; + + emit ValidatorProven(validatorStatus.nodeOperator, _witness.pubkey, address(validatorStatus.stakingVault), _withdrawalCredentials); } // node operator accounting @@ -287,6 +325,8 @@ contract PredepositGuarantee is CLProofVerifier { error PredepositDepositAmountInvalid(bytes validatorPubkey, uint256 depositAmount); error MustBeNewValidatorPubkey(bytes validatorPubkey, BondStatus bondStatus); error NotEnoughUnlockedCollateralToPredeposit(uint256 unlockedCollateral, uint256 totalDepositAmount); + error PredepositValueNotMultipleOfPrediposit(); + error stakingVaultWithdrawalCredentialsMismatch(); // depositing errors error DepositToUnprovenValidator(bytes validatorPubkey, BondStatus bondStatus); @@ -294,7 +334,8 @@ contract PredepositGuarantee is CLProofVerifier { error ValidatorNotPreDeposited(bytes validatorPubkey, BondStatus bondStatus); // prove - error WithdrawalCredentialsAreInvalid(bytes validatorPubkey); + error WithdrawalCredentialsAreInvalid(); + error WithdrawalCredentialsAreValid(); // withdrawal proven error NotEnoughUnlockedCollateralToWithdraw(uint256 unlockedCollateral, uint256 amount); diff --git a/test/0.8.25/vaults/predeposit-guarantee/cl-proof-verifyer.test.ts b/test/0.8.25/vaults/predeposit-guarantee/cl-proof-verifyer.test.ts index 87dad2156..ef283d214 100644 --- a/test/0.8.25/vaults/predeposit-guarantee/cl-proof-verifyer.test.ts +++ b/test/0.8.25/vaults/predeposit-guarantee/cl-proof-verifyer.test.ts @@ -1,15 +1,18 @@ +import { expect } from "chai"; import { hexlify, parseUnits, randomBytes } from "ethers"; import { ethers } from "hardhat"; import { CLProofVerifier__Harness, SSZMerkleTree } from "typechain-types"; -import { ValidatorStruct } from "typechain-types/contracts/0.8.25/predeposit_guarantee/PredepositGuarantee"; -import { ValidatorWitnessStruct } from "typechain-types/contracts/0.8.25/vaults/predeposit_guarantee/PredepositGuarantee"; -export const generateValidator = (customWC?: string, customPukey?: string): ValidatorStruct => { - const randomInt = (max: number): number => Math.floor(Math.random() * max); - const randomBytes32 = (): string => hexlify(randomBytes(32)); - const randomValidatorPubkey = (): string => hexlify(randomBytes(96)); +import { impersonate } from "lib"; +import { Snapshot } from "test/suite"; + +const randomBytes32 = (): string => hexlify(randomBytes(32)); +const randomInt = (max: number): number => Math.floor(Math.random() * max); +const randomValidatorPubkey = (): string => hexlify(randomBytes(48)); + +export const generateValidator = (customWC?: string, customPukey?: string) => { return { pubkey: customPukey ?? randomValidatorPubkey(), withdrawalCredentials: customWC ?? randomBytes32(), @@ -22,12 +25,28 @@ export const generateValidator = (customWC?: string, customPukey?: string): Vali }; }; +export const generateBeaconHeader = (stateRoot: string) => { + return { + slot: randomInt(1743359), + proposerIndex: randomInt(1337), + parentRoot: randomBytes32(), + stateRoot, + bodyRoot: randomBytes32(), + }; +}; + // CSM "borrowed" prefab validator object with mocked state root const STATIC_VALIDATOR = { - // state root mocked as block root because CSM proves against state and verifies block root separately - root: "0x21205c716572ae05692c0f8a4c64fd84e504cbb1a16fa0371701adbab756dd72", + blockRoot: "0x56073a5bf24e8a3ea2033ad10a5039a7a7a6884086b67053c90d38f104ae89cf", // pack(0x560000000000, 40) gIFirstValidator: "0x0000000000000000000000000000000000000000000000000056000000000028", + beaconBlockHeader: { + slot: 1743359, + proposerIndex: 1337, + parentRoot: "0x5db6dfb2b5e735bafb437a76b9e525e958d2aef589649e862bfbc02964edf5ab", + stateRoot: "0x21205c716572ae05692c0f8a4c64fd84e504cbb1a16fa0371701adbab756dd72", + bodyRoot: "0x459390eed4479eb49b71efadcc3b540bbc60073f196e0409588d6cc9eafbe5fa", + }, witness: { validatorIndex: 1551477n, beaconBlockTimestamp: 42, @@ -90,20 +109,55 @@ const STATIC_VALIDATOR = { "0xbb2952772995323016b98233c26e96e5c54955fda62e643cb56981da6aab7365", "0xda5ca7afba0d19d345e85d2825fc3078eefdd76ead776b108fe0eac9aa96e5e6", ], - } as ValidatorWitnessStruct, + }, }; -// random number integer generator - describe("CLProofVerifier.sol", () => { let CLProofVerifier: CLProofVerifier__Harness; let sszMerkleTree: SSZMerkleTree; + let BEACON_ROOTS: string; + let snapshotState: string; + let setBeaconBlockRoot: (root: string) => Promise; before(async () => { sszMerkleTree = await ethers.deployContract("SSZMerkleTree", {}); await sszMerkleTree.addValidatorLeaf(generateValidator()); const gIFirstValidator = await sszMerkleTree.getGeneralizedIndex(0n); + + // populate merkle tree with validators + for (let i = 1; i < 100; i++) { + await sszMerkleTree.addValidatorLeaf(generateValidator()); + } + CLProofVerifier = await ethers.deployContract("CLProofVerifier__Harness", [gIFirstValidator], {}); + BEACON_ROOTS = await CLProofVerifier.BEACON_ROOTS(); + + const systemSigner = await impersonate("0xfffffffffffffffffffffffffffffffffffffffe", 999999999999999999999999999n); + + setBeaconBlockRoot = async (root: string) => { + const block = await systemSigner + .sendTransaction({ + to: BEACON_ROOTS, + value: 0, + data: root, + }) + .then((tx) => tx.getBlock()); + if (!block) throw new Error("ivariant"); + return block.timestamp; + }; + + // test mocker + const mockRoot = randomBytes32(); + const timestamp = await setBeaconBlockRoot(mockRoot); + expect(await CLProofVerifier.TEST_getParentBlockRoot(timestamp)).to.equal(mockRoot); + }); + + beforeEach(async () => { + snapshotState = await Snapshot.take(); + }); + + afterEach(async () => { + await Snapshot.restore(snapshotState); }); it("should verify precalclulated validator object in merkle tree", async () => { @@ -112,26 +166,82 @@ describe("CLProofVerifier.sol", () => { [STATIC_VALIDATOR.gIFirstValidator], {}, ); - await StaticCLProofVerifier.setRoot(STATIC_VALIDATOR.root); - await StaticCLProofVerifier.TEST_validateWCProof(STATIC_VALIDATOR.witness); + + const validatorMerkle = await sszMerkleTree.getValidatorPubkeyWCParentProof(STATIC_VALIDATOR.witness.validator); + const beaconHeaderMerkle = await sszMerkleTree.getBeaconBlockHeaderProof(STATIC_VALIDATOR.beaconBlockHeader); + const validatorGIndex = await StaticCLProofVerifier.TEST_getValidatorGI(STATIC_VALIDATOR.witness.validatorIndex); + + // raw proof verification with same input as CSM + await sszMerkleTree.verifyProof( + STATIC_VALIDATOR.witness.proof, + STATIC_VALIDATOR.beaconBlockHeader.stateRoot, + validatorMerkle.root, + validatorGIndex, + ); + + // concatentate all proofs to match PG style + const concatenatedProof = [ + ...validatorMerkle.proof, + ...STATIC_VALIDATOR.witness.proof, + ...beaconHeaderMerkle.proof, + ]; + + const timestamp = await setBeaconBlockRoot(STATIC_VALIDATOR.blockRoot); + + // PG style proof verification from PK+WC to BeaconBlockRoot + await StaticCLProofVerifier.TEST_validatePubKeyWCProof( + { + proof: concatenatedProof, + pubkey: STATIC_VALIDATOR.witness.validator.pubkey, + validatorIndex: STATIC_VALIDATOR.witness.validatorIndex, + childBlockTimestamp: timestamp, + }, + STATIC_VALIDATOR.witness.validator.withdrawalCredentials, + ); }); it("can verify against dynamic merkle tree", async () => { const validator = generateValidator(); + const validatorMerkle = await sszMerkleTree.getValidatorPubkeyWCParentProof(validator); - await sszMerkleTree.addValidatorLeaf(validator); + // verify just the validator container tree from PK+WC node + await sszMerkleTree.verifyProof( + [...validatorMerkle.proof], + validatorMerkle.root, + validatorMerkle.parentNode, + validatorMerkle.parentIndex, + ); + // add validator to CL state merkle tree + await sszMerkleTree.addValidatorLeaf(validator); + const stateRoot = await sszMerkleTree.getMerkleRoot(); const validatorIndex = (await sszMerkleTree.leafCount()) - 1n; - const proof = await sszMerkleTree.getMerkleProof(validatorIndex); - const root = await sszMerkleTree.getMerkleRoot(); + const stateProof = await sszMerkleTree.getMerkleProof(validatorIndex); + const validatorGIndex = await sszMerkleTree.getGeneralizedIndex(validatorIndex); + + // verify just the state tree + await sszMerkleTree.verifyProof([...stateProof], stateRoot, validatorMerkle.root, validatorGIndex); - await CLProofVerifier.setRoot(root); + const beaconHeader = generateBeaconHeader(stateRoot); + const beaconMerkle = await sszMerkleTree.getBeaconBlockHeaderProof(beaconHeader); + // verify just the beacon tree + await sszMerkleTree.verifyProof([...beaconMerkle.proof], beaconMerkle.root, stateRoot, beaconMerkle.index); - await CLProofVerifier.TEST_validateWCProof({ - validatorIndex, - proof: [...proof], - validator, - beaconBlockTimestamp: 1n, - }); + const timestamp = await setBeaconBlockRoot(beaconMerkle.root); + + const proof = [...validatorMerkle.proof, ...stateProof, ...beaconMerkle.proof]; + await CLProofVerifier.TEST_validatePubKeyWCProof( + { + validatorIndex, + proof: [...proof], + pubkey: validator.pubkey, + childBlockTimestamp: timestamp, + }, + validator.withdrawalCredentials, + ); }); + + /* + TODO: negative tests + */ }); diff --git a/test/0.8.25/vaults/predeposit-guarantee/contracts/CLProofVerifier__harness.sol b/test/0.8.25/vaults/predeposit-guarantee/contracts/CLProofVerifier__harness.sol index 9521cac8e..795b38bf3 100644 --- a/test/0.8.25/vaults/predeposit-guarantee/contracts/CLProofVerifier__harness.sol +++ b/test/0.8.25/vaults/predeposit-guarantee/contracts/CLProofVerifier__harness.sol @@ -3,102 +3,24 @@ pragma solidity 0.8.25; -import {pack} from "contracts/0.8.25/lib/GIndex.sol"; - -import {CLProofVerifier, Validator, SSZ, ValidatorWitness, GIndex} from "contracts/0.8.25/vaults/predeposit_guarantee/CLProofVerifier.sol"; +import {pack, concat} from "contracts/0.8.25/lib/GIndex.sol"; +import {CLProofVerifier, SSZ, GIndex} from "contracts/0.8.25/vaults/predeposit_guarantee/CLProofVerifier.sol"; contract CLProofVerifier__Harness is CLProofVerifier { - bytes32 public MOCK_ROOT; - constructor(GIndex _gIFirstValidator) CLProofVerifier(_gIFirstValidator) {} - function setRoot(bytes32 _root) public { - MOCK_ROOT = _root; - } - - function _getParentBlockRoot(uint64) internal view override returns (bytes32) { - return MOCK_ROOT; - } - - function TEST_validateWCProof(ValidatorWitness calldata _witness) public view { - SSZ.verifyProof({ - proof: _witness.proof, - root: _getParentBlockRoot(_witness.beaconBlockTimestamp), - leaf: SSZ.hashTreeRoot(_witness.validator), - gIndex: _getValidatorGI(_witness.validatorIndex) - }); - } -} - -contract SSZMerkleTree { - uint256 public constant TREE_DEPTH = 32; // Adjustable tree depth (16 leaves max) - uint256 public leafCount = 0; // Number of leaves in the tree - mapping(uint256 => bytes32) public nodes; // Merkle tree nodes mapping - - /// @notice Adds a new leaf to the tree - /// @param leaf The leaf value (hashed data) - /// @return index The index of the added leaf - function addLeaf(bytes32 leaf) public returns (uint256) { - require(leafCount < (1 << TREE_DEPTH), "Tree is full"); - - uint256 index = (1 << TREE_DEPTH) + leafCount; // Compute SSZ generalized index - nodes[index] = leaf; - leafCount++; - - _updateTree(index); // Update the Merkle tree structure - - return index; + function TEST_validatePubKeyWCProof( + ValidatorWitness calldata _witness, + bytes32 _withdrawalCredentials + ) public view { + _validatePubKeyWCProof(_witness, _withdrawalCredentials); } - /// @notice Computes the Merkle root of the tree - /// @return root The computed root hash - function getMerkleRoot() public view returns (bytes32) { - return nodes[1]; // The root of the tree - } - - /// @notice Computes and returns the Merkle proof for a given leaf index - /// @param leafIndex The index of the leaf in the tree - /// @return proof The array of proof hashes - function getMerkleProof(uint256 leafIndex) public view returns (bytes32[] memory) { - require(leafIndex < leafCount, "Invalid leaf index"); - - uint256 index = (1 << TREE_DEPTH) + leafIndex; - bytes32[] memory proof = new bytes32[](TREE_DEPTH); - - for (uint256 i = 0; i < TREE_DEPTH; i++) { - uint256 siblingIndex = index % 2 == 0 ? index + 1 : index - 1; - proof[i] = nodes[siblingIndex]; - index /= 2; - } - return proof; - } - - /// @notice Returns the SSZ generalized index of a given leaf position - /// @param position The position of the leaf (0-based) - /// @return generalizedIndex The SSZ generalized index - function getGeneralizedIndex(uint256 position) public pure returns (GIndex) { - require(position < (1 << TREE_DEPTH), "Invalid position"); - - return pack((1 << TREE_DEPTH) + position, uint8(TREE_DEPTH)); - } - - /// @dev Updates the tree after adding a leaf - /// @param index The index of the new leaf - function _updateTree(uint256 index) internal { - while (index > 1) { - uint256 parentIndex = index / 2; - uint256 siblingIndex = index % 2 == 0 ? index + 1 : index - 1; - - bytes32 left = nodes[index % 2 == 0 ? index : siblingIndex]; - bytes32 right = nodes[index % 2 == 0 ? siblingIndex : index]; - - nodes[parentIndex] = sha256(abi.encodePacked(left, right)); - - index = parentIndex; - } + function TEST_getParentBlockRoot(uint64 parentBlockTimestamp) public view returns (bytes32) { + return _getParentBlockRoot(parentBlockTimestamp); } - function addValidatorLeaf(Validator calldata validator) public returns (uint256) { - return addLeaf(SSZ.hashTreeRoot(validator)); + function TEST_getValidatorGI(uint256 offset) public view returns (GIndex) { + return GI_FIRST_VALIDATOR.shr(offset); } } diff --git a/test/0.8.25/vaults/predeposit-guarantee/contracts/SSZHelpers.sol b/test/0.8.25/vaults/predeposit-guarantee/contracts/SSZHelpers.sol new file mode 100644 index 000000000..c5f5d3893 --- /dev/null +++ b/test/0.8.25/vaults/predeposit-guarantee/contracts/SSZHelpers.sol @@ -0,0 +1,334 @@ +// SPDX-License-Identifier: UNLICENSED +// for testing purposes only + +pragma solidity 0.8.25; + +import {GIndex, pack, concat} from "contracts/0.8.25/lib/GIndex.sol"; +import {SSZ} from "contracts/0.8.25/lib/SSZ.sol"; + +// As defined in phase0/beacon-chain.md:159 +type Slot is uint64; + +function unwrap(Slot slot) pure returns (uint64) { + return Slot.unwrap(slot); +} + +function gt(Slot lhs, Slot rhs) pure returns (bool) { + return lhs.unwrap() > rhs.unwrap(); +} + +function lt(Slot lhs, Slot rhs) pure returns (bool) { + return lhs.unwrap() < rhs.unwrap(); +} + +using {unwrap, lt as <, gt as >} for Slot global; + +/* + Complement to in-contract SSZ library with methods usefull for testing + original: https://github.com/lidofinance/community-staking-module/blob/7071c2096983a7780a5f147963aaa5405c0badb1/src/lib/SSZ.sol +*/ +contract SSZHelpers { + // As defined in phase0/beacon-chain.md:356 + struct Validator { + bytes pubkey; + bytes32 withdrawalCredentials; + uint64 effectiveBalance; + bool slashed; + uint64 activationEligibilityEpoch; + uint64 activationEpoch; + uint64 exitEpoch; + uint64 withdrawableEpoch; + } + + // As defined in phase0/beacon-chain.md:436 + struct BeaconBlockHeader { + Slot slot; + uint64 proposerIndex; + bytes32 parentRoot; + bytes32 stateRoot; + bytes32 bodyRoot; + } + + // canonical implementation from original SSZ + function validatorHashTreeRootCalldata(Validator calldata validator) public view returns (bytes32 root) { + bytes32 pubkeyRoot; + + assembly { + // In calldata, a dynamic field is encoded as an offset (relative to the start + // of the struct’s calldata) followed by its contents. The first 32 bytes of + // `validator` is the offset for `pubkey`. (Remember that `pubkey` is expected + // to be exactly 48 bytes long.) + let pubkeyOffset := calldataload(validator) + // The pubkey’s actual data is encoded at: + // validator + pubkeyOffset + 32 + // because the first word at that location is the length. + // Copy 48 bytes of pubkey data into memory at 0x00. + calldatacopy(0x00, add(validator, add(pubkeyOffset, 32)), 48) + // Zero the remaining 16 bytes to form a 64‐byte block. + // (0x30 = 48, so mstore at 0x30 will zero 32 bytes covering addresses 48–79; + // only bytes 48–63 matter for our 64-byte input.) + mstore(0x30, 0) + // Call the SHA‑256 precompile (at address 0x02) with the 64-byte block. + if iszero(staticcall(gas(), 0x02, 0x00, 0x40, 0x00, 0x20)) { + revert(0, 0) + } + pubkeyRoot := mload(0x00) + } + + bytes32[8] memory nodes = [ + pubkeyRoot, + validator.withdrawalCredentials, + toLittleEndian(validator.effectiveBalance), + toLittleEndian(validator.slashed), + toLittleEndian(validator.activationEligibilityEpoch), + toLittleEndian(validator.activationEpoch), + toLittleEndian(validator.exitEpoch), + toLittleEndian(validator.withdrawableEpoch) + ]; + + /// @solidity memory-safe-assembly + assembly { + // Count of nodes to hash + let count := 8 + + // Loop over levels + // prettier-ignore + for { } 1 { } { + // Loop over nodes at the given depth + + // Initialize `offset` to the offset of `proof` elements in memory. + let target := nodes + let source := nodes + let end := add(source, shl(5, count)) + + // prettier-ignore + for { } 1 { } { + // Read next two hashes to hash + mcopy(0x00, source, 0x40) + + // Call sha256 precompile + let result := staticcall( + gas(), + 0x02, + 0x00, + 0x40, + 0x00, + 0x20 + ) + + if iszero(result) { + // Precompiles returns no data on OutOfGas error. + revert(0, 0) + } + + // Store the resulting hash at the target location + mstore(target, mload(0x00)) + + // Advance the pointers + target := add(target, 0x20) + source := add(source, 0x40) + + if iszero(lt(source, end)) { + break + } + } + + count := shr(1, count) + if eq(count, 1) { + root := mload(0x00) + break + } + } + } + } + + // stupid direct hardcode to build merkle tree, proof and index for validator container for proving pubkey+wc node + function getValidatorPubkeyWCParentProof( + Validator calldata validator + ) public view returns (bytes32[] memory proof, bytes32 root, bytes32 parentNode, GIndex parentIndex) { + bytes32 pubkeyRoot = SSZ.pubkeyRoot(validator.pubkey); + + // Validator struct depth (8 -> 4 -> 2 -> 1) + bytes32[8] memory ValidatorL1 = [ + pubkeyRoot, + validator.withdrawalCredentials, + toLittleEndian(validator.effectiveBalance), + toLittleEndian(validator.slashed), + toLittleEndian(validator.activationEligibilityEpoch), + toLittleEndian(validator.activationEpoch), + toLittleEndian(validator.exitEpoch), + toLittleEndian(validator.withdrawableEpoch) + ]; + + bytes32[4] memory ValidatorL2 = [ + SSZ.sha256Pair(ValidatorL1[0], ValidatorL1[1]), + SSZ.sha256Pair(ValidatorL1[2], ValidatorL1[3]), + SSZ.sha256Pair(ValidatorL1[4], ValidatorL1[5]), + SSZ.sha256Pair(ValidatorL1[6], ValidatorL1[7]) + ]; + + parentNode = ValidatorL2[0]; + + bytes32[2] memory ValidatorL3 = [ + SSZ.sha256Pair(ValidatorL2[0], ValidatorL2[1]), + SSZ.sha256Pair(ValidatorL2[2], ValidatorL2[3]) + ]; + + root = SSZ.sha256Pair(ValidatorL3[0], ValidatorL3[1]); + // validates this hardcode against canonical implementation + require(root == validatorHashTreeRootCalldata(validator), "root mismatch"); + + uint8 proofDepth = 2; + proof = new bytes32[](proofDepth); + proof[0] = ValidatorL2[1]; + proof[1] = ValidatorL3[1]; + + // This is the parent node of `pubkey` and `withdrawalCredentials` GIndex + // it's on the start of second level from leaf level + // it's constant for all validators + uint256 VALIDATOR_TREE_DEPTH = 2; + uint256 PARENT_POSITION = 0; + parentIndex = pack((1 << VALIDATOR_TREE_DEPTH) + PARENT_POSITION, uint8(VALIDATOR_TREE_DEPTH)); + return (proof, root, parentNode, parentIndex); + } + + // canonical implementation from original SSZ + function beaconBlockHeaderHashTreeRoot(BeaconBlockHeader memory header) public view returns (bytes32 root) { + bytes32[8] memory headerNodes = [ + toLittleEndian(header.slot.unwrap()), + toLittleEndian(header.proposerIndex), + header.parentRoot, + header.stateRoot, + header.bodyRoot, + bytes32(0), + bytes32(0), + bytes32(0) + ]; + + /// @solidity memory-safe-assembly + assembly { + // Count of nodes to hash + let count := 8 + + // Loop over levels + // prettier-ignore + for { } 1 { } { + // Loop over nodes at the given depth + + // Initialize `offset` to the offset of `proof` elements in memory. + let target := headerNodes + let source := headerNodes + let end := add(source, shl(5, count)) + + // prettier-ignore + for { } 1 { } { + // Read next two hashes to hash + mcopy(0x00, source, 0x40) + + // Call sha256 precompile + let result := staticcall( + gas(), + 0x02, + 0x00, + 0x40, + 0x00, + 0x20 + ) + + if iszero(result) { + // Precompiles returns no data on OutOfGas error. + revert(0, 0) + } + + // Store the resulting hash at the target location + mstore(target, mload(0x00)) + + // Advance the pointers + target := add(target, 0x20) + source := add(source, 0x40) + + if iszero(lt(source, end)) { + break + } + } + + count := shr(1, count) + if eq(count, 1) { + root := mload(0x00) + break + } + } + } + } + + // stupid direct hardcode to build merkle tree, proof and index for validator container for proving pubkey+wc node + function getBeaconBlockHeaderProof( + BeaconBlockHeader memory header + ) public view returns (bytes32[] memory proof, bytes32 root, bytes32 leaf, GIndex index) { + // stupid hardcode to build tree for block header + bytes32[8] memory BlockHeaderL1 = [ + toLittleEndian(header.slot.unwrap()), + toLittleEndian(header.proposerIndex), + header.parentRoot, + header.stateRoot, // target leaf at position 3 + header.bodyRoot, + bytes32(0), + bytes32(0), + bytes32(0) + ]; + + bytes32[4] memory BlockHeaderL2 = [ + SSZ.sha256Pair(BlockHeaderL1[0], BlockHeaderL1[1]), + SSZ.sha256Pair(BlockHeaderL1[2], BlockHeaderL1[3]), + SSZ.sha256Pair(BlockHeaderL1[4], BlockHeaderL1[5]), + SSZ.sha256Pair(BlockHeaderL1[6], BlockHeaderL1[7]) + ]; + + bytes32[2] memory BlockHeaderL3 = [ + SSZ.sha256Pair(BlockHeaderL2[0], BlockHeaderL2[1]), + SSZ.sha256Pair(BlockHeaderL2[2], BlockHeaderL2[3]) + ]; + + root = SSZ.sha256Pair(BlockHeaderL3[0], BlockHeaderL3[1]); + leaf = header.stateRoot; + + // validates this hardcode against canonical implementation + require(root == beaconBlockHeaderHashTreeRoot(header), "root mismatch"); + + // all siblings on the way from the leaf to the root + uint256 HEADER_TREE_DEPTH = 3; + proof = new bytes32[](HEADER_TREE_DEPTH); + proof[0] = BlockHeaderL1[2]; + proof[1] = BlockHeaderL2[0]; + proof[2] = BlockHeaderL3[1]; + + uint256 PARENT_POSITION = 3; + index = pack((1 << HEADER_TREE_DEPTH) + PARENT_POSITION, uint8(HEADER_TREE_DEPTH)); + } + + // See https://github.com/succinctlabs/telepathy-contracts/blob/5aa4bb7/src/libraries/SimpleSerialize.sol#L17-L28 + function toLittleEndian(uint256 v) public pure returns (bytes32) { + v = + ((v & 0xFF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00) >> 8) | + ((v & 0x00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF00FF) << 8); + v = + ((v & 0xFFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000) >> 16) | + ((v & 0x0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF0000FFFF) << 16); + v = + ((v & 0xFFFFFFFF00000000FFFFFFFF00000000FFFFFFFF00000000FFFFFFFF00000000) >> 32) | + ((v & 0x00000000FFFFFFFF00000000FFFFFFFF00000000FFFFFFFF00000000FFFFFFFF) << 32); + v = + ((v & 0xFFFFFFFFFFFFFFFF0000000000000000FFFFFFFFFFFFFFFF0000000000000000) >> 64) | + ((v & 0x0000000000000000FFFFFFFFFFFFFFFF0000000000000000FFFFFFFFFFFFFFFF) << 64); + v = (v >> 128) | (v << 128); + return bytes32(v); + } + + function verifyProof(bytes32[] calldata proof, bytes32 root, bytes32 leaf, GIndex gIndex) public view { + SSZ.verifyProof(proof, root, leaf, gIndex); + } + + function toLittleEndian(bool v) public pure returns (bytes32) { + return bytes32(v ? 1 << 248 : 0); + } +} diff --git a/test/0.8.25/vaults/predeposit-guarantee/contracts/SSZMerkleTree.sol b/test/0.8.25/vaults/predeposit-guarantee/contracts/SSZMerkleTree.sol new file mode 100644 index 000000000..e5e87ed66 --- /dev/null +++ b/test/0.8.25/vaults/predeposit-guarantee/contracts/SSZMerkleTree.sol @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: UNLICENSED +// for testing purposes only + +pragma solidity 0.8.25; + +import {GIndex, pack, concat} from "contracts/0.8.25/lib/GIndex.sol"; +import {SSZ} from "contracts/0.8.25/lib/SSZ.sol"; + +import {SSZHelpers} from "./SSZHelpers.sol"; + +/// Merkle tree Implementation that aligns with CL implementation +/// NOT gas optimized, for testing proposes only +contract SSZMerkleTree is SSZHelpers { + uint256 public constant TREE_DEPTH = 40; // Adjustable tree depth + uint256 public leafCount = 0; // Number of leaves in the tree + mapping(uint256 => bytes32) public nodes; // Merkle tree nodes mapping + + /// @notice Adds a new leaf to the tree + /// @param leaf The leaf value (hashed data) + /// @return index The index of the added leaf + function addLeaf(bytes32 leaf) public returns (uint256) { + require(leafCount < (1 << TREE_DEPTH), "Tree is full"); + + uint256 index = (1 << TREE_DEPTH) + leafCount; // Compute SSZ generalized index + nodes[index] = leaf; + leafCount++; + + _updateTree(index); // Update the Merkle tree structure + + return index; + } + + /// @notice Computes the Merkle root of the tree + /// @return root The computed root hash + function getMerkleRoot() public view returns (bytes32) { + return nodes[1]; // The root of the tree + } + + /// @notice Computes and returns the Merkle proof for a given leaf index + /// @param leafIndex The index of the leaf in the tree + /// @return proof The array of proof hashes + function getMerkleProof(uint256 leafIndex) public view returns (bytes32[] memory) { + require(leafIndex < leafCount, "Invalid leaf index"); + + uint256 index = (1 << TREE_DEPTH) + leafIndex; + bytes32[] memory proof = new bytes32[](TREE_DEPTH); + + for (uint256 i = 0; i < TREE_DEPTH; i++) { + uint256 siblingIndex = index % 2 == 0 ? index + 1 : index - 1; + proof[i] = nodes[siblingIndex]; + index /= 2; + } + return proof; + } + + /// @notice Returns the SSZ generalized index of a given leaf position + /// @param position The position of the leaf (0-based) + /// @return generalizedIndex The SSZ generalized index + function getGeneralizedIndex(uint256 position) public pure returns (GIndex) { + require(position < (1 << TREE_DEPTH), "Invalid position"); + + return pack((1 << TREE_DEPTH) + position, uint8(TREE_DEPTH)); + } + + /// @dev Updates the tree after adding a leaf + /// @param index The index of the new leaf + function _updateTree(uint256 index) internal { + while (index > 1) { + uint256 parentIndex = index / 2; + uint256 siblingIndex = index % 2 == 0 ? index + 1 : index - 1; + + bytes32 left = nodes[index % 2 == 0 ? index : siblingIndex]; + bytes32 right = nodes[index % 2 == 0 ? siblingIndex : index]; + + nodes[parentIndex] = sha256(abi.encodePacked(left, right)); + + index = parentIndex; + } + } + + function addValidatorLeaf(SSZHelpers.Validator calldata validator) public returns (uint256) { + return addLeaf(validatorHashTreeRootCalldata(validator)); + } +}