Add stake on behalf of account functionality

.gitignore

@@ -10,3 +10,4 @@ artifacts/
 .vscode
 typechain-types
 scripts/safe/*.json
+scripts/deployStakingRewards/*testnet*.json

README.md

@@ -189,4 +189,4 @@ All upgrades go through the Timelock (2-day minimum delay):
 
 ## Security Audits
 1. [HALBORN](https://www.halborn.com/audits) has performed a security audit of `BillionsNetworkToken` smart contract and compiled report on Dec 5, 2025: [billions-token-7661d8](https://github.com/BillionsNetwork/billions-token/blob/main/audits/Billions_Token_SSC.pdf).  
-2. [HALBORN](https://www.halborn.com/audits) has performed a security audit of `StakingRewards` smart contract and compiled report on Jan 5, 2026: [staking-rewards-0b472f](https://github.com/BillionsNetwork/billions-token/blob/main/audits/Staking_Rewards_SSC.pdf).
+2. [HALBORN](https://www.halborn.com/audits) has performed a security audit of `StakingRewards` smart contract and compiled report on Jan 5, 2026: [staking-rewards-0b472f](https://github.com/BillionsNetwork/billions-token/blob/main/audits/Staking_Rewards_SSC.pdf).

contracts/staking/StakingRewards.sol

@@ -6,7 +6,7 @@ import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol
 import {Ownable2StepUpgradeable} from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
 import {ReentrancyGuardUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardUpgradeable.sol";
 import {PausableUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/PausableUpgradeable.sol";
-
+import {AccessControlUpgradeable} from "@openzeppelin/contracts-upgradeable/access/AccessControlUpgradeable.sol";
 import {IStakingRewards} from "./interfaces/IStakingRewards.sol";
 
 /**
@@ -76,10 +76,13 @@ import {IStakingRewards} from "./interfaces/IStakingRewards.sol";
 contract StakingRewards is
     IStakingRewards,
     Ownable2StepUpgradeable,
+    AccessControlUpgradeable,
     ReentrancyGuardUpgradeable,
     PausableUpgradeable
 {
     using SafeERC20 for IERC20;
+    bytes32 public constant STAKER_ON_BEHALF_ROLE = keccak256("STAKER_ON_BEHALF_ROLE");
+    bytes32 public constant REWARDS_DISTRIBUTOR_ROLE = keccak256("REWARDS_DISTRIBUTOR_ROLE");
 
     /* ========== STATE VARIABLES ========== */
 
@@ -90,7 +93,6 @@ contract StakingRewards is
     uint256 public rewardsDuration;
     uint256 public lastUpdateTime;
     uint256 public rewardPerTokenStored;
-    address public rewardsDistribution;
 
     mapping(address => uint256) public userRewardPerTokenPaid;
     mapping(address => uint256) public rewards;
@@ -102,6 +104,13 @@ contract StakingRewards is
 
     mapping(address => LockedStake) public addressToLockedStake;
 
+    uint256 public initialLockPeriodStart;
+    uint256 public initialLockPeriodDuration;
+
+    /// @dev Reserved storage gap for future upgrades. Reduces the gap by 1 for each new
+    ///      state variable added to this contract in subsequent versions.
+    uint256[50] private __gap;
+
     /* ========== CONSTRUCTOR ========== */
 
     /// @custom:oz-upgrades-unsafe-allow constructor
@@ -125,18 +134,22 @@ contract StakingRewards is
         uint256 _rewardsDuration
     ) external initializer {
         require(_owner != address(0), "Owner cannot be zero address");
+        require(_rewardsDistribution != address(0), "RewardsDistribution cannot be zero address");
         require(_rewardsToken != address(0), "RewardsToken cannot be zero address");
         require(_stakingToken != address(0), "StakingToken cannot be zero address");
 
         // Initialize inherited OZ contracts
         __Ownable_init(_owner);
+        __AccessControl_init();
         __ReentrancyGuard_init();
         __Pausable_init();
 
+        _grantRole(DEFAULT_ADMIN_ROLE, _owner);
+
         rewardsToken = IERC20(_rewardsToken);
         stakingToken = IERC20(_stakingToken);
 
-        _setRewardsDistribution(_rewardsDistribution);
+        _grantRole(REWARDS_DISTRIBUTOR_ROLE, _rewardsDistribution);
         _setRewardsDuration(_rewardsDuration);
     }
 
