Restrict distributeRewards Function to Grouping Module in EvenSplitGr…

…oupPool (storyprotocol#268)
kingster-will · Oct 12, 2024 · 71d9dd0 · 71d9dd0
1 parent 52e9f84
commit 71d9dd0
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/contracts/modules/grouping/EvenSplitGroupPool.sol b/contracts/modules/grouping/EvenSplitGroupPool.sol
@@ -124,7 +124,7 @@ contract EvenSplitGroupPool is IGroupRewardPool, ProtocolPausableUpgradeable, UU
         address groupId,
         address token,
         address[] calldata ipIds
-    ) external whenNotPaused returns (uint256[] memory rewards) {
+    ) external whenNotPaused onlyGroupingModule returns (uint256[] memory rewards) {
         rewards = _getAvailableReward(groupId, token, ipIds);
         uint256 totalRewards = 0;
         for (uint256 i = 0; i < ipIds.length; i++) {

diff --git a/test/foundry/modules/grouping/EvenSplitGroupPool.t.sol b/test/foundry/modules/grouping/EvenSplitGroupPool.t.sol
@@ -10,6 +10,7 @@ import { PILFlavors } from "../../../../contracts/lib/PILFlavors.sol";
 import { EvenSplitGroupPool } from "../../../../contracts/modules/grouping/EvenSplitGroupPool.sol";
 import { MockERC721 } from "../../mocks/token/MockERC721.sol";
 import { BaseTest } from "../../utils/BaseTest.t.sol";
+import { Errors } from "../../../../contracts/lib/Errors.sol";
 
 contract EvenSplitGroupPoolTest is BaseTest {
     using Strings for *;
@@ -188,6 +189,7 @@ contract EvenSplitGroupPoolTest is BaseTest {
         uint256[] memory rewards = rewardPool.getAvailableReward(group1, address(erc20), ipIds);
         assertEq(rewards[0], 100);
 
+        vm.prank(address(groupingModule));
         rewards = rewardPool.distributeRewards(group1, address(erc20), ipIds);
         assertEq(rewards[0], 100);
 
@@ -198,4 +200,24 @@ contract EvenSplitGroupPoolTest is BaseTest {
 
         vm.stopPrank();
     }
+
+    function test_EvenSplitGroupPool_revert_Only_GroupingModule() public {
+        vm.startPrank(address(0x123));
+        vm.expectRevert(
+            abi.encodeWithSelector(Errors.EvenSplitGroupPool__CallerIsNotGroupingModule.selector, address(0x123))
+        );
+        rewardPool.removeIp(group1, ipId1);
+
+        vm.expectRevert(
+            abi.encodeWithSelector(Errors.EvenSplitGroupPool__CallerIsNotGroupingModule.selector, address(0x123))
+        );
+        rewardPool.addIp(group1, ipId1);
+
+        vm.expectRevert(
+            abi.encodeWithSelector(Errors.EvenSplitGroupPool__CallerIsNotGroupingModule.selector, address(0x123))
+        );
+        rewardPool.distributeRewards(group1, address(erc20), new address[](0));
+
+        vm.stopPrank();
+    }
 }