fix #8 (#9)

also add gas-report command
Sol-DAO · Oct 19, 2022 · 68c5db6 · 68c5db6
1 parent 09a69c6
commit 68c5db6
Show file tree

Hide file tree

Showing 8 changed files with 94 additions and 25 deletions.
diff --git a/foundry.toml b/foundry.toml
@@ -8,5 +8,7 @@ remappings = [
   "@solbase=lib/solbase/src/"
 ]
 
+gas_reports = ["Helios", "HeliosReference"]
+
 [profile.intense.fuzz]
-runs = 10000
+runs = 10000
diff --git a/package.json b/package.json
@@ -21,6 +21,7 @@
       "test": "forge clean && forge test",
       "cov": "forge clean && forge coverage --report lcov && genhtml -o report lcov.info",
       "snapshot": "forge clean && forge snapshot",
+      "gas-report": "forge clean && forge test --gas-report",
       "lint": "prettier --print-width 120 --write src/**/*.sol && prettier --print-width 120 --write src/*.sol"
   }
 }
diff --git a/src/Helios.sol b/src/Helios.sol
@@ -6,11 +6,12 @@ import {OwnedThreeStep} from "@solbase/auth/OwnedThreeStep.sol";
 import {SafeTransferLib} from "@solbase/utils/SafeTransferLib.sol";
 import {FixedPointMathLib} from "@solbase/utils/FixedPointMathLib.sol";
 import {SafeMulticallable} from "@solbase/utils/SafeMulticallable.sol";
+import {ReentrancyGuard} from "@solbase/utils/ReentrancyGuard.sol";
 import {ERC1155, ERC1155TokenReceiver} from "@solbase/tokens/ERC1155.sol";
 
 /// @notice ERC1155 vault with router and liquidity pools.
 /// @author z0r0z.eth (SolDAO)
-contract Helios is OwnedThreeStep(tx.origin), SafeMulticallable, ERC1155, ERC1155TokenReceiver {
+contract Helios is OwnedThreeStep(tx.origin), ReentrancyGuard, SafeMulticallable, ERC1155, ERC1155TokenReceiver {
     constructor() payable {} // Clean deployment.
 
     /// -----------------------------------------------------------------------
@@ -171,7 +172,7 @@ contract Helios is OwnedThreeStep(tx.origin), SafeMulticallable, ERC1155, ERC115
         });
 
         // Swapper dictates output LP.
-        liq = swapper.addLiquidity(id, token0amount, token1amount);
+        (liq, , ) = swapper.addLiquidity(id, token0amount, token1amount);
 
         _mint(to, id, liq, data);
 
@@ -195,24 +196,27 @@ contract Helios is OwnedThreeStep(tx.origin), SafeMulticallable, ERC1155, ERC115
         uint256 token0amount,
         uint256 token1amount,
         bytes calldata data
