Added trade protection to swap functions

pooltogether · May 11, 2022 · 090a4ce · 090a4ce
1 parent bd9e309
commit 090a4ce
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 9 deletions.
diff --git a/contracts/PrizePoolLiquidator.sol b/contracts/PrizePoolLiquidator.sol
@@ -104,29 +104,56 @@ contract PrizePoolLiquidator {
     return poolLiquidatorStates[_prizePool].computeExactAmountOutAtTime(_availableStreamHaveBalance(_prizePool), amountIn, currentTime);
   }
 
-  function swapExactAmountIn(IPrizePool _prizePool, uint256 amountIn) public returns (uint256) {
-    return swapExactAmountInAtTime(_prizePool, amountIn, block.timestamp);
+  function swapExactAmountIn(IPrizePool _prizePool, uint256 amountIn, uint256 amountOutMin) external returns (uint256) {
+    return _swapExactAmountInAtTime(_prizePool, amountIn, amountOutMin, block.timestamp);
   }
 
-  function swapExactAmountInAtTime(IPrizePool _prizePool, uint256 amountIn, uint256 currentTime) public returns (uint256) {
+  function _swapExactAmountInAtTime(IPrizePool _prizePool, uint256 amountIn, uint256 amountOutMin, uint256 currentTime) internal returns (uint256) {
     uint256 availableBalance = _availableStreamHaveBalance(_prizePool);
     uint256 amountOut = poolLiquidatorStates[_prizePool].swapExactAmountInAtTime(
       availableBalance, amountIn, currentTime
     );
 
-    Target storage target = poolTargets[_prizePool];
+    require(amountOut <= availableBalance, "Whoops! have exceeds available");
+    require(amountOut >= amountOutMin, "trade does not meet min");
+
+    _swap(_prizePool, msg.sender, amountOut, amountIn);
 
+    return amountOut;
+  }
+
+  function swapExactAmountOut(IPrizePool _prizePool, uint256 amountOut, uint256 amountInMax) external returns (uint256) {
+    return _swapExactAmountOutAtTime(_prizePool, amountOut, amountInMax, block.timestamp);
+  }
+
+  function _swapExactAmountOutAtTime(
+    IPrizePool _prizePool,
+    uint256 amountOut,
+    uint256 amountInMax,
+    uint256 currentTime
+  ) internal returns (uint256) {
+    uint256 availableBalance = _availableStreamHaveBalance(_prizePool);
+    uint256 amountIn = poolLiquidatorStates[_prizePool].swapExactAmountOutAtTime(
+      availableBalance, amountOut, currentTime
+    );
+
+    require(amountIn <= amountInMax, "trade does not meet min");
     require(amountOut <= availableBalance, "Whoops! have exceeds available");
 
-    _prizePool.award(msg.sender, amountOut);
-    target.want.transferFrom(msg.sender, target.target, amountIn);
+    _swap(_prizePool, msg.sender, amountOut, amountIn);
+
+    return amountIn;
+  }
 
+  function _swap(IPrizePool _prizePool, address _account, uint256 _amountOut, uint256 _amountIn) internal {
+    Target storage target = poolTargets[_prizePool];
+    IERC20 want = target.want;
+    _prizePool.award(_account, _amountOut);
+    want.transferFrom(_account, target.target, _amountIn);
     IPrizePoolLiquidatorListener _listener = listener;
     if (address(_listener) != address(0)) {
-      _listener.afterSwap(_prizePool, _prizePool.getTicket(), amountOut, target.want, amountIn);
+      _listener.afterSwap(_prizePool, _prizePool.getTicket(), _amountOut, want, _amountIn);
     }
-
-    return amountOut;
   }
 
   function getLiquidationState(IPrizePool _prizePool) external view returns (

diff --git a/contracts/libraries/LiquidatorLib.sol b/contracts/libraries/LiquidatorLib.sol
@@ -111,6 +111,35 @@ library LiquidatorLib {
     return amountOut;
   }
 
+  function swapExactAmountOutAtTime(
+    State storage liquidationState,
+    uint256 availableBalance,
+    uint256 amountOut,
+    uint256 currentTime
+  ) internal returns (uint256) {
+    PRBMath.SD59x18 memory newExchangeRate = _increaseExchangeRateByDeltaTime(
+        liquidationState.exchangeRate,
+        liquidationState.deltaRatePerSecond,
+        liquidationState.lastSaleTime,
+        currentTime
+    );
+    require(availableBalance > 0, "Whoops! no funds available");
+    VirtualCpmmLib.Cpmm memory cpmm = VirtualCpmmLib.newCpmm(
+      liquidationState.maxSlippage, newExchangeRate, PRBMathSD59x18Typed.fromInt(availableBalance.toInt256())
+    );
+
+    uint256 amountIn = VirtualCpmmLib.getAmountIn(amountOut, cpmm.want, cpmm.have);
+    cpmm.want += amountIn;
+    cpmm.have -= amountOut;
+
+    require(amountOut <= availableBalance, "Whoops! have exceeds available");
+
+    liquidationState.lastSaleTime = currentTime;
+    liquidationState.exchangeRate = _cpmmToExchangeRate(cpmm);
+
+    return amountIn;
+  }
+
   function _cpmmToExchangeRate(VirtualCpmmLib.Cpmm memory cpmm) internal pure returns (PRBMath.SD59x18 memory) {
     return PRBMathSD59x18Typed.fromInt(int256(cpmm.have)).div(PRBMathSD59x18Typed.fromInt(int256(cpmm.want)));
   }

diff --git a/contracts/test/PrizePoolLiquidatorHarness.sol b/contracts/test/PrizePoolLiquidatorHarness.sol
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: GPL-3.0
+
+pragma solidity 0.8.6;
+
+import "../PrizePoolLiquidator.sol";
+
+contract PrizePoolLiquidatorHarness is PrizePoolLiquidator {
+
+  function swapExactAmountInAtTime(IPrizePool _prizePool, uint256 amountIn, uint256 amountOutMin, uint256 currentTime) external returns (uint256) {
+    return _swapExactAmountInAtTime(_prizePool, amountIn, amountOutMin, currentTime);
+  }
+
+  function swapExactAmountOutAtTime(
+    IPrizePool _prizePool,
+    uint256 amountOut,
+    uint256 amountInMax,
+    uint256 currentTime
+  ) external returns (uint256) {
+    return _swapExactAmountOutAtTime(_prizePool, amountOut, amountInMax, currentTime);
+  }
+}