diff --git a/.gitignore b/.gitignore index 31c9ad7d..649796e3 100644 --- a/.gitignore +++ b/.gitignore @@ -25,4 +25,5 @@ tmp/* .worktrees/ .claude/ +CLAUDE.local.md Agents.md diff --git a/contracts/src/interfaces/IRewardsRegistry.sol b/contracts/src/interfaces/IRewardsRegistry.sol index 2cd06196..3857db2c 100644 --- a/contracts/src/interfaces/IRewardsRegistry.sol +++ b/contracts/src/interfaces/IRewardsRegistry.sol @@ -79,45 +79,57 @@ interface IRewardsRegistry is IRewardsRegistryErrors, IRewardsRegistryEvents { ) external; /** - * @notice Claim rewards for an operator from a specific merkle root index + * @notice Claim rewards for an operator from a specific merkle root index using Substrate/Snowbridge positional Merkle proofs. * @param operatorAddress Address of the operator to receive rewards * @param rootIndex Index of the merkle root to claim from * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) * @dev Only callable by the AVS (Service Manager) */ function claimRewards( address operatorAddress, uint256 rootIndex, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) external; /** - * @notice Claim rewards for an operator from the latest merkle root + * @notice Claim rewards for an operator from the latest merkle root using Substrate/Snowbridge positional Merkle proofs. * @param operatorAddress Address of the operator to receive rewards * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) * @dev Only callable by the AVS (Service Manager) */ function claimLatestRewards( address operatorAddress, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) external; /** - * @notice Claim rewards for an operator from multiple merkle root indices + * @notice Claim rewards for an operator from multiple merkle root indices using Substrate/Snowbridge positional Merkle proofs. * @param operatorAddress Address of the operator to receive rewards * @param rootIndices Array of merkle root indices to claim from * @param operatorPoints Array of points earned by the operator for each root - * @param proofs Array of merkle proofs to validate the operator's rewards + * @param numberOfLeaves Array with the total number of leaves for each Merkle tree + * @param leafIndices Array of leaf indices for the operator in each Merkle tree + * @param proofs Array of positional Merkle proofs for each claim * @dev Only callable by the AVS (Service Manager) */ function claimRewardsBatch( address operatorAddress, uint256[] calldata rootIndices, uint256[] calldata operatorPoints, + uint256[] calldata numberOfLeaves, + uint256[] calldata leafIndices, bytes32[][] calldata proofs ) external; diff --git a/contracts/src/interfaces/IServiceManager.sol b/contracts/src/interfaces/IServiceManager.sol index fa62d1bf..d9273ec1 100644 --- a/contracts/src/interfaces/IServiceManager.sol +++ b/contracts/src/interfaces/IServiceManager.sol @@ -134,42 +134,54 @@ interface IServiceManager is IServiceManagerUI, IServiceManagerErrors, IServiceM function setRewardsRegistry(uint32 operatorSetId, IRewardsRegistry rewardsRegistry) external; /** - * @notice Claim rewards for an operator from a specific merkle root index + * @notice Claim rewards for an operator from a specific merkle root index using Substrate/Snowbridge positional Merkle proofs * @param operatorSetId The ID of the operator set * @param rootIndex Index of the merkle root to claim from * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) */ function claimOperatorRewards( uint32 operatorSetId, uint256 rootIndex, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) external; /** - * @notice Claim rewards for an operator from the latest merkle root + * @notice Claim rewards for an operator from the latest merkle root using Substrate/Snowbridge positional Merkle proofs * @param operatorSetId The ID of the operator set * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) */ function claimLatestOperatorRewards( uint32 operatorSetId, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) external; /** - * @notice Claim rewards for an operator from multiple merkle root indices + * @notice Claim rewards for an operator from multiple merkle root indices using Substrate/Snowbridge positional Merkle proofs * @param operatorSetId The ID of the operator set * @param rootIndices Array of merkle root indices to claim from * @param operatorPoints Array of points earned by the operator for each root - * @param proofs Array of merkle proofs to validate the operator's rewards + * @param numberOfLeaves Array with the total number of leaves for each Merkle tree + * @param leafIndices Array of leaf indices for the operator in each Merkle tree + * @param proofs Array of positional Merkle proofs for each claim */ function claimOperatorRewardsBatch( uint32 operatorSetId, uint256[] calldata rootIndices, uint256[] calldata operatorPoints, + uint256[] calldata numberOfLeaves, + uint256[] calldata leafIndices, bytes32[][] calldata proofs ) external; diff --git a/contracts/src/middleware/RewardsRegistry.sol b/contracts/src/middleware/RewardsRegistry.sol index e062e5b8..f8b9b354 100644 --- a/contracts/src/middleware/RewardsRegistry.sol +++ b/contracts/src/middleware/RewardsRegistry.sol @@ -1,7 +1,9 @@ // SPDX-License-Identifier: BUSL-1.1 pragma solidity ^0.8.27; -import {MerkleProof} from "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol"; +import {SubstrateMerkleProof} from "snowbridge/src/utils/SubstrateMerkleProof.sol"; +import {ScaleCodec} from "snowbridge/src/utils/ScaleCodec.sol"; +import {IDataHavenServiceManager} from "../interfaces/IDataHavenServiceManager.sol"; import {RewardsRegistryStorage} from "./RewardsRegistryStorage.sol"; /** @@ -69,123 +71,148 @@ contract RewardsRegistry is RewardsRegistryStorage { } /** - * @notice Claim rewards for an operator from a specific merkle root index + * @notice Claim rewards for an operator from a specific merkle root index using Substrate/Snowbridge positional Merkle proofs. * @param operatorAddress Address of the operator to receive rewards * @param rootIndex Index of the merkle root to claim from * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) * @dev Only callable by the AVS (Service Manager) */ function claimRewards( address operatorAddress, uint256 rootIndex, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) external override onlyAVS { - // Validate the claim and calculate rewards - uint256 rewardsAmount = _validateClaim(operatorAddress, rootIndex, operatorPoints, proof); + uint256 rewardsAmount = _validateClaim( + operatorAddress, rootIndex, operatorPoints, numberOfLeaves, leafIndex, proof + ); _transferRewards(operatorAddress, rewardsAmount); - // Emit the corresponding event emit RewardsClaimedForIndex(operatorAddress, rootIndex, operatorPoints, rewardsAmount); } /** - * @notice Claim rewards for an operator from the latest merkle root + * @notice Claim rewards for an operator from the latest merkle root using Substrate/Snowbridge positional Merkle proofs. * @param operatorAddress Address of the operator to receive rewards * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) * @dev Only callable by the AVS (Service Manager) */ function claimLatestRewards( address operatorAddress, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) external override onlyAVS { - // Check that we have at least one merkle root if (merkleRootHistory.length == 0) { revert RewardsMerkleRootNotSet(); } - - // Claim from the latest root index uint256 latestIndex = merkleRootHistory.length - 1; - uint256 rewardsAmount = _validateClaim(operatorAddress, latestIndex, operatorPoints, proof); + uint256 rewardsAmount = _validateClaim( + operatorAddress, latestIndex, operatorPoints, numberOfLeaves, leafIndex, proof + ); _transferRewards(operatorAddress, rewardsAmount); - // Emit the corresponding event emit RewardsClaimedForIndex(operatorAddress, latestIndex, operatorPoints, rewardsAmount); } /** - * @notice Claim rewards for an operator from multiple merkle root indices + * @notice Claim rewards for an operator from multiple merkle root indices using Substrate/Snowbridge positional Merkle proofs. * @param operatorAddress Address of the operator to receive rewards * @param rootIndices Array of merkle root indices to claim from * @param operatorPoints Array of points earned by the operator for each root - * @param proofs Array of merkle proofs to validate the operator's rewards + * @param numberOfLeaves Array with the total number of leaves for each Merkle tree + * @param leafIndices Array of leaf indices for the operator in each Merkle tree + * @param proofs Array of positional Merkle proofs for each claim * @dev Only callable by the AVS (Service Manager) */ function claimRewardsBatch( address operatorAddress, uint256[] calldata rootIndices, uint256[] calldata operatorPoints, + uint256[] calldata numberOfLeaves, + uint256[] calldata leafIndices, bytes32[][] calldata proofs ) external override onlyAVS { - // Check that the arrays have the same length - if (rootIndices.length != operatorPoints.length || rootIndices.length != proofs.length) { + if ( + rootIndices.length != operatorPoints.length || rootIndices.length != proofs.length + || rootIndices.length != numberOfLeaves.length + || rootIndices.length != leafIndices.length + ) { revert ArrayLengthMismatch(); } - // Validate all claims and accumulate the total rewards uint256 totalRewards = 0; for (uint256 i = 0; i < rootIndices.length; i++) { - totalRewards += - _validateClaim(operatorAddress, rootIndices[i], operatorPoints[i], proofs[i]); + totalRewards += _validateClaim( + operatorAddress, + rootIndices[i], + operatorPoints[i], + numberOfLeaves[i], + leafIndices[i], + proofs[i] + ); } - // Transfer the total rewards in a single transaction _transferRewards(operatorAddress, totalRewards); - // Emit the corresponding event emit RewardsBatchClaimedForIndices( operatorAddress, rootIndices, operatorPoints, totalRewards ); } /** - * @notice Internal function to validate a claim and calculate rewards + * @notice Internal function to validate a claim and calculate rewards using Substrate/Snowbridge positional Merkle proofs. * @param operatorAddress Address of the operator to receive rewards * @param rootIndex Index of the merkle root to claim from * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) * @return rewardsAmount The amount of rewards calculated */ function _validateClaim( address operatorAddress, uint256 rootIndex, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) internal returns (uint256 rewardsAmount) { - // Check that the root index to claim from exists if (rootIndex >= merkleRootHistory.length) { revert InvalidMerkleRootIndex(); } - // Check if operator has already claimed for this merkle root index if (operatorClaimedByIndex[operatorAddress][rootIndex]) { revert RewardsAlreadyClaimedForIndex(); } - // Verify the merkle proof - bytes32 leaf = keccak256(abi.encode(operatorAddress, operatorPoints)); - if (!MerkleProof.verify(proof, merkleRootHistory[rootIndex], leaf)) { - revert InvalidMerkleProof(); + // Compute Substrate-compatible leaf: keccak256(SCALE(accountId || u32LE points)) + // For DataHaven, AccountId comes from the AVS mapping (validatorEthAddressToSolochainAddress) if set. + address leafAccount = operatorAddress; + address mappedSolochain = + IDataHavenServiceManager(avs).validatorEthAddressToSolochainAddress(operatorAddress); + if (mappedSolochain != address(0)) { + leafAccount = mappedSolochain; } + bytes memory preimage = + abi.encodePacked(leafAccount, ScaleCodec.encodeU32(uint32(operatorPoints))); + bytes32 substrateLeaf = keccak256(preimage); + + bool ok = SubstrateMerkleProof.verify( + merkleRootHistory[rootIndex], substrateLeaf, leafIndex, numberOfLeaves, proof + ); + if (!ok) revert InvalidMerkleProof(); - // Calculate rewards - currently 1 point = 1 wei (placeholder) - // TODO: Update the reward calculation formula with the proper relationship rewardsAmount = operatorPoints; - - // Mark as claimed for this specific index operatorClaimedByIndex[operatorAddress][rootIndex] = true; } diff --git a/contracts/src/middleware/ServiceManagerBase.sol b/contracts/src/middleware/ServiceManagerBase.sol index 390e8591..853a1763 100644 --- a/contracts/src/middleware/ServiceManagerBase.sol +++ b/contracts/src/middleware/ServiceManagerBase.sol @@ -303,79 +303,82 @@ abstract contract ServiceManagerBase is ServiceManagerBaseStorage, IAVSRegistrar } /** - * @notice Claim rewards for an operator from a specific merkle root index + * @notice Claim rewards for an operator from a specific merkle root index using Substrate/Snowbridge positional Merkle proofs * @param operatorSetId The ID of the operator set * @param rootIndex Index of the merkle root to claim from * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) */ function claimOperatorRewards( uint32 operatorSetId, uint256 rootIndex, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) external virtual override { - // Get the rewards registry for this operator set IRewardsRegistry rewardsRegistry = operatorSetToRewardsRegistry[operatorSetId]; if (address(rewardsRegistry) == address(0)) { revert NoRewardsRegistryForOperatorSet(); } - - // Ensure the operator is part of the operator set _ensureOperatorIsPartOfOperatorSet(msg.sender, operatorSetId); - - // Forward the claim to the rewards registry - rewardsRegistry.claimRewards(msg.sender, rootIndex, operatorPoints, proof); + rewardsRegistry.claimRewards( + msg.sender, rootIndex, operatorPoints, numberOfLeaves, leafIndex, proof + ); } /** - * @notice Claim rewards for an operator from the latest merkle root + * @notice Claim rewards for an operator from the latest merkle root using Substrate/Snowbridge positional Merkle proofs * @param operatorSetId The ID of the operator set * @param operatorPoints Points earned by the operator - * @param proof Merkle proof to validate the operator's rewards + * @param numberOfLeaves The total number of leaves in the Merkle tree + * @param leafIndex The index of the operator's leaf in the Merkle tree + * @param proof Positional Merkle proof (from leaf to root) */ function claimLatestOperatorRewards( uint32 operatorSetId, uint256 operatorPoints, + uint256 numberOfLeaves, + uint256 leafIndex, bytes32[] calldata proof ) external virtual override { - // Get the rewards registry for this operator set IRewardsRegistry rewardsRegistry = operatorSetToRewardsRegistry[operatorSetId]; if (address(rewardsRegistry) == address(0)) { revert NoRewardsRegistryForOperatorSet(); } - - // Ensure the operator is part of the operator set _ensureOperatorIsPartOfOperatorSet(msg.sender, operatorSetId); - - // Forward the claim to the rewards registry - rewardsRegistry.claimLatestRewards(msg.sender, operatorPoints, proof); + rewardsRegistry.claimLatestRewards( + msg.sender, operatorPoints, numberOfLeaves, leafIndex, proof + ); } /** - * @notice Claim rewards for an operator from multiple merkle root indices + * @notice Claim rewards for an operator from multiple merkle root indices using Substrate/Snowbridge positional Merkle proofs * @param operatorSetId The ID of the operator set * @param rootIndices Array of merkle root indices to claim from * @param operatorPoints Array of points earned by the operator for each root - * @param proofs Array of merkle proofs to validate the operator's rewards + * @param numberOfLeaves Array with the total number of leaves for each Merkle tree + * @param leafIndices Array of leaf indices for the operator in each Merkle tree + * @param proofs Array of positional Merkle proofs for each claim */ function claimOperatorRewardsBatch( uint32 operatorSetId, uint256[] calldata rootIndices, uint256[] calldata operatorPoints, + uint256[] calldata numberOfLeaves, + uint256[] calldata leafIndices, bytes32[][] calldata proofs ) external virtual override { - // Get the rewards registry for this operator set IRewardsRegistry rewardsRegistry = operatorSetToRewardsRegistry[operatorSetId]; if (address(rewardsRegistry) == address(0)) { revert NoRewardsRegistryForOperatorSet(); } - - // Ensure the operator is part of the operator set _ensureOperatorIsPartOfOperatorSet(msg.sender, operatorSetId); - - // Forward the claim to the rewards registry - rewardsRegistry.claimRewardsBatch(msg.sender, rootIndices, operatorPoints, proofs); + rewardsRegistry.claimRewardsBatch( + msg.sender, rootIndices, operatorPoints, numberOfLeaves, leafIndices, proofs + ); } /** @@ -419,7 +422,7 @@ abstract contract ServiceManagerBase is ServiceManagerBaseStorage, IAVSRegistrar } /** - * @notice Returns the list of strategies that the operator has potentially restaked on the AVS + * @notice Returns the list of strategies that an operator has potentially restaked on the AVS * @param operator The address of the operator to get restaked strategies for * @dev This function is intended to be called off-chain * @dev No guarantee is made on whether the operator has shares for a strategy in a quorum or uniqueness diff --git a/contracts/test/RewardsRegistry.t.sol b/contracts/test/RewardsRegistry.t.sol index 09d065fc..34b7e988 100644 --- a/contracts/test/RewardsRegistry.t.sol +++ b/contracts/test/RewardsRegistry.t.sol @@ -4,11 +4,11 @@ pragma solidity ^0.8.13; /* solhint-disable func-name-mixedcase */ import {Test, console, stdError} from "forge-std/Test.sol"; -import {MerkleProof} from "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol"; import {AVSDeployer} from "./utils/AVSDeployer.sol"; import {RewardsRegistry} from "../src/middleware/RewardsRegistry.sol"; import {IRewardsRegistry, IRewardsRegistryErrors} from "../src/interfaces/IRewardsRegistry.sol"; +import {ScaleCodec} from "snowbridge/src/utils/ScaleCodec.sol"; contract RewardsRegistryTest is AVSDeployer { address public nonRewardsAgent; @@ -18,6 +18,8 @@ contract RewardsRegistryTest is AVSDeployer { bytes32 public merkleRoot; bytes32 public newMerkleRoot; uint256 public operatorPoints; + uint256 public leafIndex; + uint256 public numberOfLeaves; bytes32[] public validProof; bytes32[] public invalidProof; @@ -45,31 +47,33 @@ contract RewardsRegistryTest is AVSDeployer { // Set up test data operatorPoints = 100; + leafIndex = 0; // Position of our leaf in the tree + numberOfLeaves = 2; // Simple tree with 2 leaves - // For testing MerkleProof verification, we'll use the simplest case: - // A binary tree with just our target leaf and a sibling leaf + // For Substrate-compatible Merkle proofs, we need to use SCALE encoding // Our leaf (the one we want to prove exists in the tree) - bytes32 leaf = keccak256(abi.encode(operatorAddress, operatorPoints)); + bytes memory preimage = + abi.encodePacked(operatorAddress, ScaleCodec.encodeU32(uint32(operatorPoints))); + bytes32 leaf = keccak256(preimage); // Sibling leaf (another element in the Merkle tree) - bytes32 siblingLeaf = keccak256(abi.encodePacked("sibling")); + bytes memory siblingPreimage = + abi.encodePacked(address(0x1234), ScaleCodec.encodeU32(uint32(50))); + bytes32 siblingLeaf = keccak256(siblingPreimage); - // Sort leaves to follow the canonical order used by most Merkle tree libraries - (bytes32 leftLeaf, bytes32 rightLeaf) = - leaf < siblingLeaf ? (leaf, siblingLeaf) : (siblingLeaf, leaf); - - // Calculate parent node (this will be the Merkle root for our simple tree) - merkleRoot = keccak256(abi.encodePacked(leftLeaf, rightLeaf)); + // For Substrate positional merkle proof, we construct the root based on position + // Since leafIndex = 0, our leaf is on the left + merkleRoot = keccak256(abi.encodePacked(leaf, siblingLeaf)); // The proof to verify our leaf is just the sibling leaf validProof = new bytes32[](1); validProof[0] = siblingLeaf; // For tests that need a second Merkle root - bytes32 newSiblingLeaf = keccak256(abi.encodePacked("new sibling")); - (leftLeaf, rightLeaf) = - leaf < newSiblingLeaf ? (leaf, newSiblingLeaf) : (newSiblingLeaf, leaf); - newMerkleRoot = keccak256(abi.encodePacked(leftLeaf, rightLeaf)); + bytes memory newSiblingPreimage = + abi.encodePacked(address(0x5678), ScaleCodec.encodeU32(uint32(75))); + bytes32 newSiblingLeaf = keccak256(newSiblingPreimage); + newMerkleRoot = keccak256(abi.encodePacked(leaf, newSiblingLeaf)); // An invalid proof invalidProof = new bytes32[](1); @@ -77,10 +81,23 @@ contract RewardsRegistryTest is AVSDeployer { } // Helper to test our proof construction - function test_verifyProofConstruction() public view { - bytes32 leaf = keccak256(abi.encode(operatorAddress, operatorPoints)); - bool result = MerkleProof.verify(validProof, merkleRoot, leaf); - assertTrue(result, "Proof verification should succeed"); + function test_verifyProofConstruction() public { + // Test that our proof construction is valid using the contract's internal validation + vm.prank(mockRewardsAgent); + rewardsRegistry.updateRewardsMerkleRoot(merkleRoot); + + vm.deal(address(rewardsRegistry), 1000 ether); + + vm.prank(address(serviceManager)); + // This should not revert if the proof is valid + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, validProof + ); + + assertTrue( + rewardsRegistry.hasClaimedByIndex(operatorAddress, 0), + "Proof verification should succeed" + ); } /** @@ -179,7 +196,9 @@ contract RewardsRegistryTest is AVSDeployer { vm.expectEmit(true, true, true, true); emit RewardsClaimedForIndex(operatorAddress, 0, operatorPoints, operatorPoints); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, validProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, validProof + ); // Verify state changes assertTrue( @@ -201,7 +220,9 @@ contract RewardsRegistryTest is AVSDeployer { vm.expectRevert(abi.encodeWithSelector(IRewardsRegistryErrors.OnlyAVS.selector)); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, validProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_claimLatestRewards_AlreadyClaimed() public { @@ -214,14 +235,18 @@ contract RewardsRegistryTest is AVSDeployer { // First claim succeeds vm.prank(address(serviceManager)); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, validProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, validProof + ); // Second claim fails vm.prank(address(serviceManager)); vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.RewardsAlreadyClaimedForIndex.selector) ); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, validProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_claimLatestRewards_InvalidProof() public { @@ -230,7 +255,9 @@ contract RewardsRegistryTest is AVSDeployer { vm.prank(address(serviceManager)); vm.expectRevert(abi.encodeWithSelector(IRewardsRegistryErrors.InvalidMerkleProof.selector)); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, invalidProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, invalidProof + ); } function test_claimLatestRewards_NoMerkleRoot() public { @@ -239,7 +266,9 @@ contract RewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.RewardsMerkleRootNotSet.selector) ); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, validProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_claimLatestRewards_DifferentRoot() public { @@ -252,7 +281,9 @@ contract RewardsRegistryTest is AVSDeployer { // First claim succeeds vm.prank(address(serviceManager)); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, validProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, validProof + ); // Update to new merkle root vm.prank(mockRewardsAgent); @@ -260,12 +291,16 @@ contract RewardsRegistryTest is AVSDeployer { // Create a new valid proof for the new root bytes32[] memory newProof = new bytes32[](1); - bytes32 newSiblingLeaf = keccak256(abi.encodePacked("new sibling")); + bytes memory newSiblingPreimage = + abi.encodePacked(address(0x5678), ScaleCodec.encodeU32(uint32(75))); + bytes32 newSiblingLeaf = keccak256(newSiblingPreimage); newProof[0] = newSiblingLeaf; // Operator can claim again with new merkle root vm.prank(address(serviceManager)); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, newProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, newProof + ); // Verify both indices are now claimed assertTrue( @@ -290,7 +325,9 @@ contract RewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.RewardsTransferFailed.selector) ); - rewardsRegistry.claimLatestRewards(operatorAddress, operatorPoints, validProof); + rewardsRegistry.claimLatestRewards( + operatorAddress, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_receive() public { @@ -426,7 +463,9 @@ contract RewardsRegistryTest is AVSDeployer { vm.expectEmit(true, true, true, true); emit RewardsClaimedForIndex(operatorAddress, 0, operatorPoints, operatorPoints); - rewardsRegistry.claimRewards(operatorAddress, 0, operatorPoints, validProof); + rewardsRegistry.claimRewards( + operatorAddress, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); // Verify state changes assertTrue( @@ -447,7 +486,9 @@ contract RewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.InvalidMerkleRootIndex.selector) ); - rewardsRegistry.claimRewards(operatorAddress, 0, operatorPoints, validProof); + rewardsRegistry.claimRewards( + operatorAddress, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_claimRewards_AlreadyClaimed() public { @@ -459,14 +500,18 @@ contract RewardsRegistryTest is AVSDeployer { // First claim succeeds vm.prank(address(serviceManager)); - rewardsRegistry.claimRewards(operatorAddress, 0, operatorPoints, validProof); + rewardsRegistry.claimRewards( + operatorAddress, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); // Second claim fails vm.prank(address(serviceManager)); vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.RewardsAlreadyClaimedForIndex.selector) ); - rewardsRegistry.claimRewards(operatorAddress, 0, operatorPoints, validProof); + rewardsRegistry.claimRewards( + operatorAddress, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_hasClaimedByIndex() public { @@ -484,7 +529,9 @@ contract RewardsRegistryTest is AVSDeployer { // Claim vm.prank(address(serviceManager)); - rewardsRegistry.claimRewards(operatorAddress, 0, operatorPoints, validProof); + rewardsRegistry.claimRewards( + operatorAddress, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); // Now claimed assertTrue( @@ -521,7 +568,9 @@ contract RewardsRegistryTest is AVSDeployer { // Create proof for second root bytes32[] memory newProof = new bytes32[](1); - bytes32 newSiblingLeaf = keccak256(abi.encodePacked("new sibling")); + bytes memory newSiblingPreimage = + abi.encodePacked(address(0x5678), ScaleCodec.encodeU32(uint32(75))); + bytes32 newSiblingLeaf = keccak256(newSiblingPreimage); newProof[0] = newSiblingLeaf; proofs[1] = newProof; @@ -529,11 +578,17 @@ contract RewardsRegistryTest is AVSDeployer { // Batch claim vm.prank(address(serviceManager)); - vm.expectEmit(true, true, true, true); emit RewardsBatchClaimedForIndices(operatorAddress, rootIndices, points, operatorPoints * 2); - - rewardsRegistry.claimRewardsBatch(operatorAddress, rootIndices, points, proofs); + uint256[] memory widths = new uint256[](2); + widths[0] = numberOfLeaves; + widths[1] = numberOfLeaves; + uint256[] memory leafIdxs = new uint256[](2); + leafIdxs[0] = leafIndex; + leafIdxs[1] = leafIndex; + rewardsRegistry.claimRewardsBatch( + operatorAddress, rootIndices, points, widths, leafIdxs, proofs + ); // Verify both indices are claimed assertTrue( @@ -560,7 +615,11 @@ contract RewardsRegistryTest is AVSDeployer { vm.prank(address(serviceManager)); vm.expectRevert(abi.encodeWithSelector(IRewardsRegistryErrors.ArrayLengthMismatch.selector)); - rewardsRegistry.claimRewardsBatch(operatorAddress, rootIndices, points, proofs); + uint256[] memory widths = new uint256[](2); + uint256[] memory leafIdxs = new uint256[](2); + rewardsRegistry.claimRewardsBatch( + operatorAddress, rootIndices, points, widths, leafIdxs, proofs + ); } function test_claimRewardsBatch_PartialClaimFailure() public { @@ -575,7 +634,9 @@ contract RewardsRegistryTest is AVSDeployer { // Claim from index 0 first vm.prank(address(serviceManager)); - rewardsRegistry.claimRewards(operatorAddress, 0, operatorPoints, validProof); + rewardsRegistry.claimRewards( + operatorAddress, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); // Now try batch claim that includes already claimed index 0 uint256[] memory rootIndices = new uint256[](2); @@ -590,7 +651,9 @@ contract RewardsRegistryTest is AVSDeployer { proofs[0] = validProof; bytes32[] memory newProof = new bytes32[](1); - bytes32 newSiblingLeaf = keccak256(abi.encodePacked("new sibling")); + bytes memory newSiblingPreimage = + abi.encodePacked(address(0x5678), ScaleCodec.encodeU32(uint32(75))); + bytes32 newSiblingLeaf = keccak256(newSiblingPreimage); newProof[0] = newSiblingLeaf; proofs[1] = newProof; @@ -599,6 +662,14 @@ contract RewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.RewardsAlreadyClaimedForIndex.selector) ); - rewardsRegistry.claimRewardsBatch(operatorAddress, rootIndices, points, proofs); + uint256[] memory widths = new uint256[](2); + widths[0] = numberOfLeaves; + widths[1] = numberOfLeaves; + uint256[] memory leafIdxs = new uint256[](2); + leafIdxs[0] = leafIndex; + leafIdxs[1] = leafIndex; + rewardsRegistry.claimRewardsBatch( + operatorAddress, rootIndices, points, widths, leafIdxs, proofs + ); } } diff --git a/contracts/test/ServiceManagerRewardsRegistry.t.sol b/contracts/test/ServiceManagerRewardsRegistry.t.sol index bbd77929..b6f07c09 100644 --- a/contracts/test/ServiceManagerRewardsRegistry.t.sol +++ b/contracts/test/ServiceManagerRewardsRegistry.t.sol @@ -4,7 +4,6 @@ pragma solidity ^0.8.13; /* solhint-disable func-name-mixedcase */ import {Test, console, stdError} from "forge-std/Test.sol"; -import {MerkleProof} from "@openzeppelin/contracts/utils/cryptography/MerkleProof.sol"; import {IAllocationManager} from "eigenlayer-contracts/src/contracts/interfaces/IAllocationManager.sol"; @@ -13,6 +12,7 @@ import {RewardsRegistry} from "../src/middleware/RewardsRegistry.sol"; import {IRewardsRegistry, IRewardsRegistryErrors} from "../src/interfaces/IRewardsRegistry.sol"; import {ServiceManagerMock} from "./mocks/ServiceManagerMock.sol"; import {IServiceManager, IServiceManagerErrors} from "../src/interfaces/IServiceManager.sol"; +import {ScaleCodec} from "snowbridge/src/utils/ScaleCodec.sol"; contract ServiceManagerRewardsRegistryTest is AVSDeployer { // Test addresses @@ -27,6 +27,8 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { uint256 public operatorPoints; uint256 public secondOperatorPoints; uint256 public thirdOperatorPoints; + uint256 public leafIndex; + uint256 public numberOfLeaves; bytes32[] public validProof; bytes32[] public secondValidProof; bytes32[] public thirdValidProof; @@ -58,6 +60,8 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { operatorPoints = 100; secondOperatorPoints = 200; thirdOperatorPoints = 150; + leafIndex = 0; // Position of our leaf in the tree + numberOfLeaves = 2; // Simple tree with 2 leaves // Create multiple merkle trees for comprehensive batch testing _createFirstMerkleTree(); @@ -83,34 +87,50 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { } function _createFirstMerkleTree() internal { - // Create first merkle tree - bytes32 leaf = keccak256(abi.encode(operatorAddress, operatorPoints)); - bytes32 siblingLeaf = keccak256(abi.encodePacked("sibling1")); - (bytes32 leftLeaf, bytes32 rightLeaf) = - leaf < siblingLeaf ? (leaf, siblingLeaf) : (siblingLeaf, leaf); - merkleRoot = keccak256(abi.encodePacked(leftLeaf, rightLeaf)); + // Create first merkle tree with Substrate-compatible SCALE encoding + bytes memory preimage = + abi.encodePacked(operatorAddress, ScaleCodec.encodeU32(uint32(operatorPoints))); + bytes32 leaf = keccak256(preimage); + + bytes memory siblingPreimage = + abi.encodePacked(address(0x1111), ScaleCodec.encodeU32(uint32(50))); + bytes32 siblingLeaf = keccak256(siblingPreimage); + + // For Substrate positional merkle proof, we construct the root based on position + // Since leafIndex = 0, our leaf is on the left + merkleRoot = keccak256(abi.encodePacked(leaf, siblingLeaf)); validProof = new bytes32[](1); validProof[0] = siblingLeaf; } function _createSecondMerkleTree() internal { - // Create second merkle tree with different points - bytes32 leaf = keccak256(abi.encode(operatorAddress, secondOperatorPoints)); - bytes32 siblingLeaf = keccak256(abi.encodePacked("sibling2")); - (bytes32 leftLeaf, bytes32 rightLeaf) = - leaf < siblingLeaf ? (leaf, siblingLeaf) : (siblingLeaf, leaf); - secondMerkleRoot = keccak256(abi.encodePacked(leftLeaf, rightLeaf)); + // Create second merkle tree with different points using SCALE encoding + bytes memory preimage = + abi.encodePacked(operatorAddress, ScaleCodec.encodeU32(uint32(secondOperatorPoints))); + bytes32 leaf = keccak256(preimage); + + bytes memory siblingPreimage = + abi.encodePacked(address(0x2222), ScaleCodec.encodeU32(uint32(75))); + bytes32 siblingLeaf = keccak256(siblingPreimage); + + // Since leafIndex = 0, our leaf is on the left + secondMerkleRoot = keccak256(abi.encodePacked(leaf, siblingLeaf)); secondValidProof = new bytes32[](1); secondValidProof[0] = siblingLeaf; } function _createThirdMerkleTree() internal { - // Create third merkle tree with different points - bytes32 leaf = keccak256(abi.encode(operatorAddress, thirdOperatorPoints)); - bytes32 siblingLeaf = keccak256(abi.encodePacked("sibling3")); - (bytes32 leftLeaf, bytes32 rightLeaf) = - leaf < siblingLeaf ? (leaf, siblingLeaf) : (siblingLeaf, leaf); - thirdMerkleRoot = keccak256(abi.encodePacked(leftLeaf, rightLeaf)); + // Create third merkle tree with different points using SCALE encoding + bytes memory preimage = + abi.encodePacked(operatorAddress, ScaleCodec.encodeU32(uint32(thirdOperatorPoints))); + bytes32 leaf = keccak256(preimage); + + bytes memory siblingPreimage = + abi.encodePacked(address(0x3333), ScaleCodec.encodeU32(uint32(60))); + bytes32 siblingLeaf = keccak256(siblingPreimage); + + // Since leafIndex = 0, our leaf is on the left + thirdMerkleRoot = keccak256(abi.encodePacked(leaf, siblingLeaf)); thirdValidProof = new bytes32[](1); thirdValidProof[0] = siblingLeaf; } @@ -162,7 +182,7 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { emit RewardsClaimedForIndex(operatorAddress, 2, thirdOperatorPoints, thirdOperatorPoints); serviceManager.claimLatestOperatorRewards( - operatorSetId, thirdOperatorPoints, thirdValidProof + operatorSetId, thirdOperatorPoints, numberOfLeaves, leafIndex, thirdValidProof ); assertEq( @@ -180,7 +200,9 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { abi.encodeWithSelector(IServiceManagerErrors.NoRewardsRegistryForOperatorSet.selector) ); - serviceManager.claimLatestOperatorRewards(invalidSetId, operatorPoints, validProof); + serviceManager.claimLatestOperatorRewards( + invalidSetId, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_claimLatestOperatorRewards_AlreadyClaimed() public { @@ -193,7 +215,7 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { // First claim (uses latest merkle root - index 2) vm.prank(operatorAddress); serviceManager.claimLatestOperatorRewards( - operatorSetId, thirdOperatorPoints, thirdValidProof + operatorSetId, thirdOperatorPoints, numberOfLeaves, leafIndex, thirdValidProof ); // Second claim should fail @@ -203,7 +225,7 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { ); serviceManager.claimLatestOperatorRewards( - operatorSetId, thirdOperatorPoints, thirdValidProof + operatorSetId, thirdOperatorPoints, numberOfLeaves, leafIndex, thirdValidProof ); } @@ -220,7 +242,9 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { vm.expectEmit(true, true, true, true); emit RewardsClaimedForIndex(operatorAddress, 0, operatorPoints, operatorPoints); - serviceManager.claimOperatorRewards(operatorSetId, 0, operatorPoints, validProof); + serviceManager.claimOperatorRewards( + operatorSetId, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); assertEq( operatorAddress.balance, @@ -241,7 +265,7 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { // Claim from index 1 (second merkle root) vm.prank(operatorAddress); serviceManager.claimOperatorRewards( - operatorSetId, 1, secondOperatorPoints, secondValidProof + operatorSetId, 1, secondOperatorPoints, numberOfLeaves, leafIndex, secondValidProof ); assertEq( @@ -252,7 +276,9 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { // Claim from index 2 (third merkle root) vm.prank(operatorAddress); - serviceManager.claimOperatorRewards(operatorSetId, 2, thirdOperatorPoints, thirdValidProof); + serviceManager.claimOperatorRewards( + operatorSetId, 2, thirdOperatorPoints, numberOfLeaves, leafIndex, thirdValidProof + ); assertEq( operatorAddress.balance, @@ -283,7 +309,9 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.InvalidMerkleRootIndex.selector) ); - serviceManager.claimOperatorRewards(operatorSetId, 999, operatorPoints, validProof); + serviceManager.claimOperatorRewards( + operatorSetId, 999, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_claimOperatorRewards_AlreadyClaimed() public { @@ -295,14 +323,23 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { // First claim vm.prank(operatorAddress); - serviceManager.claimOperatorRewards(operatorSetId, 0, operatorPoints, validProof); + serviceManager.claimOperatorRewards( + operatorSetId, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); // Second claim should fail vm.prank(operatorAddress); vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.RewardsAlreadyClaimedForIndex.selector) ); - serviceManager.claimOperatorRewards(operatorSetId, 0, operatorPoints, validProof); + serviceManager.claimOperatorRewards( + operatorSetId, + 0, + operatorPoints, + 2, // numberOfLeaves (operator + sibling) + 0, // leafIndex (assuming operator leaf comes first) + validProof + ); } function test_claimOperatorRewardsBatch() public { @@ -337,7 +374,18 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { operatorAddress, rootIndices, points, expectedTotalRewards ); - serviceManager.claimOperatorRewardsBatch(operatorSetId, rootIndices, points, proofs); + uint256[] memory widths = new uint256[](3); + widths[0] = 2; + widths[1] = 2; + widths[2] = 2; + uint256[] memory leafIdxs = new uint256[](3); + leafIdxs[0] = 0; + leafIdxs[1] = 0; + leafIdxs[2] = 0; + + serviceManager.claimOperatorRewardsBatch( + operatorSetId, rootIndices, points, widths, leafIdxs, proofs + ); // Verify final balance includes all rewards assertEq( @@ -385,7 +433,15 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { ); vm.prank(operatorAddress); - serviceManager.claimOperatorRewardsBatch(operatorSetId, rootIndices, points, proofs); + uint256[] memory widths2 = new uint256[](2); + widths2[0] = 2; + widths2[1] = 2; + uint256[] memory leafIdxs2 = new uint256[](2); + leafIdxs2[0] = 0; + leafIdxs2[1] = 0; + serviceManager.claimOperatorRewardsBatch( + operatorSetId, rootIndices, points, widths2, leafIdxs2, proofs + ); // Verify balance and claim status assertEq( @@ -421,7 +477,20 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { vm.prank(operatorAddress); vm.expectRevert(abi.encodeWithSelector(IRewardsRegistryErrors.ArrayLengthMismatch.selector)); - serviceManager.claimOperatorRewardsBatch(operatorSetId, rootIndices, points, proofs); + + uint256[] memory numberOfLeaves = new uint256[](3); + numberOfLeaves[0] = 2; + numberOfLeaves[1] = 2; + numberOfLeaves[2] = 2; + + uint256[] memory leafIndices = new uint256[](3); + leafIndices[0] = 0; + leafIndices[1] = 0; + leafIndices[2] = 0; + + serviceManager.claimOperatorRewardsBatch( + operatorSetId, rootIndices, points, numberOfLeaves, leafIndices, proofs + ); } function test_claimOperatorRewardsBatch_AlreadyClaimedIndex() public { @@ -434,7 +503,12 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { vm.prank(operatorAddress); serviceManager.claimOperatorRewards( - operatorSetId, 1, secondOperatorPoints, secondValidProof + operatorSetId, + 1, + secondOperatorPoints, + 2, // numberOfLeaves (operator + sibling) + 0, // leafIndex (assuming operator leaf comes first) + secondValidProof ); // Now try to batch claim including the already claimed index 1 @@ -454,7 +528,18 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IRewardsRegistryErrors.RewardsAlreadyClaimedForIndex.selector) ); - serviceManager.claimOperatorRewardsBatch(operatorSetId, rootIndices, points, proofs); + + uint256[] memory numberOfLeaves = new uint256[](2); + numberOfLeaves[0] = 2; + numberOfLeaves[1] = 2; + + uint256[] memory leafIndices = new uint256[](2); + leafIndices[0] = 0; + leafIndices[1] = 0; + + serviceManager.claimOperatorRewardsBatch( + operatorSetId, rootIndices, points, numberOfLeaves, leafIndices, proofs + ); } function test_claimOperatorRewardsBatch_EmptyBatch() public { @@ -470,8 +555,13 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { uint256 initialBalance = operatorAddress.balance; + uint256[] memory numberOfLeaves = new uint256[](0); + uint256[] memory leafIndices = new uint256[](0); + vm.prank(operatorAddress); - serviceManager.claimOperatorRewardsBatch(operatorSetId, rootIndices, points, proofs); + serviceManager.claimOperatorRewardsBatch( + operatorSetId, rootIndices, points, numberOfLeaves, leafIndices, proofs + ); // Balance should remain unchanged assertEq( @@ -493,13 +583,18 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { secondOperatorSetId, IRewardsRegistry(address(secondRegistry)) ); - // Create a different merkle root for the second registry - bytes32 secondLeaf = keccak256(abi.encode(operatorAddress, operatorPoints)); - bytes32 secondSiblingLeaf = keccak256(abi.encodePacked("second sibling")); - (bytes32 leftLeaf, bytes32 rightLeaf) = secondLeaf < secondSiblingLeaf - ? (secondLeaf, secondSiblingLeaf) - : (secondSiblingLeaf, secondLeaf); - bytes32 secondRegistryMerkleRoot = keccak256(abi.encodePacked(leftLeaf, rightLeaf)); + // Create a different merkle root for the second registry using SCALE encoding + bytes memory secondLeafPreimage = + abi.encodePacked(operatorAddress, ScaleCodec.encodeU32(uint32(operatorPoints))); + bytes32 secondLeaf = keccak256(secondLeafPreimage); + + bytes memory secondSiblingPreimage = + abi.encodePacked(address(0x4444), ScaleCodec.encodeU32(uint32(80))); + bytes32 secondSiblingLeaf = keccak256(secondSiblingPreimage); + + // Since leafIndex = 0, our leaf is on the left + bytes32 secondRegistryMerkleRoot = + keccak256(abi.encodePacked(secondLeaf, secondSiblingLeaf)); // Set the merkle root in the second registry vm.prank(mockRewardsAgent); @@ -521,7 +616,7 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { ); vm.prank(operatorAddress); serviceManager.claimLatestOperatorRewards( - operatorSetId, thirdOperatorPoints, thirdValidProof + operatorSetId, thirdOperatorPoints, numberOfLeaves, leafIndex, thirdValidProof ); // Use latest root // Verify balance after first claim @@ -534,7 +629,9 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { // Claim from second registry uint256 balanceAfterFirstClaim = operatorAddress.balance; vm.prank(operatorAddress); - serviceManager.claimLatestOperatorRewards(secondOperatorSetId, operatorPoints, secondProof); + serviceManager.claimLatestOperatorRewards( + secondOperatorSetId, operatorPoints, numberOfLeaves, leafIndex, secondProof + ); // Verify balance after second claim assertEq( @@ -555,7 +652,9 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IServiceManagerErrors.OperatorNotInOperatorSet.selector) ); - serviceManager.claimLatestOperatorRewards(operatorSetId, operatorPoints, validProof); + serviceManager.claimLatestOperatorRewards( + operatorSetId, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_claimOperatorRewards_NotInOperatorSet() public { @@ -569,7 +668,9 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IServiceManagerErrors.OperatorNotInOperatorSet.selector) ); - serviceManager.claimOperatorRewards(operatorSetId, 0, operatorPoints, validProof); + serviceManager.claimOperatorRewards( + operatorSetId, 0, operatorPoints, numberOfLeaves, leafIndex, validProof + ); } function test_claimOperatorRewardsBatch_NotInOperatorSet() public { @@ -590,6 +691,12 @@ contract ServiceManagerRewardsRegistryTest is AVSDeployer { vm.expectRevert( abi.encodeWithSelector(IServiceManagerErrors.OperatorNotInOperatorSet.selector) ); - serviceManager.claimOperatorRewardsBatch(operatorSetId, rootIndices, points, proofs); + uint256[] memory widths3 = new uint256[](1); + widths3[0] = 2; + uint256[] memory leafIdxs3 = new uint256[](1); + leafIdxs3[0] = 0; + serviceManager.claimOperatorRewardsBatch( + operatorSetId, rootIndices, points, widths3, leafIdxs3, proofs + ); } } diff --git a/contracts/test/SnowbridgeIntegration.t.sol b/contracts/test/SnowbridgeIntegration.t.sol index 476540e9..d71b2070 100644 --- a/contracts/test/SnowbridgeIntegration.t.sol +++ b/contracts/test/SnowbridgeIntegration.t.sol @@ -24,6 +24,7 @@ import { IRewardsRegistryEvents, IRewardsRegistryErrors } from "../src/interfaces/IRewardsRegistry.sol"; import {SnowbridgeAndAVSDeployer} from "./utils/SnowbridgeAndAVSDeployer.sol"; +import {ScaleCodec} from "snowbridge/src/utils/ScaleCodec.sol"; import "forge-std/Test.sol"; contract SnowbridgeIntegrationTest is SnowbridgeAndAVSDeployer { @@ -105,7 +106,7 @@ contract SnowbridgeIntegrationTest is SnowbridgeAndAVSDeployer { _validatorAddresses[0], 0, _validatorPoints[0], uint256(_validatorPoints[0]) ); serviceManager.claimLatestOperatorRewards( - 0, _validatorPoints[0], rewardsProofFirstValidator + 0, _validatorPoints[0], 10, 0, rewardsProofFirstValidator ); vm.stopPrank(); @@ -126,7 +127,9 @@ contract SnowbridgeIntegrationTest is SnowbridgeAndAVSDeployer { emit IRewardsRegistryEvents.RewardsClaimedForIndex( _validatorAddresses[9], 0, _validatorPoints[9], uint256(_validatorPoints[9]) ); - serviceManager.claimLatestOperatorRewards(0, _validatorPoints[9], rewardsProofLastValidator); + serviceManager.claimLatestOperatorRewards( + 0, _validatorPoints[9], 10, 9, rewardsProofLastValidator + ); vm.stopPrank(); // Check that the last validator has received the rewards. @@ -353,11 +356,14 @@ contract SnowbridgeIntegrationTest is SnowbridgeAndAVSDeployer { bytes32[] memory leaves = new bytes32[](validators.length); for (uint256 i = 0; i < validators.length; i++) { - leaves[i] = keccak256(abi.encode(validators[i], points[i])); + // Use SCALE encoding for Substrate compatibility + bytes memory preimage = + abi.encodePacked(validators[i], ScaleCodec.encodeU32(uint32(points[i]))); + leaves[i] = keccak256(preimage); } - // We calculate the merkle root by sorting the pair before hashing (See Open Zeppelin Merkle Tree lib). - return MerkleUtils.calculateMerkleRoot(leaves, true); + // We calculate the merkle root without sorting for Substrate positional merkle tree. + return MerkleUtils.calculateMerkleRoot(leaves, false); } function _buildValidatorPointsProof( @@ -372,10 +378,13 @@ contract SnowbridgeIntegrationTest is SnowbridgeAndAVSDeployer { bytes32[] memory leaves = new bytes32[](validators.length); for (uint256 i = 0; i < validators.length; i++) { - leaves[i] = keccak256(abi.encode(validators[i], points[i])); + // Use SCALE encoding for Substrate compatibility + bytes memory preimage = + abi.encodePacked(validators[i], ScaleCodec.encodeU32(uint32(points[i]))); + leaves[i] = keccak256(preimage); } - return MerkleUtils.buildMerkleProof(leaves, leafIndex, true); + return MerkleUtils.buildMerkleProof(leaves, leafIndex, false); } function _buildMessagesMerkleTree( diff --git a/test/contract-bindings/generated.ts b/test/contract-bindings/generated.ts index d1965921..add0dab5 100644 --- a/test/contract-bindings/generated.ts +++ b/test/contract-bindings/generated.ts @@ -2103,6 +2103,8 @@ export const dataHavenServiceManagerAbi = [ inputs: [ { name: 'operatorSetId', internalType: 'uint32', type: 'uint32' }, { name: 'operatorPoints', internalType: 'uint256', type: 'uint256' }, + { name: 'numberOfLeaves', internalType: 'uint256', type: 'uint256' }, + { name: 'leafIndex', internalType: 'uint256', type: 'uint256' }, { name: 'proof', internalType: 'bytes32[]', type: 'bytes32[]' }, ], name: 'claimLatestOperatorRewards', @@ -2115,6 +2117,8 @@ export const dataHavenServiceManagerAbi = [ { name: 'operatorSetId', internalType: 'uint32', type: 'uint32' }, { name: 'rootIndex', internalType: 'uint256', type: 'uint256' }, { name: 'operatorPoints', internalType: 'uint256', type: 'uint256' }, + { name: 'numberOfLeaves', internalType: 'uint256', type: 'uint256' }, + { name: 'leafIndex', internalType: 'uint256', type: 'uint256' }, { name: 'proof', internalType: 'bytes32[]', type: 'bytes32[]' }, ], name: 'claimOperatorRewards', @@ -2127,6 +2131,8 @@ export const dataHavenServiceManagerAbi = [ { name: 'operatorSetId', internalType: 'uint32', type: 'uint32' }, { name: 'rootIndices', internalType: 'uint256[]', type: 'uint256[]' }, { name: 'operatorPoints', internalType: 'uint256[]', type: 'uint256[]' }, + { name: 'numberOfLeaves', internalType: 'uint256[]', type: 'uint256[]' }, + { name: 'leafIndices', internalType: 'uint256[]', type: 'uint256[]' }, { name: 'proofs', internalType: 'bytes32[][]', type: 'bytes32[][]' }, ], name: 'claimOperatorRewardsBatch', @@ -8135,6 +8141,8 @@ export const rewardsRegistryAbi = [ inputs: [ { name: 'operatorAddress', internalType: 'address', type: 'address' }, { name: 'operatorPoints', internalType: 'uint256', type: 'uint256' }, + { name: 'numberOfLeaves', internalType: 'uint256', type: 'uint256' }, + { name: 'leafIndex', internalType: 'uint256', type: 'uint256' }, { name: 'proof', internalType: 'bytes32[]', type: 'bytes32[]' }, ], name: 'claimLatestRewards', @@ -8147,6 +8155,8 @@ export const rewardsRegistryAbi = [ { name: 'operatorAddress', internalType: 'address', type: 'address' }, { name: 'rootIndex', internalType: 'uint256', type: 'uint256' }, { name: 'operatorPoints', internalType: 'uint256', type: 'uint256' }, + { name: 'numberOfLeaves', internalType: 'uint256', type: 'uint256' }, + { name: 'leafIndex', internalType: 'uint256', type: 'uint256' }, { name: 'proof', internalType: 'bytes32[]', type: 'bytes32[]' }, ], name: 'claimRewards', @@ -8159,6 +8169,8 @@ export const rewardsRegistryAbi = [ { name: 'operatorAddress', internalType: 'address', type: 'address' }, { name: 'rootIndices', internalType: 'uint256[]', type: 'uint256[]' }, { name: 'operatorPoints', internalType: 'uint256[]', type: 'uint256[]' }, + { name: 'numberOfLeaves', internalType: 'uint256[]', type: 'uint256[]' }, + { name: 'leafIndices', internalType: 'uint256[]', type: 'uint256[]' }, { name: 'proofs', internalType: 'bytes32[][]', type: 'bytes32[][]' }, ], name: 'claimRewardsBatch', diff --git a/test/suites/rewards-message.test.ts b/test/suites/rewards-message.test.ts new file mode 100644 index 00000000..5180c4b5 --- /dev/null +++ b/test/suites/rewards-message.test.ts @@ -0,0 +1,520 @@ +import { beforeAll, describe, expect, it } from "bun:test"; +import { logger } from "utils"; +import { + type Address, + BaseError, + ContractFunctionRevertedError, + decodeErrorResult, + decodeEventLog, + type Hex, + isAddressEqual, + padHex +} from "viem"; +import { BaseTestSuite } from "../framework"; +import { getContractInstance, parseRewardsInfoFile } from "../utils/contracts"; +import { waitForEthereumEvent } from "../utils/events"; +import * as rewardsHelpers from "../utils/rewards-helpers"; + +// Test configuration constants +const TEST_CONFIG = { + TIMEOUTS: { + ERA_END_WAIT: 600000, // 10 minutes - increased for era transitions + MESSAGE_EXECUTION: 120000, // 2 minutes + ROOT_UPDATE: 180000, // 3 minutes + CLAIM_EVENT: 30000, // 30 seconds - increased for reliability + OVERALL_TEST: 900000 // 15 minutes - increased for full suite + }, + DELAYS: { + RELAYER_INIT: 10000 // 10 seconds + } +} as const; + +class RewardsMessageTestSuite extends BaseTestSuite { + constructor() { + super({ + suiteName: "rewards-message" + }); + + this.setupHooks(); + } +} + +const suite = new RewardsMessageTestSuite(); + +let rewardsRegistry!: any; +let serviceManager!: any; +let gateway!: any; +let publicClient!: any; +let dhApi!: any; +let eraIndex!: number; +let messageId!: Hex; +let merkleRoot!: Hex; +let totalPoints!: bigint; +let newRootIndex!: bigint; +let validatorProofs!: Map; +// Persisted state from first successful claim for double-claim test +let claimedOperatorAddress!: Address; +let claimedProofData!: rewardsHelpers.ValidatorProofData; +let firstClaimGasUsed!: bigint; +let firstClaimBlockNumber!: bigint; + +describe("Rewards Message Flow", () => { + beforeAll(async () => { + logger.info("Starting rewards message flow tests"); + + // Get test connectors once for all tests + const connectors = suite.getTestConnectors(); + publicClient = connectors.publicClient; + dhApi = connectors.dhApi; + + // Acquire core contracts once for all tests + [rewardsRegistry, serviceManager, gateway] = await Promise.all([ + getContractInstance("RewardsRegistry"), + getContractInstance("ServiceManager"), + getContractInstance("Gateway") + ]); + }); + + describe("Infrastructure Setup", () => { + it("should verify rewards infrastructure deployment", async () => { + // Fetch rewards info + const rewardsInfo = await parseRewardsInfoFile(); + + expect(rewardsRegistry.address).toBeDefined(); + expect(rewardsInfo.RewardsAgent).toBeDefined(); + expect(gateway.address).toBeDefined(); + + // Validate configuration + const [agentAddress, avsAddress] = await Promise.all([ + publicClient.readContract({ + address: rewardsRegistry.address, + abi: rewardsRegistry.abi, + functionName: "rewardsAgent", + args: [] + }) as Promise
, + publicClient.readContract({ + address: rewardsRegistry.address, + abi: rewardsRegistry.abi, + functionName: "avs", + args: [] + }) as Promise
+ ]); + + expect(isAddressEqual(agentAddress, rewardsInfo.RewardsAgent as Address)).toBe(true); + expect(isAddressEqual(avsAddress, serviceManager.address as Address)).toBe(true); + + // Check DataHaven connectivity + const currentBlock = await dhApi.query.System.Number.getValue(); + expect(currentBlock > 0).toBe(true); + + logger.success("Rewards infrastructure verified"); + }); + }); + + describe("Era Transition and Message Emission", () => { + it( + "should wait for era end and capture rewards message", + async () => { + // Track current era and blocks until era end + const [currentBlock, currentEra, blocksUntilEraEnd] = await Promise.all([ + dhApi.query.System.Number.getValue(), + rewardsHelpers.getCurrentEra(dhApi), + rewardsHelpers.getBlocksUntilEraEnd(dhApi) + ]); + + logger.info("Era transition tracking:"); + logger.info(` Current block: ${currentBlock}`); + logger.info(` Current era: ${currentEra}`); + logger.info(` Blocks until era end: ${blocksUntilEraEnd}`); + + // Wait for era to end and capture the rewards message event + logger.info("⏳ Waiting for era to end and rewards message to be sent..."); + + const timeout = blocksUntilEraEnd * 6000 + TEST_CONFIG.DELAYS.RELAYER_INIT * 3; + const rewardsMessageEvent = await rewardsHelpers.waitForRewardsMessageSent( + dhApi, + currentEra, + timeout + ); + + expect(rewardsMessageEvent).not.toBeNull(); + if (!rewardsMessageEvent) throw new Error("Expected rewards message event to be defined"); + + // Store event data + messageId = rewardsMessageEvent.messageId as Hex; + merkleRoot = rewardsMessageEvent.merkleRoot as Hex; + totalPoints = rewardsMessageEvent.totalPoints; + eraIndex = rewardsMessageEvent.eraIndex; + + // Validate event data + expect(messageId).toBeDefined(); + expect(merkleRoot).toBeDefined(); + expect(totalPoints > 0n).toBe(true); + + logger.success(`Rewards message emitted for era ${eraIndex}`); + }, + TEST_CONFIG.TIMEOUTS.ERA_END_WAIT + ); + }); + + describe("Cross-Chain Message Execution", () => { + it( + "should execute rewards message on Ethereum via Gateway", + async () => { + logger.info("⏳ Waiting for message execution on Gateway..."); + + // Start watching from current block to avoid matching historical events + const fromBlock = await publicClient.getBlockNumber(); + + const executedEvent = await waitForEthereumEvent({ + client: publicClient, + address: gateway.address, + abi: gateway.abi, + eventName: "MessageExecuted", + fromBlock, + timeout: TEST_CONFIG.TIMEOUTS.MESSAGE_EXECUTION + }); + + expect(executedEvent.log).not.toBeNull(); + if (!executedEvent.log) throw new Error("Expected log to be defined"); + const log = executedEvent.log; + const _decoded = decodeEventLog({ + abi: gateway.abi, + data: log.data, + topics: log.topics, + eventName: "MessageExecuted" + }) as any; + + logger.success("Message executed on Ethereum:"); + logger.info(` Block: ${log.blockNumber}`); + logger.info(` Transaction: ${log.transactionHash}`); + }, + TEST_CONFIG.TIMEOUTS.MESSAGE_EXECUTION + ); + }); + + describe("Merkle Root Update", () => { + it( + "should update RewardsRegistry with new merkle root", + async () => { + const expectedRoot: Hex = padHex(merkleRoot, { size: 32 }); + const fromBlock = await publicClient.getBlockNumber(); + + logger.info("⏳ Waiting for merkle root update in RewardsRegistry..."); + + const rootUpdatedEvent = await waitForEthereumEvent({ + client: publicClient, + address: rewardsRegistry.address, + abi: rewardsRegistry.abi, + eventName: "RewardsMerkleRootUpdated", + args: { newRoot: expectedRoot }, + fromBlock, + timeout: TEST_CONFIG.TIMEOUTS.ROOT_UPDATE + }); + + expect(rootUpdatedEvent.log).not.toBeNull(); + if (!rootUpdatedEvent.log) throw new Error("Expected log to be defined"); + const rootLog = rootUpdatedEvent.log; + const rootDecoded = decodeEventLog({ + abi: rewardsRegistry.abi, + data: rootLog.data, + topics: rootLog.topics + }) as { args: { oldRoot: Hex; newRoot: Hex; newRootIndex: bigint } }; + const updateArgs = rootDecoded.args; + + // Store the new root index for claiming tests + newRootIndex = updateArgs.newRootIndex; + + logger.success("Merkle root updated:"); + logger.info(` Index: ${updateArgs.newRootIndex}`); + logger.info(` Old root: ${updateArgs.oldRoot}`); + logger.info(` New root: ${updateArgs.newRoot}`); + + // Verify the stored root matches the expected root + const storedRoot: Hex = (await publicClient.readContract({ + address: rewardsRegistry.address, + abi: rewardsRegistry.abi, + functionName: "merkleRootHistory", + args: [updateArgs.newRootIndex] + })) as Hex; + + expect(storedRoot.toLowerCase()).toEqual(updateArgs.newRoot.toLowerCase()); + expect(storedRoot.toLowerCase()).toEqual(expectedRoot.toLowerCase()); + }, + TEST_CONFIG.TIMEOUTS.ROOT_UPDATE + ); + }); + + describe("Merkle Proof Generation", () => { + it("should generate valid merkle proofs for all validators", async () => { + logger.info(`📊 Generating merkle proofs for era ${eraIndex}...`); + + // Get era reward points and generate proofs in parallel + const [eraPoints, proofMap] = await Promise.all([ + rewardsHelpers.getEraRewardPoints(dhApi, eraIndex), + rewardsHelpers.generateMerkleProofsForEra(dhApi, eraIndex) + ]); + + expect(eraPoints).toBeDefined(); + if (!eraPoints) throw new Error("Expected era points to be defined"); + expect(eraPoints.total > 0).toBe(true); + expect(proofMap.size > 0).toBe(true); + + // Store proofs for claiming tests + validatorProofs = proofMap; + + logger.success("Generated merkle proofs"); + + // Validate proof data structure (spot check) + const firstProofMaybe = validatorProofs.values().next().value; + expect(firstProofMaybe).toBeDefined(); + if (!firstProofMaybe) throw new Error("Expected first proof to be defined"); + const firstProof = firstProofMaybe; + expect(firstProof.proof).toBeDefined(); + expect(firstProof.points > 0).toBe(true); + expect(firstProof.numberOfLeaves > 0).toBe(true); + }); + }); + + describe("Rewards Claiming", () => { + it("should fund RewardsRegistry for payouts", async () => { + logger.info("💰 Funding RewardsRegistry for reward payouts..."); + + const { walletClient: fundingWallet } = suite.getTestConnectors(); + const fundingAmount = totalPoints; + + const fundingTx = await fundingWallet.sendTransaction({ + to: rewardsRegistry.address as Address, + value: fundingAmount, + chain: null + }); + + const fundingReceipt = await publicClient.waitForTransactionReceipt({ hash: fundingTx }); + expect(fundingReceipt.status).toBe("success"); + + // Verify contract balance + const contractBalance = await publicClient.getBalance({ + address: rewardsRegistry.address + }); + + expect(contractBalance > 0n).toBe(true); + + logger.success("RewardsRegistry funded:"); + logger.info(` Amount: ${fundingAmount} wei`); + logger.info(` Transaction: ${fundingTx}`); + logger.info(` Contract balance: ${contractBalance} wei`); + }); + + it( + "should successfully claim rewards for validator", + async () => { + logger.info("🎯 Claiming rewards for validator..."); + + // Ensure prerequisites + expect(validatorProofs).toBeDefined(); + expect(newRootIndex).toBeDefined(); + if (newRootIndex === undefined) { + throw new Error("Merkle root not updated yet; cannot claim rewards"); + } + + // Select first validator to claim + const firstEntry = validatorProofs.entries().next(); + expect(firstEntry.value).toBeDefined(); + if (!firstEntry.value) throw new Error("Expected entry to be defined"); + const entry = firstEntry.value; + const [, proofData] = entry; + + // Get validator credentials and create operator wallet + const factory = suite.getConnectorFactory(); + const credentials = rewardsHelpers.getValidatorCredentials(proofData.validatorAccount); + expect(credentials.privateKey).toBeDefined(); + if (!credentials.privateKey) throw new Error("missing validator private key"); + const operatorWallet = factory.createWalletClient(credentials.privateKey as `0x${string}`); + const resolvedOperator: Address = operatorWallet.account.address; + + // Record initial balance for validation + const balanceBefore = await publicClient.getBalance({ address: resolvedOperator }); + + // Submit claim transaction + const claimTx = await operatorWallet.writeContract({ + address: serviceManager.address as Address, + abi: serviceManager.abi, + functionName: "claimOperatorRewards", + chain: null, + args: [ + 0, // strategy index + newRootIndex, + BigInt(proofData.points), + BigInt(proofData.numberOfLeaves), + BigInt(proofData.leafIndex), + proofData.proof as readonly Hex[] + ] + }); + + logger.info(`📝 Claim transaction submitted: ${claimTx}`); + + // Wait for transaction confirmation + const claimReceipt = await publicClient.waitForTransactionReceipt({ hash: claimTx }); + expect(claimReceipt.status).toBe("success"); + + // Persist state for the double-claim test + claimedOperatorAddress = resolvedOperator; + claimedProofData = proofData; + firstClaimGasUsed = claimReceipt.gasUsed; + firstClaimBlockNumber = claimReceipt.blockNumber; + + // Wait for and validate claim event + const claimEvent = await waitForEthereumEvent({ + client: publicClient, + address: rewardsRegistry.address, + abi: rewardsRegistry.abi, + eventName: "RewardsClaimedForIndex", + fromBlock: claimReceipt.blockNumber - 1n, + timeout: TEST_CONFIG.TIMEOUTS.CLAIM_EVENT + }); + + expect(claimEvent.log).toBeDefined(); + if (!claimEvent.log) throw new Error("Expected log to be defined"); + const claimLog = claimEvent.log; + const claimDecoded = decodeEventLog({ + abi: rewardsRegistry.abi, + data: claimLog.data, + topics: claimLog.topics + }) as { + args: { + operatorAddress: Address; + rootIndex: bigint; + points: bigint; + rewardsAmount: bigint; + }; + }; + const claimArgs = claimDecoded.args; + + // Validate claim event data + expect(isAddressEqual(claimArgs.operatorAddress, resolvedOperator)).toBe(true); + expect(claimArgs.rootIndex).toEqual(newRootIndex); + expect(claimArgs.points).toEqual(BigInt(proofData.points)); + expect(claimArgs.rewardsAmount > 0n).toBe(true); + + logger.success("Rewards claimed successfully:"); + logger.info(` Operator: ${resolvedOperator}`); + logger.info(` Points: ${claimArgs.points}`); + logger.info(` Rewards: ${claimArgs.rewardsAmount} wei`); + logger.info(` Root index: ${claimArgs.rootIndex}`); + + // Validate balance change accounting for gas costs + const balanceAfter = await publicClient.getBalance({ address: resolvedOperator }); + const actualBalanceIncrease = balanceAfter - balanceBefore; + const gasUsedWei = claimReceipt.gasUsed * claimReceipt.effectiveGasPrice; + const adjustedIncrease = actualBalanceIncrease + gasUsedWei; + + logger.info("💰 Balance validation:"); + logger.info(` Gas used: ${gasUsedWei} wei`); + logger.info(` Adjusted balance increase: ${adjustedIncrease} wei`); + + expect(BigInt(adjustedIncrease)).toEqual(claimArgs.rewardsAmount); + expect(claimArgs.rewardsAmount).toEqual(BigInt(proofData.points)); + }, + TEST_CONFIG.TIMEOUTS.CLAIM_EVENT + ); + + it( + "should prevent double claiming of rewards", + async () => { + logger.info("🚫 Testing double-claim prevention (on-chain revert)..."); + + // Preconditions from previous test + expect(claimedProofData).toBeDefined(); + expect(claimedOperatorAddress).toBeDefined(); + expect(firstClaimGasUsed).toBeDefined(); + expect(firstClaimBlockNumber).toBeDefined(); + expect(newRootIndex).toBeDefined(); + if (newRootIndex === undefined) throw new Error("Merkle root not updated yet"); + + const factory = suite.getConnectorFactory(); + const credentials = rewardsHelpers.getValidatorCredentials( + claimedProofData.validatorAccount + ); + if (!credentials.privateKey) throw new Error("missing validator private key"); + const operatorWallet = factory.createWalletClient(credentials.privateKey as `0x${string}`); + + // Send a real transaction expected to revert. Provide explicit gas to avoid estimation/simulation. + const gasLimit = firstClaimGasUsed + 100_000n; + + const revertTxHash = await operatorWallet.writeContract({ + address: serviceManager.address as Address, + abi: serviceManager.abi, + functionName: "claimOperatorRewards", + args: [ + 0, + newRootIndex, + BigInt(claimedProofData.points), + BigInt(claimedProofData.numberOfLeaves), + BigInt(claimedProofData.leafIndex), + claimedProofData.proof as readonly Hex[] + ], + gas: gasLimit, + chain: null + }); + + const revertReceipt = await publicClient.waitForTransactionReceipt({ hash: revertTxHash }); + expect(revertReceipt.status).toBe("reverted"); + + // Verify custom error using eth_call at the same block + let decodedErrorName = ""; + try { + await publicClient.simulateContract({ + account: operatorWallet.account, + address: serviceManager.address as Address, + abi: serviceManager.abi, + functionName: "claimOperatorRewards", + args: [ + 0, + newRootIndex, + BigInt(claimedProofData.points), + BigInt(claimedProofData.numberOfLeaves), + BigInt(claimedProofData.leafIndex), + claimedProofData.proof as readonly Hex[] + ], + blockNumber: revertReceipt.blockNumber + }); + throw new Error("Expected simulateContract to revert"); + } catch (err: any) { + if (err instanceof BaseError) { + const revertError = err.walk((e) => e instanceof ContractFunctionRevertedError); + if (revertError instanceof ContractFunctionRevertedError) { + // First try viem's decoded data (only works if ABI included the error) + decodedErrorName = revertError.data?.errorName ?? ""; + // Fallback: decode the raw revert data using an ABI that includes the custom error + if (!decodedErrorName) { + const rawData = revertError.raw as Hex | undefined; + if (rawData) { + try { + const unionAbi = [ + ...(serviceManager.abi as any[]), + ...(rewardsRegistry.abi as any[]) + ]; + const decoded = decodeErrorResult({ abi: unionAbi as any, data: rawData }); + decodedErrorName = decoded.errorName; + } catch (_e) { + // ignore secondary decode errors + } + } + } + } else { + throw err; + } + } else { + throw err; + } + } + expect(decodedErrorName).toBe("RewardsAlreadyClaimedForIndex"); + + logger.success( + "Double-claim prevention verified (on-chain revert and correct custom error)" + ); + }, + TEST_CONFIG.TIMEOUTS.CLAIM_EVENT + ); + }); +}); diff --git a/test/utils/events.ts b/test/utils/events.ts index 8fdfbbfb..7318f295 100644 --- a/test/utils/events.ts +++ b/test/utils/events.ts @@ -1,6 +1,7 @@ import { firstValueFrom, of } from "rxjs"; -import { catchError, take, tap, timeout } from "rxjs/operators"; +import { catchError, map, filter as rxFilter, take, tap, timeout } from "rxjs/operators"; import type { Abi, Address, Log, PublicClient } from "viem"; +import { decodeEventLog } from "viem"; import { logger } from "./logger"; import type { DataHavenApi } from "./papi"; @@ -22,6 +23,8 @@ export interface DataHavenEventResult { event: string; /** Event data payload (null if timeout or error) */ data: T | null; + /** Metadata about when/where event was emitted */ + meta: any | null; } /** @@ -45,7 +48,7 @@ export interface WaitForDataHavenEventOptions { /** * Wait for a specific event on the DataHaven chain * @param options - Options for event waiting - * @returns Event result with pallet, event name, and data + * @returns Event result with pallet, event name, and converted data */ export async function waitForDataHavenEvent( options: WaitForDataHavenEventOptions @@ -55,18 +58,33 @@ export async function waitForDataHavenEvent( const eventWatcher = (api.event as any)?.[pallet]?.[event]; if (!eventWatcher?.watch) { logger.warn(`Event ${pallet}.${event} not found`); - return { pallet, event, data: null }; + return { pallet, event, data: null, meta: null }; } - let data: T | null; + let meta: any = null; + let data: T | null = null; + try { - data = await firstValueFrom( - eventWatcher.watch(filter).pipe( - tap((eventData: T) => { - logger.debug(`Event ${pallet}.${event} received`); - onEvent?.(eventData); + const matched: any = await firstValueFrom( + eventWatcher.watch().pipe( + // Log every raw emission from the watcher + tap(() => { + logger.debug(`Event ${pallet}.${event} received (raw)`); }), - take(1), // Always stop on first event + // Normalize to a consistent shape { payload, meta } + map((raw: any) => ({ payload: raw?.payload ?? raw, meta: raw?.meta ?? null })), + // Apply the optional filter BEFORE taking the first item + rxFilter(({ payload }) => { + if (!filter) return true; + try { + return filter(payload as T); + } catch { + return false; + } + }), + // Stop on the first matching event + take(1), + // Enforce an overall timeout while waiting for a matching event timeout({ first: timeoutMs, with: () => { @@ -80,11 +98,20 @@ export async function waitForDataHavenEvent( }) ) ); - } catch { + + if (matched) { + meta = matched.meta; + data = matched.payload as T; + if (data) { + onEvent?.(data); + } + } + } catch (error) { + logger.error(`Unexpected error waiting for event ${pallet}.${event}: ${error}`); data = null; } - return { pallet, event, data }; + return { pallet, event, data, meta }; } // ================== Ethereum Event Utilities ================== @@ -161,14 +188,48 @@ export async function waitForEthereumEvent( args, fromBlock, onLogs: (logs) => { - if (logs.length > 0) { - matchedLog = logs[0]; + logger.debug(`Ethereum event ${eventName} received: ${logs.length} logs`); + + // If args include non-indexed fields, viem cannot pre-filter them. + // Post-filter by decoding logs and matching provided args if any. + let selected: Log | null = null; + if (args && Object.keys(args).length > 0) { + for (const candidate of logs) { + try { + const decoded = decodeEventLog({ + abi, + eventName: eventName as any, + data: candidate.data, + topics: candidate.topics + }); + const decodedArgs = (decoded as any).args ?? {}; + const allMatch = Object.entries(args as Record).every( + ([key, value]) => decodedArgs?.[key] === value + ); + if (allMatch) { + selected = candidate; + break; + } + } catch { + // Ignore decode errors and continue scanning + } + } + } + + if (!selected && (!args || Object.keys(args).length === 0) && logs.length > 0) { + // Only fallback to first log when no args filter provided + selected = logs[0]; + } + + if (selected) { + matchedLog = selected; if (onEvent) { onEvent(matchedLog); } cleanup(); resolve(matchedLog); } + // If no selected log matched, keep watching until timeout }, onError: (error: unknown) => { // Log and continue; transient watcher errors shouldn't abort the wait diff --git a/test/utils/rewards-helpers.ts b/test/utils/rewards-helpers.ts new file mode 100644 index 00000000..d37c7870 --- /dev/null +++ b/test/utils/rewards-helpers.ts @@ -0,0 +1,228 @@ +import validatorSet from "../configs/validator-set.json"; +import { waitForDataHavenEvent } from "./events"; +import { logger } from "./logger"; +import type { DataHavenApi } from "./papi"; + +// Small hex helper +const toHex = (x: unknown): `0x${string}` => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const anyX: any = x as any; + if (anyX?.asHex) return anyX.asHex(); + const s = anyX?.toString?.() ?? ""; + return `0x${s}` as `0x${string}`; +}; + +// External Validators Rewards Events (normalized) +export interface RewardsMessageSent { + messageId: `0x${string}`; + merkleRoot: `0x${string}`; + eraIndex: number; + totalPoints: bigint; + inflation: bigint; +} + +// Era tracking utilities +export async function getCurrentEra(dhApi: DataHavenApi): Promise { + // Get the active era from ExternalValidators pallet + const activeEra = await dhApi.query.ExternalValidators.ActiveEra.getValue(); + + // ActiveEra can be null at chain genesis + if (!activeEra) { + return 0; + } + + return activeEra.index; +} + +export function getEraLengthInBlocks(dhApi: DataHavenApi): number { + // Read constants directly from runtime metadata + const consts: any = (dhApi as unknown as { consts?: unknown }).consts ?? {}; + const epochDuration = Number(consts?.Babe?.EpochDuration ?? 10); // blocks per session + const sessionsPerEra = Number(consts?.ExternalValidators?.SessionsPerEra ?? 1); + return epochDuration * sessionsPerEra; +} + +export async function getBlocksUntilEraEnd(dhApi: DataHavenApi): Promise { + const currentBlock = await dhApi.query.System.Number.getValue(); + const eraLength = getEraLengthInBlocks(dhApi) || 10; + const mod = currentBlock % eraLength; + return mod === 0 ? eraLength : eraLength - mod; +} + +// Validator monitoring and rewards data +export interface EraRewardPoints { + total: number; + individual: Map; +} + +export async function getEraRewardPoints( + dhApi: DataHavenApi, + eraIndex: number +): Promise { + try { + const rewardPoints = + await dhApi.query.ExternalValidatorsRewards.RewardPointsForEra.getValue(eraIndex); + + if (!rewardPoints) { + return null; + } + + // Convert the storage format to our interface + const individual = new Map(); + for (const [account, points] of rewardPoints.individual) { + individual.set(account.toString(), points); + } + + return { + total: rewardPoints.total, + individual + }; + } catch (error) { + logger.error(`Failed to get era reward points for era ${eraIndex}: ${error}`); + return null; + } +} + +// Merkle proof generation using DataHaven runtime API +export interface ValidatorProofData { + validatorAccount: string; + operatorAddress: string; + points: number; + proof: string[]; + leaf: string; + numberOfLeaves: number; + leafIndex: number; +} + +export async function generateMerkleProofForValidator( + dhApi: DataHavenApi, + validatorAccount: string, + eraIndex: number +): Promise<{ proof: string[]; leaf: string; numberOfLeaves: number; leafIndex: number } | null> { + try { + // Call the runtime API to generate merkle proof + const merkleProof = await dhApi.apis.ExternalValidatorsRewardsApi.generate_rewards_merkle_proof( + validatorAccount, + eraIndex + ); + + if (!merkleProof) { + logger.debug( + `No merkle proof available for validator ${validatorAccount} in era ${eraIndex}` + ); + return null; + } + + // Convert the proof to hex strings + const proof = merkleProof.proof.map((node: unknown) => toHex(node)); + + const leaf = toHex(merkleProof.leaf); + + const numberOfLeaves = Number(merkleProof.number_of_leaves as bigint); + const leafIndex = Number(merkleProof.leaf_index as bigint); + + return { proof, leaf, numberOfLeaves, leafIndex }; + } catch (error) { + logger.error(`Failed to generate merkle proof for validator ${validatorAccount}: ${error}`); + return null; + } +} + +/** + * Validator credentials containing operator address and private key + */ +export interface ValidatorCredentials { + operatorAddress: `0x${string}`; + privateKey: `0x${string}` | null; +} + +/** + * Gets validator credentials (operator address and private key) by solochain address + * @param validatorAccount The validator's solochain address + * @returns The validator's credentials including operator address and private key + */ +export function getValidatorCredentials(validatorAccount: string): ValidatorCredentials { + const normalizedAccount = validatorAccount.toLowerCase(); + + // Find matching validator by solochain address + const match = validatorSet.validators.find( + (v) => v.solochainAddress.toLowerCase() === normalizedAccount + ); + + if (match) { + return { + operatorAddress: match.publicKey as `0x${string}`, + privateKey: match.privateKey as `0x${string}` + }; + } + + // Fallback: assume the input is already an Ethereum address, but no private key available + logger.debug(`No mapping found for ${validatorAccount}, using as-is without private key`); + return { + operatorAddress: validatorAccount as `0x${string}`, + privateKey: null + }; +} + +// Generate merkle proofs for all validators in an era +export async function generateMerkleProofsForEra( + dhApi: DataHavenApi, + eraIndex: number +): Promise> { + // Get era reward points + const eraPoints = await getEraRewardPoints(dhApi, eraIndex); + if (!eraPoints) { + logger.warn(`No reward points found for era ${eraIndex}`); + return new Map(); + } + + const entries = await Promise.all( + [...eraPoints.individual].map(async ([validatorAccount, points]) => { + const merkleData = await generateMerkleProofForValidator(dhApi, validatorAccount, eraIndex); + if (!merkleData) return null; + const credentials = getValidatorCredentials(validatorAccount); + const value: ValidatorProofData = { + validatorAccount, + operatorAddress: credentials.operatorAddress, + points, + proof: merkleData.proof, + leaf: merkleData.leaf, + numberOfLeaves: merkleData.numberOfLeaves, + leafIndex: merkleData.leafIndex + }; + return [credentials.operatorAddress, value] as const; + }) + ); + + const filtered = entries.filter(Boolean) as [string, ValidatorProofData][]; + const proofs = new Map(filtered); + logger.info(`Generated ${proofs.size} merkle proofs for era ${eraIndex}`); + return proofs; +} + +// Rewards message event -> normalized return + +export async function waitForRewardsMessageSent( + dhApi: DataHavenApi, + expectedEra?: number, + timeout = 120000 +): Promise { + const result = await waitForDataHavenEvent({ + api: dhApi, + pallet: "ExternalValidatorsRewards", + event: "RewardsMessageSent", + filter: expectedEra !== undefined ? (event: any) => event.era_index === expectedEra : undefined, + timeout + }); + + if (!result?.data) return null; + + const data: any = result.data; + return { + messageId: data.message_id.asHex(), + merkleRoot: data.rewards_merkle_root.asHex(), + eraIndex: data.era_index, + totalPoints: data.total_points, + inflation: data.inflation_amount + }; +}