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

nucypher · Sep 12, 2024 · c379e45 · c379e45
1 parent 48e71c8
commit c379e45
Show file tree

Hide file tree

Showing 2 changed files with 475 additions and 441 deletions.
diff --git 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,10 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
         );
         emit RitualExtended(ritualId, ritual.endTimestamp);
     }
+
+    function withdrawAllTokens(IERC20 token) external onlyRole(TREASURY_ROLE) {
+        uint256 tokenBalance = token.balanceOf(msg.sender);
+        require(tokenBalance > 0, "Insufficient balance");
+        token.safeTransferFrom(address(this), msg.sender, tokenBalance);
+    }
 }