Merge pull request #7 from VenusProtocol/fix/certik-reaudit

Author: web3rover

contracts/PositionSwapper/PositionSwapper.sol

@@ -97,6 +97,7 @@ contract PositionSwapper is Ownable2StepUpgradeable, ReentrancyGuardUpgradeable
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor(address _comptroller, address _nativeMarket) {
         if (_comptroller == address(0)) revert ZeroAddress();
+        if (_nativeMarket == address(0)) revert ZeroAddress();
 
         COMPTROLLER = IComptroller(_comptroller);
         NATIVE_MARKET = _nativeMarket;

contracts/PositionSwapper/WBNBSwapHelper.sol

@@ -3,14 +3,17 @@ pragma solidity 0.8.25;
 
 import { IWBNB } from "../Interfaces.sol";
 import { ISwapHelper } from "./ISwapHelper.sol";
-import { IERC20Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/utils/SafeERC20Upgradeable.sol";
+import { SafeERC20Upgradeable, IERC20Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/utils/SafeERC20Upgradeable.sol";
+import { Ownable2Step } from "@openzeppelin/contracts/access/Ownable2Step.sol";
 
 /**
  * @title WBNBSwapHelper
  * @notice Swap helper that wraps or unwraps native BNB into WBNB for PositionSwapper.
  * @dev Only supports native token (BNB) wrapping into WBNB and unwrapping WBNB into BNB. Meant to be used only by the PositionSwapper.
  */
-contract WBNBSwapHelper is ISwapHelper {
+contract WBNBSwapHelper is Ownable2Step, ISwapHelper {
+    using SafeERC20Upgradeable for IERC20Upgradeable;
+
     /// @notice Address of the authorized PositionSwapper contract
     address public immutable POSITION_SWAPPER;
 
@@ -29,6 +32,21 @@ contract WBNBSwapHelper is ISwapHelper {
      */
     event SwappedToBNB(uint256 amount);
 
+    /**
+     * @notice Emitted after the owner sweeps leftover ERC-20 tokens from the contract
+     * @param token The token that was swept.
+     * @param receiver The address that received the swept tokens.
+     * @param amount The amount of tokens that were swept.
+     */
+    event SweepToken(address indexed token, address indexed receiver, uint256 amount);
+
+    /**
+     * @notice Emitted after the owner sweeps leftover native tokens (e.g., BNB) from the contract
+     * @param receiver The address that received the swept native tokens.
+     * @param amount The amount of native tokens that were swept.
+     */
+    event SweepNative(address indexed receiver, uint256 amount);
+
     /// @notice Error thrown when caller is not the authorized PositionSwapper
     error Unauthorized();
 
@@ -50,13 +68,39 @@ contract WBNBSwapHelper is ISwapHelper {
         _;
     }
 
-    constructor(address _wbnb, address _swapper) {
+    constructor(address _wbnb, address _swapper) Ownable2Step() {
         if (_wbnb == address(0) || _swapper == address(0)) revert ZeroAddress();
 
         WBNB = IWBNB(_wbnb);
         POSITION_SWAPPER = _swapper;
     }
 
+    /**
+     * @notice Allows the owner to sweep leftover ERC-20 tokens from the contract.
+     * @param token The token to sweep.
+     * @custom:event Emits SweepToken event.
+     */
+    function sweepToken(IERC20Upgradeable token) external onlyOwner {
+        uint256 balance = token.balanceOf(address(this));
+        if (balance > 0) {
+            token.safeTransfer(owner(), balance);
+            emit SweepToken(address(token), owner(), balance);
+        }
+    }
+
+    /**
+     * @notice Allows the owner to sweep leftover native tokens (e.g., BNB) from the contract.
+     * @custom:event Emits SweepToken event.
+     */
+    function sweepNative() external onlyOwner {
+        uint256 balance = address(this).balance;
+        if (balance > 0) {
+            (bool success, ) = payable(owner()).call{ value: balance }("");
+            if (!success) revert TransferFailed();
+            emit SweepNative(owner(), balance);
+        }
+    }
+
     /// @notice Allows this contract to receive native BNB
     receive() external payable {}
 
@@ -80,7 +124,7 @@ contract WBNBSwapHelper is ISwapHelper {
             WBNB.transfer(msg.sender, amount);
             emit SwappedToWBNB(amount);
         } else {
-            IERC20Upgradeable(WBNB).transferFrom(msg.sender, address(this), amount);
+            WBNB.transferFrom(msg.sender, address(this), amount);
             WBNB.withdraw(amount);
 
             (bool success, ) = payable(msg.sender).call{ value: amount }("");