vacp2p/staking-reward-streamer
1
0
Fork 0
mirror of https://github.com/vacp2p/staking-reward-streamer.git synced 2026-01-08 20:48:00 -05:00
Code Issues Packages Projects Releases Wiki Activity

refactor(StakeVault): Restore selective trust functionality for core operations

Browse Source 
Co-authored-by: 0x-r4bbit <445106+0x-r4bbit@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
2025-08-01 14:46:18 +00:00
parent 02900cc599
commit d9abb5a490
2 changed files with 79 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
src/StakeVault.sol
View File

@@ -29,8 +29,10 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
error StakeVault__UnstakingFailed();
/// @notice Emitted when not allowed to exit the system
error StakeVault__NotAllowedToExit();
/// @notice Emitted when not allowed to leave the system
error StakeVault__NotAllowedToLeave();
/// @notice Emitted when the configured stake manager is not trusted
error StakeVault__StakeManagerImplementationNotTrusted();
/// @notice Emitted when migration failed
error StakeVault__MigrationFailed();
/// @notice Emitted when the caller is not the owner of the vault
@@ -46,6 +48,8 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
IERC20 public immutable STAKING_TOKEN;
/// @notice Stake manager proxy contract
IStakeManagerProxy public stakeManager;
/// @notice Address of the trusted stake manager implementation
address public stakeManagerImplementationAddress;
/// @notice Timestamp until the funds are locked
uint256 public lockUntil;
@@ -57,6 +61,13 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
_;
}
modifier onlyTrustedStakeManager() {
if (!_stakeManagerImplementationTrusted()) {
revert StakeVault__StakeManagerImplementationNotTrusted();
}
_;
}
/*//////////////////////////////////////////////////////////////////////////
@@ -80,6 +91,7 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
/**
* @notice Initializes the contract with the owner and the stake manager.
* @dev Ensures that the stake manager implementation is trusted.
* @dev Initialization is done on proxy clones.
* @param _owner The address of the owner.
* @param _stakeManager The address of the StakeManager contract.
@@ -87,6 +99,7 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
function initialize(address _owner, address _stakeManager) public initializer {
_transferOwnership(_owner);
stakeManager = IStakeManagerProxy(_stakeManager);
stakeManagerImplementationAddress = stakeManager.implementation();
}
@@ -99,14 +112,24 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
stakeManager.registerVault();
}
/**
* @notice Trust a new stake manager implementation.
* @dev This function is only callable by the owner.
* @param stakeManagerAddress The address of the new stake manager implementation.
*/
function trustStakeManager(address stakeManagerAddress) external onlyOwner {
stakeManagerImplementationAddress = stakeManagerAddress;
}
/**
* @notice Stake tokens for a specified time.
* @dev This function is only callable by the owner.
* @dev Can only be called if the stake manager is trusted.
* @dev Reverts if the staking token transfer fails.
* @param _amount The amount of tokens to stake.
* @param _seconds The time period to stake for.
*/
function stake(uint256 _amount, uint256 _seconds) external onlyOwner {
function stake(uint256 _amount, uint256 _seconds) external onlyOwner onlyTrustedStakeManager {
_stake(_amount, _seconds, msg.sender);
}
@@ -114,31 +137,34 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
* @notice Stake tokens from a specified address for a specified time.
* @dev Overloads the `stake` function to allow staking from a specified address.
* @dev This function is only callable by the owner.
* @dev Can only be called if the stake manager is trusted.
* @dev Reverts if the staking token transfer fails.
* @param _amount The amount of tokens to stake.
* @param _seconds The time period to stake for.
* @param _from The address from which tokens will be transferred.
*/
function stake(uint256 _amount, uint256 _seconds, address _from) external onlyOwner {
function stake(uint256 _amount, uint256 _seconds, address _from) external onlyOwner onlyTrustedStakeManager {
_stake(_amount, _seconds, _from);
}
/**
* @notice Lock the staked amount for a specified time.
* @dev This function is only callable by the owner.
* @dev Can only be called if the stake manager is trusted.
* @param _seconds The time period to lock the staked amount for.
*/
function lock(uint256 _seconds) external onlyOwner {
function lock(uint256 _seconds) external onlyOwner onlyTrustedStakeManager {
stakeManager.lock(_seconds);
}
/**
* @notice Unstake a specified amount of tokens and send to the owner.
* @dev This function is only callable by the owner.
* @dev Can only be called if the stake manager is trusted.
* @dev Reverts if the staking token transfer fails.
* @param _amount The amount of tokens to unstake.
*/
function unstake(uint256 _amount) external onlyOwner {
function unstake(uint256 _amount) external onlyOwner onlyTrustedStakeManager {
_unstake(_amount, msg.sender);
}
@@ -146,6 +172,7 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
* @notice Unstake a specified amount of tokens and send to a destination address.
* @dev Overloads the `unstake` function to allow unstaking to a specified address.
* @dev This function is only callable by the owner.
* @dev Can only be called if the stake manager is trusted.
* @dev Reverts if the staking token transfer fails.
* @param _amount The amount of tokens to unstake.
* @param _destination The address to receive the unstaked tokens.
@@ -157,6 +184,7 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
external
onlyOwner
validDestination(_destination)
onlyTrustedStakeManager
{
_unstake(_amount, _destination);
}
@@ -168,29 +196,31 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
* @param _destination The address to receive the funds.
*/
function leave(address _destination) external onlyOwner validDestination(_destination) {
// Check if funds are locked before attempting to leave
if (lockUntil > block.timestamp) {
revert StakeVault__FundsLocked();
}
// Try to properly notify the stake manager that we're leaving the system.
// We use try/catch in case the upgrade was malicious and `leave()`
// either doesn't exist on the new stake manager or reverts for some reason.
// If it was a benign upgrade, it will cause the stake manager to properly update
// its internal accounting before we move the funds out.
try stakeManager.leave() {
if (lockUntil <= block.timestamp) {
STAKING_TOKEN.transfer(_destination, STAKING_TOKEN.balanceOf(address(this)));
}
STAKING_TOKEN.transfer(_destination, STAKING_TOKEN.balanceOf(address(this)));
} catch {
if (lockUntil <= block.timestamp) {
STAKING_TOKEN.transfer(_destination, STAKING_TOKEN.balanceOf(address(this)));
}
STAKING_TOKEN.transfer(_destination, STAKING_TOKEN.balanceOf(address(this)));
}
}
/**
* @notice Migrate all funds to a new vault.
* @dev This function is only callable by the owner.
* @dev Can only be called if the stake manager is trusted.
* @dev Reverts when the stake manager reverts or the funds can't be transferred.
* @param migrateTo The address of the new vault.
*/
function migrateToVault(address migrateTo) external onlyOwner {
function migrateToVault(address migrateTo) external onlyOwner onlyTrustedStakeManager {
stakeManager.migrateToVault(migrateTo);
bool success = STAKING_TOKEN.transfer(migrateTo, STAKING_TOKEN.balanceOf(address(this)));
if (!success) {
@@ -340,6 +370,13 @@ contract StakeVault is IStakeVault, Initializable, OwnableUpgradeable {
}
}
/**
* @notice Checks if the stake manager implementation is trusted.
* @return True if the stake manager implementation is trusted.
*/
function _stakeManagerImplementationTrusted() internal view virtual returns (bool) {
return stakeManagerImplementationAddress == stakeManager.implementation();
}
/*//////////////////////////////////////////////////////////////////////////

36
test/RewardsStreamerMP.t.sol
View File

@@ -2365,21 +2365,43 @@ contract LeaveTest is StakeManagerTest {
}
function test_UpgradeStakeManager() public {
// first, upgrade to new stake manager
_upgradeStakeManager();
// after upgrade, vault functions should work normally (trust functionality removed)
// first, stake with current stake manager before upgrade
_stake(alice, 100e18, 0);
// verify vault is working before upgrade
StakeVault vault = StakeVault(vaults[alice]);
assertEq(vault.amountStaked(), 100e18, "Should have 100e18 staked before upgrade");
// upgrade to new stake manager
_upgradeStakeManager();
// after upgrade, vault functions should require trusting the new stake manager
// trying to stake/unstake without trusting should revert
vm.expectRevert(StakeVault.StakeVault__StakeManagerImplementationNotTrusted.selector);
_stake(alice, 50e18, 0);
// trust the new stake manager implementation
vm.prank(alice);
vault.trustStakeManager(streamer.implementation());
// now vault functions should work with the trusted new stake manager
_stake(alice, 50e18, 0);
assertEq(vault.amountStaked(), 150e18, "Should have 150e18 staked after upgrade");
// lock the funds
vm.prank(alice);
vault.lock(365 days);
// functions should work normally without trust checks
// unstaking should fail while funds are locked (StakeManager enforces this)
// The exact error depends on StakeManager implementation, but it should revert
vm.expectRevert(); // Generic revert expectation since we don't know the exact error
_unstake(alice, 100e18);
// leave function should work based on lockUntil only
// (we would need to wait for lock to expire or manipulate time for this to work)
// leave function should work based on lockUntil only (no trust required)
// but should still fail due to time lock
vm.prank(alice);
vm.expectRevert(StakeVault.StakeVault__FundsLocked.selector);
vault.leave(alice);
}
}