fix: withdraw

allo-protocol · Aug 20, 2024 · 1c3932c · 1c3932c
1 parent 74b2ffa
commit 1c3932c
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 0 deletions.
diff --git a/contracts/core/interfaces/IBaseStrategy.sol b/contracts/core/interfaces/IBaseStrategy.sol
@@ -20,6 +20,9 @@ interface IBaseStrategy {
     /// @notice Error when the pool ID is invalid
     error BaseStrategy_INVALID_POOL_ID();
 
+    /// @notice Error when the withdraw amount leaves the pool with insufficient funds
+    error BaseStrategy_WITHDRAW_MORE_THAN_POOL_AMOUNT();
+
     /// ======================
     /// ======= Events =======
     /// ======================

diff --git a/contracts/strategies/CoreBaseStrategy.sol b/contracts/strategies/CoreBaseStrategy.sol
@@ -6,6 +6,7 @@ import "../core/interfaces/IBaseStrategy.sol";
 
 /// Libraries
 import {Transfer} from "./../core/libraries/Transfer.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
 /// @title BaseStrategy Contract
 /// @notice This contract is the base contract for all strategies
@@ -107,6 +108,16 @@ abstract contract CoreBaseStrategy is IBaseStrategy, Transfer {
         onlyPoolManager(msg.sender)
     {
         _beforeWithdraw(_token, _amount, _recipient);
+        // If the token is the pool token, revert if the amount is greater than the pool amount
+        if (_token == allo.getPool(poolId).token) {
+            if (_token == NATIVE) {
+                if (address(this).balance - _amount < poolAmount) revert BaseStrategy_WITHDRAW_MORE_THAN_POOL_AMOUNT();
+            } else {
+                if (IERC20(_token).balanceOf(address(this)) - _amount < poolAmount) {
+                    revert BaseStrategy_WITHDRAW_MORE_THAN_POOL_AMOUNT();
+                }
+            }
+        }
         _transferAmount(_token, _recipient, _amount);
         _afterWithdraw(_token, _amount, _recipient);
 

diff --git a/test/foundry/strategies/CoreBaseStrategy.t.sol b/test/foundry/strategies/CoreBaseStrategy.t.sol
@@ -1,6 +1,7 @@
 pragma solidity 0.8.19;
 
 import "forge-std/Test.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 
 // Test libraries
 import {AlloSetup} from "../shared/AlloSetup.sol";
@@ -43,4 +44,33 @@ contract CoreBaseStrategyTest is Test, AlloSetup {
         strategy.increasePoolAmount(100);
         assertEq(strategy.getPoolAmount(), 100);
     }
+
+    function test_withdraw() public {
+        vm.mockCall(address(allo()), abi.encodeWithSelector(IAllo.isPoolManager.selector), abi.encode(true));
+
+        /// Increase pool amount
+        vm.prank(address(allo()));
+        strategy.increasePoolAmount(100);
+
+        address _token = allo().getPool(0).token;
+
+        vm.mockCall(_token, abi.encodeWithSelector(IERC20.balanceOf.selector, address(strategy)), abi.encode(150));
+        strategy.withdraw(_token, 50, address(this));
+
+        assertEq(strategy.getPoolAmount(), 100);
+    }
+
+    function testRevert_withdrawMoreThanPoolAmount() public {
+        vm.mockCall(address(allo()), abi.encodeWithSelector(IAllo.isPoolManager.selector), abi.encode(true));
+
+        /// Increase pool amount
+        vm.prank(address(allo()));
+        strategy.increasePoolAmount(100);
+
+        address _token = allo().getPool(0).token;
+
+        vm.mockCall(_token, abi.encodeWithSelector(IERC20.balanceOf.selector, address(strategy)), abi.encode(100));
+        vm.expectRevert(IBaseStrategy.BaseStrategy_WITHDRAW_MORE_THAN_POOL_AMOUNT.selector);
+        strategy.withdraw(_token, 50, address(this));
+    }
 }