-    ) external payable returns (uint256 liq) {
+    ) external payable nonReentrant returns (uint256 liq) {
         require(id <= totalSupply, "Helios: PAIR_DOESNT_EXIST");
 
         Pair storage pair = pairs[id];
 
         // If base is address(0), assume native token and overwrite amount.
-        if (pair.token0 == address(0)) {
-            token0amount = msg.value;
+        if (pair.token0 == address(0)) token0amount = msg.value;
+
+        // Swapper dictates output LP and adjustments to token amounts
+        (liq, token0amount, token1amount) = pair.swapper.addLiquidity(id, token0amount, token1amount);
 
-            pair.token1.safeTransferFrom(msg.sender, address(this), token1amount);
+        // We know the exact amounts, now transfer and/or refund any ETH
+        if (pair.token0 == address(0)) {
+            if (msg.value > token0amount) {
+                (bool success, ) = payable(msg.sender).call{value: msg.value - token0amount}("");
+                require(success, "Helios: REFUND_FAILED");
+            }
         } else {
             pair.token0.safeTransferFrom(msg.sender, address(this), token0amount);
-
-            pair.token1.safeTransferFrom(msg.sender, address(this), token1amount);
         }
-
-        // Swapper dictates output LP.
-        liq = pair.swapper.addLiquidity(id, token0amount, token1amount);
+        pair.token1.safeTransferFrom(msg.sender, address(this), token1amount);
 
         require(liq != 0, "Helios: INSUFFICIENT_LIQUIDITY_MINTED");
 

diff --git a/src/HeliosReference.sol b/src/HeliosReference.sol
@@ -6,12 +6,19 @@ import {OwnedThreeStep} from "@solbase/auth/OwnedThreeStep.sol";
 import {SafeTransferLib} from "@solbase/utils/SafeTransferLib.sol";
 import {FixedPointMathLib} from "@solbase/utils/FixedPointMathLib.sol";
 import {SafeMulticallable} from "@solbase/utils/SafeMulticallable.sol";
+import {ReentrancyGuard} from "@solbase/utils/ReentrancyGuard.sol";
 import {ERC1155, ERC1155TokenReceiver} from "@solbase/tokens/ERC1155.sol";
 
 /// @notice ERC1155 vault with router and liquidity pools.
 /// @dev Reference implementation (emphasizes clarity, deemphasizes gas cost)
 /// @author z0r0z.eth (SolDAO)
-contract HeliosReference is OwnedThreeStep(tx.origin), SafeMulticallable, ERC1155, ERC1155TokenReceiver {
+contract HeliosReference is
+    OwnedThreeStep(tx.origin),
+    ReentrancyGuard,
+    SafeMulticallable,
+    ERC1155,
+    ERC1155TokenReceiver
+{
     constructor() payable {} // Clean deployment.
 
     /// -----------------------------------------------------------------------
@@ -172,7 +179,7 @@ contract HeliosReference is OwnedThreeStep(tx.origin), SafeMulticallable, ERC115
         });
 
         // Swapper dictates output LP.
-        liq = swapper.addLiquidity(id, token0amount, token1amount);
+        (liq, , ) = swapper.addLiquidity(id, token0amount, token1amount);
 
         _mint(to, id, liq, data);
 
@@ -196,24 +203,27 @@ contract HeliosReference is OwnedThreeStep(tx.origin), SafeMulticallable, ERC115
         uint256 token0amount,
         uint256 token1amount,
         bytes calldata data
-    ) external payable returns (uint256 liq) {
+    ) external payable nonReentrant returns (uint256 liq) {
         require(id <= totalSupply, "Helios: PAIR_DOESNT_EXIST");
 
         Pair storage pair = pairs[id];
 
         // If base is address(0), assume native token and overwrite amount.
-        if (pair.token0 == address(0)) {
-            token0amount = msg.value;
+        if (pair.token0 == address(0)) token0amount = msg.value;
+
+        // Swapper dictates output LP and adjustments to token amounts
+        (liq, token0amount, token1amount) = pair.swapper.addLiquidity(id, token0amount, token1amount);
 
-            pair.token1.safeTransferFrom(msg.sender, address(this), token1amount);
+        // We know the exact amounts, now transfer and/or refund any ETH
+        if (pair.token0 == address(0)) {
+            if (msg.value > token0amount) {
+                (bool success, ) = payable(msg.sender).call{value: msg.value - token0amount}("");
+                require(success, "Helios: REFUND_FAILED");
+            }
         } else {
             pair.token0.safeTransferFrom(msg.sender, address(this), token0amount);
-
-            pair.token1.safeTransferFrom(msg.sender, address(this), token1amount);
         }
-
-        // Swapper dictates output LP.
-        liq = pair.swapper.addLiquidity(id, token0amount, token1amount);
+        pair.token1.safeTransferFrom(msg.sender, address(this), token1amount);
 
         require(liq != 0, "Helios: INSUFFICIENT_LIQUIDITY_MINTED");
 

diff --git a/src/interfaces/IHelios.sol b/src/interfaces/IHelios.sol
@@ -20,7 +20,13 @@ interface IHelios {
         uint256 id,
         uint256 token0amount,
         uint256 token1amount
-    ) external returns (uint256 liq);
+    )
+        external
+        returns (
+            uint256 liq,
+            uint256 amount0,
+            uint256 amount1
+        );
 
     function removeLiquidity(uint256 id, uint256 liq) external returns (uint256 amount0out, uint256 amount1out);
 

diff --git a/src/swappers/XYKswapper.sol b/src/swappers/XYKswapper.sol
@@ -38,7 +38,15 @@ contract XYKswapper is ReentrancyGuard {
         uint256 id,
         uint256 token0amount,
         uint256 token1amount
-    ) external nonReentrant returns (uint256 liq) {
+    )
+        external
+        nonReentrant
+        returns (
+            uint256 liq,
+            uint256 amount0,
+            uint256 amount1
+        )
+    {
         IHelios.Pair memory pair = IHelios(msg.sender).pairs(id);
 
         uint256 reserve0 = pair.reserve0;
@@ -49,8 +57,12 @@ contract XYKswapper is ReentrancyGuard {
 
         if (totalSupply == 0) {
             liq = (token0amount * token1amount).sqrt() - MIN_LP;
+            amount0 = token0amount;
+            amount1 = token1amount;
         } else {
             liq = min((token0amount * totalSupply) / reserve0, (token1amount * totalSupply) / reserve1);
+            amount0 = (reserve0 * liq) / totalSupply;
+            amount1 = (reserve1 * liq) / totalSupply;
         }
 
         require(liq != 0, "XYKswapper: INSUFFICIENT_LIQUIDITY_MINTED");

diff --git a/src/test/Helios.t.sol b/src/test/Helios.t.sol
@@ -76,6 +76,23 @@ contract HeliosTest is ERC1155TokenReceiver, Test {
         helios.createPair(address(this), token0, token1, 1_000 ether, 1_000 ether, xykSwapper, 1, "");
     }
 
+    function testAddLiquidity() public payable {
+        (uint256 id, uint256 c) = helios.createPair(
+            address(this),
+            token0,
+            token1,
+            1_000 ether,
+            1_000 ether,
+            xykSwapper,
+            1,
+            ""
+        );
+        uint256 b = helios.addLiquidity(address(this), id, 1_000 ether, 1_000 ether, "");
+        uint256 a = helios.addLiquidity(address(this), id, 1_000 ether, 1_000 ether, "");
+        require(a == b, "a!=b");
+        require(b == c, "b!=c");
+    }
+
     function testXYKpairSwap(uint256 amountIn) public payable {
         uint256 b0 = MockERC20(token0).balanceOf(address(this));
         vm.assume(amountIn > 100000 && amountIn < b0);

diff --git a/src/test/HeliosReference.t.sol b/src/test/HeliosReference.t.sol
@@ -61,6 +61,23 @@ contract HeliosReferenceTest is ERC1155TokenReceiver, Test {
         helios.createPair(address(this), token0, token1, 1_000 ether, 1_000 ether, xykSwapper, 1, "");
     }
 
+    function testAddLiquidity() public payable {
+        (uint256 id, uint256 c) = helios.createPair(
+            address(this),
+            token0,
+            token1,
+            1_000 ether,
+            1_000 ether,
+            xykSwapper,
+            1,
+            ""
+        );
+        uint256 b = helios.addLiquidity(address(this), id, 1_000 ether, 1_000 ether, "");
+        uint256 a = helios.addLiquidity(address(this), id, 1_000 ether, 1_000 ether, "");
+        require(a == b, "a!=b");
+        require(b == c, "b!=c");
+    }
+
     function testXYKpairSwap(uint256 amountIn) public payable {
         uint256 b0 = MockERC20(token0).balanceOf(address(this));
         vm.assume(amountIn > 100000 && amountIn < b0);