Merge pull request #75 from yieldnest/chad/sc-319/audit-8-1-fixes

Fix audit issues 8.1
yieldnest · Apr 1, 2024 · 2716e0b · 2716e0b
2 parents b87f183 + b7dee5e
commit 2716e0b
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 19 deletions.
diff --git a/src/StakingNodesManager.sol b/src/StakingNodesManager.sol
@@ -231,8 +231,8 @@ contract StakingNodesManager is
         uint256 totalDepositAmount = newValidators.length * DEFAULT_VALIDATOR_STAKE;
         ynETH.withdrawETH(totalDepositAmount); // Withdraw ETH from depositPool
 
-        uint256 validatorsLength = newValidators.length;
-        for (uint256 i = 0; i < validatorsLength; i++) {
+        uint256 newValidatorCount = newValidators.length;
+        for (uint256 i = 0; i < newValidatorCount; i++) {
 
             ValidatorData calldata validator = newValidators[i];
             if (usedValidators[validator.publicKey]) {
@@ -250,10 +250,12 @@ contract StakingNodesManager is
      */
     function validateNodes(ValidatorData[] calldata newValidators) public view {
 
+        uint256 nodeCount = nodes.length;
+
         for (uint256 i = 0; i < newValidators.length; i++) {
             uint256 nodeId = newValidators[i].nodeId;
 
-            if (nodeId >= nodes.length) {
+            if (nodeId >= nodeCount) {
                 revert InvalidNodeId(nodeId);
             }
         }

diff --git a/src/ynLSD.sol b/src/ynLSD.sol
@@ -289,22 +289,26 @@ contract ynLSD is IynLSD, ynBase, ReentrancyGuardUpgradeable, IynLSDEvents {
         view
         returns (uint256[] memory assetBalances)
     {
-        assetBalances = new uint256[](assets.length);
-        IStrategy[] memory assetStrategies = new IStrategy[](assets.length);
+        uint256 assetsCount = assets.length;
+
+        assetBalances = new uint256[](assetsCount);
+        IStrategy[] memory assetStrategies = new IStrategy[](assetsCount);
 
         // Add balances for funds held directly in ynLSD.
-        for (uint256 i = 0; i < assets.length; i++) {
-            assetStrategies[i] = strategies[assets[i]];
+        for (uint256 i = 0; i < assetsCount; i++) {
+            IERC20 asset = assets[i];
+            assetStrategies[i] = strategies[asset];
 
-            uint256 balanceThis = assets[i].balanceOf(address(this));
+            uint256 balanceThis = asset.balanceOf(address(this));
             assetBalances[i] += balanceThis;
         }
 
         // Add balances contained in each LSDStakingNode, including those managed by strategies.
-        for (uint256 i; i < nodes.length; i++ ) {
+        uint256 nodesCount = nodes.length;
+        for (uint256 i; i < nodesCount; i++ ) {
 
             ILSDStakingNode node = nodes[i];
-            for (uint256 j = 0; j < assets.length; j++) {
+            for (uint256 j = 0; j < assetsCount; j++) {
 
                 IERC20 asset = assets[j];
                 uint256 balanceNode = asset.balanceOf(address(node));
@@ -355,15 +359,16 @@ contract ynLSD is IynLSD, ynBase, ReentrancyGuardUpgradeable, IynLSDEvents {
         onlyRole(LSD_STAKING_NODE_CREATOR_ROLE)
         returns (ILSDStakingNode) {
 
-        if (nodes.length >= maxNodeCount) {
+        uint256 nodeId = nodes.length;
+
+        if (nodeId >= maxNodeCount) {
             revert TooManyStakingNodes(maxNodeCount);
         }
 
         BeaconProxy proxy = new BeaconProxy(address(upgradeableBeacon), "");
         ILSDStakingNode node = ILSDStakingNode(payable(proxy));
 
-        uint256 nodeId = nodes.length;
-        initializeLSDStakingNode(node);
+        initializeLSDStakingNode(node, nodeId);
 
         nodes.push(node);
 
@@ -377,12 +382,12 @@ contract ynLSD is IynLSD, ynBase, ReentrancyGuardUpgradeable, IynLSDEvents {
      * @dev This function checks the current initialized version of the node and performs initialization if it hasn't been done.
      * For future versions, additional conditional blocks should be added to handle version-specific initialization.
      * @param node The ILSDStakingNode instance to be initialized.
+     * @param nodeId The ID of the staking node.
      */
-    function initializeLSDStakingNode(ILSDStakingNode node) virtual internal {
+    function initializeLSDStakingNode(ILSDStakingNode node, uint256 nodeId) virtual internal {
 
          uint64 initializedVersion = node.getInitializedVersion();
          if (initializedVersion == 0) {
-             uint256 nodeId = nodes.length;
              node.initialize(
                ILSDStakingNode.Init(IynLSD(address(this)), nodeId)
              );
@@ -438,7 +443,7 @@ contract ynLSD is IynLSD, ynBase, ReentrancyGuardUpgradeable, IynLSDEvents {
 
         // Reinitialize all nodes to ensure compatibility with the new implementation.
         for (uint256 i = 0; i < nodeCount; i++) {
-            initializeLSDStakingNode(nodes[i]);
+            initializeLSDStakingNode(nodes[i], nodeCount);
         }
 
         emit UpgradedStakingNodeImplementationContract(address(_implementationContract), nodeCount);
@@ -490,4 +495,3 @@ contract ynLSD is IynLSD, ynBase, ReentrancyGuardUpgradeable, IynLSDEvents {
         _;
     }
 }
-
diff --git a/test/foundry/mocks/TestYnLSDV2.sol b/test/foundry/mocks/TestYnLSDV2.sol
@@ -6,11 +6,10 @@ import "../../../src/interfaces/ILSDStakingNode.sol";
 import "./TestLSDStakingNodeV2.sol";
 
 contract TestYnLSDV2 is ynLSD {
-    function initializeLSDStakingNode(ILSDStakingNode node) override internal {
+    function initializeLSDStakingNode(ILSDStakingNode node, uint256 nodeId) override internal {
 
          uint64 initializedVersion = node.getInitializedVersion();
          if (initializedVersion == 0) {
-             uint256 nodeId = nodes.length;
              node.initialize(
                 ILSDStakingNode.Init(IynLSD(address(this)), nodeId)
              );