Merge pull request #4 from GenerationSoftware/gen-447-c4-22-vaultboos…

…ter-users-tokens-will-be-stuck-if-they Added deposit protection
GenerationSoftware · Aug 27, 2023 · 745c9fc · 745c9fc
2 parents 3736815 + f447cb9
commit 745c9fc
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/src/VaultBooster.sol b/src/VaultBooster.sol
@@ -31,6 +31,10 @@ error VaultZeroAddress();
 /// @notice Emitted when the owner is set to the zero address.
 error OwnerZeroAddress();
 
+/// @notice Emitted when someone tries to deposit when no boost has been set for a token
+/// @param token The token that was attempted to be deposited
+error CannotDepositWithoutBoost(IERC20 token);
+
 /// @notice Struct that holds the boost data
 struct Boost {
   address liquidationPair;
@@ -179,7 +183,7 @@ contract VaultBooster is Ownable, ILiquidationSource {
   /// @dev Useful because it ensures `accrue` is called before depositing
   /// @param _token The token to deposit
   /// @param _amount The amount to deposit
-  function deposit(IERC20 _token, uint256 _amount) external {
+  function deposit(IERC20 _token, uint256 _amount) onlyBoosted(_token) external {
     if (0 == _amount) revert ZeroAmountDeposit();
     _accrue(_token);
     _token.safeTransferFrom(msg.sender, address(this), _amount);
@@ -319,4 +323,11 @@ contract VaultBooster is Ownable, ILiquidationSource {
     _;
   }
 
+  modifier onlyBoosted(IERC20 _token) {
+    if (_boosts[_token].liquidationPair == address(0)) {
+      revert CannotDepositWithoutBoost(_token);
+    }
+    _;
+  }
+
 }
diff --git a/test/VaultBooster.t.sol b/test/VaultBooster.t.sol
@@ -14,7 +14,9 @@ import {
   ZeroAmountWithdraw,
   ZeroAmountDeposit,
   VaultZeroAddress,
-  OwnerZeroAddress
+  OwnerZeroAddress,
+  CannotDepositWithoutBoost,
+  InsufficientAvailableBalance
 } from "../src/VaultBooster.sol";
 
 import { IFlashSwapCallback } from "pt-v5-liquidator-interfaces/interfaces/IFlashSwapCallback.sol";
@@ -136,7 +138,7 @@ contract VaultBoosterTest is Test {
     assertEq(boost.available, 0.5e18);
   }
 
-  function testDeposit() public {
+  function testDeposit_success() public {
     mockBoostTokenBalance(1e18);
     booster.setBoost(boostToken, liquidationPair, UD2x18.wrap(0), 0, 1e18);
     vm.mockCall(address(boostToken), abi.encodeWithSelector(IERC20.transferFrom.selector, address(this), address(booster), 2e18), abi.encode(true));
@@ -160,6 +162,11 @@ contract VaultBoosterTest is Test {
     booster.deposit(boostToken, 0); // zero amount
   }
 
+  function testDeposit_CannotDepositWithoutBoost() public {
+    vm.expectRevert(abi.encodeWithSelector(CannotDepositWithoutBoost.selector, boostToken));
+    booster.deposit(boostToken, 2e18);
+  }
+
   function testAccrue_tokensPerSecond() public {
     vm.warp(0);
     mockBoostTokenBalance(1e18);