Resolve L-02: enable harvest to be called before setting cometRewards

brritoxyz · Jan 12, 2024 · a9fecd6 · a9fecd6
1 parent 8dcd945
commit a9fecd6
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 11 deletions.
diff --git a/src/BrrETH.sol b/src/BrrETH.sol
@@ -45,7 +45,7 @@ contract BrrETH is Ownable, ERC4626 {
         uint256 supplyAssets,
         uint256 fees
     );
-    event SetCometRewards(address);
+    event SetCometRewards(address, bool);
     event SetRouter(address);
     event SetRewardFee(uint256);
     event SetProtocolFeeReceiver(address);
@@ -196,7 +196,7 @@ contract BrrETH is Ownable, ERC4626 {
     }
 
     /// @notice Claim rewards and convert them into the vault asset.
-    function harvest() external {
+    function harvest() public {
         cometRewards.claim(_COMET, address(this), true);
 
         ICometRewards.RewardConfig memory rewardConfig = cometRewards
@@ -257,13 +257,18 @@ contract BrrETH is Ownable, ERC4626 {
     /**
      * @notice Set the Comet Rewards contract.
      * @param  _cometRewards  address  Comet Rewards contract address.
+     * @param  shouldHarvest  bool     Whether to call `harvest` before setting `cometRewards`.
      */
-    function setCometRewards(address _cometRewards) external onlyOwner {
+    function setCometRewards(
+        address _cometRewards,
+        bool shouldHarvest
+    ) external onlyOwner {
         if (_cometRewards == address(0)) revert InvalidCometRewards();
+        if (shouldHarvest) harvest();
 
         cometRewards = ICometRewards(_cometRewards);
 
-        emit SetCometRewards(_cometRewards);
+        emit SetCometRewards(_cometRewards, shouldHarvest);
     }
 
     /**

diff --git a/test/BrrETH.t.sol b/test/BrrETH.t.sol
@@ -499,50 +499,91 @@ contract BrrETHTest is Helper {
     function testCannotSetCometRewardsUnauthorized() external {
         address msgSender = address(0);
         address cometRewards = address(0xbeef);
+        bool shouldHarvest = false;
 
         assertTrue(msgSender != vault.owner());
 
         vm.prank(msgSender);
         vm.expectRevert(Ownable.Unauthorized.selector);
 
-        vault.setCometRewards(cometRewards);
+        vault.setCometRewards(cometRewards, shouldHarvest);
     }
 
     function testCannotSetCometRewardsInvalidCometRewards() external {
         address cometRewards = address(0);
+        bool shouldHarvest = false;
 
         vm.expectRevert(BrrETH.InvalidCometRewards.selector);
 
-        vault.setCometRewards(cometRewards);
+        vault.setCometRewards(cometRewards, shouldHarvest);
     }
 
     function testSetCometRewards() external {
         address cometRewards = address(0xbeef);
+        bool shouldHarvest = false;
 
         assertTrue(cometRewards != address(vault.cometRewards()));
 
         vm.expectEmit(true, true, true, true, address(vault));
 
-        emit BrrETH.SetCometRewards(cometRewards);
+        emit BrrETH.SetCometRewards(cometRewards, shouldHarvest);
 
-        vault.setCometRewards(cometRewards);
+        vault.setCometRewards(cometRewards, shouldHarvest);
 
         assertEq(cometRewards, address(vault.cometRewards()));
     }
 
-    function testSetCometRewardsFuzz(address cometRewards) external {
+    function testSetCometRewardsShouldHarvest() external {
+        address cometRewards = address(0xbeef);
+        bool shouldHarvest = true;
+
+        assertTrue(cometRewards != address(vault.cometRewards()));
+
+        // Deposit and accrue enough time to ensure `harvest` is called (i.e. emits `Harvest` event).
+        vault.deposit{value: 1 ether}(address(this));
+
+        skip(1 days);
+
+        // Event members are unchecked, we just need to know that `harvest` was called.
+        vm.expectEmit(false, false, false, false, address(vault));
+
+        emit BrrETH.Harvest(_COMP, 0, 0, 0);
+
+        vm.expectEmit(true, true, true, true, address(vault));
+
+        emit BrrETH.SetCometRewards(cometRewards, shouldHarvest);
+
+        vault.setCometRewards(cometRewards, shouldHarvest);
+
+        assertEq(cometRewards, address(vault.cometRewards()));
+    }
+
+    function testSetCometRewardsFuzz(
+        address cometRewards,
+        bool shouldHarvest
+    ) external {
         vm.assume(
             cometRewards != address(0) &&
                 cometRewards != address(vault.cometRewards())
         );
 
         assertTrue(cometRewards != address(vault.cometRewards()));
 
+        if (shouldHarvest) {
+            vault.deposit{value: 1 ether}(address(this));
+
+            skip(1 days);
+
+            vm.expectEmit(false, false, false, false, address(vault));
+
+            emit BrrETH.Harvest(_COMP, 0, 0, 0);
+        }
+
         vm.expectEmit(true, true, true, true, address(vault));
 
-        emit BrrETH.SetCometRewards(cometRewards);
+        emit BrrETH.SetCometRewards(cometRewards, shouldHarvest);
 
-        vault.setCometRewards(cometRewards);
+        vault.setCometRewards(cometRewards, shouldHarvest);
 
         assertEq(cometRewards, address(vault.cometRewards()));
     }