From 0a29e38d929be188895da6a3305197ce4591938c Mon Sep 17 00:00:00 2001
From: Jun Kim <64379343+junkim012@users.noreply.github.com>
Date: Sun, 12 May 2024 01:19:37 -0400
Subject: [PATCH] test: show that depositing the allocation cap diff or supply
 cap diff cannot violate either max caps

---
 test/helpers/VaultSharedSetup.sol |  6 +++
 test/unit/fuzz/vault/Vault.t.sol  | 73 +++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/test/helpers/VaultSharedSetup.sol b/test/helpers/VaultSharedSetup.sol
index 39618c79..bbf85f5f 100644
--- a/test/helpers/VaultSharedSetup.sol
+++ b/test/helpers/VaultSharedSetup.sol
@@ -284,4 +284,10 @@ contract VaultSharedSetup is IonPoolSharedSetup {
         IonPoolExposed(address(rsEthIonPool)).setSupplyFactor(7.1336673e27);
         IonPoolExposed(address(rswEthIonPool)).setSupplyFactor(10.1336673e27);
     }
+
+    function _zeroFloorSub(uint256 x, uint256 y) internal pure returns (uint256 z) {
+        assembly {
+            z := mul(gt(x, y), sub(x, y))
+        }
+    }
 }
diff --git a/test/unit/fuzz/vault/Vault.t.sol b/test/unit/fuzz/vault/Vault.t.sol
index 32cc524f..5e8d4a7f 100644
--- a/test/unit/fuzz/vault/Vault.t.sol
+++ b/test/unit/fuzz/vault/Vault.t.sol
@@ -10,6 +10,7 @@ import { console2 } from "forge-std/console2.sol";
 import { IERC20 } from "openzeppelin-contracts/contracts/interfaces/IERC20.sol";
 
 using Math for uint256;
+using WadRayMath for uint256;
 
 contract Vault_Fuzz is VaultSharedSetup {
     function setUp() public override {
@@ -55,6 +56,78 @@ contract Vault_Fuzz is VaultSharedSetup {
 
         assertEq(normalizedAmt, sharesToBurn);
     }
+
+    // NOTE Supplying the diff can revert if the normalized mint amount
+    // truncates to zero. Otherwise, it should be impossible to supply the
+    // 'diff' and end up violating the supply cap.
+
+    function testFuzz_DepositToFillSupplyCap(uint256 assets, uint256 supplyFactor) public {
+        supplyFactor = bound(supplyFactor, 1e27, 10e27);
+        IonPoolExposed(address(weEthIonPool)).setSupplyFactor(supplyFactor);
+
+        uint256 supplyCap = bound(assets, 100e18, type(uint128).max);
+        weEthIonPool.updateSupplyCap(supplyCap);
+
+        uint256 initialDeposit = bound(assets, 1e18, supplyCap - 10e18);
+        supply(address(this), weEthIonPool, initialDeposit);
+        uint256 initialTotalNormalized = weEthIonPool.totalSupplyUnaccrued();
+
+        uint256 supplyCapDiff = _zeroFloorSub(supplyCap, weEthIonPool.getTotalUnderlyingClaims());
+
+        // `IonPool.supply` math
+        uint256 amountScaled = supplyCapDiff.rayDivDown(supplyFactor);
+        uint256 resultingTotalNormalized = initialTotalNormalized + amountScaled;
+
+        uint256 resultingTotalClaim = resultingTotalNormalized.rayMulDown(supplyFactor);
+
+        supply(address(this), weEthIonPool, supplyCapDiff);
+
+        assertEq(
+            resultingTotalClaim, weEthIonPool.getTotalUnderlyingClaims(), "resulting should be the same as calculated"
+        );
+
+        // Is it possible that depositing this supplyCapDiff results in a revert?
+        // `IonPool` compares `getTotalUnderlyingClaims > _supplyCap`
+        assertLe(resultingTotalClaim, supplyCap, "supply cap reached");
+        assertLe(weEthIonPool.getTotalUnderlyingClaims(), supplyCap, "supply cap reached");
+    }
+
+    // Supplying the diff in the allocation cap should never end up violating
+    // the allocation cap.
+    // Is it possible that the `maxDeposit` returns more than the allocation cap?
+    function testFuzz_DepositToFillAllocationCap(uint256 assets, uint256 supplyFactor) public {
+        supplyFactor = bound(supplyFactor, 1e27, 10e27);
+        IonPoolExposed(address(weEthIonPool)).setSupplyFactor(supplyFactor);
+
+        uint256 allocationCap = bound(assets, 100e18, type(uint128).max);
+        updateAllocationCaps(vault, allocationCap, type(uint128).max, 0);
+
+        // Deposit, but leave some room below the allocation cap.
+        uint256 depositAmt = bound(assets, 1e18, allocationCap - 10e18);
+        setERC20Balance(address(BASE_ASSET), address(this), depositAmt);
+        vault.deposit(depositAmt, address(this));
+
+        uint256 initialTotalNormalized = weEthIonPool.totalSupplyUnaccrued();
+
+        uint256 allocationCapDiff = _zeroFloorSub(allocationCap, weEthIonPool.getUnderlyingClaimOf(address(vault)));
+
+        uint256 amountScaled = allocationCapDiff.rayDivDown(supplyFactor);
+        uint256 resultingTotalNormalized = initialTotalNormalized + amountScaled;
+        uint256 resultingTotalClaim = resultingTotalNormalized.rayMulDown(supplyFactor);
+
+        // Try to deposit a little more than the first allocation cap would
+        // allow, then check whether it's possible to violate the first
+        // allocation cap.
+
+        setERC20Balance(address(BASE_ASSET), address(this), allocationCapDiff + 123e18);
+        vault.deposit(allocationCapDiff + 123e18, address(this));
+
+        uint256 actualTotalClaim = weEthIonPool.getUnderlyingClaimOf(address(vault));
+        assertEq(resultingTotalClaim, actualTotalClaim, "expected and actual must be equal");
+
+        assertLe(resultingTotalClaim, allocationCap, "expected claim le to allocation cap");
+        assertLe(actualTotalClaim, allocationCap, "actual claim le to allocation cap");
+    }
 }
 
 contract VaultWithYieldAndFee_Fuzz is VaultSharedSetup {