Fixing: Security Vulnerabilities in VaultEngine

src/VaultEngine.sol

@@ -32,6 +32,9 @@ contract VaultEngine is ReentrancyGuard, IVaultEngine {
     using PriceOracle for AggregatorV3Interface;
     using VaultMath for uint256;
 
+    uint256 constant MAX_SUPPORTED_TOKENS = 50; // Reasonable limit
+
+
     // State Variables
     VaultStablecoin private immutable i_vaultStablecoin;
 
@@ -80,8 +83,9 @@ contract VaultEngine is ReentrancyGuard, IVaultEngine {
             revert VaultErrors.Vault__ZeroAddress();
         }
 
+        uint256 length = tokenAddresses.length; // Cache length
         // Initialize supported tokens and price feeds
-        for (uint256 i = 0; i < tokenAddresses.length; i++) {
+        for (uint256 i = 0; i < length;) {
             if (tokenAddresses[i] == address(0) || priceFeedAddresses[i] == address(0)) {
                 revert VaultErrors.Vault__ZeroAddress();
             }
@@ -100,15 +104,40 @@ contract VaultEngine is ReentrancyGuard, IVaultEngine {
      * @param amountCollateral Amount of collateral to deposit
      * @param amountStablecoinToMint Amount of stablecoins to mint
      */
+
     function depositCollateralAndMintStablecoin(
         address tokenCollateralAddress,
         uint256 amountCollateral,
         uint256 amountStablecoinToMint
-    ) external {
-        depositCollateral(tokenCollateralAddress, amountCollateral);
-        mintStablecoin(amountStablecoinToMint);
+    ) external nonReentrant {  
+        if (amountCollateral == 0 || amountStablecoinToMint == 0) {
+            revert VaultErrors.Vault__ZeroAmount();
+        }
+        if (s_priceFeeds[tokenCollateralAddress] == address(0)) {
+            revert VaultErrors.Vault__TokenNotSupported();
+        }
+
+        // Update state before external calls (CEI pattern)
+        s_collateralDeposited[msg.sender][tokenCollateralAddress] += amountCollateral;
+        s_stablecoinMinted[msg.sender] += amountStablecoinToMint;
+
+        // Validate health factor
+        _revertIfHealthFactorIsBroken(msg.sender);
+
+        // External calls last
+        bool collateralSuccess = IERC20(tokenCollateralAddress).transferFrom(
+            msg.sender, address(this), amountCollateral
+        );
+        require(collateralSuccess, "Collateral transfer failed");
+
+        bool mintSuccess = i_vaultStablecoin.mint(msg.sender, amountStablecoinToMint);
+        require(mintSuccess, "Mint failed");
+
+        emit CollateralDeposited(msg.sender, tokenCollateralAddress, amountCollateral);
+        emit StablecoinMinted(msg.sender, amountStablecoinToMint);
     }
 
+
     /**
      * @notice Redeems collateral and burns stablecoins in one transaction
      * @param tokenCollateralAddress Address of the collateral token
@@ -240,6 +269,8 @@ contract VaultEngine is ReentrancyGuard, IVaultEngine {
         _revertIfHealthFactorIsBroken(msg.sender); // This should never hit
     }
 
+
+
     // Private Functions
 
     function _redeemCollateral(
@@ -260,10 +291,11 @@ contract VaultEngine is ReentrancyGuard, IVaultEngine {
     function _burnStablecoin(uint256 amountStablecoinToBurn, address onBehalfOf, address stablecoinFrom) private {
         s_stablecoinMinted[onBehalfOf] -= amountStablecoinToBurn;
 
-        bool success = i_vaultStablecoin.transferFrom(stablecoinFrom, address(this), amountStablecoinToBurn);
-        if (!success) {
-            revert VaultErrors.Vault__TransferFailed();
-        }
+        bool success = i_vaultStablecoin.transferFrom(onBehalfOf, address(this), amountStablecoinToBurn);
+        // if (!success) {
+        //     revert VaultErrors.Vault__TransferFailed();
+        // } 
+        require(success, "Transfer failed");
 
         i_vaultStablecoin.burn(amountStablecoinToBurn);
         emit StablecoinBurned(onBehalfOf, amountStablecoinToBurn);
@@ -282,7 +314,7 @@ contract VaultEngine is ReentrancyGuard, IVaultEngine {
 
     function _healthFactor(address user) private view returns (uint256) {
         (uint256 totalStablecoinMinted, uint256 collateralValueInUsd) = _getAccountInformation(user);
-        return VaultMath.calculateHealthFactor(totalStablecoinMinted, collateralValueInUsd);
+        return VaultMath.calculateHealthFactor(collateralValueInUsd, totalStablecoinMinted);
     }
 
     function _getTokenPrice(address token) private view returns (uint256) {
@@ -292,7 +324,7 @@ contract VaultEngine is ReentrancyGuard, IVaultEngine {
 
     function _getUsdValue(address token, uint256 amount) private view returns (uint256) {
         uint256 price = _getTokenPrice(token);
-        return VaultMath.getUsdValue(amount, price);
+        return VaultMath.getUsdValue(amount, tokenPriceInUsd);
     }
 
     function _revertIfHealthFactorIsBroken(address user) internal view {
@@ -316,13 +348,21 @@ contract VaultEngine is ReentrancyGuard, IVaultEngine {
         return _getAccountInformation(user);
     }
 
-    function getCollateralValue(address user) public view returns (uint256 totalCollateralValueInUsd) {
-        for (uint256 i = 0; i < s_collateralTokens.length; i++) {
+    function getCollateralValue(address user, uint256 startIndex, uint256 maxTokens) public view returns (uint256 totalValue, uint256 nextIndex) {
+        uint256 endIndex = startIndex + maxTokens;
+        if (endIndex > s_collateralTokens.length) {
+            endIndex = s_collateralTokens.length;
+        }
+
+        for (uint256 i = startIndex; i < endIndex; i++) {
             address token = s_collateralTokens[i];
             uint256 amount = s_collateralDeposited[user][token];
-            totalCollateralValueInUsd += _getUsdValue(token, amount);
+            if (amount > 0) {  // Skip zero balances
+                totalValue += _getUsdValue(token, amount);
+            }
         }
-        return totalCollateralValueInUsd;
+
+        nextIndex = endIndex < s_collateralTokens.length ? endIndex : 0;
     }
 
     function getTokenAmountFromUsd(address token, uint256 usdAmountInWei) external view returns (uint256) {

src/VaultStablecoin.sol

@@ -15,6 +15,10 @@ import "./libraries/VaultErrors.sol";
  *      Only the VaultEngine contract can mint and burn tokens.
  */
 contract VaultStablecoin is ERC20Burnable, Ownable, IVaultStablecoin {
+    uint256 public constant MAX_DAILY_MINT = 1000000e18; // 1M vUSD daily limit
+    uint256 public lastMintTimestamp;
+    uint256 public dailyMintAmount;
+
     /**
      * @notice Initializes the Vault Stablecoin
      * @param initialOwner Address that will own this contract (should be VaultEngine)
@@ -42,6 +46,20 @@ contract VaultStablecoin is ERC20Burnable, Ownable, IVaultStablecoin {
         if (amount == 0) {
             revert VaultErrors.Vault__ZeroAmount();
         }
+
+        // Reset daily counter if needed
+        if (block.timestamp > lastMintTimestamp + 1 days) {
+            dailyMintAmount = 0;
+            lastMintTimestamp = block.timestamp;
+        }
+
+        // Check daily mint limit
+        require(dailyMintAmount + amount <= MAX_DAILY_MINT, "Daily mint limit exceeded");
+        dailyMintAmount += amount;
+
+        // Check total supply growth rate
+        uint256 currentSupply = totalSupply();
+        require(amount <= currentSupply / 100, "Cannot mint more than 1% of supply at once");
 
         _mint(to, amount);
         return true;

src/libraries/PriceOracle.sol

@@ -3,6 +3,7 @@ pragma solidity ^0.8.20;
 
 import "lib/chainlink-brownie-contracts/contracts/src/v0.8/shared/interfaces/AggregatorV3Interface.sol";
 import "./VaultErrors.sol";
+import "./VaultMath.sol";
 
 /**
  * @title PriceOracle
@@ -12,50 +13,37 @@ import "./VaultErrors.sol";
  */
 library PriceOracle {
     uint256 private constant TIMEOUT = 3 hours;
+    uint256 constant STALE_BLOCK_THRESHOLD = 240; // ~1 hour at 15s blocks
+
 
     /**
      * @notice Gets the latest price data with staleness validation
      * @param priceFeed Chainlink price feed interface
-     * @return roundId The round ID
-     * @return price The asset price
-     * @return startedAt Timestamp when the round started
-     * @return updatedAt Timestamp when the round was last updated
-     * @return answeredInRound The round ID of the round in which the answer was computed
+     * @return price The latest asset price with additional feed precision
      */
-    function getLatestPrice(AggregatorV3Interface priceFeed)
-        internal
-        view
-        returns (uint80 roundId, int256 price, uint256 startedAt, uint256 updatedAt, uint80 answeredInRound)
-    {
-        (roundId, price, startedAt, updatedAt, answeredInRound) = priceFeed.latestRoundData();
+    function getLatestPrice(AggregatorV3Interface priceFeed) internal view returns (uint256) {
+        (uint80 roundId, int256 price, uint256 startedAt, uint256 updatedAt, uint80 answeredInRound) = priceFeed.latestRoundData();
+        require(price > 0, "Invalid price");
+        require(answeredInRound >= roundId, "Stale price");
 
-        // Check for stale price data
-        if (updatedAt == 0 || answeredInRound < roundId) {
-            revert VaultErrors.Vault__StalePrice();
-        }
-
-        // Check if price is too old
+        // Use both timestamp AND block-based staleness checks
         uint256 secondsSinceUpdate = block.timestamp - updatedAt;
-        if (secondsSinceUpdate > TIMEOUT) {
-            revert VaultErrors.Vault__StalePrice();
-        }
+        require(secondsSinceUpdate <= TIMEOUT, "Price too stale (time)");
 
-        // Validate price is positive
-        if (price <= 0) {
-            revert VaultErrors.Vault__InvalidPriceData();
-        }
+        // Additional block-based check for extra security
+        require(block.number - updatedAt <= STALE_BLOCK_THRESHOLD, "Price too stale (blocks)");
 
-        return (roundId, price, startedAt, updatedAt, answeredInRound);
+        return uint256(price) * VaultMath.ADDITIONAL_FEED_PRECISION;
     }
 
+
     /**
      * @notice Gets only the price from the latest round data
      * @param priceFeed Chainlink price feed interface
      * @return price The latest price
      */
     function getPrice(AggregatorV3Interface priceFeed) internal view returns (uint256 price) {
-        (, int256 rawPrice,,,) = getLatestPrice(priceFeed);
-        return uint256(rawPrice);
+        return getLatestPrice(priceFeed);
     }
 
     /**

src/libraries/VaultMath.sol

@@ -33,10 +33,10 @@ library VaultMath{
     {
         if (totalStablecoinMinted == 0) return type(uint256).max;
 
-        uint256 collateralAdjustedForThreshold = 
+        uint256 healthFactor = 
             (collateralValueInUsd * LIQUIDATION_THRESHOLD) / LIQUIDATION_PRECISION;
 
-        return (collateralAdjustedForThreshold * PRECISION) / totalStablecoinMinted;
+        return healthFactor;
     }
 
     /**
@@ -55,16 +55,20 @@ library VaultMath{
 
     /**
      * @notice Converts token amount to USD value using price feed data
-     * @param tokenAmount Amount of tokens
-     * @param tokenPriceInUsd Token price in USD (8 decimals from Chainlink)
-     * @return usdValue Equivalent USD value
+     * @param amount Amount of tokens
+     * @param token Token price in USD (8 decimals from Chainlink)
      */
-    function getUsdValue(uint256 tokenAmount, uint256 tokenPriceInUsd)
-        internal
+    function getUsdValue(address token, uint256 amount, uint256 tokenPriceInUsd)
+        public
         pure
-        returns (uint256 usdValue)
+        returns (uint256)
     {
-        return (tokenAmount * tokenPriceInUsd * ADDITIONAL_FEED_PRECISION) / PRECISION;
+        require(amount <= type(uint256).max / tokenPriceInUsd, "Amount too large");
+        require(tokenPriceInUsd <= type(uint256).max / ADDITIONAL_FEED_PRECISION, "Price too large");
+        // Safe calculation with overflow protection
+        uint256 intermediateResult = amount * tokenPriceInUsd;
+        require(intermediateResult <= type(uint256).max / ADDITIONAL_FEED_PRECISION, "Calculation overflow");
+        return (intermediateResult * ADDITIONAL_FEED_PRECISION) / PRECISION;
     }
 
     /**
@@ -81,4 +85,33 @@ library VaultMath{
         uint256 collateralAmount = getTokenAmountFromUsd(debtToCover, tokenPriceInUsd);
         return (collateralAmount * LIQUIDATION_BONUS) / LIQUIDATION_PRECISION;
     }
+
+    // Add minimum amount validation and proper precision handling
+    function calculateCollateralAdjustedForThreshold(uint256 collateralValueInUsd)
+        internal pure returns (uint256) {
+
+        // Ensure minimum collateral value to prevent precision issues
+        require(collateralValueInUsd >= LIQUIDATION_PRECISION, "Collateral too small");
+
+        // Use higher precision arithmetic
+        uint256 threshold = (collateralValueInUsd * LIQUIDATION_THRESHOLD) / LIQUIDATION_PRECISION;
+
+        // Ensure result is never zero if input is non-zero
+        require(threshold > 0 || collateralValueInUsd == 0, "Threshold calculation error");
+
+        return threshold;
+    }
+
+    // Alternative: Use SafeMath with scaling
+    function calculateCollateralAdjustedForThresholdSafe(uint256 collateralValueInUsd)
+        internal pure returns (uint256) {
+
+        if (collateralValueInUsd == 0) return 0;
+
+        // Scale up calculation to maintain precision
+        uint256 scaledResult = (collateralValueInUsd * LIQUIDATION_THRESHOLD * 1e18) / (LIQUIDATION_PRECISION * 1e18);
+
+        return scaledResult;
+    }
+
 }

src/mocks/MockERC20.sol

@@ -31,8 +31,8 @@ contract MockERC20 is ERC20 {
         _mint(account, amount);
     }
 
-    function burn(address account, uint256 amount) external {
-        _burn(account, amount);
+    function burn(uint256 amount) external {
+        _burn(msg.sender, amount);
     }
 
     function faucet(uint256 amount) external {