Upgrade sjoe

.gitignore

@@ -18,3 +18,7 @@ contracts/.deps/
 .idea
 .openzeppelin
 .vscode
+
+out
+broadcast
+forge-cache/*

.gitmodules

@@ -0,0 +1,4 @@
+[submodule "lib/forge-std"]
+	path = lib/forge-std
+	url = https://github.com/foundry-rs/forge-std
+	branch = v1.5.5

contracts/StableJoeStaking.sol

@@ -15,7 +15,7 @@ import "@openzeppelin/contracts-upgradeable/token/ERC20/SafeERC20Upgradeable.sol
  * harvests. Users deposit JOE and receive a share of what has been sent by MoneyMaker based on their participation of
  * the total deposited JOE. It is similar to a MasterChef, but we allow for claiming of different reward tokens
  * (in case at some point we wish to change the stablecoin rewarded).
- * Every time `updateReward(token)` is called, We distribute the balance of that tokens as rewards to users that are
+ * Every time `_updateReward(token)` is called, We distribute the balance of that tokens as rewards to users that are
  * currently staking inside this contract, and they can claim it using `withdraw(0)`
  */
 contract StableJoeStaking is Initializable, OwnableUpgradeable {
@@ -40,28 +40,47 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
          */
     }
 
-    IERC20Upgradeable public joe;
+    // @dev gap to keep the storage ordering, replace `IERC20Upgradeable public joe;`
+    uint256[1] private __gap0;
+
+    /// @notice The address of the JOE token
+    IERC20Upgradeable public immutable joe;
 
     /// @dev Internal balance of JOE, this gets updated on user deposits / withdrawals
     /// this allows to reward users with JOE
     uint256 public internalJoeBalance;
+
     /// @notice Array of tokens that users can claim
     IERC20Upgradeable[] public rewardTokens;
+
+    /// @notice Mapping to check if a token is a reward token
     mapping(IERC20Upgradeable => bool) public isRewardToken;
+
     /// @notice Last reward balance of `token`
     mapping(IERC20Upgradeable => uint256) public lastRewardBalance;
 
+    /// @notice The address where deposit fees will be sent
     address public feeCollector;
+    /// @notice Reentrancy guard
+    bool public reentrant;
 
     /// @notice The deposit fee, scaled to `DEPOSIT_FEE_PERCENT_PRECISION`
     uint256 public depositFeePercent;
+
+    /// @dev gap to keep the storage ordering, replace `uint256 public DEPOSIT_FEE_PERCENT_PRECISION;`
+    uint256[1] private __gap1;
+
     /// @notice The precision of `depositFeePercent`
-    uint256 public DEPOSIT_FEE_PERCENT_PRECISION;
+    uint256 public constant DEPOSIT_FEE_PERCENT_PRECISION = 1e18;
 
     /// @notice Accumulated `token` rewards per share, scaled to `ACC_REWARD_PER_SHARE_PRECISION`
     mapping(IERC20Upgradeable => uint256) public accRewardPerShare;
+
+    /// @dev gap to keep the storage ordering, replace `uint256 public ACC_REWARD_PER_SHARE_PRECISION;`
+    uint256[1] private __gap3;
+
     /// @notice The precision of `accRewardPerShare`
-    uint256 public ACC_REWARD_PER_SHARE_PRECISION;
+    uint256 public constant ACC_REWARD_PER_SHARE_PRECISION = 1e24;
 
     /// @dev Info of each user that stakes JOE
     mapping(address => UserInfo) private userInfo;
@@ -87,42 +106,58 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
     /// @notice Emitted when owner removes a token from the reward tokens list
     event RewardTokenRemoved(address token);
 
