[VEN-1567]: add sweep token in ReserveHelpers

contracts/RiskFund/ProtocolShareReserve.sol

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: BSD-3-Clause
 pragma solidity 0.8.13;
 
-import { Ownable2StepUpgradeable } from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import { IERC20Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/IERC20Upgradeable.sol";
 import { SafeERC20Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/utils/SafeERC20Upgradeable.sol";
 
@@ -11,12 +10,7 @@ import { ReserveHelpers } from "./ReserveHelpers.sol";
 import { IRiskFund } from "./IRiskFund.sol";
 import { ensureNonzeroAddress } from "../lib/validators.sol";
 
-/**
- * @title ProtocolShareReserve
- * @author Venus
- * @notice Contract used to store and distribute the reserves generated in the markets.
- */
-contract ProtocolShareReserve is Ownable2StepUpgradeable, ExponentialNoError, ReserveHelpers, IProtocolShareReserve {
+contract ProtocolShareReserve is ExponentialNoError, ReserveHelpers, IProtocolShareReserve {
     using SafeERC20Upgradeable for IERC20Upgradeable;
 
     address private protocolIncome;

contracts/RiskFund/ReserveHelpers.sol

@@ -1,21 +1,21 @@
 // SPDX-License-Identifier: BSD-3-Clause
 pragma solidity 0.8.13;
 
+import { Ownable2StepUpgradeable } from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import { IERC20Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/IERC20Upgradeable.sol";
 import { SafeERC20Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/utils/SafeERC20Upgradeable.sol";
 
 import { ensureNonzeroAddress } from "../lib/validators.sol";
 import { ComptrollerInterface } from "../ComptrollerInterface.sol";
 import { PoolRegistryInterface } from "../Pool/PoolRegistryInterface.sol";
 
-/**
- * @title ReserveHelpers
- * @author Venus
- * @notice Contract with basic features to track/hold different assets for different Comptrollers.
- */
-contract ReserveHelpers {
+contract ReserveHelpers is Ownable2StepUpgradeable {
     using SafeERC20Upgradeable for IERC20Upgradeable;
 
+    uint256 private constant NOT_ENTERED = 1;
+
+    uint256 private constant ENTERED = 2;
+
     // Store the previous state for the asset transferred to ProtocolShareReserve combined(for all pools).
     mapping(address => uint256) internal assetsReserves;
 
@@ -26,6 +26,11 @@ contract ReserveHelpers {
     // Address of pool registry contract
     address internal poolRegistry;
 
+    /**
+     * @dev Guard variable for re-entrancy checks
+     */
+    uint256 internal status;
+
     /**
      * @dev This empty reserved space is put in place to allow future versions to add new
      * variables without shifting down storage in the inheritance chain.
@@ -38,6 +43,40 @@ contract ReserveHelpers {
     /// @param amount An amount by which the reserves have increased
     event AssetsReservesUpdated(address indexed comptroller, address indexed asset, uint256 amount);
 
+    /// @notice event emitted on sweep token success
+    event SweepToken(address indexed token, address indexed to, uint256 amount);
+
+    /**
+     * @dev Prevents a contract from calling itself, directly or indirectly.
+     */
+    modifier nonReentrant() {
+        require(status != ENTERED, "re-entered");
+        status = ENTERED;
+        _;
+        status = NOT_ENTERED;
+    }
+
+    /**
+     * @notice A public function to sweep accidental BEP-20 transfers to this contract. Tokens are sent to the address `to`, provided in input
+     * @param _token The address of the BEP-20 token to sweep
+     * @param _to Recipient of the output tokens.
+     * @custom:error ZeroAddressNotAllowed is thrown when asset address is zero
+     * @custom:access Only Owner
+     */
+    function sweepToken(address _token, address _to) external onlyOwner nonReentrant {
+        ensureNonzeroAddress(_to);
+        uint256 balanceDfference_;
+        uint256 balance_ = IERC20Upgradeable(_token).balanceOf(address(this));
+
+        require(balance_ > assetsReserves[_token], "ReserveHelpers: Zero surplus tokens");
+        unchecked {
+            balanceDfference_ = balance_ - assetsReserves[_token];
+        }
+
+        IERC20Upgradeable(_token).safeTransfer(_to, balanceDfference_);
+        emit SweepToken(_token, _to, balanceDfference_);
+    }
+
     /**
      * @notice Get the Amount of the asset in the risk fund for the specific pool.
      * @param comptroller  Comptroller address(pool).

contracts/RiskFund/RiskFund.sol

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: BSD-3-Clause
 pragma solidity 0.8.13;
 
-import { Ownable2StepUpgradeable } from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import { IERC20Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/IERC20Upgradeable.sol";
 import { SafeERC20Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/utils/SafeERC20Upgradeable.sol";
 import { AccessControlledV8 } from "@venusprotocol/governance-contracts/contracts/Governance/AccessControlledV8.sol";
@@ -24,14 +23,7 @@ import { ensureNonzeroAddress } from "../lib/validators.sol";
  * @notice Contract with basic features to track/hold different assets for different Comptrollers.
  * @dev This contract does not support BNB.
  */
-contract RiskFund is
-    Ownable2StepUpgradeable,
-    AccessControlledV8,
-    ExponentialNoError,
-    ReserveHelpers,
-    MaxLoopsLimitHelper,
-    IRiskFund
-{
+contract RiskFund is AccessControlledV8, ExponentialNoError, ReserveHelpers, MaxLoopsLimitHelper, IRiskFund {
     using SafeERC20Upgradeable for IERC20Upgradeable;
 
     address private pancakeSwapRouter;

tests/hardhat/ProtocolShareReserve.ts

@@ -1,5 +1,6 @@
 import { FakeContract, smock } from "@defi-wonderland/smock";
 import { loadFixture } from "@nomicfoundation/hardhat-network-helpers";
+import { SignerWithAddress } from "@nomiclabs/hardhat-ethers/signers";
 import { expect } from "chai";
 import { constants } from "ethers";
 import { ethers, upgrades } from "hardhat";
@@ -43,9 +44,10 @@ describe("ProtocolShareReserve: Tests", function () {
   /**
    * Deploying required contracts along with the poolRegistry.
    */
-
+  let signer2: SignerWithAddress;
   before(async function () {
     await loadFixture(fixture);
+    [, signer2] = await ethers.getSigners();
   });
 
   it("Revert on invalid asset address.", async function () {
@@ -99,4 +101,57 @@ describe("ProtocolShareReserve: Tests", function () {
     expect(liquidatedShareBal).equal(convertToUnit(63, 18));
     expect(protocolShareReserveBal).equal(convertToUnit(10, 18));
   });
+
+  it("Revert if try to sweep tokens by non admin", async function () {
+    await expect(protocolShareReserve.connect(signer2).sweepToken(mockDAI.address, signer2.address)).to.be.revertedWith(
+      "Ownable: caller is not the owner",
+    );
+  });
+
+  it("Revert if recipient address is zero", async function () {
+    await expect(
+      protocolShareReserve.sweepToken(mockDAI.address, ethers.constants.AddressZero),
+    ).to.be.revertedWithCustomError(protocolShareReserve, "ZeroAddressNotAllowed");
+  });
+
+  it("Revert if there are no surplus tokens to sweep", async function () {
+    await mockDAI.transfer(protocolShareReserve.address, convertToUnit(100, 18));
+
+    fakeComptroller.isComptroller.returns(true);
+    // Update assetReserves with all available balance
+    await protocolShareReserve.updateAssetsState(
+      fakeComptroller.address, // Mock comptroller address
+      mockDAI.address,
+    );
+
+    await expect(protocolShareReserve.sweepToken(mockDAI.address, signer2.address)).to.be.revertedWith(
+      "ReserveHelpers: Zero surplus tokens",
+    );
+  });
+
+  it("Success on sweep tokens", async function () {
+    const amount = convertToUnit(100, 18);
+    const excessAmount = convertToUnit(50, 18);
+    const protocolShareReserveBalPrev = await mockDAI.balanceOf(protocolShareReserve.address);
+    await mockDAI.transfer(protocolShareReserve.address, amount);
+
+    fakeComptroller.isComptroller.returns(true);
+    // Update assetReserves with all available balance
+    await protocolShareReserve.updateAssetsState(
+      fakeComptroller.address, // Mock comptroller address
+      mockDAI.address,
+    );
+    let protocolShareReserveBal = await mockDAI.balanceOf(protocolShareReserve.address);
+    expect(protocolShareReserveBal.sub(protocolShareReserveBalPrev)).equal(amount);
+
+    // Sending some extra funds but not updating assetReserves
+    await mockDAI.transfer(protocolShareReserve.address, excessAmount);
+
+    await protocolShareReserve.sweepToken(mockDAI.address, signer2.address);
+    protocolShareReserveBal = await mockDAI.balanceOf(protocolShareReserve.address);
+    expect(protocolShareReserveBal).equal(protocolShareReserveBalPrev.add(amount));
+
+    const recipientBal = await mockDAI.balanceOf(signer2.address);
+    expect(recipientBal).equal(excessAmount);
+  });
 });