Spearbit post-audit fixes

src/ConstantSum/ConstantSum.sol

@@ -2,14 +2,18 @@
 pragma solidity 0.8.22;
 
 import {
-    FixedPointMathLib, computeTradingFunction
+    FixedPointMathLib,
+    computeTradingFunction,
+    computeSwapDeltaLiquidity,
+    computeDeltaLiquidity
 } from "./ConstantSumMath.sol";
 import {
     decodePriceUpdate,
     decodeFeeUpdate,
     decodeControllerUpdate
 } from "./ConstantSumUtils.sol";
 import { PairStrategy, IStrategy, Pool } from "src/PairStrategy.sol";
+import { EPSILON } from "src/lib/StrategyLib.sol";
 
 struct InternalParams {
     uint256 price;
@@ -33,6 +37,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 +52,7 @@ contract ConstantSum is PairStrategy {
     function init(
         address,
         uint256 poolId,
-        Pool calldata,
+        Pool calldata pool,
         bytes calldata data
     )
         public
@@ -58,21 +65,106 @@ contract ConstantSum is PairStrategy {
         )
     {
         ConstantSumParams memory params;
-        (reserves, totalLiquidity, params) =
-            abi.decode(data, (uint256[], uint256, ConstantSumParams));
+
+        (reserves, params) = abi.decode(data, (uint256[], ConstantSumParams));
+        totalLiquidity =
+            computeDeltaLiquidity(reserves[0], reserves[1], params.price);
+
+        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 =
             tradingFunction(reserves, totalLiquidity, abi.encode(params));
 
-        valid = invariant >= 0;
+        valid = invariant >= 0 && invariant <= EPSILON;
 
         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 +197,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 +221,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 +230,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
+        );
     }
 }

src/ConstantSum/ConstantSumMath.sol

@@ -6,15 +6,16 @@ import { ConstantSumParams } from "src/ConstantSum/ConstantSum.sol";
 import { ONE } from "src/lib/StrategyLib.sol";
 
 using FixedPointMathLib for uint256;
-using FixedPointMathLib for int256;
 
 function computeTradingFunction(
     uint256[] memory reserves,
     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(
@@ -24,77 +25,28 @@ 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[] memory reserves = new uint256[](2);
     reserves[0] = rx;
     reserves[1] = ry;
-    return abi.encode(reserves, L, params);
+    return abi.encode(reserves, 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);
+    }
 }

src/ConstantSum/ConstantSumSolver.sol

@@ -12,14 +12,14 @@ import {
 import {
     ONE,
     computeInitialPoolData,
-    FixedPointMathLib
+    FixedPointMathLib,
+    computeSwapDeltaLiquidity
 } from "./ConstantSumMath.sol";
 
 contract ConstantSumSolver {
     error NotEnoughLiquidity();
 
     using FixedPointMathLib for uint256;
-    using FixedPointMathLib for int256;
 
     struct Reserves {
         uint256 rx;
@@ -59,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
             );
@@ -68,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);
 
@@ -81,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(

src/ConstantSum/ConstantSumUtils.sol

@@ -20,7 +20,6 @@ function encodeControllerUpdate(address controller)
 
 function decodeFeeUpdate(bytes memory data) pure returns (uint256 swapFee) {
     (, swapFee) = abi.decode(data, (UpdateCode, uint256));
-    return swapFee;
 }
 
 function decodePriceUpdate(bytes memory data)