+    /// @notice Emitted when owner sweeps a token
+    event TokenSwept(address token, address to, uint256 amount);
+
+    /**
+     * @notice Reentrancy guard
+     */
+    modifier nonReentrant() {
+        require(!reentrant, "StableJoeStaking: reentrant call");
+        reentrant = true;
+        _;
+        reentrant = false;
+    }
+
+    /**
+     * @notice Construct a new StableJoeStaking contract
+     * @param _joe The address of the JOE token
+     */
+    constructor(IERC20Upgradeable _joe) initializer {
+        require(address(_joe) != address(0), "StableJoeStaking: joe can't be address(0)");
+
+        joe = _joe;
+    }
+
     /**
      * @notice Initialize a new StableJoeStaking contract
      * @dev This contract needs to receive an ERC20 `_rewardToken` in order to distribute them
-     * (with MoneyMaker in our case)
      * @param _rewardToken The address of the ERC20 reward token
-     * @param _joe The address of the JOE token
      * @param _feeCollector The address where deposit fees will be sent
      * @param _depositFeePercent The deposit fee percent, scalled to 1e18, e.g. 3% is 3e16
      */
     function initialize(
         IERC20Upgradeable _rewardToken,
-        IERC20Upgradeable _joe,
         address _feeCollector,
         uint256 _depositFeePercent
     ) external initializer {
         __Ownable_init();
         require(address(_rewardToken) != address(0), "StableJoeStaking: reward token can't be address(0)");
-        require(address(_joe) != address(0), "StableJoeStaking: joe can't be address(0)");
         require(_feeCollector != address(0), "StableJoeStaking: fee collector can't be address(0)");
         require(_depositFeePercent <= 5e17, "StableJoeStaking: max deposit fee can't be greater than 50%");
 
-        joe = _joe;
         depositFeePercent = _depositFeePercent;
         feeCollector = _feeCollector;
 
         isRewardToken[_rewardToken] = true;
         rewardTokens.push(_rewardToken);
-        DEPOSIT_FEE_PERCENT_PRECISION = 1e18;
-        ACC_REWARD_PER_SHARE_PRECISION = 1e24;
     }
 
     /**
      * @notice Deposit JOE for reward token allocation
      * @param _amount The amount of JOE to deposit
      */
-    function deposit(uint256 _amount) external {
+    function deposit(uint256 _amount) external nonReentrant {
         UserInfo storage user = userInfo[_msgSender()];
 
         uint256 _fee = _amount.mul(depositFeePercent).div(DEPOSIT_FEE_PERCENT_PRECISION);
@@ -135,7 +170,7 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
         uint256 _len = rewardTokens.length;
         for (uint256 i; i < _len; i++) {
             IERC20Upgradeable _token = rewardTokens[i];
-            updateReward(_token);
+            _updateReward(_token);
 
             uint256 _previousRewardDebt = user.rewardDebt[_token];
             user.rewardDebt[_token] = _newAmount.mul(accRewardPerShare[_token]).div(ACC_REWARD_PER_SHARE_PRECISION);
@@ -146,15 +181,17 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
                     .div(ACC_REWARD_PER_SHARE_PRECISION)
                     .sub(_previousRewardDebt);
                 if (_pending != 0) {
-                    safeTokenTransfer(_token, _msgSender(), _pending);
+                    _safeTokenTransfer(_token, _msgSender(), _pending);
                     emit ClaimReward(_msgSender(), address(_token), _pending);
                 }
             }
         }
 
         internalJoeBalance = internalJoeBalance.add(_amountMinusFee);
-        joe.safeTransferFrom(_msgSender(), feeCollector, _fee);
-        joe.safeTransferFrom(_msgSender(), address(this), _amountMinusFee);
+
+        if (_fee > 0) joe.safeTransferFrom(_msgSender(), feeCollector, _fee);
+        if (_amountMinusFee > 0) joe.safeTransferFrom(_msgSender(), address(this), _amountMinusFee);
+
         emit Deposit(_msgSender(), _amountMinusFee, _fee);
     }
 
@@ -188,9 +225,11 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
             "StableJoeStaking: token can't be added"
         );
         require(rewardTokens.length < 25, "StableJoeStaking: list of token too big");
+        require(accRewardPerShare[_rewardToken] == 0, "StableJoeStaking: reward token can't be re-added");
+
         rewardTokens.push(_rewardToken);
         isRewardToken[_rewardToken] = true;
-        updateReward(_rewardToken);
+
         emit RewardTokenAdded(address(_rewardToken));
     }
 
