Merge pull request #71 from kphed/update-set-flywheelrewards

Transfer balance from old the FlywheelRewards only if it's not the zero address
fei-protocol · Nov 13, 2023 · 379c738 · 379c738
2 parents c9b4003 + ebb7d36
commit 379c738
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/src/FlywheelCore.sol b/src/FlywheelCore.sol
@@ -40,6 +40,8 @@ contract FlywheelCore is Auth {
     /// @notice optional booster module for calculating virtual balances on strategies
     IFlywheelBooster public flywheelBooster;
 
+    error InvalidAddress();
+
     constructor(
         ERC20 _rewardToken,
         IFlywheelRewards _flywheelRewards,
@@ -163,9 +165,12 @@ contract FlywheelCore is Auth {
 
     /// @notice swap out the flywheel rewards contract
     function setFlywheelRewards(IFlywheelRewards newFlywheelRewards) external requiresAuth {
-        uint256 oldRewardBalance = rewardToken.balanceOf(address(flywheelRewards));
-        if (oldRewardBalance > 0) {
-            rewardToken.safeTransferFrom(address(flywheelRewards), address(newFlywheelRewards), oldRewardBalance);
+        if (address(newFlywheelRewards) == address(0)) revert InvalidAddress();
+        if (address(flywheelRewards) != address(0)) {
+            uint256 oldRewardBalance = rewardToken.balanceOf(address(flywheelRewards));
+
+            if (oldRewardBalance != 0)
+                rewardToken.safeTransferFrom(address(flywheelRewards), address(newFlywheelRewards), oldRewardBalance);
         }
 
         flywheelRewards = newFlywheelRewards;