Skip to content

Commit da2dba7

Browse files
tmigonetmigone
committed
feat: add fn to force withdraw legacy stake and delegation
Signed-off-by: Tomás Migone <tomas@edgeandnode.com>
1 parent a5bbbf8 commit da2dba7

File tree

4 files changed

+327
-26
lines changed

4 files changed

+327
-26
lines changed

‎packages/horizon/contracts/staking/HorizonStaking.sol‎

Lines changed: 49 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,11 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
111111
_withdraw(msg.sender);
112112
}
113113

114+
/// @inheritdoc IHorizonStakingMain
115+
function forceWithdraw(address serviceProvider) external override notPaused {
116+
_withdraw(serviceProvider);
117+
}
118+
114119
/*
115120
* PROVISIONS
116121
*/
@@ -322,33 +327,15 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
322327
address serviceProvider,
323328
address // deprecated - kept for backwards compatibility
324329
) external override notPaused returns (uint256) {
325-
// Get the delegation pool of the indexer
326-
address delegator = msg.sender;
327-
DelegationPoolInternal storage pool = _legacyDelegationPools[serviceProvider];
328-
DelegationInternal storage delegation = pool.delegators[delegator];
329-
330-
// Validation
331-
uint256 tokensToWithdraw = 0;
332-
uint256 currentEpoch = _graphEpochManager().currentEpoch();
333-
if (
334-
delegation.__DEPRECATED_tokensLockedUntil > 0 && currentEpoch >= delegation.__DEPRECATED_tokensLockedUntil
335-
) {
336-
tokensToWithdraw = delegation.__DEPRECATED_tokensLocked;
337-
}
338-
require(tokensToWithdraw > 0, HorizonStakingNothingToWithdraw());
339-
340-
// Reset lock
341-
delegation.__DEPRECATED_tokensLocked = 0;
342-
delegation.__DEPRECATED_tokensLockedUntil = 0;
343-
344-
emit StakeDelegatedWithdrawn(serviceProvider, delegator, tokensToWithdraw);
345-
346-
// -- Interactions --
347-
348-
// Return tokens to the delegator
349-
_graphToken().pushTokens(delegator, tokensToWithdraw);
330+
return _withdrawDelegatedLegacy(serviceProvider, msg.sender);
331+
}
350332

351-
return tokensToWithdraw;
333+
/// @inheritdoc IHorizonStakingMain
334+
function forceWithdrawDelegated(
335+
address serviceProvider,
336+
address delegator
337+
) external override notPaused returns (uint256) {
338+
return _withdrawDelegatedLegacy(serviceProvider, delegator);
352339
}
353340

