From ab7478e1f490fa92e57ed98069737cde64ff0bd5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Tue, 20 Jun 2023 23:38:33 +0200
Subject: [PATCH] Treat fees as pending until ritual is final

Either ritual is successful and fees can be used by the protocol, or ritual is failed and fees are refunded to initiator (consider partial refund)
---
 .../contracts/coordination/Coordinator.sol    | 42 +++++++++++++++++--
 1 file changed, 39 insertions(+), 3 deletions(-)

diff --git a/contracts/contracts/coordination/Coordinator.sol b/contracts/contracts/coordination/Coordinator.sol
index b7b423db..efbe71c2 100644
--- a/contracts/contracts/coordination/Coordinator.sol
+++ b/contracts/contracts/coordination/Coordinator.sol
@@ -70,6 +70,8 @@ contract Coordinator is AccessControlDefaultAdminRules {
     bool public isInitiationPublic;
     IFeeModel feeModel;  // TODO: Consider making feeModel specific to each ritual
     IReimbursementPool reimbursementPool;
+    uint256 public totalPendingFees;
+    mapping(uint256 => uint256) public pendingFees;
 
     constructor(
         IAccessControlApplication _stakes,
@@ -162,9 +164,10 @@ contract Coordinator is AccessControlDefaultAdminRules {
         require(2 <= length && length <= maxDkgSize, "Invalid number of nodes");
         require(duration > 0, "Invalid ritual duration");  // TODO: We probably want to restrict it more
 
-        processRitualPayment(providers, duration);
-
         uint32 id = uint32(rituals.length);
+
+        processRitualPayment(id, providers, duration);
+
         Ritual storage ritual = rituals.push();
         ritual.initiator = msg.sender;
         ritual.authority = authority;
@@ -293,6 +296,8 @@ contract Coordinator is AccessControlDefaultAdminRules {
                 // TODO: Consider including public key in event
             }
         }
+
+        processPendingFee(ritualId);
         processReimbursement(initialGasLeft);
     }
 
@@ -318,14 +323,45 @@ contract Coordinator is AccessControlDefaultAdminRules {
         return getParticipantFromProvider(rituals[ritualID], provider);
     }
 
-    function processRitualPayment(address[] calldata providers, uint32 duration) internal {
+    function processRitualPayment(uint256 ritualID, address[] calldata providers, uint32 duration) internal {
         uint256 ritualCost = feeModel.getRitualInitiationCost(providers, duration);
         if (ritualCost > 0){
+            totalPendingFees += ritualCost;
+            assert(pendingFees[ritualID] == 0);  // TODO: This is an invariant, not sure if actually needed
+            pendingFees[ritualID] += ritualCost;
             IERC20 currency = IERC20(feeModel.currency());
             currency.transferFrom(msg.sender, address(this), ritualCost);
             // TODO: Define methods to manage these funds
         }
     }
+
+    function processPendingFee(uint256 ritualID) public {
+        Ritual storage ritual = rituals[ritualID];
+        RitualState state = getRitualState(ritual);
+        require(
+            state == RitualState.TIMEOUT ||
+            state == RitualState.INVALID ||
+            state == RitualState.FINALIZED,
+            "Ritual is not ended"
+        );
+        uint256 pending = pendingFees[ritualID];
+        require(pending > 0, "No pending fees for this ritual");
+
+        // Finalize fees for this ritual
+        totalPendingFees -= pending;
+        delete pendingFees[ritualID];
+        // Transfer fees back to initiator if failed
+        if(state == RitualState.TIMEOUT || state == RitualState.INVALID){
+            // Amount to refund depends on how much work nodes did for the ritual.
+            // TODO: Validate if this is enough to remove griefing attacks
+            uint256 executedTransactions = ritual.totalTranscripts + ritual.totalAggregations;
+            uint256 expectedTransactions = 2 * ritual.dkgSize;
+            uint256 consumedFee = pending * executedTransactions / expectedTransactions;
+            uint256 refundableFee = pending - consumedFee;
+            IERC20 currency = IERC20(feeModel.currency());
+            currency.transferFrom(address(this), ritual.initiator, refundableFee);
+        }
+    }
     
     function processReimbursement(uint256 initialGasLeft) internal {
         if(address(reimbursementPool) != address(0)){ // TODO: Consider defining a method