Merge branch 'fix/spearbit-audit' into fix/minor-issues

primitivefinance · Apr 1, 2024 · af6ac26 · af6ac26
2 parents 28315ce + 86b6b0b
commit af6ac26
Show file tree

Hide file tree

Showing 33 changed files with 831 additions and 492 deletions.
diff --git a/src/ConstantSum/ConstantSum.sol b/src/ConstantSum/ConstantSum.sol
@@ -2,7 +2,10 @@
 pragma solidity 0.8.22;
 
 import {
-    FixedPointMathLib, computeTradingFunction
+    FixedPointMathLib,
+    computeTradingFunction,
+    computeSwapDeltaLiquidity,
+    computeDeltaLiquidity
 } from "./ConstantSumMath.sol";
 import {
     decodePriceUpdate,
@@ -33,6 +36,9 @@ enum UpdateCode {
 contract ConstantSum is PairStrategy {
     using FixedPointMathLib for uint256;
 
+    /// @notice Thrown when the expected liquidity is not met.
+    error InvalidDeltaLiquidity();
+
     /// @inheritdoc IStrategy
     string public constant name = "ConstantSum";
 
@@ -45,7 +51,7 @@ contract ConstantSum is PairStrategy {
     function init(
         address,
         uint256 poolId,
-        Pool calldata,
+        Pool calldata pool,
         bytes calldata data
     )
         public
@@ -61,8 +67,13 @@ contract ConstantSum is PairStrategy {
         (reserves, totalLiquidity, params) =
             abi.decode(data, (uint256[], uint256, ConstantSumParams));
 
+        if (pool.reserves.length != 2 || reserves.length != 2) {
+            revert InvalidReservesLength();
+        }
+
         internalParams[poolId].price = params.price;
         internalParams[poolId].swapFee = params.swapFee;
+        internalParams[poolId].controller = params.controller;
 
         // Get the trading function and check this is valid
         invariant =
@@ -73,6 +84,84 @@ contract ConstantSum is PairStrategy {
         return (valid, invariant, reserves, totalLiquidity);
     }
 
+    function validateAllocate(
+        address,
+        uint256 poolId,
+        Pool memory pool,
+        bytes calldata data
+    )
+        external
+        view
+        override
+        returns (
+            bool valid,
+            int256 invariant,
+            uint256[] memory deltas,
+            uint256 deltaLiquidity
+        )
+    {
+        (uint256 deltaX, uint256 deltaY, uint256 minDeltaL) =
+            abi.decode(data, (uint256, uint256, uint256));
+
+        deltaLiquidity =
+            computeDeltaLiquidity(deltaX, deltaY, internalParams[poolId].price);
+        if (deltaLiquidity < minDeltaL) revert InvalidDeltaLiquidity();
+
+        deltas = new uint256[](2);
+        deltas[0] = deltaX;
+        deltas[1] = deltaY;
+
+        pool.reserves[0] += deltaX;
+        pool.reserves[1] += deltaY;
+
+        invariant = tradingFunction(
+            pool.reserves,
+            pool.totalLiquidity + deltaLiquidity,
+            getPoolParams(poolId)
+        );
+
+        valid = invariant >= 0;
+    }
+
+    function validateDeallocate(
+        address,
+        uint256 poolId,
+        Pool memory pool,
+        bytes calldata data
+    )
+        external
+        view
+        override
+        returns (
+            bool valid,
+            int256 invariant,
+            uint256[] memory deltas,
+            uint256 deltaLiquidity
+        )
+    {
+        (uint256 deltaX, uint256 deltaY, uint256 maxDeltaL) =
+            abi.decode(data, (uint256, uint256, uint256));
+
+        deltaLiquidity =
+            computeDeltaLiquidity(deltaX, deltaY, internalParams[poolId].price);
+        if (maxDeltaL > deltaLiquidity) revert InvalidDeltaLiquidity();
+
+        deltas = new uint256[](2);
+        deltas[0] = deltaX;
+        deltas[1] = deltaY;
+
+        pool.reserves[0] -= deltaX;
+        pool.reserves[1] -= deltaY;
+
+        invariant = tradingFunction(
+            pool.reserves,
+            pool.totalLiquidity - deltaLiquidity,
+            getPoolParams(poolId)
+        );
+
+        valid = invariant >= 0;
+    }
+
     /// @inheritdoc IStrategy
     function update(
         address sender,
@@ -105,6 +194,7 @@ contract ConstantSum is PairStrategy {
 
         params.price = internalParams[poolId].price;
         params.swapFee = internalParams[poolId].swapFee;
+        params.controller = internalParams[poolId].controller;
 
         return abi.encode(params);
     }
@@ -128,7 +218,7 @@ contract ConstantSum is PairStrategy {
         Pool memory,
         bytes memory
     ) internal pure override returns (uint256[] memory) {
-        return new uint256[](0);
+        return new uint256[](2);
     }
 
     /// @inheritdoc PairStrategy
@@ -137,6 +227,20 @@ contract ConstantSum is PairStrategy {
         Pool memory,
         bytes memory
     ) internal pure override returns (uint256[] memory) {
-        return new uint256[](0);
+        return new uint256[](2);
+    }
+
+    /// @inheritdoc PairStrategy
+    function _computeSwapDeltaLiquidity(
+        Pool memory,
+        bytes memory params,
+        uint256 tokenInIndex,
+        uint256,
+        uint256 amountIn,
+        uint256
+    ) internal pure override returns (uint256) {
+        return computeSwapDeltaLiquidity(
+            amountIn, abi.decode(params, (ConstantSumParams)), tokenInIndex == 0
+        );
     }
 }
diff --git a/src/ConstantSum/ConstantSumMath.sol b/src/ConstantSum/ConstantSumMath.sol
@@ -12,8 +12,10 @@ function computeTradingFunction(
     uint256 totalLiquidity,
     uint256 price
 ) pure returns (int256) {
-    return int256(reserves[0].divWadUp(totalLiquidity))
-        + int256(reserves[1].divWadUp(totalLiquidity.mulWadUp(price))) - int256(ONE);
+    return int256(
+        price.mulWadUp(reserves[0].divWadUp(totalLiquidity))
+            + reserves[1].divWadUp(totalLiquidity)
+    ) - int256(ONE);
 }
 
 function computeInitialPoolData(
@@ -23,77 +25,29 @@ function computeInitialPoolData(
 ) pure returns (bytes memory) {
     // The pool can be initialized with any non-negative amount of rx, and ry.
     // so we have to allow a user to pass an amount of both even if one is zero.
-    uint256 L = rx + ry.divWadUp(params.price);
+    uint256 L = rx.mulWadDown(params.price) + ry;
     uint256[] memory reserves = new uint256[](2);
     reserves[0] = rx;
     reserves[1] = ry;
     return abi.encode(reserves, L, params);
 }
 
-function computeDeallocateGivenDeltaX(
+function computeDeltaLiquidity(
     uint256 deltaX,
-    uint256 rX,
-    uint256 rY,
-    uint256 totalLiquidity
-) pure returns (uint256 deltaY, uint256 deltaL) {
-    uint256 a = deltaX.divWadDown(rX);
-    if (rY > 0) {
-        deltaY = a.mulWadDown(rY);
-    }
-    deltaL = a.mulWadDown(totalLiquidity);
-}
-
-function computeDeallocateGivenDeltaY(
-    uint256 deltaY,
-    uint256 rX,
-    uint256 rY,
-    uint256 totalLiquidity
-) pure returns (uint256 deltaX, uint256 deltaL) {
-    uint256 a = deltaY.divWadDown(rY);
-    if (rX > 0) {
-        deltaX = a.mulWadDown(rX);
-    }
-    deltaL = a.mulWadDown(totalLiquidity);
-}
-
-function computeAllocateGivenDeltaX(
-    uint256 deltaX,
-    uint256 rX,
-    uint256 rY,
-    uint256 totalLiquidity
-) pure returns (uint256 deltaY, uint256 deltaL) {
-    uint256 a = deltaX.divWadUp(rX);
-    if (rY > 0) {
-        deltaY = a.mulWadUp(rY);
-    }
-    deltaL = a.mulWadUp(totalLiquidity);
-}
-
-function computeAllocateGivenDeltaY(
     uint256 deltaY,
-    uint256 rX,
-    uint256 rY,
-    uint256 totalLiquidity
-) pure returns (uint256 deltaX, uint256 deltaL) {
-    uint256 a = deltaY.divWadUp(rY);
-    if (rX > 0) {
-        deltaX = a.mulWadUp(rX);
-    }
-    deltaL = a.mulWadUp(totalLiquidity);
-}
-
-function computeDeltaGivenDeltaLRoundUp(
-    uint256 reserve,
-    uint256 deltaLiquidity,
-    uint256 totalLiquidity
+    uint256 price
 ) pure returns (uint256) {
-    return reserve.mulWadUp(deltaLiquidity.divWadUp(totalLiquidity));
+    return price.mulWadUp(deltaX) + deltaY;
 }
 
-function computeDeltaGivenDeltaLRoundDown(
-    uint256 reserve,
-    uint256 deltaLiquidity,
-    uint256 totalLiquidity
+function computeSwapDeltaLiquidity(
+    uint256 delta,
+    ConstantSumParams memory params,
+    bool isSwapXForY
 ) pure returns (uint256) {
-    return reserve.mulWadDown(deltaLiquidity.divWadDown(totalLiquidity));
+    if (isSwapXForY) {
+        return (params.swapFee).mulWadUp(delta);
+    } else {
+        return (params.swapFee).mulDivUp(delta, params.price);
+    }
 }
diff --git a/src/ConstantSum/ConstantSumSolver.sol b/src/ConstantSum/ConstantSumSolver.sol
@@ -12,7 +12,8 @@ import {
 import {
     ONE,
     computeInitialPoolData,
-    FixedPointMathLib
+    FixedPointMathLib,
+    computeSwapDeltaLiquidity
 } from "./ConstantSumMath.sol";
 
 contract ConstantSumSolver {
@@ -58,7 +59,7 @@ contract ConstantSumSolver {
         SimulateSwapState memory state;
 
         if (swapXIn) {
-            state.deltaLiquidity = amountIn.mulWadUp(poolParams.swapFee);
+            computeSwapDeltaLiquidity(amountIn, poolParams, true);
             state.amountOut = amountIn.mulWadDown(poolParams.price).mulWadDown(
                 ONE - poolParams.swapFee
             );
@@ -67,8 +68,7 @@ contract ConstantSumSolver {
                 revert NotEnoughLiquidity();
             }
         } else {
-            state.deltaLiquidity =
-                amountIn.mulWadUp(poolParams.swapFee).divWadUp(poolParams.price);
+            computeSwapDeltaLiquidity(amountIn, poolParams, false);
             state.amountOut = (ONE - poolParams.swapFee).mulWadDown(amountIn)
                 .divWadDown(poolParams.price);
 
@@ -80,13 +80,9 @@ contract ConstantSumSolver {
         bytes memory swapData;
 
         if (swapXIn) {
-            swapData = abi.encode(
-                0, 1, amountIn, state.amountOut, state.deltaLiquidity
-            );
+            swapData = abi.encode(0, 1, amountIn, state.amountOut);
         } else {
-            swapData = abi.encode(
-                1, 0, amountIn, state.amountOut, state.deltaLiquidity
-            );
+            swapData = abi.encode(1, 0, amountIn, state.amountOut);
         }
 
         (bool valid,,,,,,) = IStrategy(strategy).validateSwap(