354341
/*
@@ -1122,6 +1109,42 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
11221109
emit OperatorSet(msg.sender, _verifier, _operator, _allowed);
11231110
}
11241111

1112+
/**
1113+
* @notice Withdraw legacy undelegated tokens for a delegator.
1114+
* @dev This function handles pre-Horizon undelegations where tokens are locked
1115+
* in the legacy delegation pool.
1116+
* @param _serviceProvider The service provider address
1117+
* @param _delegator The delegator address
1118+
* @return The amount of tokens withdrawn
1119+
*/
1120+
function _withdrawDelegatedLegacy(address _serviceProvider, address _delegator) private returns (uint256) {
1121+
DelegationPoolInternal storage pool = _legacyDelegationPools[_serviceProvider];
1122+
DelegationInternal storage delegation = pool.delegators[_delegator];
1123+
1124+
// Validation
1125+
uint256 tokensToWithdraw = 0;
1126+
uint256 currentEpoch = _graphEpochManager().currentEpoch();
1127+
if (
1128+
delegation.__DEPRECATED_tokensLockedUntil > 0 && currentEpoch >= delegation.__DEPRECATED_tokensLockedUntil
1129+
) {
1130+
tokensToWithdraw = delegation.__DEPRECATED_tokensLocked;
1131+
}
1132+
require(tokensToWithdraw > 0, HorizonStakingNothingToWithdraw());
1133+
1134+
// Reset lock
1135+
delegation.__DEPRECATED_tokensLocked = 0;
1136+
delegation.__DEPRECATED_tokensLockedUntil = 0;
1137+
1138+
emit StakeDelegatedWithdrawn(_serviceProvider, _delegator, tokensToWithdraw);
1139+
1140+
// -- Interactions --
1141+
1142+
// Return tokens to the delegator
1143+
_graphToken().pushTokens(_delegator, tokensToWithdraw);
1144+
1145+
return tokensToWithdraw;
1146+
}
1147+
11251148
/**
11261149
* @notice Check if an operator is authorized for the caller on a specific verifier / data service.
11271150
* @dev Note that this function handles the special case where the verifier is the subgraph data service,

‎packages/horizon/test/unit/staking/delegation/forceWithdrawDelegated.t.sol‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity 0.8.27;
3+
4+
import "forge-std/Test.sol";
5+
6+
import { IHorizonStakingMain } from "@graphprotocol/interfaces/contracts/horizon/internal/IHorizonStakingMain.sol";
7+
import { IHorizonStakingTypes } from "@graphprotocol/interfaces/contracts/horizon/internal/IHorizonStakingTypes.sol";
8+
9+
import { HorizonStakingTest } from "../HorizonStaking.t.sol";
10+
11+
contract HorizonStakingForceWithdrawDelegatedTest is HorizonStakingTest {
12+
/*
13+
* MODIFIERS
14+
*/
15+
16+
modifier useDelegator() {
17+
resetPrank(users.delegator);
18+
_;
19+
}
20+
21+
/*
22+
* HELPERS
23+
*/
24+
25+
function _setLegacyDelegation(
26+
address _indexer,
27+
address _delegator,
28+
uint256 _shares,
29+
uint256 __DEPRECATED_tokensLocked,
30+
uint256 __DEPRECATED_tokensLockedUntil
31+
) public {
32+
// Calculate the base storage slot for the serviceProvider in the mapping
33+
bytes32 baseSlot = keccak256(abi.encode(_indexer, uint256(20)));
34+
35+
// Calculate the slot for the delegator's DelegationInternal struct
36+
bytes32 delegatorSlot = keccak256(abi.encode(_delegator, bytes32(uint256(baseSlot) + 4)));
37+
38+
// Use vm.store to set each field of the struct
39+
vm.store(address(staking), bytes32(uint256(delegatorSlot)), bytes32(_shares));
40+
vm.store(address(staking), bytes32(uint256(delegatorSlot) + 1), bytes32(__DEPRECATED_tokensLocked));
41+
vm.store(address(staking), bytes32(uint256(delegatorSlot) + 2), bytes32(__DEPRECATED_tokensLockedUntil));
42+
}
43+
44+
/*
45+
* ACTIONS
46+
*/
47+
48+
function _forceWithdrawDelegated(address _indexer, address _delegator) internal {
49+
IHorizonStakingTypes.DelegationPool memory pool = staking.getDelegationPool(
50+
_indexer,
51+
subgraphDataServiceLegacyAddress
52+
);
53+
uint256 beforeStakingBalance = token.balanceOf(address(staking));
54+
uint256 beforeDelegatorBalance = token.balanceOf(_delegator);
55+
56+
vm.expectEmit(address(staking));
57+
emit IHorizonStakingMain.StakeDelegatedWithdrawn(_indexer, _delegator, pool.tokens);
58+
staking.forceWithdrawDelegated(_indexer, _delegator);
59+
60+
uint256 afterStakingBalance = token.balanceOf(address(staking));
61+
uint256 afterDelegatorBalance = token.balanceOf(_delegator);
62+
63+
assertEq(afterStakingBalance, beforeStakingBalance - pool.tokens);
64+
assertEq(afterDelegatorBalance - pool.tokens, beforeDelegatorBalance);
65+
66+
DelegationInternal memory delegation = _getStorage_Delegation(
67+
_indexer,
68+
subgraphDataServiceLegacyAddress,
69+
_delegator,
70+
true
71+
);
72+
assertEq(delegation.shares, 0);
73+
assertEq(delegation.__DEPRECATED_tokensLocked, 0);
74+
assertEq(delegation.__DEPRECATED_tokensLockedUntil, 0);
75+
}
76+
77+
/*
78+
* TESTS
79+
*/
80+
81+
function testForceWithdrawDelegated_Tokens(uint256 tokensLocked) public useDelegator {
82+
vm.assume(tokensLocked > 0);
83+
84+
_setStorage_DelegationPool(users.indexer, tokensLocked, 0, 0);
85+
_setLegacyDelegation(users.indexer, users.delegator, 0, tokensLocked, 1);
86+
token.transfer(address(staking), tokensLocked);
87+
88+
// switch to a third party (not the delegator)
89+
resetPrank(users.operator);
90+
91+
_forceWithdrawDelegated(users.indexer, users.delegator);
92+
}
93+
94+
function testForceWithdrawDelegated_CalledByDelegator(uint256 tokensLocked) public useDelegator {
95+
vm.assume(tokensLocked > 0);
96+
97+
_setStorage_DelegationPool(users.indexer, tokensLocked, 0, 0);
98+
_setLegacyDelegation(users.indexer, users.delegator, 0, tokensLocked, 1);
99+
token.transfer(address(staking), tokensLocked);
100+
101+
// delegator can also call forceWithdrawDelegated on themselves
102+
_forceWithdrawDelegated(users.indexer, users.delegator);
103+
}
104+
105+
function testForceWithdrawDelegated_RevertWhen_NoTokens() public useDelegator {
106+
_setStorage_DelegationPool(users.indexer, 0, 0, 0);
107+
_setLegacyDelegation(users.indexer, users.delegator, 0, 0, 0);
108+
109+
// switch to a third party
110+
resetPrank(users.operator);
111+
112+
bytes memory expectedError = abi.encodeWithSignature("HorizonStakingNothingToWithdraw()");
113+
vm.expectRevert(expectedError);
114+
staking.forceWithdrawDelegated(users.indexer, users.delegator);
115+
}
116+
117+
function testForceWithdrawDelegated_RevertWhen_StillLocked(uint256 tokensLocked) public useDelegator {
118+
vm.assume(tokensLocked > 0);
119+
120+
// Set a future epoch for tokensLockedUntil
121+
uint256 futureEpoch = 1000;
122+
_setStorage_DelegationPool(users.indexer, tokensLocked, 0, 0);
123+
_setLegacyDelegation(users.indexer, users.delegator, 0, tokensLocked, futureEpoch);
124+
token.transfer(address(staking), tokensLocked);
125+
126+
// switch to a third party
127+
resetPrank(users.operator);
128+
129+
// Should revert because tokens are still locked (current epoch < futureEpoch)
130+
bytes memory expectedError = abi.encodeWithSignature("HorizonStakingNothingToWithdraw()");
131+
vm.expectRevert(expectedError);
132+
staking.forceWithdrawDelegated(users.indexer, users.delegator);
133+
}
134+
}

