From 606f43274b8d04f728e26036e221113ecfb2f2fb Mon Sep 17 00:00:00 2001 From: Dasari Manoj <122810482+Manudasari265@users.noreply.github.com> Date: Fri, 12 Jul 2024 23:33:09 +0530 Subject: [PATCH 1/2] Audited the Staking Contract --- .../Manoj/Session03/StakingContractAudit.txt | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 Bootcampers/Manoj/Session03/StakingContractAudit.txt diff --git a/Bootcampers/Manoj/Session03/StakingContractAudit.txt b/Bootcampers/Manoj/Session03/StakingContractAudit.txt new file mode 100644 index 0000000..9548cd4 --- /dev/null +++ b/Bootcampers/Manoj/Session03/StakingContractAudit.txt @@ -0,0 +1,110 @@ +1. The Lockup period and Reward Amount should be immutable, as their values are constant even after the contract is deployed. + Currently, they are not set as immutable. + +2. In the User struct, the stack id is initialized with uint8, which should be uint256 to avoid limiting the number of stakes. + The same issue is present in the code with { mapping(address => uint8) public userStakeCount; }. + +3. The constructor constructor(IERC20 _token) incorrectly uses Ownable(msg.sender), + as there is no need for Ownable to take any parameters in the constructor. + +4. The logic + if (rewardPool >= REWARD_AMOUNT) { + user.stakeAmount += REWARD_AMOUNT; + rewardPool -= REWARD_AMOUNT; + } + needs clarification. Rewards are added to the stake amount only if the reward pool is sufficient, + which might cause confusion and potential fund manipulation. + +5. In the initializeUser function, setting userStakeData[msg.sender][0] to a zero User struct is redundant. + +Final Code : + +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import "@openzeppelin/contracts/access/Ownable.sol"; + +/** + * @title StakingContract + * @dev A contract that allows users to stake ERC20 tokens, earn rewards, and unstake after a lockup period. + */ +contract BRBStaking is Ownable { + IERC20 public token; + uint256 public totalStaked; + uint256 public rewardPool; + uint256 public immutable LOCKUP_PERIOD; + uint256 public immutable REWARD_AMOUNT; + + /** + * @dev Struct to represent a user's staking information. + */ + struct User { + address userAddress; + uint256 stakeAmount; + bool initialized; + uint256 timeStamp; + uint256 stakeID; + } + + mapping(address => mapping(uint256 => User)) public userStakeData; + mapping(address => uint256) public userStakeCount; + + event UserInitialized(address indexed user); + event TokensStaked(address indexed user, uint256 amount, uint256 stakeID); + event TokensUnstaked(address indexed user, uint256 amount, uint256 stakeID); + event RewardsAdded(uint256 amount); + + constructor(IERC20 _token, uint256 _lockupPeriod, uint256 _rewardAmount) { + token = _token; + LOCKUP_PERIOD = _lockupPeriod; + REWARD_AMOUNT = _rewardAmount; + } + + function initializeUser() external { + require(!userStakeData[msg.sender][0].initialized, "User already initialized"); + User memory user = User(msg.sender, 0, true, 0, 0); + userStakeData[msg.sender][0] = user; + emit UserInitialized(msg.sender); + } + + function stake(uint256 _amount) external { + require(userStakeData[msg.sender][0].initialized, "User not initialized"); + require(token.transferFrom(msg.sender, address(this), _amount), "Token transfer failed"); + + uint256 stakeID = userStakeCount[msg.sender]; + User memory user = User(msg.sender, _amount, true, block.timestamp, stakeID); + userStakeData[msg.sender][stakeID] = user; + + userStakeCount[msg.sender]++; + totalStaked += _amount; + + emit TokensStaked(msg.sender, _amount, stakeID); + } + + function unstake(uint256 _stakeID) external { + User storage user = userStakeData[msg.sender][_stakeID]; + require(user.initialized, "Stake not found"); + require(block.timestamp >= user.timeStamp + LOCKUP_PERIOD, "Lockup period not completed"); + + uint256 amountToTransfer = user.stakeAmount; + if (rewardPool >= REWARD_AMOUNT) { + amountToTransfer += REWARD_AMOUNT; + rewardPool -= REWARD_AMOUNT; + } + + totalStaked -= user.stakeAmount; + delete userStakeData[msg.sender][_stakeID]; + + require(token.transfer(msg.sender, amountToTransfer), "Token transfer failed"); + + emit TokensUnstaked(msg.sender, user.stakeAmount, _stakeID); + } + + function addReward(uint256 _amount) external onlyOwner { + require(token.transferFrom(msg.sender, address(this), _amount), "Token transfer failed"); + rewardPool += _amount; + + emit RewardsAdded(_amount); + } +} \ No newline at end of file From abc010ca92e0184901bdbe0fa066ed194f171b01 Mon Sep 17 00:00:00 2001 From: Dasari Manoj <122810482+Manudasari265@users.noreply.github.com> Date: Sat, 13 Jul 2024 12:06:05 +0530 Subject: [PATCH 2/2] Added test cases for Staking contract --- .../.github/workflows/test.yml | 34 ++ .../Session03/brbStakingContract/.gitignore | 14 + .../Session03/brbStakingContract/.gitmodules | 6 + .../brbStakingContract/.vscode/settings.json | 3 + .../Session03/brbStakingContract/README.md | 77 +++++ .../Session03/brbStakingContract/foundry.toml | 6 + .../brbStakingContract/package-lock.json | 6 + .../brbStakingContract/remapping.txt | 5 + .../brbStakingContract/script/Counter.s.sol | 12 + .../brbStakingContract/src/BRBStaking.sol | 173 +++++++++++ .../brbStakingContract/src/BRBToken.sol | 14 + .../brbStakingContract/test/BRBStaking.t.sol | 290 ++++++++++++++++++ 12 files changed, 640 insertions(+) create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/.github/workflows/test.yml create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/.gitignore create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/.gitmodules create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/.vscode/settings.json create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/README.md create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/foundry.toml create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/package-lock.json create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/remapping.txt create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/script/Counter.s.sol create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/src/BRBStaking.sol create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/src/BRBToken.sol create mode 100644 Bootcampers/Manoj/Session03/brbStakingContract/test/BRBStaking.t.sol diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/.github/workflows/test.yml b/Bootcampers/Manoj/Session03/brbStakingContract/.github/workflows/test.yml new file mode 100644 index 0000000..9282e82 --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/.github/workflows/test.yml @@ -0,0 +1,34 @@ +name: test + +on: workflow_dispatch + +env: + FOUNDRY_PROFILE: ci + +jobs: + check: + strategy: + fail-fast: true + + name: Foundry project + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Install Foundry + uses: foundry-rs/foundry-toolchain@v1 + with: + version: nightly + + - name: Run Forge build + run: | + forge --version + forge build --sizes + id: build + + - name: Run Forge tests + run: | + forge test -vvv + id: test diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/.gitignore b/Bootcampers/Manoj/Session03/brbStakingContract/.gitignore new file mode 100644 index 0000000..85198aa --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/.gitignore @@ -0,0 +1,14 @@ +# Compiler files +cache/ +out/ + +# Ignores development broadcast logs +!/broadcast +/broadcast/*/31337/ +/broadcast/**/dry-run/ + +# Docs +docs/ + +# Dotenv file +.env diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/.gitmodules b/Bootcampers/Manoj/Session03/brbStakingContract/.gitmodules new file mode 100644 index 0000000..690924b --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/.gitmodules @@ -0,0 +1,6 @@ +[submodule "lib/forge-std"] + path = lib/forge-std + url = https://github.com/foundry-rs/forge-std +[submodule "lib/openzeppelin-contracts"] + path = lib/openzeppelin-contracts + url = https://github.com/OpenZeppelin/openzeppelin-contracts diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/.vscode/settings.json b/Bootcampers/Manoj/Session03/brbStakingContract/.vscode/settings.json new file mode 100644 index 0000000..bf47ba9 --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "solidity.compileUsingRemoteVersion": "v0.8.26+commit.8a97fa7a" +} \ No newline at end of file diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/README.md b/Bootcampers/Manoj/Session03/brbStakingContract/README.md new file mode 100644 index 0000000..0cc0439 --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/README.md @@ -0,0 +1,77 @@ +## BRB Staking with Foundry + +### Details of Staking Contract: +* Initialization Requirement: “Users must be able to initialize their staking profile.” +* Staking Tokens: “Users should be able to stake ERC20 tokens in the contract.” +* Lockup Period: “A 7-day lockup period must be enforced before users can unstake their tokens.” +* Unstaking Tokens: “Users should be able to unstake their tokens after the 7-day lockup period.” +* Fixed Reward: “Upon successful unstaking after the lockup period, users should receive a fixed reward of 100 tokens.” +* Reward Addition by Admin: “The admin should be able to add tokens to the reward pool.” + + + +### Details on Foundry Framework +**Foundry is a blazing fast, portable and modular toolkit for Ethereum application development written in Rust.** + +Foundry consists of: + +- **Forge**: Ethereum testing framework (like Truffle, Hardhat and DappTools). +- **Cast**: Swiss army knife for interacting with EVM smart contracts, sending transactions and getting chain data. +- **Anvil**: Local Ethereum node, akin to Ganache, Hardhat Network. +- **Chisel**: Fast, utilitarian, and verbose solidity REPL. + +## Documentation + +https://book.getfoundry.sh/ + +## Usage + +### Build + +```shell +$ forge build +``` + +### Test + +```shell +$ forge test +``` + +### Format + +```shell +$ forge fmt +``` + +### Gas Snapshots + +```shell +$ forge snapshot +``` + +### Anvil + +```shell +$ anvil +``` + +### Deploy + +```shell +$ forge script script/Counter.s.sol:CounterScript --rpc-url --private-key +``` + +### Cast + +```shell +$ cast +``` + +### Help + +```shell +$ forge --help +$ anvil --help +$ cast --help +``` diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/foundry.toml b/Bootcampers/Manoj/Session03/brbStakingContract/foundry.toml new file mode 100644 index 0000000..25b918f --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/foundry.toml @@ -0,0 +1,6 @@ +[profile.default] +src = "src" +out = "out" +libs = ["lib"] + +# See more config options https://github.com/foundry-rs/foundry/blob/master/crates/config/README.md#all-options diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/package-lock.json b/Bootcampers/Manoj/Session03/brbStakingContract/package-lock.json new file mode 100644 index 0000000..d057ccf --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/package-lock.json @@ -0,0 +1,6 @@ +{ + "name": "brbStakingContract", + "lockfileVersion": 3, + "requires": true, + "packages": {} +} diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/remapping.txt b/Bootcampers/Manoj/Session03/brbStakingContract/remapping.txt new file mode 100644 index 0000000..2c2800d --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/remapping.txt @@ -0,0 +1,5 @@ +@openzeppelin/contracts/=lib/openzeppelin-contracts/contracts/ +ds-test/=lib/openzeppelin-contracts/lib/forge-std/lib/ds-test/src/ +erc4626-tests/=lib/openzeppelin-contracts/lib/erc4626-tests/ +forge-std/=lib/forge-std/src/ +openzeppelin-contracts/=lib/openzeppelin-contracts/ diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/script/Counter.s.sol b/Bootcampers/Manoj/Session03/brbStakingContract/script/Counter.s.sol new file mode 100644 index 0000000..df9ee8b --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/script/Counter.s.sol @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import {Script, console} from "forge-std/Script.sol"; + +contract CounterScript is Script { + function setUp() public {} + + function run() public { + vm.broadcast(); + } +} diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/src/BRBStaking.sol b/Bootcampers/Manoj/Session03/brbStakingContract/src/BRBStaking.sol new file mode 100644 index 0000000..4703613 --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/src/BRBStaking.sol @@ -0,0 +1,173 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import "@openzeppelin/contracts/access/Ownable.sol"; +import "forge-std/console.sol"; + +/** + * @title StakingContract + * @dev A contract that allows users to stake ERC20 tokens, earn rewards, and unstake after a lockup period. + */ +contract BRBStaking is Ownable { + IERC20 public token; + uint256 public totalStaked; + uint256 public rewardPool; + uint256 public constant LOCKUP_PERIOD = 7 days; + uint256 public constant REWARD_AMOUNT = 100 ether; + + /** + * @dev Struct to represent a user's staking information. + */ + struct User { + address userAddress; + bool initialized; + uint8 stakeID; + uint256 stakeAmount; + uint256 timeStamp; + } + + //mapping(address => mapping(uint256 => User)) public userStakeData; + mapping(address => mapping(uint256 => User)) public userStakeData; + mapping(address => uint8) public userStakeCount; + + /** + * @dev Event emitted when a user initializes their staking profile. + * @param user The address of the user. + */ + event UserInitialized(address indexed user); + + /** + * @dev Event emitted when a user stakes tokens. + * @param user The address of the user. + * @param amount The amount of tokens staked. + * @param stakeID The ID of the stake. + */ + event TokensStaked(address indexed user, uint256 amount, uint256 stakeID); + + /** + * @dev Event emitted when a user unstakes tokens. + * @param user The address of the user. + * @param amount The amount of tokens unstaked. + * @param stakeID The ID of the stake. + */ + event TokensUnstaked(address indexed user, uint256 amount, uint256 stakeID); + + /** + * @dev Event emitted when the admin adds rewards to the pool. + * @param amount The amount of tokens added to the reward pool. + */ + event RewardsAdded(uint256 amount); + + /** + * @dev Constructor to initialize the staking contract with the token address. + * @param _token The address of the ERC20 token to be staked. + */ + constructor(IERC20 _token) Ownable(msg.sender) { + token = _token; + } + + /** + * @notice Initializes the user's staking profile. + */ + function initializeUser() external { + require( + !userStakeData[msg.sender][0].initialized, + "User already initialized" + ); + + // User memory user; + // user.userAddress = msg.sender; + // user.stakeAmount = 0; + // user.timeStamp = 0; + // user.stakeID = 0; + // user.initialized = true; + + userStakeData[msg.sender][0] = User({ + userAddress: msg.sender, + initialized: true, + stakeID: 0, + stakeAmount: 0, + timeStamp: 0 + }); + emit UserInitialized(msg.sender); + } + + /** + * @notice Allows a user to stake tokens. + * @param _amount The amount of tokens to stake. + */ + function stake(uint256 _amount) external payable { + require( + userStakeData[msg.sender][0].initialized, + "User not initialized" + ); + require( + token.transferFrom(msg.sender, address(this), _amount), + "Token transfer failed" + ); + + uint8 stakeID = userStakeCount[msg.sender] + 1; + // userStakeCount[msg.sender]++; + // User memory user; + // user.stakeAmount = _amount; + // user.userAddress = msg.sender; + // user.timeStamp = block.timestamp; + // user.stakeID = stakeID; + // user.initialized = true; + + userStakeData[msg.sender][stakeID] = User({ + stakeAmount: _amount, + userAddress: msg.sender, + timeStamp: block.timestamp, + stakeID: stakeID, + initialized: true + }); + + // optimized + // userStakeCount[msg.sender]++; + totalStaked += _amount; + + emit TokensStaked(msg.sender, _amount, stakeID); + } + + /** + * @notice Allows a user to unstake tokens after the lockup period. + * @param _stakeID The ID of the stake to unstake. + */ + function unstake(uint256 _stakeID) external { + User memory user = userStakeData[msg.sender][_stakeID]; + require(user.initialized, "User not Initialized"); + + require( + block.timestamp >= user.timeStamp + LOCKUP_PERIOD, + "Lockup period not completed" + ); + + totalStaked -= user.stakeAmount; + if (rewardPool >= REWARD_AMOUNT) { + user.stakeAmount += REWARD_AMOUNT; + token.transfer(msg.sender, user.stakeAmount); + rewardPool -= REWARD_AMOUNT; + } + + emit TokensUnstaked(msg.sender, user.stakeAmount, _stakeID); + + // delete userStakeData[msg.sender][_stakeID]; + } + + /** + * @notice Adds rewards to the reward pool. Only callable by the admin. + * @param _amount The amount of tokens to add to the reward pool. + */ + function addReward(uint256 _amount) external onlyOwner { + require( + token.transferFrom(msg.sender, address(this), _amount), + "Token transfer failed" + ); + require(_amount == 100 ether, "Amount must be 100 ether"); + rewardPool += _amount; + + emit RewardsAdded(_amount); + } +} \ No newline at end of file diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/src/BRBToken.sol b/Bootcampers/Manoj/Session03/brbStakingContract/src/BRBToken.sol new file mode 100644 index 0000000..5fc44d3 --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/src/BRBToken.sol @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT +// Compatible with OpenZeppelin Contracts ^5.0.0 +pragma solidity ^0.8.24; + +import "@openzeppelin/contracts/token/ERC20/ERC20.sol"; + +contract BRBToken is ERC20 { + constructor(address initialOwner) + ERC20("BRBToken", "BRB") + { + _mint(initialOwner , 10000 ether); + + } +} diff --git a/Bootcampers/Manoj/Session03/brbStakingContract/test/BRBStaking.t.sol b/Bootcampers/Manoj/Session03/brbStakingContract/test/BRBStaking.t.sol new file mode 100644 index 0000000..51d3034 --- /dev/null +++ b/Bootcampers/Manoj/Session03/brbStakingContract/test/BRBStaking.t.sol @@ -0,0 +1,290 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.24; + +import "forge-std/Test.sol"; +import "../src/BRBStaking.sol"; +import "../src/BRBToken.sol"; + +contract StakingContractTest is Test { + BRBStaking stakingContract; + BRBToken token; + + event TokensStaked(address indexed user, uint256 amount, uint256 stakeID); + event TokensUnstaked(address indexed user, uint256 amount, uint256 stakeID); + event RewardsAdded(uint256 amount); + + address owner; + address user; + address user2; + + function setUp() public { + owner = address(0x111); + user = address(0x222); + user2 = address(0x333); + + vm.startPrank(owner); + token = new BRBToken(owner); + stakingContract = new BRBStaking(token); + token.transfer(user, 1000 ether); + token.transfer(user2, 1000 ether); + vm.stopPrank(); + } + + function testInitializeUserTwice() public { + vm.startPrank(user); + stakingContract.initializeUser(); + + vm.expectRevert("User already initialized"); + stakingContract.initializeUser(); + vm.stopPrank(); + } + + function testInitializeUser() public { + vm.prank(user); + stakingContract.initializeUser(); + + (uint256 stakeAmount, uint256 timeStamp, address userAddress, uint8 stakeID, bool initialized) = stakingContract.userStakeData(user, 0); + assertEq(userAddress, user); + assertTrue(initialized); + assertEq(stakeAmount, 0); + assertEq(timeStamp, 0); + assertEq(stakeID, 0); + } + + function testStake() public { + uint256 stakeAmount = 100 ether; + + vm.startPrank(user); + token.approve(address(stakingContract), stakeAmount); + stakingContract.initializeUser(); + stakingContract.stake(stakeAmount); + vm.stopPrank(); + + (uint256 storedStakeAmount, uint256 timeStamp, address userAddress, uint8 stakeID, bool initialized) = stakingContract.userStakeData(user, 1); + assertEq(userAddress, user); + assertEq(storedStakeAmount, stakeAmount); + assertEq(stakeID, 1); + assertTrue(initialized); + assertGt(timeStamp, 0); + } + + function testUnstakeFunction() public { + uint256 stakeAmount = 100 ether; + uint256 rewardAmount = 100 ether; + + vm.startPrank(user); + token.approve(address(stakingContract), stakeAmount); + stakingContract.initializeUser(); + stakingContract.stake(stakeAmount); + vm.stopPrank(); + + vm.startPrank(owner); + token.approve(address(stakingContract), rewardAmount); + stakingContract.addReward(rewardAmount); + vm.stopPrank(); + + vm.warp(block.timestamp + 7 days); + + uint256 userBalanceBefore = token.balanceOf(user); + + vm.prank(user); + stakingContract.unstake(1); + + uint256 userBalanceAfter = token.balanceOf(user); + assertEq(userBalanceAfter, userBalanceBefore + stakeAmount + rewardAmount); + } + + function testAddReward() public { + uint256 rewardAmount = 100 ether; + vm.startPrank(owner); + token.approve(address(stakingContract), rewardAmount); + stakingContract.addReward(rewardAmount); + vm.stopPrank(); + + assertEq(stakingContract.rewardPool(), rewardAmount); + } + + function testUnstakeBeforeTime() public { + uint256 stakeAmount = 100 ether; + + vm.startPrank(user); + token.approve(address(stakingContract), stakeAmount); + stakingContract.initializeUser(); + stakingContract.stake(stakeAmount); + + vm.expectRevert("Lockup period not completed"); + stakingContract.unstake(1); + vm.stopPrank(); + } + + function testEmitStakeEvent() public { + vm.startPrank(user); + token.approve(address(stakingContract), 100 ether); + stakingContract.initializeUser(); + + vm.expectEmit(true, true, true, true); + emit TokensStaked(user, 100 ether, 1); + stakingContract.stake(100 ether); + vm.stopPrank(); + } + + function testEmitUnstakeEvent() public { + uint256 stakeAmount = 100 ether; + uint256 rewardAmount = 100 ether; + + vm.startPrank(user); + token.approve(address(stakingContract), stakeAmount); + stakingContract.initializeUser(); + stakingContract.stake(stakeAmount); + vm.stopPrank(); + + vm.startPrank(owner); + token.approve(address(stakingContract), rewardAmount); + stakingContract.addReward(rewardAmount); + vm.stopPrank(); + + vm.warp(block.timestamp + 7 days); + + vm.startPrank(user); + vm.expectEmit(true, true, true, true); + emit TokensUnstaked(user, stakeAmount + rewardAmount, 1); + stakingContract.unstake(1); + vm.stopPrank(); + } + + function testEmitRewardEvent() public { + uint256 rewardAmount = 100 ether; + vm.startPrank(owner); + token.approve(address(stakingContract), rewardAmount); + vm.expectEmit(true, true, true, true); + emit RewardsAdded(rewardAmount); + stakingContract.addReward(rewardAmount); + vm.stopPrank(); + } + + // 1. Only owner should be able to call the addRewards() + function testOnlyOwnerCanAddRewards() public { + uint256 rewardAmount = 100 ether; + vm.prank(user); + vm.expectRevert("Ownable: caller is not the owner"); + stakingContract.addReward(rewardAmount); + } + + // 2. Uninitialized User should not be able to STAKE + function testUninitializedUserCannotStake() public { + uint256 stakeAmount = 100 ether; + vm.startPrank(user); + token.approve(address(stakingContract), stakeAmount); + vm.expectRevert("User not initialized"); + stakingContract.stake(stakeAmount); + vm.stopPrank(); + } + + // 2. Uninitialized User should not be able to UNSTAKE + function testUninitializedUserCannotUnstake() public { + vm.prank(user); + vm.expectRevert("User not Initialized"); + stakingContract.unstake(1); + } + + // 3. Reward should exactly be 100 tokens for all stakers + + function testRewardDistribution() public { + uint256 stakeAmount = 100 ether; + uint256 rewardAmount = 100 ether; + + // User 1 stakes + vm.startPrank(user); + token.approve(address(stakingContract), stakeAmount); + stakingContract.initializeUser(); + stakingContract.stake(stakeAmount); + vm.stopPrank(); + + // User 2 stakes + vm.startPrank(user2); + token.approve(address(stakingContract), stakeAmount); + stakingContract.initializeUser(); + stakingContract.stake(stakeAmount); + vm.stopPrank(); + + // Add reward + vm.startPrank(owner); + token.approve(address(stakingContract), rewardAmount); + stakingContract.addReward(rewardAmount); + vm.stopPrank(); + + vm.warp(block.timestamp + 7 days); + + // User 1 unstakes + uint256 user1BalanceBefore = token.balanceOf(user); + vm.prank(user); + stakingContract.unstake(1); + uint256 user1BalanceAfter = token.balanceOf(user); + uint256 user1Reward = user1BalanceAfter - user1BalanceBefore - stakeAmount; + + // User 2 unstakes + uint256 user2BalanceBefore = token.balanceOf(user2); + vm.prank(user2); + stakingContract.unstake(1); + uint256 user2BalanceAfter = token.balanceOf(user2); + uint256 user2Reward = user2BalanceAfter - user2BalanceBefore - stakeAmount; + + // Check that rewards are distributed equally + assertEq(user1Reward, rewardAmount / 2); + assertEq(user2Reward, rewardAmount / 2); + } + + // 4. If a Staker tries to UNSTAKE before 7 days, it should revert + function testUnstakeBeforeLockupPeriod() public { + uint256 stakeAmount = 100 ether; + + vm.startPrank(user); + token.approve(address(stakingContract), stakeAmount); + stakingContract.initializeUser(); + stakingContract.stake(stakeAmount); + + // Try to unstake before the lockup period (7 days) + vm.expectRevert("Lockup period not completed"); + stakingContract.unstake(1); + + vm.stopPrank(); + } + + // 5. Event Emission of Stake, Unstake, and RewardAdded should be accurately tested + function testEventEmission() public { + uint256 stakeAmount = 100 ether; + uint256 rewardAmount = 50 ether; + + vm.startPrank(user); + token.approve(address(stakingContract), stakeAmount); + stakingContract.initializeUser(); + + // Stake and check for event + vm.expectEmit(true, true, true, true); + emit TokensStaked(user, stakeAmount, 1); + stakingContract.stake(stakeAmount); + + vm.stopPrank(); + + vm.startPrank(owner); + token.approve(address(stakingContract), rewardAmount); + + // Add reward and check for event + vm.expectEmit(true, true, true, true); + emit RewardsAdded(rewardAmount); + stakingContract.addReward(rewardAmount); + + vm.stopPrank(); + + vm.warp(block.timestamp + 7 days); + + vm.startPrank(user); + + // Unstake and check for event + vm.expectEmit(true, true, true, true); + emit TokensUnstaked(user, stakeAmount + rewardAmount, 1); // Assuming all reward goes to this user + stakingContract.unstake(1); + + vm.stopPrank(); + } +} \ No newline at end of file