diff --git a/src/Pectra.sol b/src/Pectra.sol index 3e43a7a..397823e 100644 --- a/src/Pectra.sol +++ b/src/Pectra.sol @@ -64,8 +64,14 @@ contract Pectra { _; } - function getFee(address target) public view returns (uint256 fee) { - (bool readOK, bytes memory feeData) = target.staticcall(""); + function getConsolidationFee() public view returns (uint256 fee) { + (bool readOK, bytes memory feeData) = consolidationTarget.staticcall(""); + if (!readOK) return MIN_FEE; + fee = uint256(bytes32(feeData)); + } + + function getExitFee() public view returns (uint256 fee) { + (bool readOK, bytes memory feeData) = exitTarget.staticcall(""); if (!readOK) return MIN_FEE; fee = uint256(bytes32(feeData)); } @@ -82,7 +88,7 @@ contract Pectra { revert InvalidTargetPubkeyLength(targetPubkey); } - uint256 consolidationFee = getFee(consolidationTarget); + uint256 consolidationFee = getConsolidationFee(); require(msg.value >= batchSize * consolidationFee, InsufficientFeePerValidator()); for (uint256 i = 0; i < batchSize; ++i) { @@ -105,7 +111,7 @@ contract Pectra { require(batchSize >= MIN_VALIDATORS, MinimumValidatorRequired()); require(batchSize <= MAX_VALIDATORS, TooManyValidators()); - uint256 switchFee = getFee(consolidationTarget); + uint256 switchFee = getConsolidationFee(); require(msg.value >= batchSize * switchFee, InsufficientFeePerValidator()); for (uint256 i = 0; i < batchSize; ++i) { @@ -128,7 +134,7 @@ contract Pectra { require(batchSize >= MIN_VALIDATORS, MinimumValidatorRequired()); require(batchSize <= MAX_VALIDATORS, TooManyValidators()); - uint256 exitFee = getFee(exitTarget); + uint256 exitFee = getExitFee(); require(msg.value >= batchSize * exitFee, InsufficientFeePerValidator()); for (uint256 i = 0; i < batchSize; ++i) { diff --git a/test/Pectra.t.sol b/test/Pectra.t.sol index 01947eb..a2fa7ec 100644 --- a/test/Pectra.t.sol +++ b/test/Pectra.t.sol @@ -40,19 +40,19 @@ contract PectraTest is Test { // ───────────────────────────────────────────────────────────────────────────── function testGetFee_ConsolidationTarget() public view { - uint256 fee = pectra.getFee(consolidationTarget); + uint256 fee = pectra.getConsolidationFee(); assertEq(fee, 1 wei, "Fee from consolidationTarget should be 1 wei"); } function testGetFee_ExitTarget() public view { - uint256 fee = pectra.getFee(exitTarget); + uint256 fee = pectra.getExitFee(); assertEq(fee, 1 wei, "Fee from exitTarget should be 1 wei"); } function testGetFee_FailedCall() public { // Temporarily set target to revert code vm.etch(consolidationTarget, revertCode); - uint256 fee = pectra.getFee(consolidationTarget); + uint256 fee = pectra.getConsolidationFee(); assertEq(fee, pectra.MIN_FEE(), "Fee should default to MIN_FEE when call fails"); // Reset back to fee code vm.etch(consolidationTarget, feeCode); @@ -186,7 +186,7 @@ contract PectraTest is Test { bytes memory target = validPubkey(); // Get the fee from the target - uint256 fee = pectra.getFee(consolidationTarget); + uint256 fee = pectra.getConsolidationFee(); assertEq(fee, 1 wei, "Fee should be 1 wei"); uint256 totalValue = count * fee; @@ -258,7 +258,7 @@ contract PectraTest is Test { } // Get the fee from the target - uint256 fee = pectra.getFee(consolidationTarget); + uint256 fee = pectra.getConsolidationFee(); assertEq(fee, 1 wei, "Fee should be 1 wei"); uint256 totalValue = count * fee; @@ -390,7 +390,7 @@ contract PectraTest is Test { } // Get the fee from the target - uint256 fee = pectra.getFee(exitTarget); + uint256 fee = pectra.getExitFee(); assertEq(fee, 1 wei, "Fee should be 1 wei"); uint256 totalValue = count * fee; @@ -412,7 +412,7 @@ contract PectraTest is Test { } // Get the fee from the target - uint256 fee = pectra.getFee(exitTarget); + uint256 fee = pectra.getExitFee(); assertEq(fee, 1 wei, "Fee should be 1 wei"); uint256 totalValue = count * fee; @@ -438,7 +438,7 @@ contract PectraTest is Test { bytes memory target = validPubkey(); // Get the fee from the target - uint256 fee = pectra.getFee(consolidationTarget); + uint256 fee = pectra.getConsolidationFee(); assertEq(fee, 1 wei, "Fee should be 1 wei"); uint256 totalValue = count * fee; @@ -464,7 +464,7 @@ contract PectraTest is Test { } // Get the fee from the target - uint256 fee = pectra.getFee(consolidationTarget); + uint256 fee = pectra.getConsolidationFee(); assertEq(fee, 1 wei, "Fee should be 1 wei"); uint256 totalValue = count * fee; @@ -490,7 +490,7 @@ contract PectraTest is Test { } // Get the fee from the target - uint256 fee = pectra.getFee(exitTarget); + uint256 fee = pectra.getExitFee(); assertEq(fee, 1 wei, "Fee should be 1 wei"); uint256 totalValue = count * fee;