@@ -200,7 +239,6 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
      */
     function removeRewardToken(IERC20Upgradeable _rewardToken) external onlyOwner {
         require(isRewardToken[_rewardToken], "StableJoeStaking: token can't be removed");
-        updateReward(_rewardToken);
         isRewardToken[_rewardToken] = false;
         uint256 _len = rewardTokens.length;
         for (uint256 i; i < _len; i++) {
@@ -253,7 +291,7 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
      * @notice Withdraw JOE and harvest the rewards
      * @param _amount The amount of JOE to withdraw
      */
-    function withdraw(uint256 _amount) external {
+    function withdraw(uint256 _amount) external nonReentrant {
         UserInfo storage user = userInfo[_msgSender()];
         uint256 _previousAmount = user.amount;
         require(_amount <= _previousAmount, "StableJoeStaking: withdraw amount exceeds balance");
@@ -264,7 +302,7 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
         if (_previousAmount != 0) {
             for (uint256 i; i < _len; i++) {
                 IERC20Upgradeable _token = rewardTokens[i];
-                updateReward(_token);
+                _updateReward(_token);
 
                 uint256 _pending = _previousAmount
                     .mul(accRewardPerShare[_token])
@@ -273,7 +311,7 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
                 user.rewardDebt[_token] = _newAmount.mul(accRewardPerShare[_token]).div(ACC_REWARD_PER_SHARE_PRECISION);
 
                 if (_pending != 0) {
-                    safeTokenTransfer(_token, _msgSender(), _pending);
+                    _safeTokenTransfer(_token, _msgSender(), _pending);
                     emit ClaimReward(_msgSender(), address(_token), _pending);
                 }
             }
@@ -287,10 +325,13 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
     /**
      * @notice Withdraw without caring about rewards. EMERGENCY ONLY
      */
-    function emergencyWithdraw() external {
+    function emergencyWithdraw() external nonReentrant {
         UserInfo storage user = userInfo[_msgSender()];
 
         uint256 _amount = user.amount;
+
+        require(_amount > 0, "StableJoeStaking: can't withdraw 0");
+
         user.amount = 0;
         uint256 _len = rewardTokens.length;
         for (uint256 i; i < _len; i++) {
@@ -303,11 +344,11 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
     }
 
     /**
-     * @notice Update reward variables
+     * @dev Update reward variables
+     * Needs to be called before any deposit or withdrawal
      * @param _token The address of the reward token
-     * @dev Needs to be called before any deposit or withdrawal
      */
-    function updateReward(IERC20Upgradeable _token) public {
+    function _updateReward(IERC20Upgradeable _token) internal {
         require(isRewardToken[_token], "StableJoeStaking: wrong reward token");
 
         uint256 _totalJoe = internalJoeBalance;
@@ -329,13 +370,30 @@ contract StableJoeStaking is Initializable, OwnableUpgradeable {
     }
 
     /**
-     * @notice Safe token transfer function, just in case if rounding error
+     * @notice Sweep token to the `_to` address
+     * @param _token The address of the token to sweep
+     * @param _to The address that will receive `_token` balance
+     */
+    function sweep(IERC20Upgradeable _token, address _to) external onlyOwner {
+        require(!isRewardToken[_token] && address(_token) != address(joe), "StableJoeStaking: token can't be swept");
+
+        uint256 _balance = _token.balanceOf(address(this));
+
+        require(_balance > 0, "StableJoeStaking: can't sweep 0");
+
+        _token.safeTransfer(_to, _balance);
+
+        emit TokenSwept(address(_token), _to, _balance);
+    }
+
+    /**
+     * @dev Safe token transfer function, just in case if rounding error
      * causes pool to not have enough reward tokens
      * @param _token The address of then token to transfer
      * @param _to The address that will receive `_amount` `rewardToken`
      * @param _amount The amount to send to `_to`
      */
-    function safeTokenTransfer(
+    function _safeTokenTransfer(
         IERC20Upgradeable _token,
         address _to,
         uint256 _amount

foundry.toml

@@ -0,0 +1,30 @@
+[profile.default]
+src = 'contracts'
+out = 'out'
+libs = ["node_modules", "lib"]
+test = 'test/foundry'
+cache_path  = 'forge-cache'
+
+remappings = [
+    '@ensdomains/=node_modules/@ensdomains/',
+    '@openzeppelin/=node_modules/@openzeppelin/',
+    '@solidity-parser/=node_modules/truffle-flattener/node_modules/@solidity-parser/',
+    '@uniswap/=node_modules/@uniswap/',
+    'eth-gas-reporter/=node_modules/eth-gas-reporter/',
+    'hardhat-deploy/=node_modules/hardhat-deploy/',
+    'hardhat/=node_modules/hardhat/',
+]
+
+[rpc_endpoints]
+avalanche = "https://api.avax.network/ext/bc/C/rpc"
+fuji = "https://api.avax-test.network/ext/bc/C/rpc"
+arbitrum = "https://arb1.arbitrum.io/rpc"
+bsc = "https://bscrpc.com"
+
+[etherscan]
+arbitrum = { key = "${ARBISCAN_API_KEY}", chain = 42161 }
+avalanche = { key = "${SNOWTRACE_API_KEY}", chain = 43114 }
+arbitrum_goerli = { key = "${ARBISCAN_API_KEY}", chain = 421613 }
+fuji = { key = "${SNOWTRACE_API_KEY}", chain = 43113 }
+
+# See more config options https://github.com/foundry-rs/foundry/tree/master/config