Merge pull request sushi-labs#81 from sushiswap/mockUpgrade

Added Origin Check in Rewarder
CodeXSky · Apr 19, 2021 · 787ca14 · 787ca14
2 parents 6f84e57 + 913c2f3
commit 787ca14
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/contracts/mocks/RewarderMock.sol b/contracts/mocks/RewarderMock.sol
@@ -12,13 +12,15 @@ contract RewarderMock is IRewarder {
     uint256 private immutable rewardMultiplier;
     IERC20 private immutable rewardToken;
     uint256 private constant REWARD_TOKEN_DIVISOR = 1e18;
+    address private immutable MASTERCHEF_V2;
 
-    constructor (uint256 _rewardMultiplier, IERC20 _rewardToken) public {
+    constructor (uint256 _rewardMultiplier, IERC20 _rewardToken, address _MASTERCHEF_V2) public {
         rewardMultiplier = _rewardMultiplier;
         rewardToken = _rewardToken;
+        MASTERCHEF_V2 = _MASTERCHEF_V2;
     }
 
-    function onSushiReward (uint256, address user, address to, uint256 sushiAmount, uint256) override external {
+    function onSushiReward (uint256, address user, address to, uint256 sushiAmount, uint256) onlyMCV2 override external {
         uint256 pendingReward = sushiAmount.mul(rewardMultiplier) / REWARD_TOKEN_DIVISOR;
         uint256 rewardBal = rewardToken.balanceOf(address(this));
         if (pendingReward > rewardBal) {
@@ -35,5 +37,13 @@ contract RewarderMock is IRewarder {
         _rewardAmounts[0] = sushiAmount.mul(rewardMultiplier) / REWARD_TOKEN_DIVISOR;
         return (_rewardTokens, _rewardAmounts);
     }
+
+    modifier onlyMCV2 {
+        require(
+            msg.sender == MASTERCHEF_V2,
+            "Only MCV2 can call this function."
+        );
+        _;
+    }
 
 }
diff --git a/test/MasterChefV2.test.ts b/test/MasterChefV2.test.ts
@@ -31,7 +31,7 @@ describe("MasterChefV2", function () {
         ["rlp", this.ERC20Mock, ["LP", "rLPT", getBigNumber(10)]],
         ["r", this.ERC20Mock, ["Reward", "RewardT", getBigNumber(100000)]],
     ])
-    await deploy(this, [["rewarder", this.RewarderMock, [getBigNumber(1), this.r.address]]])
+    await deploy(this, [["rewarder", this.RewarderMock, [getBigNumber(1), this.r.address, this.chef2.address]]])
     await this.dummy.approve(this.chef2.address, getBigNumber(10))
     await this.chef2.init(this.dummy.address)
     await this.rlp.transfer(this.bob.address, getBigNumber(1))