From 71d9dd0b3df59f6cf19cb18244946dec52bc6502 Mon Sep 17 00:00:00 2001
From: Kingter <83567446+kingster-will@users.noreply.github.com>
Date: Fri, 11 Oct 2024 21:31:17 -0700
Subject: [PATCH] Restrict distributeRewards Function to Grouping Module in
 EvenSplitGroupPool (#268)

---
 .../modules/grouping/EvenSplitGroupPool.sol   |  2 +-
 .../modules/grouping/EvenSplitGroupPool.t.sol | 22 +++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/contracts/modules/grouping/EvenSplitGroupPool.sol b/contracts/modules/grouping/EvenSplitGroupPool.sol
index afad10e0..c03e438b 100644
--- 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
index 7d584b5c..332e3dca 100644
--- 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();
+    }
 }