From 0cb829df842d90fe1087ab3d2dab5f3d255f4660 Mon Sep 17 00:00:00 2001
From: Akshat Mittal <itsreallyakshat@gmail.com>
Date: Thu, 9 May 2024 20:09:11 +0530
Subject: [PATCH] Emergency Actions (#2)

* Emergency actions

* Block token when removing

* Remove redundant check
---
 .../rewards/GenericMultiRewardsVault.sol      | 55 ++++++++++++++++++-
 contracts/rewards/definitions.sol             |  2 +
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/contracts/rewards/GenericMultiRewardsVault.sol b/contracts/rewards/GenericMultiRewardsVault.sol
index 4468cb7..8a980ab 100644
--- a/contracts/rewards/GenericMultiRewardsVault.sol
+++ b/contracts/rewards/GenericMultiRewardsVault.sol
@@ -90,7 +90,7 @@ contract GenericMultiRewardsVault is ERC4626, Ownable {
      * @param _rewardTokens Array of rewardTokens for which rewards should be claimed.
      * @dev This function will revert if any of the rewardTokens have zero rewards accrued.
      */
-    function claimRewards(address user, IERC20[] memory _rewardTokens) external accrueRewards(msg.sender, user) {
+    function claimRewards(address user, IERC20[] memory _rewardTokens) external accrueRewards(_msgSender(), user) {
         for (uint8 i; i < _rewardTokens.length; i++) {
             uint256 rewardAmount = accruedRewards[user][_rewardTokens[i]];
 
@@ -112,6 +112,7 @@ contract GenericMultiRewardsVault is ERC4626, Ownable {
     mapping(IERC20 rewardToken => RewardInfo rewardInfo) public rewardInfos;
     mapping(IERC20 rewardToken => address distributor) public distributorInfo;
     mapping(IERC20 rewardToken => uint256 excessRewards) public leftoverRewards;
+    mapping(IERC20 rewardToken => bool isBlocked) public isRewardTokenBlocked;
 
     mapping(address user => mapping(IERC20 rewardToken => uint256 rewardIndex)) public userIndex; // {qRewardTok}
     mapping(address user => mapping(IERC20 rewardToken => uint256 accruedRewards)) public accruedRewards; // {qRewardTok}
@@ -135,6 +136,9 @@ contract GenericMultiRewardsVault is ERC4626, Ownable {
         if (asset() == address(rewardToken)) {
             revert Errors.RewardTokenCanNotBeStakingToken();
         }
+        if (isRewardTokenBlocked[rewardToken]) {
+            revert Errors.RewardTokenBlocked(rewardToken);
+        }
 
         RewardInfo memory rewards = rewardInfos[rewardToken];
         if (rewards.lastUpdatedTimestamp > 0) {
@@ -146,7 +150,7 @@ contract GenericMultiRewardsVault is ERC4626, Ownable {
                 revert Errors.ZeroRewardsSpeed();
             }
 
-            SafeERC20.safeTransferFrom(rewardToken, msg.sender, address(this), amount);
+            SafeERC20.safeTransferFrom(rewardToken, _msgSender(), address(this), amount);
         }
 
         rewardTokens.push(rewardToken);
@@ -171,6 +175,51 @@ contract GenericMultiRewardsVault is ERC4626, Ownable {
         emit Events.RewardInfoUpdate(rewardToken, rewardsPerSecond, rewardsEndTimestamp);
     }
 
+    /**
+     * @notice Removes a rewardToken. Caller must be owner. EMERGENCY ONLY.
+     * @param rewardToken Token that can be earned by staking.
+     * @dev This action is destructive and will disable reward accrual for this token.
+     *      ONLY CONSIDER USING DURING EMERGENCY.
+     */
+    function removeRewardToken(IERC20 rewardToken, bool _accrue) external onlyOwner {
+        RewardInfo memory rewards = rewardInfos[rewardToken];
+
+        if (rewards.lastUpdatedTimestamp == 0) {
+            revert Errors.InvalidRewardToken(rewardToken);
+        }
+        if (_accrue) {
+            _accrueRewards(rewardToken, _accrueStatic(rewards));
+        }
+
+        delete rewardInfos[rewardToken];
+        delete distributorInfo[rewardToken];
+        delete leftoverRewards[rewardToken];
+        isRewardTokenBlocked[rewardToken] = true;
+
+        uint256 totalRewardTokens = rewardTokens.length;
+
+        for (uint256 i; i < totalRewardTokens; ++i) {
+            if (rewardTokens[i] == rewardToken) {
+                rewardTokens[i] = rewardTokens[totalRewardTokens - 1];
+                rewardTokens.pop();
+
+                break;
+            }
+        }
+    }
+    /**
+     * @notice Ragequit the system. EMERGENCY ONLY.
+     * @dev Allows the user to withdraw their tokens without accruing or claiming rewards.
+     *      ONLY CONSIDER USING DURING EMERGENCY.
+     */
+    function ragequit() external {
+        uint256 shares = balanceOf(_msgSender());
+        uint256 assets = previewRedeem(shares);
+
+        totalDeposited -= assets;
+        super._withdraw(_msgSender(), _msgSender(), _msgSender(), assets, shares);
+    }
+
     /**
      * @notice Updates distributor for `rewardToken`
      * @param rewardToken Token that can be earned by staking.
@@ -253,7 +302,7 @@ contract GenericMultiRewardsVault is ERC4626, Ownable {
             revert Errors.RewardTokenDoesNotExist(rewardToken);
         }
 
-        SafeERC20.safeTransferFrom(rewardToken, msg.sender, address(this), amount);
+        SafeERC20.safeTransferFrom(rewardToken, _msgSender(), address(this), amount);
 
         _accrueRewards(rewardToken, rewards.rewardsPerSecond == 0 ? amount : _accrueStatic(rewards));
 
diff --git a/contracts/rewards/definitions.sol b/contracts/rewards/definitions.sol
index d322323..e050df6 100644
--- a/contracts/rewards/definitions.sol
+++ b/contracts/rewards/definitions.sol
@@ -16,9 +16,11 @@ struct RewardInfo {
 
 abstract contract Errors {
     // Reward Management
+    error InvalidRewardToken(IERC20 rewardToken);
     error RewardTokenAlreadyExist(IERC20 rewardToken);
     error RewardTokenDoesNotExist(IERC20 rewardToken);
     error RewardTokenCanNotBeStakingToken();
+    error RewardTokenBlocked(IERC20 rewardToken);
     error ZeroAmount();
     error RewardsAreDynamic(IERC20 rewardToken);
     error ZeroRewardsSpeed();