From ebb7d36cde980da66131e64d17b1e0ea6a643f89 Mon Sep 17 00:00:00 2001
From: kphed <kpcrypto69@gmail.com>
Date: Mon, 13 Nov 2023 07:38:26 -0500
Subject: [PATCH] Prevent new Flywheel Rewards contracts from being the zero
 addr

---
 src/FlywheelCore.sol | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/FlywheelCore.sol b/src/FlywheelCore.sol
index 94db1f3..eeb8bca 100644
--- 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,6 +165,7 @@ contract FlywheelCore is Auth {
 
     /// @notice swap out the flywheel rewards contract
     function setFlywheelRewards(IFlywheelRewards newFlywheelRewards) external requiresAuth {
+        if (address(newFlywheelRewards) == address(0)) revert InvalidAddress();
         if (address(flywheelRewards) != address(0)) {
             uint256 oldRewardBalance = rewardToken.balanceOf(address(flywheelRewards));