New fix

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "lib/forge-std"]
+	path = lib/forge-std
+	url = https://github.com/foundry-rs/forge-std

contracts/erc4626/Vault.sol

@@ -18,6 +18,56 @@ import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
 import "../common/safe-HTS/SafeHTS.sol";
 import "../common/safe-HTS/IHederaTokenService.sol";
 
+/**
+ * @title Vault
+ *
+ * The contract which represents a custom Vault.
+ *
+ * /**
+ * @title SafeMath
+ * @dev Math operations with safety checks that throw on error
+ */
+library SafeMath {
+    /**
+     * @dev Multiplies two numbers, throws on overflow.
+     */
+    function mul(uint256 a, uint256 b) internal pure returns (uint256 c) {
+        if (a == 0) {
+            return 0;
+        }
+        c = a * b;
+        assert(c / a == b);
+        return c;
+    }
+
+    /**
+     * @dev Integer division of two numbers, truncating the quotient.
+     */
+    function div(uint256 a, uint256 b) internal pure returns (uint256) {
+        // assert(b > 0); // Solidity automatically throws when dividing by 0
+        // uint256 c = a / b;
+        // assert(a == b * c + a % b); // There is no case in which this doesn't hold
+        return a / b;
+    }
+
+    /**
+     * @dev Subtracts two numbers, throws on overflow (i.e. if subtrahend is greater than minuend).
+     */
+    function sub(uint256 a, uint256 b) internal pure returns (uint256) {
+        assert(b <= a);
+        return a - b;
+    }
+
+    /**
+     * @dev Adds two numbers, throws on overflow.
+     */
+    function add(uint256 a, uint256 b) internal pure returns (uint256 c) {
+        c = a + b;
+        assert(c >= a);
+        return c;
+    }
+}
+
 /**
  * @title Hedera Vault
  *
@@ -26,6 +76,7 @@ import "../common/safe-HTS/IHederaTokenService.sol";
 contract HederaVault is IERC4626, FeeConfiguration, TokenBalancer, Ownable, ReentrancyGuard {
     using SafeTransferLib for ERC20;
     using FixedPointMathLib for uint256;
+    using SafeMath for uint256;
     using Bits for uint256;
 
     // Staking token
@@ -461,11 +512,17 @@ contract HederaVault is IERC4626, FeeConfiguration, TokenBalancer, Ownable, Reen
         require(_vestingPeriod != 0, "Vault: Vesting period can't be zero");
 
         // Ensure the reward token is not the same as the staking token or the share token
-        require(_token != address(asset) && _token != share, "Vault: Reward and Staking tokens cannot be same");
+        require(
+            _token != address(asset) && _token != address(share),
+            "Vault: Reward and Staking tokens cannot be same"
+        );
 
         // Retrieve the reward info for the specified token
         RewardsInfo storage rewardInfo = tokensRewardInfo[_token];
 
+        // Update the vesting period even if the token already exists
+        rewardInfo.vestingPeriod = _vestingPeriod;
+
         // Get the current time for reward period calculations
         uint256 currentTime = block.timestamp;
 
@@ -483,18 +540,14 @@ contract HederaVault is IERC4626, FeeConfiguration, TokenBalancer, Ownable, Reen
         if (!tokenExists) {
             rewardTokens.push(_token);
 
-            // Set the vesting period for the new reward token
-            rewardInfo.vestingPeriod = _vestingPeriod;
-
-            // Associate the reward token with the vault
             SafeHTS.safeAssociateToken(_token, address(this));
         }
 
         // Add a new reward period for the token
         _addRewardPeriod(_token, _amount, currentTime);
 
         // Transfer the reward tokens from the sender to the vault
-        ERC20(_token).safeTransferFrom(msg.sender, address(this), _amount);
+        ERC20(_token).transferFrom(msg.sender, address(this), _amount);
 
         // Emit an event indicating that the reward has been added
         emit RewardAdded(_token, _amount);
@@ -648,6 +701,31 @@ contract HederaVault is IERC4626, FeeConfiguration, TokenBalancer, Ownable, Reen
      * @param _currentTime The current block timestamp.
      */
     function _addRewardPeriod(address _token, uint256 _amount, uint256 _currentTime) internal {
+        // // Ensure the token address is not zero (an invalid address)
+        // require(_token != address(0), "Vault: Token address can't be zero");
+        // // Ensure the amount is not zero
+        // require(_amount != 0, "Vault: Amount can't be zero");
+        // // Ensure the current time is not zero
+        // require(_currentTime != 0, "Vault: Current time can't be zero");
+
+        // // Retrieve the rewards information for the specified token
+        // RewardsInfo storage rewardInfo = tokensRewardInfo[_token];
+        // // Get the number of existing reward periods for this token
+        // uint256 rewardPeriodsLength = rewardInfo.rewardPeriods.length;
+
+        // // If there are existing reward periods, update the end time of the last period
+        // if (rewardPeriodsLength > 0) {
+        //     rewardInfo.rewardPeriods[rewardPeriodsLength - 1].endTime = _currentTime;
+        // }
+
+        // // Calculate the reward per share for the new period
+        // uint256 rewardPerShare = _amount / assetTotalSupply;
+
+        // // Add a new reward period starting at the current time with the calculated reward per share
+        // rewardInfo.rewardPeriods.push(
+        //     RewardPeriod({startTime: _currentTime, endTime: 0, rewardPerShare: rewardPerShare})
+        // );
+
         // Ensure the token address is not zero (an invalid address)
         require(_token != address(0), "Vault: Token address can't be zero");
         // Ensure the amount is not zero
@@ -666,7 +744,8 @@ contract HederaVault is IERC4626, FeeConfiguration, TokenBalancer, Ownable, Reen
         }
 
         // Calculate the reward per share for the new period
-        uint256 rewardPerShare = _amount / assetTotalSupply;
+        uint256 rewardPerShare = _amount.div(assetTotalSupply);
+        require(rewardPerShare > 0, "Vault: rewardPerShare must be greater than zero");
 
         // Add a new reward period starting at the current time with the calculated reward per share
         rewardInfo.rewardPeriods.push(

contracts/mockErc4626/asset.sol

@@ -0,0 +1,21 @@
+// // SPDX-License-Identifier: MIT
+// pragma solidity 0.8.20;
+// pragma abicoder v2;
+
+// import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+
+// /**
+//  * @title CustomERC20
+//  * @dev Implementation of the ERC20 Token to be used as shares.
+//  */
+// contract AssetToken is ERC20 {
+//     constructor() ERC20("ASSET", "ASST") {}
+
+//     function mint(address to, uint256 amount) public {
+//         _mint(to, amount);
+//     }
+
+//     function burnFrom(address account, uint256 amount) public {
+//         _burn(account, amount);
+//     }
+// }

contracts/mockErc4626/rewardToken.sol

@@ -0,0 +1,21 @@
+// // SPDX-License-Identifier: MIT
+// pragma solidity 0.8.20;
+// pragma abicoder v2;
+
+// import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+
+// /**
+//  * @title CustomERC20
+//  * @dev Implementation of the ERC20 Token to be used as shares.
+//  */
+// contract RewardToken1 is ERC20 {
+//     constructor() ERC20("Reward1", "RT1") {}
+
+//     function mint(address to, uint256 amount) public {
+//         _mint(to, amount);
+//     }
+
+//     function burnFrom(address account, uint256 amount) public {
+//         _burn(account, amount);
+//     }
+// }