Add distributors in MerkleDistributor (#10)

* Add distributors in MerkleDistributor * Ignore macos stuff * Rework * Add gas snapshots * Fix tests * Fix tests 2 * Review fixes * Snapshots * Fix docstring --------- Co-authored-by: Dmitri Tsumak <[email protected]>
stakewise · Feb 6, 2025 · 0dae0b0 · 0dae0b0
1 parent 251e398
commit 0dae0b0
Show file tree

Hide file tree

Showing 6 changed files with 141 additions and 15 deletions.
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -32,13 +32,17 @@ EthAaveLeverageStrategyTest:test_upgradeProxy_NotRegisteredProxy() (gas: 45298)
 EthAaveLeverageStrategyTest:test_upgradeProxy_VaultUpgradeConfigSameAddress() (gas: 980507)
 EthAaveLeverageStrategyTest:test_upgradeProxy_VaultUpgradeConfigZeroAddress() (gas: 959850)
 EthAaveLeverageStrategyTest:test_upgradeProxy_WithExitingPosition() (gas: 1101835)
-MerkleDistributorTest:test_claim() (gas: 731451)
-MerkleDistributorTest:test_constructor() (gas: 1420037)
-MerkleDistributorTest:test_distributeOneTime() (gas: 334784)
-MerkleDistributorTest:test_distributePeriodically() (gas: 358162)
-MerkleDistributorTest:test_setRewardsDelay() (gas: 71336)
-MerkleDistributorTest:test_setRewardsMinOracles() (gas: 142651)
-MerkleDistributorTest:test_setRewardsRoot() (gas: 401435)
+MerkleDistributorTest:test_addDistributorNormal() (gas: 69504)
+MerkleDistributorTest:test_addDistributorUnauthorized() (gas: 37009)
+MerkleDistributorTest:test_claim() (gas: 731525)
+MerkleDistributorTest:test_constructor() (gas: 1474501)
+MerkleDistributorTest:test_distributeOneTime() (gas: 384076)
+MerkleDistributorTest:test_distributePeriodically() (gas: 407781)
+MerkleDistributorTest:test_removeDistributorNormal() (gas: 139333)
+MerkleDistributorTest:test_removeDistributorUnauthorized() (gas: 86492)
+MerkleDistributorTest:test_setRewardsDelay() (gas: 71380)
+MerkleDistributorTest:test_setRewardsMinOracles() (gas: 142827)
+MerkleDistributorTest:test_setRewardsRoot() (gas: 401545)
 StakeHelpersTest:test_calculateStake() (gas: 338892)
 StakeHelpersTest:test_calculateUnstake() (gas: 524501)
 StrategiesRegistryTest:test_addStrategyProxy() (gas: 240040)

diff --git a/.gitignore b/.gitignore
@@ -11,3 +11,7 @@ docs/
 # IDE
 .idea
 .vscode
+.history
+
+# MacOS
+.DS_Store
diff --git a/snapshots/MerkleDistributorTest.json b/snapshots/MerkleDistributorTest.json
@@ -1,9 +1,11 @@
 {
+  "MerkleDistributorTest_test_addDistributor": "51298",
   "MerkleDistributorTest_test_claim": "146428",
-  "MerkleDistributorTest_test_constructor": "1410695",
-  "MerkleDistributorTest_test_distributeOneTime": "74966",
-  "MerkleDistributorTest_test_distributePeriodically": "73090",
+  "MerkleDistributorTest_test_constructor": "1465059",
+  "MerkleDistributorTest_test_distributeOneTime": "74989",
+  "MerkleDistributorTest_test_distributePeriodically": "73113",
+  "MerkleDistributorTest_test_removeDistributor": "26886",
   "MerkleDistributorTest_test_setRewardsDelay": "30905",
-  "MerkleDistributorTest_test_setRewardsMinOracles": "36907",
+  "MerkleDistributorTest_test_setRewardsMinOracles": "36951",
   "MerkleDistributorTest_test_setRewardsRoot": "52139"
 }
diff --git a/src/MerkleDistributor.sol b/src/MerkleDistributor.sol
@@ -24,8 +24,12 @@ contract MerkleDistributor is Ownable2Step, EIP712, IMerkleDistributor {
 
     IKeeperOracles private immutable _keeper;
 
+    /// @inheritdoc IMerkleDistributor
     mapping(address token => mapping(address user => uint256 cumulativeAmount)) public claimedAmounts;
 
+    /// @inheritdoc IMerkleDistributor
+    mapping(address distributor => bool isEnabled) public distributors;
+
     /// @inheritdoc IMerkleDistributor
     bytes32 public rewardsRoot;
 
@@ -60,6 +64,14 @@ contract MerkleDistributor is Ownable2Step, EIP712, IMerkleDistributor {
         _transferOwnership(_initialOwner);
     }
 
+    /**
+     * @notice Reverts if called by any account other than an enabled distributor.
+     */
+    modifier onlyDistributor() {
+        if (!distributors[msg.sender]) revert Errors.AccessDenied();
+        _;
+    }
+
     /// @inheritdoc IMerkleDistributor
     function getNextRewardsRootUpdateTimestamp() public view returns (uint64) {
         return lastUpdateTimestamp + rewardsDelay;
@@ -119,14 +131,20 @@ contract MerkleDistributor is Ownable2Step, EIP712, IMerkleDistributor {
         emit RewardsMinOraclesUpdated(msg.sender, _rewardsMinOracles);
     }
 
+    /// @inheritdoc IMerkleDistributor
+    function setDistributor(address distributor, bool isEnabled) external onlyOwner {
+        distributors[distributor] = isEnabled;
+        emit DistributorUpdated(msg.sender, distributor, isEnabled);
+    }
+
     /// @inheritdoc IMerkleDistributor
     function distributePeriodically(
         address token,
         uint256 amount,
         uint256 delayInSeconds,
         uint256 durationInSeconds,
         bytes calldata extraData
-    ) external onlyOwner {
+    ) external onlyDistributor {
         if (amount == 0) revert InvalidAmount();
         if (durationInSeconds == 0) revert InvalidDuration();
 
@@ -140,7 +158,7 @@ contract MerkleDistributor is Ownable2Step, EIP712, IMerkleDistributor {
         uint256 amount,
         string calldata rewardsIpfsHash,
         bytes calldata extraData
-    ) external onlyOwner {
+    ) external onlyDistributor {
         if (amount == 0) revert InvalidAmount();
 
         SafeERC20.safeTransferFrom(IERC20(token), msg.sender, address(this), amount);

diff --git a/src/interfaces/IMerkleDistributor.sol b/src/interfaces/IMerkleDistributor.sol
@@ -66,6 +66,7 @@ interface IMerkleDistributor {
 
     /**
      * @notice Emitted when the rewards are claimed
+     * @param caller The address of the caller
      * @param account The address of the account
      * @param tokens The list of tokens
      * @param cumulativeAmounts The cumulative amounts of tokens
@@ -74,6 +75,14 @@ interface IMerkleDistributor {
         address indexed caller, address indexed account, address[] tokens, uint256[] cumulativeAmounts
     );
 
+    /**
+     * @notice Emitted when a distributor is added or removed
+     * @param caller The address of the caller
+     * @param distributor The address of the distributor
+     * @param isEnabled The status of the distributor
+     */
+    event DistributorUpdated(address indexed caller, address indexed distributor, bool isEnabled);
+
     /**
      * @notice Get the current rewards Merkle Tree root
      * @return The current rewards Merkle Tree root
@@ -112,6 +121,14 @@ interface IMerkleDistributor {
      */
     function claimedAmounts(address token, address user) external view returns (uint256 cumulativeAmount);
 
+    /**
+     * @notice Get the status of a distributor, is it enabled or not
+     * @param distributor The address of the distributor
+     */
+    function distributors(
+        address distributor
+    ) external view returns (bool isEnabled);
+
     /**
      * @notice Get the next rewards root update timestamp
      * @return The next rewards root update timestamp
@@ -146,6 +163,13 @@ interface IMerkleDistributor {
         uint64 newRewardsMinOracles
     ) external;
 
+    /**
+     * @notice Add or remove a distributor. Can only be called by the owner.
+     * @param distributor The address of the distributor
+     * @param isEnabled The status of the distributor, true for adding distributor, false for removing distributor
+     */
+    function setDistributor(address distributor, bool isEnabled) external;
+
     /**
      * @notice Distribute tokens every rewards delay for a specific duration
      * @param token The address of the token

diff --git a/test/MerkleDistributor.t.sol b/test/MerkleDistributor.t.sol
@@ -85,11 +85,81 @@ contract MerkleDistributorTest is Test {
         assertEq(distributor.rewardsDelay(), newDelay, 'Should correctly update rewardsDelay');
     }
 
+    function test_addDistributorUnauthorized() public {
+        address account = address(5);
+
+        // Try to add from invalid address
+        vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, address(this)));
+        distributor.setDistributor(account, true);
+
+        assertEq(distributor.distributors(account), false, 'Should not add distributor');
+    }
+
+    function test_addDistributorNormal() public {
+        address account = address(5);
+
+        // Add a valid distributor
+        vm.startPrank(owner);
+        vm.expectEmit(address(distributor));
+        emit IMerkleDistributor.DistributorUpdated(owner, account, true);
+        vm.startSnapshotGas('MerkleDistributorTest_test_addDistributor');
+        distributor.setDistributor(account, true);
+        vm.stopSnapshotGas();
+        vm.stopPrank();
+
+        assertEq(distributor.distributors(account), true, 'Should correctly add distributor');
+        assertEq(distributor.distributors(vm.addr(6)), false, 'Not every address is a distributor');
+    }
+
+    function test_removeDistributorUnauthorized() public {
+        address account = address(5);
+
+        // Add distributor
+        vm.prank(owner);
+        distributor.setDistributor(account, true);
+
+        // Try to remove from invalid address
+        vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, address(this)));
+        distributor.setDistributor(account, false);
+
+        assertEq(distributor.distributors(account), true, 'Should not remove distributor');
+    }
+
+    function test_removeDistributorNormal() public {
+        // 2 distributors, remove the first one
+        address account = address(5);
+        address account2 = address(6);
+
+        // Add distributors
+        vm.startPrank(owner);
+        distributor.setDistributor(account, true);
+        distributor.setDistributor(account2, true);
+        vm.stopPrank();
+
+        // Remove distributor
+        vm.startPrank(owner);
+        vm.expectEmit(address(distributor));
+        emit IMerkleDistributor.DistributorUpdated(owner, account, false);
+        vm.startSnapshotGas('MerkleDistributorTest_test_removeDistributor');
+        distributor.setDistributor(account, false);
+        vm.stopSnapshotGas();
+        vm.stopPrank();
+
+        assertEq(distributor.distributors(account), false, 'Should correctly remove distributor');
+        assertEq(distributor.distributors(account2), true, 'Should correctly remove distributor');
+    }
+
     function test_distributePeriodically() public {
         uint256 amount = 100 ether;
-        vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, address(this)));
+
+        // Ensure unauthorized accounts cannot call the function
+        vm.expectRevert(abi.encodeWithSelector(Errors.AccessDenied.selector));
         distributor.distributePeriodically(address(swiseToken), amount, 3600, 86_400, '');
 
+        // Add owner to distributors
+        vm.prank(owner);
+        distributor.setDistributor(owner, true);
+
         // Impersonate the owner to approve SWISE and distribute tokens
         vm.startPrank(owner);
         deal(address(swiseToken), owner, amount); // Give owner SWISE
@@ -115,9 +185,13 @@ contract MerkleDistributorTest is Test {
         uint256 amount = 100 ether;
 
         // Ensure unauthorized accounts cannot call the function
-        vm.expectRevert(abi.encodeWithSelector(Ownable.OwnableUnauthorizedAccount.selector, address(this)));
+        vm.expectRevert(abi.encodeWithSelector(Errors.AccessDenied.selector));
         distributor.distributeOneTime(address(swiseToken), amount, 'ipfsHash', 'extraData');
 
+        // Add owner to distributors
+        vm.prank(owner);
+        distributor.setDistributor(owner, true);
+
         // Impersonate the owner to approve SWISE and distribute tokens
         vm.startPrank(owner);
         deal(address(swiseToken), owner, amount); // Give owner SWISE
-Original file line number
+Diff line change
@@ Expand Up / @@ -11,3 +11,7 @@ docs/ @@
     # IDE
     .idea
     .vscode
+    .history
+    # MacOS
+    .DS_Store