diff --git a/contracts/StakeManagerV2.sol b/contracts/StakeManagerV2.sol index f925e15..b2446d8 100644 --- a/contracts/StakeManagerV2.sol +++ b/contracts/StakeManagerV2.sol @@ -22,6 +22,7 @@ contract StakeManagerV2 is ReentrancyGuardUpgradeable { using SafeERC20Upgradeable for IBnbX; + using SafeERC20Upgradeable for IERC20Upgradeable; bytes32 public constant OPERATOR_ROLE = keccak256("OPERATOR_ROLE"); bytes32 public constant MANAGER_ROLE = keccak256("MANAGER_ROLE"); @@ -47,6 +48,9 @@ contract StakeManagerV2 is uint256 public totalBnbxSupplyAtUndelegation; bool public redemptionEnabled; + bool public assetCustodied; + uint256 public sweepToCustodyTimestamp; + // @custom:oz-upgrades-unsafe-allow constructor constructor() { _disableInitializers(); @@ -174,6 +178,7 @@ contract StakeManagerV2 is /// @dev Caller must approve the contract to burn the specified BNBx amount before calling. /// @dev The exchange rate is based on totalBnbUndelegated and totalBnbxSupplyAtUndelegation snapshot. function redeemBnbxForBnb(uint256 _amountInBnbX) external override nonReentrant returns (uint256) { + if (assetCustodied) revert AssetCustodied(); if (!redemptionEnabled) revert RedemptionNotEnabled(); if (_amountInBnbX == 0) revert ZeroAmount(); if (totalBnbxSupplyAtUndelegation == 0) revert ZeroAmount(); @@ -439,6 +444,47 @@ contract StakeManagerV2 is emit SetRedemptionEnabled(_enabled); } + /// @notice Set the delay before `sweepToCustody` can be called. + /// @dev Can only be called by an address with the DEFAULT_ADMIN_ROLE. + function setCustodyDelay(uint256 _custodyDelay) external onlyRole(DEFAULT_ADMIN_ROLE) { + if (_custodyDelay == 0) revert ZeroCustodyDelay(); + sweepToCustodyTimestamp = block.timestamp + _custodyDelay; + emit SetCustodyDelay(sweepToCustodyTimestamp); + } + + /// @notice Sweep all BNB or ERC20 tokens to a custody address after the delay has elapsed. + /// @dev Can only be called by an address with the DEFAULT_ADMIN_ROLE. + /// @dev Pass `address(0)` as `_asset` to sweep native BNB. + function sweepToCustody( + address _asset, + address _custody + ) + external + nonReentrant + onlyRole(DEFAULT_ADMIN_ROLE) + { + if (_custody == address(0)) revert ZeroAddress(); + if (sweepToCustodyTimestamp == 0 || block.timestamp < sweepToCustodyTimestamp) { + revert CustodyDelayNotElapsed(); + } + + assetCustodied = true; + + uint256 bal; + if (_asset == address(0)) { + bal = address(this).balance; + if (bal == 0) revert ZeroAmount(); + (bool success,) = payable(_custody).call{ value: bal }(""); + if (!success) revert TransferFailed(); + } else { + bal = IERC20Upgradeable(_asset).balanceOf(address(this)); + if (bal == 0) revert ZeroAmount(); + IERC20Upgradeable(_asset).safeTransfer(_custody, bal); + } + + emit SweptToCustody(_asset, _custody, bal); + } + /*////////////////////////////////////////////////////////////// internal functions //////////////////////////////////////////////////////////////*/ diff --git a/contracts/interfaces/IStakeManagerV2.sol b/contracts/interfaces/IStakeManagerV2.sol index 6f2b6fd..bcfcd3b 100644 --- a/contracts/interfaces/IStakeManagerV2.sol +++ b/contracts/interfaces/IStakeManagerV2.sol @@ -33,6 +33,9 @@ interface IStakeManagerV2 { error WithdrawalBelowMinimum(); error RedemptionNotEnabled(); error InsufficientBnbBalance(); + error CustodyDelayNotElapsed(); + error ZeroCustodyDelay(); + error AssetCustodied(); function delegate(string calldata _referralId) external payable returns (uint256); function requestWithdraw(uint256 _amount, string calldata _referralId) external returns (uint256); @@ -49,6 +52,8 @@ interface IStakeManagerV2 { function pause() external; function unpause() external; function setRedemptionEnabled(bool _enabled) external; + function setCustodyDelay(uint256 _custodyDelay) external; + function sweepToCustody(address _asset, address _custody) external; function convertBnbToBnbX(uint256 _amount) external view returns (uint256); function convertBnbXToBnb(uint256 _amountInBnbX) external view returns (uint256); @@ -72,4 +77,6 @@ interface IStakeManagerV2 { event ClaimedAllBnbFromAllOperators(uint256 _totalClaimedBnb); event RedeemedBnbxForBnb(address indexed _account, uint256 _amountInBnbX, uint256 _amountInBnb); event SetRedemptionEnabled(bool _enabled); + event SetCustodyDelay(uint256 _sweepToCustodyTimestamp); + event SweptToCustody(address indexed _asset, address indexed _custody, uint256 _amount); } diff --git a/test/fork-tests/StakeManagerV2Custody.t.sol b/test/fork-tests/StakeManagerV2Custody.t.sol new file mode 100644 index 0000000..01c0646 --- /dev/null +++ b/test/fork-tests/StakeManagerV2Custody.t.sol @@ -0,0 +1,339 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity 0.8.25; + +import "forge-std/Test.sol"; + +import { ProxyAdmin } from "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol"; +import { ITransparentUpgradeableProxy } from + "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol"; + +import "contracts/StakeManagerV2.sol"; + +/// @dev Fork tests for the custody sweep mechanism added to `StakeManagerV2`. +/// +/// Standalone setup (does NOT inherit `StakeManagerV2Setup`) — the parent +/// fixture exercises `STAKE_HUB` via `_clearCurrentPendingTransactions` +/// which requires an archive-depth BSC RPC. These tests only need the +/// proxy upgraded to the new impl + funded with BNB, so we bypass that +/// machinery and stay compatible with non-archive BSC endpoints. +contract StakeManagerV2Custody is Test { + // Mainnet addresses (see StakeManagerV2Setup.t.sol). + address internal proxyAdmin = 0xF90e293D34a42CB592Be6BE6CA19A9963655673C; + address internal timelock = 0xD990A252E7e36700d47520e46cD2B3E446836488; + address internal admin = 0x79A2Ae748AC8bE4118B7a8096681B30310c3adBE; // internal multisig (DEFAULT_ADMIN_ROLE) + address internal manager = 0x79A2Ae748AC8bE4118B7a8096681B30310c3adBE; // internal multisig (MANAGER_ROLE) + + StakeManagerV2 internal stakeManagerV2 = StakeManagerV2(payable(0x3b961e83400D51e6E1AF5c450d3C7d7b80588d28)); + + address internal custody; + address internal attacker; + + event SetCustodyDelay(uint256 _sweepToCustodyTimestamp); + event SweptToCustody(address indexed _asset, address indexed _custody, uint256 _amount); + + function setUp() public { + string memory rpcUrl = vm.envString("BSC_MAINNET_RPC_URL"); + vm.createSelectFork(rpcUrl); + + custody = makeAddr("custody"); + attacker = makeAddr("attacker"); + + address newImpl = address(new StakeManagerV2()); + vm.prank(timelock); + ProxyAdmin(proxyAdmin).upgrade(ITransparentUpgradeableProxy(address(stakeManagerV2)), newImpl); + + // The live contract holds delegator BNB at the fork block; zero it so + // each test starts from a deterministic balance and can assert exact + // sweep amounts. + vm.deal(address(stakeManagerV2), 0); + } + + // -------- access control -------- + + function test_setCustodyDelay_revertsForNonAdmin() public { + vm.expectRevert(); + vm.prank(attacker); + stakeManagerV2.setCustodyDelay(1 days); + } + + function test_sweepToCustody_revertsForNonAdmin() public { + _arm(1 days); + skip(1 days); + _fundContract(1 ether); + + vm.expectRevert(); + vm.prank(attacker); + stakeManagerV2.sweepToCustody(address(0), custody); + } + + // -------- setCustodyDelay -------- + + function test_setCustodyDelay_setsTargetAndEmits() public { + uint256 delay = 7 days; + uint256 expectedTarget = block.timestamp + delay; + + vm.expectEmit(false, false, false, true); + emit SetCustodyDelay(expectedTarget); + + vm.prank(admin); + stakeManagerV2.setCustodyDelay(delay); + + assertEq(stakeManagerV2.sweepToCustodyTimestamp(), expectedTarget); + } + + function test_setCustodyDelay_revertsOnZero() public { + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.ZeroCustodyDelay.selector); + stakeManagerV2.setCustodyDelay(0); + } + + function test_setCustodyDelay_overwritesTargetOnReconfig() public { + _arm(7 days); + uint256 firstTarget = stakeManagerV2.sweepToCustodyTimestamp(); + + skip(1 days); + + _arm(3 days); + uint256 secondTarget = stakeManagerV2.sweepToCustodyTimestamp(); + + assertEq(secondTarget, block.timestamp + 3 days); + assertTrue(secondTarget != firstTarget); + } + + function test_setCustodyDelay_reconfigGatesSweep() public { + _arm(7 days); + skip(6 days); + _fundContract(1 ether); + + _arm(7 days); + + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); + stakeManagerV2.sweepToCustody(address(0), custody); + + skip(7 days); + vm.prank(admin); + stakeManagerV2.sweepToCustody(address(0), custody); + assertEq(custody.balance, 1 ether); + } + + // -------- sweepToCustody negative paths -------- + + function test_sweepToCustody_revertsBeforeArming() public { + _fundContract(1 ether); + + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); + stakeManagerV2.sweepToCustody(address(0), custody); + } + + function test_sweepToCustody_revertsBeforeTargetElapsed() public { + _arm(7 days); + _fundContract(1 ether); + skip(6 days); + + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); + stakeManagerV2.sweepToCustody(address(0), custody); + } + + function test_sweepToCustody_revertsOnZeroCustody() public { + _arm(1 days); + skip(1 days); + _fundContract(1 ether); + + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.ZeroAddress.selector); + stakeManagerV2.sweepToCustody(address(0), address(0)); + } + + function test_sweepToCustody_revertsOnZeroAmount() public { + _arm(1 days); + skip(1 days); + // contract has 0 BNB — full-balance sweep reverts + + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.ZeroAmount.selector); + stakeManagerV2.sweepToCustody(address(0), custody); + } + + function test_sweepToCustody_revertsOnFailedTransfer() public { + _arm(1 days); + skip(1 days); + _fundContract(1 ether); + + address rejector = address(new RejectETH()); + + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.TransferFailed.selector); + stakeManagerV2.sweepToCustody(address(0), rejector); + } + + // -------- sweepToCustody happy path -------- + + function test_sweepToCustody_sweepsFullBalanceAndEmits() public { + _arm(1 days); + skip(1 days); + _fundContract(5 ether); + + uint256 preCustody = custody.balance; + + vm.expectEmit(true, true, false, true); + emit SweptToCustody(address(0), custody, 5 ether); + + vm.prank(admin); + stakeManagerV2.sweepToCustody(address(0), custody); + + assertEq(custody.balance, preCustody + 5 ether); + assertEq(address(stakeManagerV2).balance, 0); + assertTrue(stakeManagerV2.assetCustodied()); + } + + function test_sweepToCustody_setsAssetCustodiedAndBlocksRedeem() public { + _arm(1 days); + skip(1 days); + _fundContract(1 ether); + + assertFalse(stakeManagerV2.assetCustodied()); + + vm.prank(admin); + stakeManagerV2.sweepToCustody(address(0), custody); + + assertTrue(stakeManagerV2.assetCustodied()); + + vm.expectRevert(IStakeManagerV2.AssetCustodied.selector); + stakeManagerV2.redeemBnbxForBnb(1 ether); + } + + // -------- sweepToCustody ERC20 path -------- + + function test_sweepToCustody_erc20_sweepsFullBalanceAndEmits() public { + _arm(1 days); + skip(1 days); + + MockERC20 token = new MockERC20(); + token.mint(address(stakeManagerV2), 1000 ether); + + vm.expectEmit(true, true, false, true); + emit SweptToCustody(address(token), custody, 1000 ether); + + vm.prank(admin); + stakeManagerV2.sweepToCustody(address(token), custody); + + assertEq(token.balanceOf(custody), 1000 ether); + assertEq(token.balanceOf(address(stakeManagerV2)), 0); + assertTrue(stakeManagerV2.assetCustodied()); + } + + function test_sweepToCustody_erc20_revertsOnZeroBalance() public { + _arm(1 days); + skip(1 days); + + MockERC20 token = new MockERC20(); + + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.ZeroAmount.selector); + stakeManagerV2.sweepToCustody(address(token), custody); + } + + function test_sweepToCustody_erc20_revertsBeforeTargetElapsed() public { + _arm(7 days); + MockERC20 token = new MockERC20(); + token.mint(address(stakeManagerV2), 1 ether); + + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); + stakeManagerV2.sweepToCustody(address(token), custody); + } + + function test_sweepToCustody_erc20_setsAssetCustodiedAndBlocksRedeem() public { + _arm(1 days); + skip(1 days); + + MockERC20 token = new MockERC20(); + token.mint(address(stakeManagerV2), 1 ether); + + vm.prank(admin); + stakeManagerV2.sweepToCustody(address(token), custody); + + assertTrue(stakeManagerV2.assetCustodied()); + + vm.expectRevert(IStakeManagerV2.AssetCustodied.selector); + stakeManagerV2.redeemBnbxForBnb(1 ether); + } + + function test_sweepToCustody_canSweepBnbThenErc20() public { + _arm(1 days); + skip(1 days); + _fundContract(2 ether); + + MockERC20 token = new MockERC20(); + token.mint(address(stakeManagerV2), 500 ether); + + vm.prank(admin); + stakeManagerV2.sweepToCustody(address(0), custody); + assertEq(custody.balance, 2 ether); + assertTrue(stakeManagerV2.assetCustodied()); + + // Second sweep, different asset — timestamp remains armed, balance > 0 + vm.prank(admin); + stakeManagerV2.sweepToCustody(address(token), custody); + assertEq(token.balanceOf(custody), 500 ether); + } + + // -------- helpers -------- + + function _arm(uint256 delay) internal { + vm.prank(admin); + stakeManagerV2.setCustodyDelay(delay); + } + + function _fundContract(uint256 amount) internal { + vm.deal(address(stakeManagerV2), amount); + } +} + +contract RejectETH { + receive() external payable { + revert("no eth"); + } +} + +contract MockERC20 { + string public name = "Mock"; + string public symbol = "MOCK"; + uint8 public decimals = 18; + uint256 public totalSupply; + mapping(address => uint256) public balanceOf; + mapping(address => mapping(address => uint256)) public allowance; + + event Transfer(address indexed from, address indexed to, uint256 value); + event Approval(address indexed owner, address indexed spender, uint256 value); + + function mint(address to, uint256 amount) external { + balanceOf[to] += amount; + totalSupply += amount; + emit Transfer(address(0), to, amount); + } + + function transfer(address to, uint256 amount) external returns (bool) { + balanceOf[msg.sender] -= amount; + balanceOf[to] += amount; + emit Transfer(msg.sender, to, amount); + return true; + } + + function transferFrom(address from, address to, uint256 amount) external returns (bool) { + allowance[from][msg.sender] -= amount; + balanceOf[from] -= amount; + balanceOf[to] += amount; + emit Transfer(from, to, amount); + return true; + } + + function approve(address spender, uint256 amount) external returns (bool) { + allowance[msg.sender][spender] = amount; + emit Approval(msg.sender, spender, amount); + return true; + } +}