add withdraw functions

smartcontractkit · Mar 12, 2024 · a4a7ce7 · a4a7ce7
1 parent b8b4ed2
commit a4a7ce7
Show file tree

Hide file tree

Showing 6 changed files with 321 additions and 314 deletions.
diff --git a/contracts/src/v0.8/automation/dev/interfaces/v2_3/IAutomationRegistryMaster2_3.sol b/contracts/src/v0.8/automation/dev/interfaces/v2_3/IAutomationRegistryMaster2_3.sol
diff --git a/contracts/src/v0.8/automation/dev/v2_3/AutomationRegistry2_3.sol b/contracts/src/v0.8/automation/dev/v2_3/AutomationRegistry2_3.sol
@@ -361,7 +361,6 @@ contract AutomationRegistry2_3 is AutomationRegistryBase2_3, OCR2Abstract, Chain
       nonce: s_storage.nonce,
       configCount: s_storage.configCount,
       latestConfigBlockNumber: s_storage.latestConfigBlockNumber
-      //      ownerLinkBalance: s_storage.ownerLinkBalance
     });
     s_fallbackGasPrice = onchainConfig.fallbackGasPrice;
     s_fallbackLinkPrice = onchainConfig.fallbackLinkPrice;

diff --git a/contracts/src/v0.8/automation/dev/v2_3/AutomationRegistryBase2_3.sol b/contracts/src/v0.8/automation/dev/v2_3/AutomationRegistryBase2_3.sol
@@ -164,6 +164,8 @@ abstract contract AutomationRegistryBase2_3 is ConfirmedOwner {
   error ValueNotChanged();
   error ZeroAddressNotAllowed();
   error OnlyFinanceAdmin();
+  error TransferFailed();
+  error InsufficientBalance(uint256 available, uint256 requested);
 
   enum MigrationPermission {
     NONE,
@@ -501,7 +503,6 @@ abstract contract AutomationRegistryBase2_3 is ConfirmedOwner {
   event FundsAdded(uint256 indexed id, address indexed from, uint96 amount);
   event FundsWithdrawn(uint256 indexed id, uint256 amount, address to);
   event InsufficientFundsUpkeepReport(uint256 indexed id, bytes trigger);
-  event OwnerFundsWithdrawn(uint96 amount);
   event Paused(address account);
   event PayeesUpdated(address[] transmitters, address[] payees);
   event PayeeshipTransferRequested(address indexed transmitter, address indexed from, address indexed to);
@@ -533,6 +534,7 @@ abstract contract AutomationRegistryBase2_3 is ConfirmedOwner {
   event Unpaused(address account);
   // Event to emit when a billing configuration is set
   event BillingConfigSet(IERC20 indexed token, BillingConfig config);
+  event FeesWithdrawn(address indexed recipient, address indexed assetAddress, uint256 amount);
 
   /**
    * @param link address of the LINK Token

diff --git a/contracts/src/v0.8/automation/dev/v2_3/AutomationRegistryLogicA2_3.sol b/contracts/src/v0.8/automation/dev/v2_3/AutomationRegistryLogicA2_3.sol
@@ -295,8 +295,7 @@ contract AutomationRegistryLogicA2_3 is AutomationRegistryBase2_3, Chainable {
       }
     }
     s_upkeep[id].balance = upkeep.balance - cancellationFee;
-    s_reserveLinkBalance = s_reserveLinkBalance + cancellationFee;
-    //    s_storage.ownerLinkBalance = s_storage.ownerLinkBalance + cancellationFee;
+    s_reserveLinkBalance = s_reserveLinkBalance - cancellationFee;
 
     emit UpkeepCanceled(id, uint64(height));
   }

diff --git a/contracts/src/v0.8/automation/dev/v2_3/AutomationRegistryLogicB2_3.sol b/contracts/src/v0.8/automation/dev/v2_3/AutomationRegistryLogicB2_3.sol
@@ -136,6 +136,36 @@ contract AutomationRegistryLogicB2_3 is AutomationRegistryBase2_3 {
     emit FundsWithdrawn(id, amountToWithdraw, to);
   }
 
+  /**
+   * @notice LINK available to withdraw by the finance team
+   */
+  function linkAvailableForPayment() public view returns (uint256) {
+    return i_link.balanceOf(address(this)) - s_reserveLinkBalance;
+  }
+
+  function withdrawLinkFees(address to, uint256 amount) external {
+    _onlyFinanceAdminAllowed();
+    if (to == ZERO_ADDRESS) revert InvalidRecipient();
+
+    uint256 available = linkAvailableForPayment();
+    if (amount > available) revert InsufficientBalance(available, amount);
+
+    i_link.transfer(to, amount);
+    emit FeesWithdrawn(to, address(i_link), amount);
+  }
+
+  function withdrawNonLinkFees(address assetAddress, address to, uint256 amount) external {
+    _onlyFinanceAdminAllowed();
+    if (to == ZERO_ADDRESS) revert InvalidRecipient();
+
+    bool transferStatus = IERC20(assetAddress).transfer(to, amount);
+    if (!transferStatus) {
+      revert TransferFailed();
+    }
+
+    emit FeesWithdrawn(to, assetAddress, amount);
+  }
+
   // ================================================================
   // |                       NODE MANAGEMENT                        |
   // ================================================================
@@ -193,25 +223,6 @@ contract AutomationRegistryLogicB2_3 is AutomationRegistryBase2_3 {
     emit UpkeepPrivilegeConfigSet(upkeepId, newPrivilegeConfig);
   }
 
-  /**
-   * @notice withdraws the owner's LINK balance
-   */
-  //  function withdrawOwnerFunds() external onlyOwner {
-  //    uint96 amount = s_storage.ownerLinkBalance;
-  //    s_reserveLinkBalance = s_reserveLinkBalance - amount;
-  //    s_storage.ownerLinkBalance = 0;
-  //    emit OwnerFundsWithdrawn(amount);
-  //    i_link.transfer(msg.sender, amount);
-  //  }
-
-  /**
-   * @notice allows the owner to withdraw any LINK accidentally sent to the contract
-   */
-  function recoverFunds() external onlyOwner {
-    uint256 total = i_link.balanceOf(address(this));
-    i_link.transfer(msg.sender, total - s_reserveLinkBalance);
-  }
-
   /**
    * @notice sets the payees for the transmitters
    */
@@ -444,7 +455,6 @@ contract AutomationRegistryLogicB2_3 is AutomationRegistryBase2_3 {
   {
     state = State({
       nonce: s_storage.nonce,
-      //      ownerLinkBalance: s_storage.ownerLinkBalance,
       expectedLinkBalance: s_reserveLinkBalance,
       totalPremium: s_hotVars.totalPremium,
       numUpkeeps: s_upkeepIDs.length(),