Merge pull request #55 from yieldnest/chad/sc-319/audit-6-6-balance-t…

…ransfers

Fix audit issue 6.6 rebasing balance transfers

src/LSDStakingNode.sol

@@ -89,12 +89,15 @@ contract LSDStakingNode is ILSDStakingNode, Initializable, ReentrancyGuardUpgrad
                 revert UnsupportedAsset(asset);
             }
 
+            uint256 balanceBefore = asset.balanceOf(address(this));
             ynLSD.retrieveAsset(nodeId, assets[i], amount);
+            uint256 balanceAfter = asset.balanceOf(address(this));
+            uint256 retrievedAmount = balanceAfter - balanceBefore;
 
-            asset.forceApprove(address(strategyManager), amount);
+            asset.forceApprove(address(strategyManager), retrievedAmount);
 
-            uint256 eigenShares = strategyManager.depositIntoStrategy(IStrategy(strategy), asset, amount);
-            emit DepositToEigenlayer(assets[i], strategy, amount, eigenShares);
+            uint256 eigenShares = strategyManager.depositIntoStrategy(IStrategy(strategy), asset, retrievedAmount);
+            emit DepositToEigenlayer(assets[i], strategy, retrievedAmount, eigenShares);
         }
     }
 

test/foundry/integration/LSDStakingNode.t.sol

