diff --git a/contracts/contracts/coordination/Coordinator.sol b/contracts/contracts/coordination/Coordinator.sol index 2672fede..2372e36c 100644 --- a/contracts/contracts/coordination/Coordinator.sol +++ b/contracts/contracts/coordination/Coordinator.sol @@ -2,6 +2,8 @@ pragma solidity ^0.8.0; +import "@openzeppelin/contracts/token/ERC20/IERC20.sol"; +import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; import "@openzeppelin-upgradeable/contracts/access/extensions/AccessControlDefaultAdminRulesUpgradeable.sol"; import "@openzeppelin-upgradeable/contracts/proxy/utils/Initializable.sol"; import "./IFeeModel.sol"; @@ -15,6 +17,8 @@ import "./IEncryptionAuthorizer.sol"; * @notice Coordination layer for Threshold Access Control (TACo 🌮) */ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable { + using SafeERC20 for IERC20; + // DKG Protocol event StartRitual(uint32 indexed ritualId, address indexed authority, address[] participants); event StartAggregationRound(uint32 indexed ritualId); @@ -642,4 +646,10 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable ); emit RitualExtended(ritualId, ritual.endTimestamp); } + + function withdrawAllTokens(IERC20 token) external onlyRole(TREASURY_ROLE) { + uint256 tokenBalance = token.balanceOf(address(this)); + require(tokenBalance > 0, "Insufficient balance"); + token.safeTransfer(msg.sender, tokenBalance); + } } diff --git a/tests/test_coordinator.py b/tests/test_coordinator.py index 0d896319..58d38bf3 100644 --- a/tests/test_coordinator.py +++ b/tests/test_coordinator.py @@ -650,3 +650,27 @@ def test_upgrade( assert ritual_struct["publicKey"] == (b"\x00" * 32, b"\x00" * 16) # publicKey assert not ritual_struct["aggregatedTranscript"] # aggregatedTranscript assert ritual_struct["feeModel"] == fee_model.address # feeModel + + +def test_withdraw_tokens(coordinator, initiator, erc20, treasury, deployer): + + # Let's send some tokens to Coordinator by mistake + erc20.transfer(coordinator.address, 42, sender=initiator) + assert erc20.balanceOf(coordinator.address) == 42 + + # Only accounts with TREASURY_ROLE can withdraw + with ape.reverts(): + coordinator.withdrawAllTokens(erc20.address, sender=treasury) + + # Treasury is granted proper role and withdraws all tokens + treasury_balance_before = erc20.balanceOf(treasury.address) + + coordinator.grantRole(coordinator.TREASURY_ROLE(), treasury, sender=deployer) + coordinator.withdrawAllTokens(erc20.address, sender=treasury) + + assert erc20.balanceOf(coordinator.address) == 0 + assert erc20.balanceOf(treasury.address) == 42 + treasury_balance_before + + # Can't withdraw when there's no tokens + with ape.reverts("Insufficient balance"): + coordinator.withdrawAllTokens(erc20.address, sender=treasury)