diff --git a/contracts/contract/RocketBase.sol b/contracts/contract/RocketBase.sol index fc447097d..f50e40f05 100644 --- a/contracts/contract/RocketBase.sol +++ b/contracts/contract/RocketBase.sol @@ -70,8 +70,13 @@ abstract contract RocketBase { _; } - - + /** + * @dev Throws if called by any sender that isn't a registered node or the their respective withdrawal address + */ + modifier onlyRegisteredNodeOrWithdrawalAddress(address _nodeAddress) { + require(getBool(keccak256(abi.encodePacked("node.exists", _nodeAddress))) || _nodeAddress == rocketStorage.getNodeWithdrawalAddress(_nodeAddress), "Invalid minipool owner"); + _; + } /*** Methods **********************************************************/ diff --git a/contracts/contract/node/RocketNodeStaking.sol b/contracts/contract/node/RocketNodeStaking.sol index ede700805..8209c1a36 100644 --- a/contracts/contract/node/RocketNodeStaking.sol +++ b/contracts/contract/node/RocketNodeStaking.sol @@ -252,10 +252,10 @@ contract RocketNodeStaking is RocketBase, RocketNodeStakingInterface { } /// @notice Withdraw staked RPL back to the node account - /// Only accepts calls from registered nodes + /// Only accepts calls from registered nodes or their respective withdrawal addresses /// Withdraws to withdrawal address if set, otherwise defaults to node address /// @param _amount The amount of RPL to withdraw - function withdrawRPL(uint256 _amount) override external onlyLatestContract("rocketNodeStaking", address(this)) onlyRegisteredNode(msg.sender) { + function withdrawRPL(uint256 _amount) override external onlyLatestContract("rocketNodeStaking", address(this)) onlyRegisteredNodeOrWithdrawalAddress(msg.sender) { // Load contracts RocketDAOProtocolSettingsRewardsInterface rocketDAOProtocolSettingsRewards = RocketDAOProtocolSettingsRewardsInterface(getContractAddress("rocketDAOProtocolSettingsRewards")); RocketVaultInterface rocketVault = RocketVaultInterface(getContractAddress("rocketVault"));