From be59344db35c73926565db093b085305c1cc85f1 Mon Sep 17 00:00:00 2001
From: Lucas Manuel <lucasmanuel.tech@gmail.com>
Date: Thu, 20 Jun 2024 08:16:05 -0400
Subject: [PATCH] feat: Add rounding testing  (SC-474) (#8)

---
 src/PSM3.sol        |   1 -
 test/Rounding.t.sol | 161 ++++++++++++++++++++++++++++++++++++++++++++
 test/Withdraw.t.sol |   1 -
 3 files changed, 161 insertions(+), 2 deletions(-)
 create mode 100644 test/Rounding.t.sol

diff --git a/src/PSM3.sol b/src/PSM3.sol
index c90e7f7..de8946a 100644
--- a/src/PSM3.sol
+++ b/src/PSM3.sol
@@ -11,7 +11,6 @@ interface IRateProviderLike {
     function getConversionRate() external view returns (uint256);
 }
 
-// TODO: Prove that we're always rounding against user
 contract PSM3 is IPSM3 {
 
     using SafeERC20 for IERC20;
diff --git a/test/Rounding.t.sol b/test/Rounding.t.sol
new file mode 100644
index 0000000..c9bd18c
--- /dev/null
+++ b/test/Rounding.t.sol
@@ -0,0 +1,161 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+pragma solidity ^0.8.13;
+
+import "forge-std/Test.sol";
+
+import { PSMTestBase } from "test/PSMTestBase.sol";
+
+import { MockERC20 } from "erc20-helpers/MockERC20.sol";
+
+contract RoundingTests is PSMTestBase {
+
+    address user = makeAddr("user");
+
+    function setUp() public override {
+        super.setUp();
+
+        // Seed the PSM with max liquidity so withdrawals can always be performed
+        _deposit(address(dai),  address(this), DAI_TOKEN_MAX);
+        _deposit(address(sDai), address(this), SDAI_TOKEN_MAX);
+        _deposit(address(usdc), address(this), USDC_TOKEN_MAX);
+
+        // Set an exchange rate that will cause rounding
+        rateProvider.__setConversionRate(1.25e27 * uint256(100) / 99);
+    }
+
+    function test_roundAgainstUser_dai() public {
+        _deposit(address(dai), address(user), 1e18);
+
+        assertEq(dai.balanceOf(address(user)), 0);
+
+        vm.prank(user);
+        psm.withdraw(address(dai), address(user), 1e18);
+
+        assertEq(dai.balanceOf(address(user)), 1e18 - 1);  // Rounds against user
+    }
+
+    function test_roundAgainstUser_usdc() public {
+        _deposit(address(usdc), address(user), 1e6);
+
+        assertEq(usdc.balanceOf(address(user)), 0);
+
+        vm.prank(user);
+        psm.withdraw(address(usdc), address(user), 1e6);
+
+        assertEq(usdc.balanceOf(address(user)), 1e6 - 1);  // Rounds against user
+    }
+
+    function test_roundAgainstUser_sDai() public {
+        _deposit(address(sDai), address(user), 1e18);
+
+        assertEq(sDai.balanceOf(address(user)), 0);
+
+        vm.prank(user);
+        psm.withdraw(address(sDai), address(user), 1e18);
+
+        assertEq(sDai.balanceOf(address(user)), 1e18 - 1);  // Rounds against user
+    }
+
+    function testFuzz_roundingAgainstUser_multiUser_dai(
+        uint256 rate1,
+        uint256 rate2,
+        uint256 amount1,
+        uint256 amount2
+    )
+        public
+    {
+        _runRoundingAgainstUsersFuzzTest(
+            dai,
+            DAI_TOKEN_MAX,
+            rate1,
+            rate2,
+            amount1,
+            amount2,
+            4
+        );
+    }
+
+    function testFuzz_roundingAgainstUser_multiUser_usdc(
+        uint256 rate1,
+        uint256 rate2,
+        uint256 amount1,
+        uint256 amount2
+    )
+        public
+    {
+        _runRoundingAgainstUsersFuzzTest(
+            usdc,
+            USDC_TOKEN_MAX,
+            rate1,
+            rate2,
+            amount1,
+            amount2,
+            1  // Lower precision so rounding errors are lower
+        );
+    }
+
+    function testFuzz_roundingAgainstUser_multiUser_sDai(
+        uint256 rate1,
+        uint256 rate2,
+        uint256 amount1,
+        uint256 amount2
+    )
+        public
+    {
+        _runRoundingAgainstUsersFuzzTest(
+            sDai,
+            SDAI_TOKEN_MAX,
+            rate1,
+            rate2,
+            amount1,
+            amount2,
+            4  // sDai has higher rounding errors that can be introduced because of rate conversion
+        );
+    }
+
+    function _runRoundingAgainstUsersFuzzTest(
+        MockERC20 asset,
+        uint256   tokenMax,
+        uint256   rate1,
+        uint256   rate2,
+        uint256   amount1,
+        uint256   amount2,
+        uint256   roundingTolerance
+    ) internal {
+        address user1 = makeAddr("user1");
+        address user2 = makeAddr("user2");
+
+        rate1 = _bound(rate1, 1e27,  10e27);
+        rate2 = _bound(rate2, rate1, 10e27);
+
+        amount1 = _bound(amount1, 1, tokenMax);
+        amount2 = _bound(amount2, 1, tokenMax);
+
+        rateProvider.__setConversionRate(rate1);
+
+        _deposit(address(asset), address(user1), amount1);
+
+        assertEq(asset.balanceOf(address(user1)), 0);
+
+        vm.prank(user1);
+        psm.withdraw(address(asset), address(user1), amount1);
+
+        // Rounds against user up to one unit, always rounding down
+        assertApproxEqAbs(asset.balanceOf(address(user1)), amount1, roundingTolerance);
+        assertLe(asset.balanceOf(address(user1)), amount1);
+
+        rateProvider.__setConversionRate(rate2);
+
+        _deposit(address(asset), address(user2), amount2);
+
+        assertEq(asset.balanceOf(address(user2)), 0);
+
+        vm.prank(user2);
+        psm.withdraw(address(asset), address(user2), amount2);
+
+        // Rounds against user up to one unit, always rounding down
+
+        assertApproxEqAbs(asset.balanceOf(address(user2)), amount2, roundingTolerance);
+        assertLe(asset.balanceOf(address(user2)), amount2);
+    }
+}
diff --git a/test/Withdraw.t.sol b/test/Withdraw.t.sol
index 37ec99b..041da77 100644
--- a/test/Withdraw.t.sol
+++ b/test/Withdraw.t.sol
@@ -456,7 +456,6 @@ contract PSMWithdrawTests is PSMTestBase {
     function _getExpectedWithdrawnAmount(MockERC20 asset, address user, uint256 amount)
         internal view returns (uint256 withdrawAmount)
     {
-        // TODO: See if convertToAssets can be used
         uint256 balance    = asset.balanceOf(address(psm));
         uint256 userAssets = psm.convertToAssetValue(psm.shares(user));