From f7bb2bae82c405cce96a0e4d517201ba3a1984ab Mon Sep 17 00:00:00 2001
From: liyaxing <liyaxing@buckyos.com>
Date: Wed, 31 Jul 2024 18:12:14 +0800
Subject: [PATCH] Add lock state for user stake in dividend contract

---
 contracts/dividend.sol | 48 +++++++++++++++++++++++++++++++++++++++---
 test/test_dividend2.ts |  4 +++-
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/contracts/dividend.sol b/contracts/dividend.sol
index 069962c..70fea05 100644
--- a/contracts/dividend.sol
+++ b/contracts/dividend.sol
@@ -64,6 +64,10 @@ contract DividendContract is Initializable, UUPSUpgradeable, ReentrancyGuardUpgr
     mapping(address => bool) private tokenWhiteList;
     address[] private tokenWhiteListArray;
 
+    // Lock state for users
+    mapping (address => uint256) lockState;
+    uint256 public lockBlocks;
+    address public proposalContract;
 
     event TokenAddedToWhitelist(address token);
     event TokenRemovedFromWhitelist(address token);
@@ -73,17 +77,20 @@ contract DividendContract is Initializable, UUPSUpgradeable, ReentrancyGuardUpgr
     event NewCycle(uint256 cycleIndex, uint256 startBlock);
     event Withdraw(address indexed user, address token, uint256 amount);
 
-    function initialize(address _stakingToken, uint256 _cycleMaxLength, address[] memory _tokenList) public initializer {
+    function initialize(address _stakingToken, uint256 _cycleMaxLength, address[] memory _tokenList, uint256 _lockBlocks, address _proposalContract) public initializer {
         __UUPSUpgradeable_init();
         __ReentrancyGuard_init();
         __Ownable_init(msg.sender);
-        __DividendContractUpgradable_init(_stakingToken, _cycleMaxLength, _tokenList);
+        __DividendContractUpgradable_init(_stakingToken, _cycleMaxLength, _tokenList, _lockBlocks, _proposalContract);
     }
 
-    function __DividendContractUpgradable_init(address _stakingToken, uint256 _cycleMaxLength, address[] memory _tokenList) public onlyInitializing {
+    function __DividendContractUpgradable_init(address _stakingToken, uint256 _cycleMaxLength, address[] memory _tokenList, uint256 _lockBlocks, address _proposalContract) public onlyInitializing {
         stakingToken = _stakingToken;
         cycleMaxLength = _cycleMaxLength;
 
+        lockBlocks = _lockBlocks;
+        proposalContract = _proposalContract;
+
         for (uint i = 0; i < _tokenList.length; i++) {
             tokenWhiteList[_tokenList[i]] = true;
             tokenWhiteListArray.push(_tokenList[i]);
@@ -208,6 +215,36 @@ contract DividendContract is Initializable, UUPSUpgradeable, ReentrancyGuardUpgr
         return tokenBalances[token];
     }
 
+    function _updateLockState(address user) internal {
+        lockState[user] = block.number;
+    }
+
+    /**
+     * Update the lock state of the user, only the proposal contract can call this function, and return the current stake amount of the user
+     * @param user the user address to check
+     * @return the current stake amount of the user in total
+     */
+    function updateLockState(address user) external returns (uint256) {
+        require(msg.sender == proposalContract, "Only proposal contract can call this function");
+
+        uint256 stakeAmount = _getStakeAmount(user, currentCycleIndex);
+
+        if (stakeAmount > 0) {
+            _updateLockState(user);
+        }
+        
+        return stakeAmount;
+    }
+
+    /**
+     * Get the lock state of the user, use can not unstake if the lock state is true
+     * @param user the user address to check
+     * @return true if the user unstake is locked, otherwise false
+     */
+    function isUserUnstakeLocked(address user) public view returns (bool) {
+        return block.number - lockState[user] <= lockBlocks;
+    }
+
     // Deposit token as rewards to the current cycle
     function _depositToken(address token, uint256 amount) internal {
         require(amount > 0, "Cannot deposit 0");
@@ -328,6 +365,7 @@ contract DividendContract is Initializable, UUPSUpgradeable, ReentrancyGuardUpgr
 
         // console.log("user stake ===> amount %d, cycle %d, user %s", amount, currentCycleIndex, msg.sender);
 
+        // The last stake record of the user is always the newest stake amount of the user in the contract
         StakeRecord[] storage stakeRecords = UserStakeRecords[msg.sender];
         if (stakeRecords.length == 0) {
             stakeRecords.push(StakeRecord(currentCycleIndex, amount));
@@ -343,6 +381,9 @@ contract DividendContract is Initializable, UUPSUpgradeable, ReentrancyGuardUpgr
         // Update the total staked amount of the contract
         totalStaked += amount;
 
+        // Update the lock state of the user
+        _updateLockState(msg.sender);
+
         // Emit the stake event
         emit Stake(msg.sender, amount);
     }
@@ -358,6 +399,7 @@ contract DividendContract is Initializable, UUPSUpgradeable, ReentrancyGuardUpgr
         require(amount > 0, "Cannot unstake 0");
         StakeRecord[] storage stakeRecords = UserStakeRecords[msg.sender];
         require(stakeRecords.length > 0, "No stake record found");
+        require(!isUserUnstakeLocked(msg.sender), "Unstake is locked");
         
         // console.log("user unstake <=== amount %d, cycle %d, user %s", amount, currentCycleIndex, msg.sender);
 
diff --git a/test/test_dividend2.ts b/test/test_dividend2.ts
index c23cb28..28a3bb0 100644
--- a/test/test_dividend2.ts
+++ b/test/test_dividend2.ts
@@ -30,10 +30,12 @@ describe("Devidend", function () {
             "0x0000000000000000000000000000000000000000",
         ];
 
+        const proposalContract = "0x0000000000000000000000000000000000000000";
+
         dividend = (await (
             await upgrades.deployProxy(
                 await ethers.getContractFactory("DividendContract"),
-                [await dmc.getAddress(), 1000, tokenWhiteList]
+                [await dmc.getAddress(), 1000, tokenWhiteList, 500, proposalContract]
             )
         ).waitForDeployment()) as unknown as DividendContract;