Only call balanceOf and perform a transfer if flywheel isn't the zero…

… addr
fei-protocol · Nov 13, 2023 · 4875d5a · 4875d5a
1 parent d493a37
commit 4875d5a
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/FlywheelCore.sol b/src/FlywheelCore.sol
@@ -163,9 +163,11 @@ contract FlywheelCore is Auth {
 
     /// @notice swap out the flywheel rewards contract
     function setFlywheelRewards(IFlywheelRewards newFlywheelRewards) external requiresAuth {
-        uint256 oldRewardBalance = rewardToken.balanceOf(address(flywheelRewards));
-        if (address(flywheelRewards) != address(0) && oldRewardBalance > 0) {
-            rewardToken.safeTransferFrom(address(flywheelRewards), address(newFlywheelRewards), oldRewardBalance);
+        if (address(flywheelRewards) != address(0)) {
+            uint256 oldRewardBalance = rewardToken.balanceOf(address(flywheelRewards));
+
+            if (oldRewardBalance != 0)
+                rewardToken.safeTransferFrom(address(flywheelRewards), address(newFlywheelRewards), oldRewardBalance);
         }
 
         flywheelRewards = newFlywheelRewards;