‎packages/horizon/test/unit/staking/stake/forceWithdraw.t.sol‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity 0.8.27;
3+
4+
import "forge-std/Test.sol";
5+
6+
import { IHorizonStakingMain } from "@graphprotocol/interfaces/contracts/horizon/internal/IHorizonStakingMain.sol";
7+
8+
import { HorizonStakingTest } from "../HorizonStaking.t.sol";
9+
10+
contract HorizonStakingForceWithdrawTest is HorizonStakingTest {
11+
/*
12+
* HELPERS
13+
*/
14+
15+
function _forceWithdraw(address _serviceProvider) internal {
16+
(, address msgSender, ) = vm.readCallers();
17+
18+
// before
19+
ServiceProviderInternal memory beforeServiceProvider = _getStorage_ServiceProviderInternal(_serviceProvider);
20+
uint256 beforeServiceProviderBalance = token.balanceOf(_serviceProvider);
21+
uint256 beforeCallerBalance = token.balanceOf(msgSender);
22+
uint256 beforeStakingBalance = token.balanceOf(address(staking));
23+
24+
// forceWithdraw
25+
vm.expectEmit(address(staking));
26+
emit IHorizonStakingMain.HorizonStakeWithdrawn(_serviceProvider, beforeServiceProvider.__DEPRECATED_tokensLocked);
27+
staking.forceWithdraw(_serviceProvider);
28+
29+
// after
30+
ServiceProviderInternal memory afterServiceProvider = _getStorage_ServiceProviderInternal(_serviceProvider);
31+
uint256 afterServiceProviderBalance = token.balanceOf(_serviceProvider);
32+
uint256 afterCallerBalance = token.balanceOf(msgSender);
33+
uint256 afterStakingBalance = token.balanceOf(address(staking));
34+
35+
// assert - tokens go to service provider, not caller
36+
assertEq(afterServiceProviderBalance - beforeServiceProviderBalance, beforeServiceProvider.__DEPRECATED_tokensLocked);
37+
assertEq(afterCallerBalance, beforeCallerBalance); // caller balance unchanged
38+
assertEq(beforeStakingBalance - afterStakingBalance, beforeServiceProvider.__DEPRECATED_tokensLocked);
39+
40+
// assert - service provider state updated
41+
assertEq(
42+
afterServiceProvider.tokensStaked,
43+
beforeServiceProvider.tokensStaked - beforeServiceProvider.__DEPRECATED_tokensLocked
44+
);
45+
assertEq(afterServiceProvider.tokensProvisioned, beforeServiceProvider.tokensProvisioned);
46+
assertEq(afterServiceProvider.__DEPRECATED_tokensAllocated, beforeServiceProvider.__DEPRECATED_tokensAllocated);
47+
assertEq(afterServiceProvider.__DEPRECATED_tokensLocked, 0);
48+
assertEq(afterServiceProvider.__DEPRECATED_tokensLockedUntil, 0);
49+
}
50+
51+
/*
52+
* TESTS
53+
*/
54+
55+
function testForceWithdraw_Tokens(uint256 tokens, uint256 tokensLocked) public useIndexer {
56+
tokens = bound(tokens, 1, MAX_STAKING_TOKENS);
57+
tokensLocked = bound(tokensLocked, 1, tokens);
58+
59+
// simulate locked tokens ready to withdraw
60+
token.transfer(address(staking), tokens);
61+
_setStorage_ServiceProvider(users.indexer, tokens, 0, tokensLocked, block.number, 0);
62+
63+
_createProvision(users.indexer, subgraphDataServiceAddress, tokens, 0, MAX_THAWING_PERIOD);
64+
65+
// switch to a different user (not the service provider)
66+
resetPrank(users.delegator);
67+
68+
_forceWithdraw(users.indexer);
69+
}
70+
71+
function testForceWithdraw_CalledByServiceProvider(uint256 tokens, uint256 tokensLocked) public useIndexer {
72+
tokens = bound(tokens, 1, MAX_STAKING_TOKENS);
73+
tokensLocked = bound(tokensLocked, 1, tokens);
74+
75+
// simulate locked tokens ready to withdraw
76+
token.transfer(address(staking), tokens);
77+
_setStorage_ServiceProvider(users.indexer, tokens, 0, tokensLocked, block.number, 0);
78+
79+
_createProvision(users.indexer, subgraphDataServiceAddress, tokens, 0, MAX_THAWING_PERIOD);
80+
81+
// before
82+
ServiceProviderInternal memory beforeServiceProvider = _getStorage_ServiceProviderInternal(users.indexer);
83+
uint256 beforeServiceProviderBalance = token.balanceOf(users.indexer);
84+
uint256 beforeStakingBalance = token.balanceOf(address(staking));
85+
86+
// service provider can also call forceWithdraw on themselves
87+
vm.expectEmit(address(staking));
88+
emit IHorizonStakingMain.HorizonStakeWithdrawn(users.indexer, beforeServiceProvider.__DEPRECATED_tokensLocked);
89+
staking.forceWithdraw(users.indexer);
90+
91+
// after
92+
ServiceProviderInternal memory afterServiceProvider = _getStorage_ServiceProviderInternal(users.indexer);
93+
uint256 afterServiceProviderBalance = token.balanceOf(users.indexer);
94+
uint256 afterStakingBalance = token.balanceOf(address(staking));
95+
96+
// assert
97+
assertEq(afterServiceProviderBalance - beforeServiceProviderBalance, beforeServiceProvider.__DEPRECATED_tokensLocked);
98+
assertEq(beforeStakingBalance - afterStakingBalance, beforeServiceProvider.__DEPRECATED_tokensLocked);
99+
assertEq(afterServiceProvider.__DEPRECATED_tokensLocked, 0);
100+
assertEq(afterServiceProvider.__DEPRECATED_tokensLockedUntil, 0);
101+
}
102+
103+
function testForceWithdraw_RevertWhen_ZeroTokens(uint256 tokens) public useIndexer {
104+
tokens = bound(tokens, 1, MAX_STAKING_TOKENS);
105+
106+
// simulate zero locked tokens
107+
token.transfer(address(staking), tokens);
108+
_setStorage_ServiceProvider(users.indexer, tokens, 0, 0, 0, 0);
109+
110+
_createProvision(users.indexer, subgraphDataServiceLegacyAddress, tokens, 0, MAX_THAWING_PERIOD);
111+
112+
// switch to a different user
113+
resetPrank(users.delegator);
114+
115+
vm.expectRevert(abi.encodeWithSelector(IHorizonStakingMain.HorizonStakingInvalidZeroTokens.selector));
116+
staking.forceWithdraw(users.indexer);
117+
}
118+
}