@@ -82,12 +82,11 @@ contract LSDStakingNodeTest is IntegrationBaseTest {
 		// 1. Obtain stETH and Deposit assets to ynLSD by User
 		IERC20 stETH = IERC20(chainAddresses.lsd.STETH_ADDRESS);
         uint256 amount = 1 ether;
-        (bool success, ) = chainAddresses.lsd.STETH_ADDRESS.call{value: amount + 1}("");
+        (bool success, ) = chainAddresses.lsd.STETH_ADDRESS.call{value: amount}("");
         require(success, "ETH transfer failed");
         uint256 balance = stETH.balanceOf(address(this));
-        assertEq(compareWithThreshold(balance, amount, 1), true, "Amount not received");
-		stETH.approve(address(ynlsd), amount);
-		ynlsd.deposit(stETH, amount, address(this));
+		stETH.approve(address(ynlsd), balance);
+		ynlsd.deposit(stETH, balance, address(this));
 
 		// 2. Deposit should fail when paused
 		IERC20[] memory assets = new IERC20[](1);

test/foundry/integration/ynLSD.t.sol

@@ -214,7 +214,10 @@ contract ynLSDAssetTest is IntegrationBaseTest {
         uint256 expectedBalance = balanceInStrategyForNode * oraclePrice / 1e18;
 
         // Assert that totalAssets reflects the deposit
-        assertEq(totalAssetsAfterDeposit - totalAssetsBeforeDeposit, expectedBalance, "Total assets do not reflect the deposit");
+        assertEq(
+            compareWithThreshold(totalAssetsAfterDeposit - totalAssetsBeforeDeposit,expectedBalance, 1), true, 
+            "Total assets do not reflect the deposit"
+        );
     }
 
     function testPreviewDeposit() public {

test/foundry/scenarios/Invariants.sol

@@ -2,35 +2,50 @@
 pragma solidity ^0.8.24;
 
 library Invariants {
+
+	/**
+	* @dev Compares two uint256 values and checks if their difference is within a specified threshold.
+	* @param value1 The first uint256 value.
+	* @param value2 The second uint256 value.
+	* @param threshold The threshold for the difference between value1 and value2.
+	* @return bool Returns true if the difference between value1 and value2 is less than or equal to the threshold.
+	*/
+	function compareWithThreshold(uint256 value1, uint256 value2, uint256 threshold) public pure returns (bool) {
+		if(value1 > value2) {
+			return (value1 - value2) <= threshold;
+		} else {
+			return (value2 - value1) <= threshold;
+		}
+	}
 
-  /// Share and Assets Invariants
-  function shareMintIntegrity (uint256 totalSupply, uint256 previousTotal, uint256 newShares) public pure {
-    require(totalSupply == previousTotal + newShares,
-      "Invariant: Total supply should be equal to previous total plus new shares"
-    );
-  }
+	/// Share and Assets Invariants
+	function shareMintIntegrity (uint256 totalSupply, uint256 previousTotal, uint256 newShares) public pure {
+		require(compareWithThreshold(totalSupply, previousTotal + newShares, 2) == true,
+			"Invariant: Total supply should be equal to previous total plus new shares"
+		);
+	}
 
-  function totalDepositIntegrity (uint256 totalDeposited, uint256 previousTotal, uint256 newDeposited) public pure {
-    require(totalDeposited == previousTotal + newDeposited,
-      "Invariant: Total deposited should be equal to previous total plus new deposited"
-    );
-  }
+	function totalDepositIntegrity (uint256 totalDeposited, uint256 previousTotal, uint256 newDeposited) public pure {
+		require(compareWithThreshold(totalDeposited, previousTotal + newDeposited, 2) == true,
+			"Invariant: Total deposited should be equal to previous total plus new deposited"
+		);
+	}
 
-  function userSharesIntegrity (uint256 userShares, uint256 previousShares, uint256 newShares) public pure {
-    require(userShares == previousShares + newShares,
-      "Invariant: User shares should be equal to previous shares plus new shares"
-    );
-  }
+	function userSharesIntegrity (uint256 userShares, uint256 previousShares, uint256 newShares) public pure {
+		require(compareWithThreshold(userShares, previousShares + newShares, 2) == true,
+			"Invariant: User shares should be equal to previous shares plus new shares"
+		);
+	}
 
-  function totalAssetsIntegrity (uint256 totalAssets, uint256 previousAssets, uint256 newAssets) public pure {
-    require(totalAssets == previousAssets + newAssets,
-      "Invariant: Total assets should be equal to previous assets plus new assets"
-    );
-  }
+	function totalAssetsIntegrity (uint256 totalAssets, uint256 previousAssets, uint256 newAssets) public pure {
+		require(compareWithThreshold(totalAssets, previousAssets + newAssets, 2) == true,
+			"Invariant: Total assets should be equal to previous assets plus new assets"
+		);
+	}
 
-  function totalBalanceIntegrity (uint256 balance, uint256 previousBalance, uint256 newBalance) public pure {
-    require(balance == previousBalance + newBalance,
-      "Invariant: Total balance should be equal to previous balance plus new balance"
-    );
-  }
+	function totalBalanceIntegrity (uint256 balance, uint256 previousBalance, uint256 newBalance) public pure {
+		require(compareWithThreshold(balance, previousBalance + newBalance, 2) == true,
+			"Invariant: Total balance should be equal to previous balance plus new balance"
+		);
+	}
 }

test/foundry/scenarios/ynLSD.spec.sol

@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: BSD-3-Clause
+pragma solidity ^0.8.24;
+
+import { IntegrationBaseTest } from "test/foundry/integration/IntegrationBaseTest.sol";
+import { Invariants } from "test/foundry/scenarios/Invariants.sol";
+import { IERC20 } from "@openzeppelin/contracts/interfaces/IERC20.sol";
+
+contract YnLSDScenarioTest1 is IntegrationBaseTest {
+
+	/**
+	Scenario 1: Successful LSD Deposit and Share Minting 
+	Objective: Test that a user can deposit stETH and receive 
+	the correct amount of shares in return.
+	*/
+
+	function test_ynLSD_Scenario_1_Fuzz(uint256 amount1, uint256 amount2, uint256 amount3) public {
+
+		/**
+		Users deposit random amounts
+		- Check the total assets of ynLSD
+		- Check the share balance of each user
+		- Check total supply of ynLSD
+		*/
+
+		address asset = chainAddresses.lsd.STETH_ADDRESS;
+
+		User_stETH_deposit(asset, amount1, address(0x01));
+		User_stETH_deposit(asset, amount2, address(0x02));
+		User_stETH_deposit(asset, amount3, address(0x03));
+	}
+
+	function User_stETH_deposit(address asset, uint256 amount, address user) public {
+
+		vm.assume(amount > 1 && amount < 10_000 ether);
+
+		uint256 previousTotalShares = ynlsd.totalSupply();
+		uint256 previousTotalAssets = ynlsd.getTotalAssets()[0];
+
+		vm.startPrank(user);
+		vm.deal(user, amount);
+		(bool success,) = asset.call{ value: amount }("");
+		require(success, "ETH transfer failed");
+		IERC20 steth = IERC20(asset);
+
+		uint256 userDeposit = IERC20(asset).balanceOf(user);
+
+		steth.approve(address(ynlsd), userDeposit);
+		ynlsd.deposit(steth, userDeposit, user);
+
+		uint256 userShares = ynlsd.balanceOf(user);
+
+		uint256 currentTotalAssets = ynlsd.getTotalAssets()[0];
+		uint256 currentTotalShares = ynlsd.totalSupply();
+
+		runInvariants(
+			user, 
+			previousTotalAssets, 
+			previousTotalShares,
+			currentTotalAssets,
+			currentTotalShares,
+			userDeposit, 
+			userShares
+		);
+	}
+
+	function runInvariants(
+		address user, 
+		uint256 previousTotalAssets,
+		uint256 previousTotalShares,
+		uint256 currentTotalAssets,
+		uint256 currentTotalShares,
+		uint256 userDeposit,
+		uint256 userShares
+	) public  view{
+		Invariants.totalAssetsIntegrity(currentTotalAssets, previousTotalAssets, userDeposit);
+		Invariants.shareMintIntegrity(currentTotalShares, previousTotalShares, userShares);
+		Invariants.userSharesIntegrity(ynlsd.balanceOf(user), 0, userShares);
+	}
+}