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

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

Hide file tree

Showing 2 changed files with 32 additions and 0 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);
+    }
 }
diff --git a/tests/test_coordinator.py b/tests/test_coordinator.py
@@ -650,3 +650,25 @@ 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):
+
+    # Let's send some tokens to Coordinator by mistake
+    erc20.transfer(coordinator.address, 42, sender=initiator)
+    assert erc20.balanceOf(coordinator.address) == 42
+
+    # Only treasury can withdraw
+    with ape.reverts():
+        coordinator.withdrawAllTokens(erc20.address, sender=initiator)
+
+    # Treasury withdraws all tokens
+    treasury_balance_before = erc20.balanceOf(treasury.address)
+    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)