‎packages/interfaces/contracts/horizon/internal/IHorizonStakingMain.sol‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -952,4 +952,30 @@ interface IHorizonStakingMain {
952952
* @return Whether the operator is authorized or not
953953
*/
954954
function isAuthorized(address serviceProvider, address verifier, address operator) external view returns (bool);
955+
956+
/**
957+
* @notice Withdraw service provider legacy locked tokens.
958+
* This is a permissionless function that allows anyone to withdraw on behalf of a service provider.
959+
* It only allows withdrawing tokens that were unstaked before the Horizon upgrade.
960+
* @dev Tokens are always sent to the service provider, not the caller.
961+
*
962+
* Emits a {HorizonStakeWithdrawn} event.
963+
*
964+
* @param serviceProvider Address of service provider to withdraw funds from
965+
*/
966+
function forceWithdraw(address serviceProvider) external;
967+
968+
/**
969+
* @notice Withdraw delegator legacy undelegated tokens.
970+
* This is a permissionless function that allows anyone to withdraw on behalf of a delegator.
971+
* It only allows withdrawing tokens that were undelegated before the Horizon upgrade.
972+
* @dev Tokens are always sent to the delegator, not the caller.
973+
*
974+
* Emits a {StakeDelegatedWithdrawn} event.
975+
*
976+
* @param serviceProvider The service provider address
977+
* @param delegator The delegator address to withdraw funds for
978+
* @return The amount of tokens withdrawn
979+
*/
980+
function forceWithdrawDelegated(address serviceProvider, address delegator) external returns (uint256);
955981
}

0 commit comments

Comments
 (0)