From 9f8a03b88301ae62dd84fcff86b52ed7648ebf0c Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Thu, 14 May 2026 15:47:15 +0530 Subject: [PATCH 1/7] add custody sweep to StakeManagerV2 --- contracts/StakeManagerV2.sol | 46 ++++++++++++++++++++++++ contracts/interfaces/IStakeManagerV2.sol | 6 ++-- 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/contracts/StakeManagerV2.sol b/contracts/StakeManagerV2.sol index f925e15..f380519 100644 --- a/contracts/StakeManagerV2.sol +++ b/contracts/StakeManagerV2.sol @@ -38,6 +38,8 @@ contract StakeManagerV2 is uint256 public firstUnprocessedUserIndex; uint256 public firstUnbondingBatchIndex; uint256 public minWithdrawableBnbx; + uint256 public custodyDelay; + uint256 public custodyConfigTimestamp; WithdrawalRequest[] private withdrawalRequests; BatchWithdrawalRequest[] private batchWithdrawalRequests; @@ -439,6 +441,50 @@ contract StakeManagerV2 is emit SetRedemptionEnabled(_enabled); } + /// @notice Set the delay (seconds) that must elapse after the last + /// reconfiguration before `sweepToCustody` is callable. + /// @dev Bumps `custodyConfigTimestamp` to restart the sweep delay clock. + /// @dev Can only be called by an address with the MANAGER_ROLE. + /// @param _custodyDelay New delay in seconds. + function setCustodyDelay(uint256 _custodyDelay) external onlyRole(MANAGER_ROLE) { + custodyDelay = _custodyDelay; + custodyConfigTimestamp = block.timestamp; + emit SetCustodyDelay(_custodyDelay, block.timestamp); + } + + /// @notice Sweep an explicit amount of BNB from the contract to a + /// caller-provided custody address. The caller picks `_amount` because + /// bnbX is an ongoing protocol — draining the full balance would brick + /// `claimWithdrawal` payouts. Manager is responsible for leaving enough + /// BNB to cover live unclaimed batch withdrawals. + /// @dev Reverts until `setCustodyDelay` has been called at least once + /// (`custodyConfigTimestamp != 0`) AND + /// `block.timestamp >= custodyConfigTimestamp + custodyDelay`. Any + /// subsequent `setCustodyDelay` restarts the window. + /// @dev Can only be called by an address with the MANAGER_ROLE. + /// @param _custody Address that receives the swept BNB. + /// @param _amount Amount of BNB (wei) to send to `_custody`. + function sweepToCustody( + address _custody, + uint256 _amount + ) + external + nonReentrant + onlyRole(MANAGER_ROLE) + { + if (_custody == address(0)) revert ZeroAddress(); + if (_amount == 0) revert ZeroAmount(); + if (custodyConfigTimestamp == 0) revert CustodyDelayNotConfigured(); + if (block.timestamp < custodyConfigTimestamp + custodyDelay) { + revert CustodyDelayNotElapsed(); + } + + (bool success,) = payable(_custody).call{ value: _amount }(""); + if (!success) revert TransferFailed(); + + emit Swept(_custody, _amount); + } + /*////////////////////////////////////////////////////////////// internal functions //////////////////////////////////////////////////////////////*/ diff --git a/contracts/interfaces/IStakeManagerV2.sol b/contracts/interfaces/IStakeManagerV2.sol index 6f2b6fd..267344c 100644 --- a/contracts/interfaces/IStakeManagerV2.sol +++ b/contracts/interfaces/IStakeManagerV2.sol @@ -33,6 +33,8 @@ interface IStakeManagerV2 { error WithdrawalBelowMinimum(); error RedemptionNotEnabled(); error InsufficientBnbBalance(); + error CustodyDelayNotConfigured(); + error CustodyDelayNotElapsed(); function delegate(string calldata _referralId) external payable returns (uint256); function requestWithdraw(uint256 _amount, string calldata _referralId) external returns (uint256); @@ -68,8 +70,4 @@ interface IStakeManagerV2 { event SetMaxActiveRequestsPerUser(uint256 _maxActiveRequestsPerUser); event SetMaxExchangeRateSlippageBps(uint256 _maxExchangeRateSlippageBps); event SetMinWithdrawableBnbx(uint256 _minWithdrawableBnbx); - event UndelegatedAllBnbFromAllOperators(uint256 _totalBnbUndelegated, uint256 _totalBnbxSupply); - event ClaimedAllBnbFromAllOperators(uint256 _totalClaimedBnb); - event RedeemedBnbxForBnb(address indexed _account, uint256 _amountInBnbX, uint256 _amountInBnb); - event SetRedemptionEnabled(bool _enabled); } From 7db111b27e34bf628eba43b77ce1801f3f90c901 Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Thu, 14 May 2026 15:47:44 +0530 Subject: [PATCH 2/7] test: add fork tests for custody sweep --- test/fork-tests/StakeManagerV2Custody.t.sol | 172 ++++++++++++++++++++ 1 file changed, 172 insertions(+) create mode 100644 test/fork-tests/StakeManagerV2Custody.t.sol diff --git a/test/fork-tests/StakeManagerV2Custody.t.sol b/test/fork-tests/StakeManagerV2Custody.t.sol new file mode 100644 index 0000000..7ec3c68 --- /dev/null +++ b/test/fork-tests/StakeManagerV2Custody.t.sol @@ -0,0 +1,172 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity 0.8.25; + +import "./StakeManagerV2Setup.t.sol"; + +import { ITransparentUpgradeableProxy } from + "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol"; + +/// @dev Fork tests for the custody sweep mechanism added to `StakeManagerV2`. +/// The production proxy still points at the pre-custody impl, so each test +/// first upgrades the proxy to a freshly deployed impl that contains the new +/// state vars + functions. +contract StakeManagerV2Custody is StakeManagerV2Setup { + address internal custody; + address internal attacker; + + event SetCustodyDelay(uint256 _custodyDelay, uint256 _custodyConfigTimestamp); + event Swept(address indexed _custody, uint256 _amount); + + function setUp() public override { + super.setUp(); + custody = makeAddr("custody"); + attacker = makeAddr("attacker"); + _upgradeToCustodyImpl(); + } + + // -------- access control -------- + + function test_setCustodyDelay_revertsForNonManager() public { + vm.expectRevert(); + vm.prank(attacker); + stakeManagerV2.setCustodyDelay(1 days); + } + + function test_sweepToCustody_revertsForNonManager() public { + _arm(1 days); + skip(1 days); + _fundContract(1 ether); + + vm.expectRevert(); + vm.prank(attacker); + stakeManagerV2.sweepToCustody(custody, 1 ether); + } + + // -------- setCustodyDelay -------- + + function test_setCustodyDelay_setsValuesAndEmits() public { + uint256 delay = 7 days; + + vm.expectEmit(false, false, false, true); + emit SetCustodyDelay(delay, block.timestamp); + + vm.prank(manager); + stakeManagerV2.setCustodyDelay(delay); + + assertEq(stakeManagerV2.custodyDelay(), delay); + assertEq(stakeManagerV2.custodyConfigTimestamp(), block.timestamp); + } + + function test_setCustodyDelay_reconfigRestartsClock() public { + _arm(7 days); + skip(6 days); + + uint256 t1 = block.timestamp; + vm.prank(manager); + stakeManagerV2.setCustodyDelay(7 days); + + assertEq(stakeManagerV2.custodyConfigTimestamp(), t1); + + _fundContract(1 ether); + + skip(6 days); + vm.prank(manager); + vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); + stakeManagerV2.sweepToCustody(custody, 1 ether); + + skip(1 days); + vm.prank(manager); + stakeManagerV2.sweepToCustody(custody, 1 ether); + assertEq(custody.balance, 1 ether); + } + + // -------- sweepToCustody negative paths -------- + + function test_sweepToCustody_revertsBeforeArming() public { + _fundContract(1 ether); + + vm.prank(manager); + vm.expectRevert(IStakeManagerV2.CustodyDelayNotConfigured.selector); + stakeManagerV2.sweepToCustody(custody, 1 ether); + } + + function test_sweepToCustody_revertsBeforeDelayElapsed() public { + _arm(7 days); + _fundContract(1 ether); + skip(6 days); + + vm.prank(manager); + vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); + stakeManagerV2.sweepToCustody(custody, 1 ether); + } + + function test_sweepToCustody_revertsOnZeroCustody() public { + _arm(1 days); + skip(1 days); + _fundContract(1 ether); + + vm.prank(manager); + vm.expectRevert(IStakeManagerV2.ZeroAddress.selector); + stakeManagerV2.sweepToCustody(address(0), 1 ether); + } + + function test_sweepToCustody_revertsOnZeroAmount() public { + _arm(1 days); + skip(1 days); + _fundContract(1 ether); + + vm.prank(manager); + vm.expectRevert(IStakeManagerV2.ZeroAmount.selector); + stakeManagerV2.sweepToCustody(custody, 0); + } + + // -------- sweepToCustody happy path -------- + + function test_sweepToCustody_movesFundsAndEmits() public { + _arm(1 days); + skip(1 days); + _fundContract(5 ether); + + uint256 preBal = address(stakeManagerV2).balance; + uint256 preCustody = custody.balance; + + vm.expectEmit(true, false, false, true); + emit Swept(custody, 2 ether); + + vm.prank(manager); + stakeManagerV2.sweepToCustody(custody, 2 ether); + + assertEq(custody.balance, preCustody + 2 ether); + assertEq(address(stakeManagerV2).balance, preBal - 2 ether); + } + + function test_sweepToCustody_partialThenRemainder() public { + _arm(1 days); + skip(1 days); + _fundContract(3 ether); + + vm.startPrank(manager); + stakeManagerV2.sweepToCustody(custody, 1 ether); + stakeManagerV2.sweepToCustody(custody, 2 ether); + vm.stopPrank(); + + assertEq(custody.balance, 3 ether); + } + + // -------- helpers -------- + + function _arm(uint256 delay) internal { + vm.prank(manager); + stakeManagerV2.setCustodyDelay(delay); + } + + function _fundContract(uint256 amount) internal { + vm.deal(address(stakeManagerV2), address(stakeManagerV2).balance + amount); + } + + function _upgradeToCustodyImpl() internal { + address newImpl = address(new StakeManagerV2()); + vm.prank(timelock); + ProxyAdmin(proxyAdmin).upgrade(ITransparentUpgradeableProxy(address(stakeManagerV2)), newImpl); + } +} From a6783c47c35a4988817c82f9216d8389565a6082 Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Thu, 14 May 2026 16:42:36 +0530 Subject: [PATCH 3/7] =?UTF-8?q?refactor:=20simplify=20custody=20sweep=20?= =?UTF-8?q?=E2=80=94=20absolute=20target=20time,=20admin-gated?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- contracts/StakeManagerV2.sol | 37 ++++---- contracts/interfaces/IStakeManagerV2.sol | 2 + test/fork-tests/StakeManagerV2Custody.t.sol | 97 +++++++++++++++------ 3 files changed, 90 insertions(+), 46 deletions(-) diff --git a/contracts/StakeManagerV2.sol b/contracts/StakeManagerV2.sol index f380519..aebd642 100644 --- a/contracts/StakeManagerV2.sol +++ b/contracts/StakeManagerV2.sol @@ -38,8 +38,6 @@ contract StakeManagerV2 is uint256 public firstUnprocessedUserIndex; uint256 public firstUnbondingBatchIndex; uint256 public minWithdrawableBnbx; - uint256 public custodyDelay; - uint256 public custodyConfigTimestamp; WithdrawalRequest[] private withdrawalRequests; BatchWithdrawalRequest[] private batchWithdrawalRequests; @@ -49,6 +47,8 @@ contract StakeManagerV2 is uint256 public totalBnbxSupplyAtUndelegation; bool public redemptionEnabled; + uint256 public sweepToCustodyTimestamp; + // @custom:oz-upgrades-unsafe-allow constructor constructor() { _disableInitializers(); @@ -441,27 +441,26 @@ contract StakeManagerV2 is emit SetRedemptionEnabled(_enabled); } - /// @notice Set the delay (seconds) that must elapse after the last - /// reconfiguration before `sweepToCustody` is callable. - /// @dev Bumps `custodyConfigTimestamp` to restart the sweep delay clock. - /// @dev Can only be called by an address with the MANAGER_ROLE. - /// @param _custodyDelay New delay in seconds. - function setCustodyDelay(uint256 _custodyDelay) external onlyRole(MANAGER_ROLE) { - custodyDelay = _custodyDelay; - custodyConfigTimestamp = block.timestamp; - emit SetCustodyDelay(_custodyDelay, block.timestamp); + /// @notice Arm `sweepToCustody` so it becomes callable at + /// `block.timestamp + _custodyDelay`. Subsequent calls overwrite the + /// target time, so admin can shorten or extend the window. + /// @dev Can only be called by an address with the DEFAULT_ADMIN_ROLE. + /// @param _custodyDelay Seconds from now until `sweepToCustody` opens. + function setCustodyDelay(uint256 _custodyDelay) external onlyRole(DEFAULT_ADMIN_ROLE) { + if (_custodyDelay == 0) revert ZeroAmount(); + sweepToCustodyTimestamp = block.timestamp + _custodyDelay; + emit SetCustodyDelay(sweepToCustodyTimestamp); } /// @notice Sweep an explicit amount of BNB from the contract to a /// caller-provided custody address. The caller picks `_amount` because /// bnbX is an ongoing protocol — draining the full balance would brick - /// `claimWithdrawal` payouts. Manager is responsible for leaving enough + /// `claimWithdrawal` payouts. Admin is responsible for leaving enough /// BNB to cover live unclaimed batch withdrawals. /// @dev Reverts until `setCustodyDelay` has been called at least once - /// (`custodyConfigTimestamp != 0`) AND - /// `block.timestamp >= custodyConfigTimestamp + custodyDelay`. Any - /// subsequent `setCustodyDelay` restarts the window. - /// @dev Can only be called by an address with the MANAGER_ROLE. + /// (`sweepToCustodyTimestamp != 0`) AND + /// `block.timestamp >= sweepToCustodyTimestamp`. + /// @dev Can only be called by an address with the DEFAULT_ADMIN_ROLE. /// @param _custody Address that receives the swept BNB. /// @param _amount Amount of BNB (wei) to send to `_custody`. function sweepToCustody( @@ -470,12 +469,12 @@ contract StakeManagerV2 is ) external nonReentrant - onlyRole(MANAGER_ROLE) + onlyRole(DEFAULT_ADMIN_ROLE) { if (_custody == address(0)) revert ZeroAddress(); if (_amount == 0) revert ZeroAmount(); - if (custodyConfigTimestamp == 0) revert CustodyDelayNotConfigured(); - if (block.timestamp < custodyConfigTimestamp + custodyDelay) { + if (sweepToCustodyTimestamp == 0) revert CustodyDelayNotConfigured(); + if (block.timestamp < sweepToCustodyTimestamp) { revert CustodyDelayNotElapsed(); } diff --git a/contracts/interfaces/IStakeManagerV2.sol b/contracts/interfaces/IStakeManagerV2.sol index 267344c..b65ce35 100644 --- a/contracts/interfaces/IStakeManagerV2.sol +++ b/contracts/interfaces/IStakeManagerV2.sol @@ -70,4 +70,6 @@ interface IStakeManagerV2 { event SetMaxActiveRequestsPerUser(uint256 _maxActiveRequestsPerUser); event SetMaxExchangeRateSlippageBps(uint256 _maxExchangeRateSlippageBps); event SetMinWithdrawableBnbx(uint256 _minWithdrawableBnbx); + event SetCustodyDelay(uint256 _sweepToCustodyTimestamp); + event Swept(address indexed _custody, uint256 _amount); } diff --git a/test/fork-tests/StakeManagerV2Custody.t.sol b/test/fork-tests/StakeManagerV2Custody.t.sol index 7ec3c68..d06b2bc 100644 --- a/test/fork-tests/StakeManagerV2Custody.t.sol +++ b/test/fork-tests/StakeManagerV2Custody.t.sol @@ -9,12 +9,13 @@ import { ITransparentUpgradeableProxy } from /// @dev Fork tests for the custody sweep mechanism added to `StakeManagerV2`. /// The production proxy still points at the pre-custody impl, so each test /// first upgrades the proxy to a freshly deployed impl that contains the new -/// state vars + functions. +/// state var + functions. Both `setCustodyDelay` and `sweepToCustody` are +/// DEFAULT_ADMIN_ROLE-gated, so tests prank `admin`. contract StakeManagerV2Custody is StakeManagerV2Setup { address internal custody; address internal attacker; - event SetCustodyDelay(uint256 _custodyDelay, uint256 _custodyConfigTimestamp); + event SetCustodyDelay(uint256 _sweepToCustodyTimestamp); event Swept(address indexed _custody, uint256 _amount); function setUp() public override { @@ -26,13 +27,21 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { // -------- access control -------- - function test_setCustodyDelay_revertsForNonManager() public { + function test_setCustodyDelay_revertsForNonAdmin() public { vm.expectRevert(); vm.prank(attacker); stakeManagerV2.setCustodyDelay(1 days); } - function test_sweepToCustody_revertsForNonManager() public { + function test_setCustodyDelay_revertsForManager() public { + // Manager used to control this; with DEFAULT_ADMIN_ROLE-only the + // manager multisig must no longer be sufficient. + vm.expectRevert(); + vm.prank(manager); + stakeManagerV2.setCustodyDelay(1 days); + } + + function test_sweepToCustody_revertsForNonAdmin() public { _arm(1 days); skip(1 days); _fundContract(1 ether); @@ -44,38 +53,54 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { // -------- setCustodyDelay -------- - function test_setCustodyDelay_setsValuesAndEmits() public { + function test_setCustodyDelay_setsTargetAndEmits() public { uint256 delay = 7 days; + uint256 expectedTarget = block.timestamp + delay; vm.expectEmit(false, false, false, true); - emit SetCustodyDelay(delay, block.timestamp); + emit SetCustodyDelay(expectedTarget); - vm.prank(manager); + vm.prank(admin); stakeManagerV2.setCustodyDelay(delay); - assertEq(stakeManagerV2.custodyDelay(), delay); - assertEq(stakeManagerV2.custodyConfigTimestamp(), block.timestamp); + assertEq(stakeManagerV2.sweepToCustodyTimestamp(), expectedTarget); } - function test_setCustodyDelay_reconfigRestartsClock() public { + function test_setCustodyDelay_revertsOnZero() public { + vm.prank(admin); + vm.expectRevert(IStakeManagerV2.ZeroAmount.selector); + stakeManagerV2.setCustodyDelay(0); + } + + function test_setCustodyDelay_overwritesTargetOnReconfig() public { _arm(7 days); - skip(6 days); + uint256 firstTarget = stakeManagerV2.sweepToCustodyTimestamp(); - uint256 t1 = block.timestamp; - vm.prank(manager); - stakeManagerV2.setCustodyDelay(7 days); + skip(1 days); - assertEq(stakeManagerV2.custodyConfigTimestamp(), t1); + _arm(3 days); + uint256 secondTarget = stakeManagerV2.sweepToCustodyTimestamp(); - _fundContract(1 ether); + assertEq(secondTarget, block.timestamp + 3 days); + // Second call should fully overwrite, not extend. + assertTrue(secondTarget != firstTarget); + } + function test_setCustodyDelay_reconfigGatesSweep() public { + _arm(7 days); skip(6 days); - vm.prank(manager); + _fundContract(1 ether); + + // 6 days in, original target is 1 day away. Reconfiguring with a + // 7-day delay must push the window further out, not let sweep fire. + _arm(7 days); + + vm.prank(admin); vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); stakeManagerV2.sweepToCustody(custody, 1 ether); - skip(1 days); - vm.prank(manager); + skip(7 days); + vm.prank(admin); stakeManagerV2.sweepToCustody(custody, 1 ether); assertEq(custody.balance, 1 ether); } @@ -85,17 +110,17 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { function test_sweepToCustody_revertsBeforeArming() public { _fundContract(1 ether); - vm.prank(manager); + vm.prank(admin); vm.expectRevert(IStakeManagerV2.CustodyDelayNotConfigured.selector); stakeManagerV2.sweepToCustody(custody, 1 ether); } - function test_sweepToCustody_revertsBeforeDelayElapsed() public { + function test_sweepToCustody_revertsBeforeTargetElapsed() public { _arm(7 days); _fundContract(1 ether); skip(6 days); - vm.prank(manager); + vm.prank(admin); vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); stakeManagerV2.sweepToCustody(custody, 1 ether); } @@ -105,7 +130,7 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { skip(1 days); _fundContract(1 ether); - vm.prank(manager); + vm.prank(admin); vm.expectRevert(IStakeManagerV2.ZeroAddress.selector); stakeManagerV2.sweepToCustody(address(0), 1 ether); } @@ -115,11 +140,23 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { skip(1 days); _fundContract(1 ether); - vm.prank(manager); + vm.prank(admin); vm.expectRevert(IStakeManagerV2.ZeroAmount.selector); stakeManagerV2.sweepToCustody(custody, 0); } + 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(rejector, 1 ether); + } + // -------- sweepToCustody happy path -------- function test_sweepToCustody_movesFundsAndEmits() public { @@ -133,7 +170,7 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { vm.expectEmit(true, false, false, true); emit Swept(custody, 2 ether); - vm.prank(manager); + vm.prank(admin); stakeManagerV2.sweepToCustody(custody, 2 ether); assertEq(custody.balance, preCustody + 2 ether); @@ -145,7 +182,7 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { skip(1 days); _fundContract(3 ether); - vm.startPrank(manager); + vm.startPrank(admin); stakeManagerV2.sweepToCustody(custody, 1 ether); stakeManagerV2.sweepToCustody(custody, 2 ether); vm.stopPrank(); @@ -156,7 +193,7 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { // -------- helpers -------- function _arm(uint256 delay) internal { - vm.prank(manager); + vm.prank(admin); stakeManagerV2.setCustodyDelay(delay); } @@ -170,3 +207,9 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { ProxyAdmin(proxyAdmin).upgrade(ITransparentUpgradeableProxy(address(stakeManagerV2)), newImpl); } } + +contract RejectETH { + receive() external payable { + revert("no eth"); + } +} From faca3ce8ec58e2aa79138fe72d3ebef6bfc05af4 Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Thu, 14 May 2026 19:59:30 +0530 Subject: [PATCH 4/7] fix: rename ZeroAmount -> ZeroCustodyDelay in setCustodyDelay --- contracts/StakeManagerV2.sol | 2 +- contracts/interfaces/IStakeManagerV2.sol | 1 + test/fork-tests/StakeManagerV2Custody.t.sol | 55 +++++++++++---------- 3 files changed, 30 insertions(+), 28 deletions(-) diff --git a/contracts/StakeManagerV2.sol b/contracts/StakeManagerV2.sol index aebd642..8dfcb53 100644 --- a/contracts/StakeManagerV2.sol +++ b/contracts/StakeManagerV2.sol @@ -447,7 +447,7 @@ contract StakeManagerV2 is /// @dev Can only be called by an address with the DEFAULT_ADMIN_ROLE. /// @param _custodyDelay Seconds from now until `sweepToCustody` opens. function setCustodyDelay(uint256 _custodyDelay) external onlyRole(DEFAULT_ADMIN_ROLE) { - if (_custodyDelay == 0) revert ZeroAmount(); + if (_custodyDelay == 0) revert ZeroCustodyDelay(); sweepToCustodyTimestamp = block.timestamp + _custodyDelay; emit SetCustodyDelay(sweepToCustodyTimestamp); } diff --git a/contracts/interfaces/IStakeManagerV2.sol b/contracts/interfaces/IStakeManagerV2.sol index b65ce35..fa170a5 100644 --- a/contracts/interfaces/IStakeManagerV2.sol +++ b/contracts/interfaces/IStakeManagerV2.sol @@ -35,6 +35,7 @@ interface IStakeManagerV2 { error InsufficientBnbBalance(); error CustodyDelayNotConfigured(); error CustodyDelayNotElapsed(); + error ZeroCustodyDelay(); function delegate(string calldata _referralId) external payable returns (uint256); function requestWithdraw(uint256 _amount, string calldata _referralId) external returns (uint256); diff --git a/test/fork-tests/StakeManagerV2Custody.t.sol b/test/fork-tests/StakeManagerV2Custody.t.sol index d06b2bc..4c5d7c7 100644 --- a/test/fork-tests/StakeManagerV2Custody.t.sol +++ b/test/fork-tests/StakeManagerV2Custody.t.sol @@ -1,28 +1,46 @@ // SPDX-License-Identifier: UNLICENSED pragma solidity 0.8.25; -import "./StakeManagerV2Setup.t.sol"; +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`. -/// The production proxy still points at the pre-custody impl, so each test -/// first upgrades the proxy to a freshly deployed impl that contains the new -/// state var + functions. Both `setCustodyDelay` and `sweepToCustody` are -/// DEFAULT_ADMIN_ROLE-gated, so tests prank `admin`. -contract StakeManagerV2Custody is StakeManagerV2Setup { +/// +/// 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 Swept(address indexed _custody, uint256 _amount); - function setUp() public override { - super.setUp(); + function setUp() public { + string memory rpcUrl = vm.envString("BSC_MAINNET_RPC_URL"); + vm.createSelectFork(rpcUrl); + custody = makeAddr("custody"); attacker = makeAddr("attacker"); - _upgradeToCustodyImpl(); + + address newImpl = address(new StakeManagerV2()); + vm.prank(timelock); + ProxyAdmin(proxyAdmin).upgrade(ITransparentUpgradeableProxy(address(stakeManagerV2)), newImpl); } // -------- access control -------- @@ -33,14 +51,6 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { stakeManagerV2.setCustodyDelay(1 days); } - function test_setCustodyDelay_revertsForManager() public { - // Manager used to control this; with DEFAULT_ADMIN_ROLE-only the - // manager multisig must no longer be sufficient. - vm.expectRevert(); - vm.prank(manager); - stakeManagerV2.setCustodyDelay(1 days); - } - function test_sweepToCustody_revertsForNonAdmin() public { _arm(1 days); skip(1 days); @@ -68,7 +78,7 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { function test_setCustodyDelay_revertsOnZero() public { vm.prank(admin); - vm.expectRevert(IStakeManagerV2.ZeroAmount.selector); + vm.expectRevert(IStakeManagerV2.ZeroCustodyDelay.selector); stakeManagerV2.setCustodyDelay(0); } @@ -82,7 +92,6 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { uint256 secondTarget = stakeManagerV2.sweepToCustodyTimestamp(); assertEq(secondTarget, block.timestamp + 3 days); - // Second call should fully overwrite, not extend. assertTrue(secondTarget != firstTarget); } @@ -91,8 +100,6 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { skip(6 days); _fundContract(1 ether); - // 6 days in, original target is 1 day away. Reconfiguring with a - // 7-day delay must push the window further out, not let sweep fire. _arm(7 days); vm.prank(admin); @@ -200,12 +207,6 @@ contract StakeManagerV2Custody is StakeManagerV2Setup { function _fundContract(uint256 amount) internal { vm.deal(address(stakeManagerV2), address(stakeManagerV2).balance + amount); } - - function _upgradeToCustodyImpl() internal { - address newImpl = address(new StakeManagerV2()); - vm.prank(timelock); - ProxyAdmin(proxyAdmin).upgrade(ITransparentUpgradeableProxy(address(stakeManagerV2)), newImpl); - } } contract RejectETH { From cfb72fc638ef2e62e29380fee7115637139371ee Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Mon, 18 May 2026 15:48:57 +0530 Subject: [PATCH 5/7] fix: restore missing events and add function sigs to interface --- contracts/StakeManagerV2.sol | 16 ++-------------- contracts/interfaces/IStakeManagerV2.sol | 6 ++++++ 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/contracts/StakeManagerV2.sol b/contracts/StakeManagerV2.sol index 8dfcb53..655d19d 100644 --- a/contracts/StakeManagerV2.sol +++ b/contracts/StakeManagerV2.sol @@ -441,28 +441,16 @@ contract StakeManagerV2 is emit SetRedemptionEnabled(_enabled); } - /// @notice Arm `sweepToCustody` so it becomes callable at - /// `block.timestamp + _custodyDelay`. Subsequent calls overwrite the - /// target time, so admin can shorten or extend the window. + /// @notice Set the delay before `sweepToCustody` can be called. /// @dev Can only be called by an address with the DEFAULT_ADMIN_ROLE. - /// @param _custodyDelay Seconds from now until `sweepToCustody` opens. function setCustodyDelay(uint256 _custodyDelay) external onlyRole(DEFAULT_ADMIN_ROLE) { if (_custodyDelay == 0) revert ZeroCustodyDelay(); sweepToCustodyTimestamp = block.timestamp + _custodyDelay; emit SetCustodyDelay(sweepToCustodyTimestamp); } - /// @notice Sweep an explicit amount of BNB from the contract to a - /// caller-provided custody address. The caller picks `_amount` because - /// bnbX is an ongoing protocol — draining the full balance would brick - /// `claimWithdrawal` payouts. Admin is responsible for leaving enough - /// BNB to cover live unclaimed batch withdrawals. - /// @dev Reverts until `setCustodyDelay` has been called at least once - /// (`sweepToCustodyTimestamp != 0`) AND - /// `block.timestamp >= sweepToCustodyTimestamp`. + /// @notice Sweep BNB to a custody address after the delay has elapsed. /// @dev Can only be called by an address with the DEFAULT_ADMIN_ROLE. - /// @param _custody Address that receives the swept BNB. - /// @param _amount Amount of BNB (wei) to send to `_custody`. function sweepToCustody( address _custody, uint256 _amount diff --git a/contracts/interfaces/IStakeManagerV2.sol b/contracts/interfaces/IStakeManagerV2.sol index fa170a5..ffdf149 100644 --- a/contracts/interfaces/IStakeManagerV2.sol +++ b/contracts/interfaces/IStakeManagerV2.sol @@ -52,6 +52,8 @@ interface IStakeManagerV2 { function pause() external; function unpause() external; function setRedemptionEnabled(bool _enabled) external; + function setCustodyDelay(uint256 _custodyDelay) external; + function sweepToCustody(address _custody, uint256 _amount) external; function convertBnbToBnbX(uint256 _amount) external view returns (uint256); function convertBnbXToBnb(uint256 _amountInBnbX) external view returns (uint256); @@ -71,6 +73,10 @@ interface IStakeManagerV2 { event SetMaxActiveRequestsPerUser(uint256 _maxActiveRequestsPerUser); event SetMaxExchangeRateSlippageBps(uint256 _maxExchangeRateSlippageBps); event SetMinWithdrawableBnbx(uint256 _minWithdrawableBnbx); + event UndelegatedAllBnbFromAllOperators(uint256 _totalBnbUndelegated, uint256 _totalBnbxSupply); + event ClaimedAllBnbFromAllOperators(uint256 _totalClaimedBnb); + event RedeemedBnbxForBnb(address indexed _account, uint256 _amountInBnbX, uint256 _amountInBnb); + event SetRedemptionEnabled(bool _enabled); event SetCustodyDelay(uint256 _sweepToCustodyTimestamp); event Swept(address indexed _custody, uint256 _amount); } From 564cc5b0561a5332fc3d11e4d59cb473a5d89c7f Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Mon, 18 May 2026 16:04:08 +0530 Subject: [PATCH 6/7] refactor: sweepToCustody takes asset+custody, sweeps full balance, sets assetCustodied --- contracts/StakeManagerV2.sol | 31 ++++++++---- contracts/interfaces/IStakeManagerV2.sol | 6 +-- test/fork-tests/StakeManagerV2Custody.t.sol | 53 +++++++++++---------- 3 files changed, 53 insertions(+), 37 deletions(-) diff --git a/contracts/StakeManagerV2.sol b/contracts/StakeManagerV2.sol index 655d19d..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,7 @@ contract StakeManagerV2 is uint256 public totalBnbxSupplyAtUndelegation; bool public redemptionEnabled; + bool public assetCustodied; uint256 public sweepToCustodyTimestamp; // @custom:oz-upgrades-unsafe-allow constructor @@ -176,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(); @@ -449,27 +452,37 @@ contract StakeManagerV2 is emit SetCustodyDelay(sweepToCustodyTimestamp); } - /// @notice Sweep BNB to a custody address after the delay has elapsed. + /// @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 _custody, - uint256 _amount + address _asset, + address _custody ) external nonReentrant onlyRole(DEFAULT_ADMIN_ROLE) { if (_custody == address(0)) revert ZeroAddress(); - if (_amount == 0) revert ZeroAmount(); - if (sweepToCustodyTimestamp == 0) revert CustodyDelayNotConfigured(); - if (block.timestamp < sweepToCustodyTimestamp) { + if (sweepToCustodyTimestamp == 0 || block.timestamp < sweepToCustodyTimestamp) { revert CustodyDelayNotElapsed(); } - (bool success,) = payable(_custody).call{ value: _amount }(""); - if (!success) revert TransferFailed(); + 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 Swept(_custody, _amount); + emit SweptToCustody(_asset, _custody, bal); } /*////////////////////////////////////////////////////////////// diff --git a/contracts/interfaces/IStakeManagerV2.sol b/contracts/interfaces/IStakeManagerV2.sol index ffdf149..bcfcd3b 100644 --- a/contracts/interfaces/IStakeManagerV2.sol +++ b/contracts/interfaces/IStakeManagerV2.sol @@ -33,9 +33,9 @@ interface IStakeManagerV2 { error WithdrawalBelowMinimum(); error RedemptionNotEnabled(); error InsufficientBnbBalance(); - error CustodyDelayNotConfigured(); 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); @@ -53,7 +53,7 @@ interface IStakeManagerV2 { function unpause() external; function setRedemptionEnabled(bool _enabled) external; function setCustodyDelay(uint256 _custodyDelay) external; - function sweepToCustody(address _custody, uint256 _amount) external; + function sweepToCustody(address _asset, address _custody) external; function convertBnbToBnbX(uint256 _amount) external view returns (uint256); function convertBnbXToBnb(uint256 _amountInBnbX) external view returns (uint256); @@ -78,5 +78,5 @@ interface IStakeManagerV2 { event RedeemedBnbxForBnb(address indexed _account, uint256 _amountInBnbX, uint256 _amountInBnb); event SetRedemptionEnabled(bool _enabled); event SetCustodyDelay(uint256 _sweepToCustodyTimestamp); - event Swept(address indexed _custody, uint256 _amount); + 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 index 4c5d7c7..8203294 100644 --- a/test/fork-tests/StakeManagerV2Custody.t.sol +++ b/test/fork-tests/StakeManagerV2Custody.t.sol @@ -29,7 +29,7 @@ contract StakeManagerV2Custody is Test { address internal attacker; event SetCustodyDelay(uint256 _sweepToCustodyTimestamp); - event Swept(address indexed _custody, uint256 _amount); + event SweptToCustody(address indexed _asset, address indexed _custody, uint256 _amount); function setUp() public { string memory rpcUrl = vm.envString("BSC_MAINNET_RPC_URL"); @@ -58,7 +58,7 @@ contract StakeManagerV2Custody is Test { vm.expectRevert(); vm.prank(attacker); - stakeManagerV2.sweepToCustody(custody, 1 ether); + stakeManagerV2.sweepToCustody(address(0), custody); } // -------- setCustodyDelay -------- @@ -104,11 +104,11 @@ contract StakeManagerV2Custody is Test { vm.prank(admin); vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); - stakeManagerV2.sweepToCustody(custody, 1 ether); + stakeManagerV2.sweepToCustody(address(0), custody); skip(7 days); vm.prank(admin); - stakeManagerV2.sweepToCustody(custody, 1 ether); + stakeManagerV2.sweepToCustody(address(0), custody); assertEq(custody.balance, 1 ether); } @@ -118,8 +118,8 @@ contract StakeManagerV2Custody is Test { _fundContract(1 ether); vm.prank(admin); - vm.expectRevert(IStakeManagerV2.CustodyDelayNotConfigured.selector); - stakeManagerV2.sweepToCustody(custody, 1 ether); + vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); + stakeManagerV2.sweepToCustody(address(0), custody); } function test_sweepToCustody_revertsBeforeTargetElapsed() public { @@ -129,7 +129,7 @@ contract StakeManagerV2Custody is Test { vm.prank(admin); vm.expectRevert(IStakeManagerV2.CustodyDelayNotElapsed.selector); - stakeManagerV2.sweepToCustody(custody, 1 ether); + stakeManagerV2.sweepToCustody(address(0), custody); } function test_sweepToCustody_revertsOnZeroCustody() public { @@ -139,17 +139,17 @@ contract StakeManagerV2Custody is Test { vm.prank(admin); vm.expectRevert(IStakeManagerV2.ZeroAddress.selector); - stakeManagerV2.sweepToCustody(address(0), 1 ether); + stakeManagerV2.sweepToCustody(address(0), address(0)); } function test_sweepToCustody_revertsOnZeroAmount() public { _arm(1 days); skip(1 days); - _fundContract(1 ether); + // contract has 0 BNB — full-balance sweep reverts vm.prank(admin); vm.expectRevert(IStakeManagerV2.ZeroAmount.selector); - stakeManagerV2.sweepToCustody(custody, 0); + stakeManagerV2.sweepToCustody(address(0), custody); } function test_sweepToCustody_revertsOnFailedTransfer() public { @@ -161,40 +161,43 @@ contract StakeManagerV2Custody is Test { vm.prank(admin); vm.expectRevert(IStakeManagerV2.TransferFailed.selector); - stakeManagerV2.sweepToCustody(rejector, 1 ether); + stakeManagerV2.sweepToCustody(address(0), rejector); } // -------- sweepToCustody happy path -------- - function test_sweepToCustody_movesFundsAndEmits() public { + function test_sweepToCustody_sweepsFullBalanceAndEmits() public { _arm(1 days); skip(1 days); _fundContract(5 ether); - uint256 preBal = address(stakeManagerV2).balance; uint256 preCustody = custody.balance; - vm.expectEmit(true, false, false, true); - emit Swept(custody, 2 ether); + vm.expectEmit(true, true, false, true); + emit SweptToCustody(address(0), custody, 5 ether); vm.prank(admin); - stakeManagerV2.sweepToCustody(custody, 2 ether); + stakeManagerV2.sweepToCustody(address(0), custody); - assertEq(custody.balance, preCustody + 2 ether); - assertEq(address(stakeManagerV2).balance, preBal - 2 ether); + assertEq(custody.balance, preCustody + 5 ether); + assertEq(address(stakeManagerV2).balance, 0); + assertTrue(stakeManagerV2.assetCustodied()); } - function test_sweepToCustody_partialThenRemainder() public { + function test_sweepToCustody_setsAssetCustodiedAndBlocksRedeem() public { _arm(1 days); skip(1 days); - _fundContract(3 ether); + _fundContract(1 ether); + + assertFalse(stakeManagerV2.assetCustodied()); + + vm.prank(admin); + stakeManagerV2.sweepToCustody(address(0), custody); - vm.startPrank(admin); - stakeManagerV2.sweepToCustody(custody, 1 ether); - stakeManagerV2.sweepToCustody(custody, 2 ether); - vm.stopPrank(); + assertTrue(stakeManagerV2.assetCustodied()); - assertEq(custody.balance, 3 ether); + vm.expectRevert(IStakeManagerV2.AssetCustodied.selector); + stakeManagerV2.redeemBnbxForBnb(1 ether); } // -------- helpers -------- From 3d25d95a8a36e409c8c8c221441fb606bef4857e Mon Sep 17 00:00:00 2001 From: blockgroot <170620375+blockgroot@users.noreply.github.com> Date: Tue, 19 May 2026 17:53:05 +0530 Subject: [PATCH 7/7] test: add ERC20 sweepToCustody tests for full balance, zero balance, and custody delay --- test/fork-tests/StakeManagerV2Custody.t.sol | 122 +++++++++++++++++++- 1 file changed, 121 insertions(+), 1 deletion(-) diff --git a/test/fork-tests/StakeManagerV2Custody.t.sol b/test/fork-tests/StakeManagerV2Custody.t.sol index 8203294..01c0646 100644 --- a/test/fork-tests/StakeManagerV2Custody.t.sol +++ b/test/fork-tests/StakeManagerV2Custody.t.sol @@ -41,6 +41,11 @@ contract StakeManagerV2Custody is Test { 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 -------- @@ -200,6 +205,82 @@ contract StakeManagerV2Custody is Test { 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 { @@ -208,7 +289,7 @@ contract StakeManagerV2Custody is Test { } function _fundContract(uint256 amount) internal { - vm.deal(address(stakeManagerV2), address(stakeManagerV2).balance + amount); + vm.deal(address(stakeManagerV2), amount); } } @@ -217,3 +298,42 @@ contract RejectETH { 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; + } +}