Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add function to withdraw all tokens from Coordinator for a token address #333

Merged
merged 1 commit into from
Sep 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions contracts/contracts/coordination/Coordinator.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand All @@ -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);
Expand Down Expand Up @@ -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);
}
}
24 changes: 24 additions & 0 deletions tests/test_coordinator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔥


# 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)
Loading