Merge pull request #333 from cygnusv/withdraw

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

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);
+    }
 }

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)