@@ -220,12 +233,18 @@ contract StakingRewards is
      * @dev Tokens are staked unlocked by default. Use lockStake() to lock staked tokens.
      * @param amount The amount of tokens to stake
      */
-    function stake(uint256 amount) public nonReentrant whenNotPaused updateReward(msg.sender) {
-        require(amount > 0, "Cannot stake 0");
-        _totalSupply = _totalSupply + amount;
-        _balances[msg.sender] = _balances[msg.sender] + amount;
-        stakingToken.safeTransferFrom(msg.sender, address(this), amount);
-        emit Staked(msg.sender, amount);
+    function stake(uint256 amount) public {
+        _stake(msg.sender, amount);
+    }
+
+    /**
+     * @notice Stakes tokens on behalf of an account
+     * @dev Tokens are staked unlocked by default. Use stakeAndLockOnBehalf to lock staked tokens.
+     * @param account The address on whose behalf to stake
+     * @param amount The amount of tokens to stake
+     */
+    function stakeOnBehalf(address account, uint256 amount) public onlyRole(STAKER_ON_BEHALF_ROLE) {
+        _stake(account, amount);
     }
 
     /**
@@ -237,35 +256,8 @@ contract StakingRewards is
      * @param amount The amount of staked tokens to lock
      * @param lockDuration The duration in seconds to lock the tokens
      */
-    function lockStake(uint256 amount, uint256 lockDuration) public whenNotPaused {
-        require(lockDuration > 0, "Lock duration must be greater than 0");
-        require(amount > 0, "Amount must be greater than 0");
-
-        // Check that user has enough unlocked staked balance to lock
-        uint256 currentLockedStakeAmount = getLockedStakeAmount(msg.sender);
-        uint256 newUnlockTimestamp = block.timestamp + lockDuration;
-
-        // Check if user has any locked tokens
-        if (currentLockedStakeAmount != 0) {
-            require(
-                amount >= currentLockedStakeAmount,
-                "Cannot reduce the amount of locked tokens"
-            );
-            require(
-                newUnlockTimestamp >= addressToLockedStake[msg.sender].unlockTimestamp,
-                "Cannot shorten the lock duration"
-            );
-        }
-
-        // Check that user has enough staked balance to lock
-        require(_balances[msg.sender] >= amount, "Not enough staked balance to lock");
-
-        // Update the locked tokens
-        addressToLockedStake[msg.sender].lockDuration = lockDuration;
-        addressToLockedStake[msg.sender].unlockTimestamp = newUnlockTimestamp;
-        addressToLockedStake[msg.sender].amount = amount;
-
-        emit StakeLocked(msg.sender, amount, lockDuration);
+    function lockStake(uint256 amount, uint256 lockDuration) public {
+        _lockStake(msg.sender, amount, lockDuration);
     }
 
     /**
@@ -280,8 +272,31 @@ contract StakingRewards is
         uint256 amountToLock,
         uint256 lockDuration
     ) public {
-        stake(amountToStake);
-        lockStake(amountToLock, lockDuration);
+        _stake(msg.sender, amountToStake);
+        _lockStake(msg.sender, amountToLock, lockDuration);
+    }
+
+    /**
+     * @notice Stakes and locks tokens on behalf of an account for a specified duration
+     * @dev Can only be called by the authorized stakerOnBehalf address. If the account already has a locked stake,
+     *      the new lock duration will be the maximum of the existing unlock timestamp and the new lock duration.
+     * @param account The address on whose behalf to stake and lock
+     * @param amount The amount of tokens to stake and lock
+     * @param lockDuration The duration in seconds to lock the tokens
+     */
+    function stakeAndLockOnBehalf(
+        address account,
+        uint256 amount,
+        uint256 lockDuration
+    ) public onlyRole(STAKER_ON_BEHALF_ROLE) {
+        uint256 currentLockedStakeAmount = getLockedStakeAmount(account);
+        if (currentLockedStakeAmount > 0) {
+            if (block.timestamp + lockDuration < addressToLockedStake[account].unlockTimestamp) {
+                lockDuration = addressToLockedStake[account].unlockTimestamp - block.timestamp;
+            }
+        }
+        _stake(account, amount);
+        _lockStake(account, currentLockedStakeAmount + amount, lockDuration);
     }
 
     /**
@@ -292,6 +307,12 @@ contract StakingRewards is
     function withdraw(uint256 amount) public nonReentrant updateReward(msg.sender) {
         require(amount > 0, "Cannot withdraw 0");
 
+        require(
+            initialLockPeriodStart == 0 ||
+                block.timestamp >= initialLockPeriodStart + initialLockPeriodDuration,
+            "Withdraw not allowed during initial lock period"
+        );
+
         // Check if there are any locked tokens
         uint256 lockedStakeAmount = getLockedStakeAmount(msg.sender);
 
@@ -311,6 +332,11 @@ contract StakingRewards is
      * @notice Claims all pending rewards for the caller
      */
     function getReward() public nonReentrant updateReward(msg.sender) {
+        require(
+            initialLockPeriodStart == 0 ||
+                block.timestamp >= initialLockPeriodStart + initialLockPeriodDuration,
+            "Get rewards not allowed during initial lock period"
+        );
         uint256 reward = rewards[msg.sender];
         if (reward > 0) {
             rewards[msg.sender] = 0;
@@ -336,7 +362,7 @@ contract StakingRewards is
      */
     function notifyRewardAmount(
         uint256 reward
-    ) external onlyRewardsDistribution updateReward(address(0)) {
+    ) external onlyRole(REWARDS_DISTRIBUTOR_ROLE) updateReward(address(0)) {
         if (block.timestamp >= periodFinish) {
             rewardRate = reward / rewardsDuration;
         } else {
@@ -362,11 +388,19 @@ contract StakingRewards is
     }
 
     /**
-     * @notice Sets the rewards distribution address
-     * @param _rewardsDistribution The address authorized to call notifyRewardAmount
+     * @notice Sets the initial period during which withdrawals and get rewards are not allowed
+     * @param duration The duration in seconds for which which withdrawals and get rewards are not allowed
      */
-    function setRewardsDistribution(address _rewardsDistribution) external onlyOwner {
-        _setRewardsDistribution(_rewardsDistribution);
+    function setInitialLockPeriod(uint256 duration) external onlyOwner {
+        require(duration > 0, "Initial lock period must be greater than 0");
+        require(
+            initialLockPeriodStart == 0 ||
+                block.timestamp > initialLockPeriodStart + initialLockPeriodDuration,
+            "Previous initial lock period must be complete before changing the duration for the new period"
+        );
+        initialLockPeriodStart = block.timestamp;
+        initialLockPeriodDuration = duration;
+        emit InitialLockPeriodUpdated(duration);
     }
 
     /**
@@ -414,13 +448,65 @@ contract StakingRewards is
     /* ========== INTERNAL FUNCTIONS ========== */
 
     /**
-     * @notice Internal function to set the rewards distribution address
-     * @param _rewardsDistribution The address authorized to call notifyRewardAmount
+     * @notice Stakes tokens
+     * @dev Tokens are staked unlocked by default. Use lockStake() to lock staked tokens.
+     * @param account The address on whose behalf to stake
+     * @param amount The amount of tokens to stake
      */
-    function _setRewardsDistribution(address _rewardsDistribution) internal {
-        require(_rewardsDistribution != address(0), "RewardsDistribution cannot be zero address");
-        rewardsDistribution = _rewardsDistribution;
-        emit RewardsDistributionUpdated(_rewardsDistribution);
+    function _stake(
+        address account,
+        uint256 amount
+    ) internal nonReentrant whenNotPaused updateReward(account) {
+        require(amount > 0, "Cannot stake 0");
+        _totalSupply = _totalSupply + amount;
+        _balances[account] = _balances[account] + amount;
+        stakingToken.safeTransferFrom(msg.sender, address(this), amount);
+        emit Staked(account, amount);
+    }
+
+    /**
+     * @notice Locks staked tokens for a specified duration
+     * @dev User must have enough staked balance to cover the new lock amount. If a lock already exists,
+     *      the new amount cannot be less than the currently locked amount and the new unlock timestamp
+     *      cannot be earlier than the current unlock timestamp. Increasing the lock amount uses additional
+     *      unlocked staked balance on top of the already locked tokens.
+     * @param amount The amount of staked tokens to lock
+     * @param account The address on whose behalf to lock the tokens
+     * @param lockDuration The duration in seconds to lock the tokens
+     */
+    function _lockStake(
+        address account,
+        uint256 amount,
+        uint256 lockDuration
+    ) internal whenNotPaused {
+        require(lockDuration > 0, "Lock duration must be greater than 0");
+        require(amount > 0, "Amount must be greater than 0");
+
+        // Check that user has enough unlocked staked balance to lock
+        uint256 currentLockedStakeAmount = getLockedStakeAmount(account);
+        uint256 newUnlockTimestamp = block.timestamp + lockDuration;
+
+        // Check if user has any locked tokens
+        if (currentLockedStakeAmount != 0) {
+            require(
+                amount >= currentLockedStakeAmount,
+                "Cannot reduce the amount of locked tokens"
+            );
+            require(
+                newUnlockTimestamp >= addressToLockedStake[account].unlockTimestamp,
+                "Cannot shorten the lock duration"
+            );
+        }
+
+        // Check that user has enough staked balance to lock
+        require(_balances[account] >= amount, "Not enough staked balance to lock");
+
+        // Update the locked tokens
+        addressToLockedStake[account].lockDuration = lockDuration;
+        addressToLockedStake[account].unlockTimestamp = newUnlockTimestamp;
+        addressToLockedStake[account].amount = amount;
+
+        emit StakeLocked(account, amount, lockDuration);
     }
 
     /**
@@ -439,11 +525,6 @@ contract StakingRewards is
 
     /* ========== MODIFIERS ========== */
 
-    modifier onlyRewardsDistribution() {
-        require(msg.sender == rewardsDistribution, "Caller is not RewardsDistribution contract");
-        _;
-    }
-
     modifier updateReward(address account) {
         rewardPerTokenStored = rewardPerToken();
         lastUpdateTime = lastTimeRewardApplicable();
@@ -464,4 +545,5 @@ contract StakingRewards is
     event Recovered(address token, uint256 amount);
     event RewardsDistributionUpdated(address indexed newRewardsDistribution);
     event StakeLocked(address indexed user, uint256 amount, uint256 lockDuration);
+    event InitialLockPeriodUpdated(uint256 newDuration);
 }

contracts/staking/interfaces/IStakingRewards.sol

@@ -8,7 +8,7 @@
  * @notice Interface for the StakingRewards contract
  * @dev Based on Synthetix StakingRewards with token locking support
  */
 interface IStakingRewards {
     /* ========== STRUCTS ========== */
 
     struct LockedStake {
@@ -19,7 +19,7 @@
 
     /* ========== VIEWS ========== */
 
     function rewardsToken() external view returns (IERC20);
 
     function stakingToken() external view returns (IERC20);
 
@@ -33,8 +33,6 @@
 
     function rewardPerTokenStored() external view returns (uint256);
 
-    function rewardsDistribution() external view returns (address);
-
     function userRewardPerTokenPaid(address account) external view returns (uint256);
 
     function rewards(address account) external view returns (uint256);
@@ -61,6 +59,8 @@
 
     function stake(uint256 amount) external;
 
+    function stakeOnBehalf(address account, uint256 amount) external;
+
     function withdraw(uint256 amount) external;
 
     function lockStake(uint256 amount, uint256 lockDuration) external;
@@ -71,6 +71,12 @@
         uint256 lockDuration
     ) external;
 
+    function stakeAndLockOnBehalf(
+        address account,
+        uint256 amountToStake,
+        uint256 lockDuration
+    ) external;
+
     function getReward() external;
 
     function exit() external;
@@ -79,13 +85,13 @@
 
     function notifyRewardAmount(uint256 reward) external;
 
-    function setRewardsDistribution(address _rewardsDistribution) external;
-
     function recoverERC20(address tokenAddress, uint256 tokenAmount) external;
 
     function setRewardsDuration(uint256 _rewardsDuration) external;
 
     function pause() external;
 
     function unpause() external;
+
+    function setInitialLockPeriod(uint256 duration) external;
 }