From a3317af9b5e1c7b29ecd1136a3bfc8ca13942b5b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Thu, 12 Sep 2024 11:54:30 +0200
Subject: [PATCH] Add function to withdraw all tokens from Coordinator for a
 token address

---
 .../contracts/coordination/Coordinator.sol      |  9 +++++++++
 tests/test_coordinator.py                       | 17 +++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/contracts/contracts/coordination/Coordinator.sol b/contracts/contracts/coordination/Coordinator.sol
index 2672fede..641d07d1 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/utils/SafeERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/IERC20.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,9 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
         );
         emit RitualExtended(ritualId, ritual.endTimestamp);
     }
+
+    function withdrawAllTokens(IERC20 token) external onlyRole(TREASURY_ROLE) {
+        uint256 tokenBalance = token.balanceOf(msg.sender);
+        token.safeTransferFrom(address(this), msg.sender, tokenBalance);
+    }
 }
diff --git a/tests/test_coordinator.py b/tests/test_coordinator.py
index 0d896319..5ed626ff 100644
--- a/tests/test_coordinator.py
+++ b/tests/test_coordinator.py
@@ -650,3 +650,20 @@ 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):
+    
+    erc20.transfer(coordinator.address, 42, sender=initiator)
+
+    assert erc20.balanceOf(coordinator.address) == 42
+
+    treasury_balance_before = erc20.balanceOf(treasury.address)
+
+    with ape.reverts("Can't withdraw pending fees"):
+        coordinator.withdrawAllTokens(erc20.address, sender=initiator)
+
+    coordinator.withdrawAllTokens(erc20.address, sender=treasury)
+
+    assert erc20.balanceOf(coordinator.address) == 0
+    assert erc20.balanceOf(treasury.address) == 42 + treasury_balance_before
\ No newline at end of file