Beefing up the Coordinator: Restrictions to initiator, ritual duration, fee model, automatic TX reimbursements, etc

ape-config.yaml

@@ -11,10 +11,10 @@ plugins:
 dependencies:
   - name: openzeppelin
     github: OpenZeppelin/openzeppelin-contracts
-    version: 4.8.1
+    version: 4.9.1
   - name: openzeppelin-upgradeable
     github: OpenZeppelin/openzeppelin-contracts-upgradeable
-    version: 4.8.1
+    version: 4.9.1
   - name: fx-portal
     github: 0xPolygon/fx-portal
     version: 1.0.5
@@ -23,13 +23,20 @@ solidity:
   version: 0.8.20
   evm_version: paris
   import_remapping:
-    - "@openzeppelin/contracts=openzeppelin/v4.8.1"
-    - "@openzeppelin-upgradeable/contracts=openzeppelin-upgradeable/v4.8.1"
+    - "@openzeppelin/contracts=openzeppelin/v4.9.1"
+    - "@openzeppelin-upgradeable/contracts=openzeppelin-upgradeable/v4.9.1"
     - "@fx-portal/contracts=fx-portal/v1.0.5"
 
 deployments:
   polygon:
+    mainnet:
+      - contract_type: DAI
+        address: '0x8f3Cf7ad23Cd3CaDbD9735AFf958023239c6A063'
     mumbai:
+      - contract_type: DAI
+        address: '0x001B3B4d0F3714Ca98ba10F6042DaEbF0B1B7b6F'
+      - contract_type: StakeInfo
+        address: '0xC1379866Fb0c100DCBFAb7b470009C4827D47DD8'
       - fx_child: '0xCf73231F28B7331BBe3124B907840A94851f9f11'
       - verify: False
   ethereum:

contracts/contracts/coordination/Coordinator.sol

@@ -2,29 +2,33 @@
 
 pragma solidity ^0.8.0;
 
-import "@openzeppelin/contracts/access/Ownable.sol";
+import "@openzeppelin/contracts/access/AccessControlDefaultAdminRules.sol";
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import "./IFeeModel.sol";
+import "./IReimbursementPool.sol";
 import "../lib/BLS12381.sol";
 import "../../threshold/IAccessControlApplication.sol";
 
 /**
 * @title Coordinator
 * @notice Coordination layer for DKG-TDec
 */
-contract Coordinator is Ownable {
+contract Coordinator is AccessControlDefaultAdminRules {
 
     // Ritual
-    event StartRitual(uint32 indexed ritualId, address indexed initiator, address[] participants);
+    event StartRitual(uint32 indexed ritualId, address indexed authority, address[] participants);
     event StartAggregationRound(uint32 indexed ritualId);
     // TODO: Do we want the public key here? If so, we want 2 events or do we reuse this event?
-    event EndRitual(uint32 indexed ritualId, address indexed initiator, bool successful);
+    event EndRitual(uint32 indexed ritualId, bool successful);
 
     // Node
     event TranscriptPosted(uint32 indexed ritualId, address indexed node, bytes32 transcriptDigest);
     event AggregationPosted(uint32 indexed ritualId, address indexed node, bytes32 aggregatedTranscriptDigest);
 
     // Admin
     event TimeoutChanged(uint32 oldTimeout, uint32 newTimeout);
-    event MaxDkgSizeChanged(uint32 oldSize, uint32 newSize);
+    event MaxDkgSizeChanged(uint16 oldSize, uint16 newSize);
 
     enum RitualState {
         NON_INITIATED,
@@ -45,26 +49,46 @@ contract Coordinator is Ownable {
     // TODO: Optimize layout
     struct Ritual {
         address initiator;
-        uint32 dkgSize;
         uint32 initTimestamp;
-        uint32 totalTranscripts;
-        uint32 totalAggregations;
-        BLS12381.G1Point publicKey;
+        uint32 endTimestamp;
+        uint16 totalTranscripts;
+        uint16 totalAggregations;
+        address authority;
+        uint16 dkgSize;
         bool aggregationMismatch;
+        BLS12381.G1Point publicKey;
         bytes aggregatedTranscript;
         Participant[] participant;
     }
 
-    Ritual[] public rituals;
+    using SafeERC20 for IERC20;
+
+    bytes32 public constant INITIATOR_ROLE = keccak256("INITIATOR_ROLE");
 
     IAccessControlApplication public immutable application;
-    uint32 public timeout;
-    uint32 public maxDkgSize;
 
-    constructor(IAccessControlApplication app, uint32 _timeout, uint32 _maxDkgSize) {
-        application = app;
+    Ritual[] public rituals;
+    uint32 public timeout;
+    uint16 public maxDkgSize;
+    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,
+        uint32 _timeout,
+        uint16 _maxDkgSize,
+        address _admin,
+        IFeeModel _feeModel
+    ) AccessControlDefaultAdminRules(0, _admin)
+    {
+        require(address(_feeModel.stakes()) == address(_stakes), "Invalid stakes for fee model");
+        application = _stakes;
         timeout = _timeout;
         maxDkgSize = _maxDkgSize;
+        feeModel = IFeeModel(_feeModel);
     }
 
     function getRitualState(uint256 ritualId) external view returns (RitualState){
@@ -95,17 +119,31 @@ contract Coordinator is Ownable {
         }
     }
 
+    function makeInitiationPublic() external onlyRole(DEFAULT_ADMIN_ROLE) {
+        isInitiationPublic = true;
+        _setRoleAdmin(INITIATOR_ROLE, bytes32(0));
+    }
 
-    function setTimeout(uint32 newTimeout) external onlyOwner {
+    function setTimeout(uint32 newTimeout) external onlyRole(DEFAULT_ADMIN_ROLE) {
         emit TimeoutChanged(timeout, newTimeout);
         timeout = newTimeout;
     }
 
-    function setMaxDkgSize(uint32 newSize) external onlyOwner {
+    function setMaxDkgSize(uint16 newSize) external onlyRole(DEFAULT_ADMIN_ROLE) {
         emit MaxDkgSizeChanged(maxDkgSize, newSize);
         maxDkgSize = newSize;
     }
 
+    function setReimbursementPool(IReimbursementPool pool) external onlyRole(DEFAULT_ADMIN_ROLE) {
+        require(
+            address(pool) == address(0) || 
+            pool.isAuthorized(address(this)),
+            "Invalid ReimbursementPool"
+        );
+        reimbursementPool = pool;
+        // TODO: Events
+    }
+
     function numberOfRituals() external view returns(uint256) {
         return rituals.length;
     }
@@ -115,16 +153,27 @@ contract Coordinator is Ownable {
         return ritual.participant;
     }
 
-    function initiateRitual(address[] calldata providers) external returns (uint32) {
+    function initiateRitual(
+        address[] calldata providers,
+        address authority,
+        uint32 duration
+    ) external returns (uint32) {
+        require(
+            isInitiationPublic || hasRole(INITIATOR_ROLE, msg.sender),
+            "Sender can't initiate ritual"
+        );
         // TODO: Validate service fees, expiration dates, threshold
         uint256 length = providers.length;
         require(2 <= length && length <= maxDkgSize, "Invalid number of nodes");
+        require(duration > 0, "Invalid ritual duration");  // TODO: We probably want to restrict it more
 
         uint32 id = uint32(rituals.length);
         Ritual storage ritual = rituals.push();
-        ritual.initiator = msg.sender;  // TODO: Consider sponsor model
-        ritual.dkgSize = uint32(length);
+        ritual.initiator = msg.sender;
+        ritual.authority = authority;
+        ritual.dkgSize = uint16(length);
         ritual.initTimestamp = uint32(block.timestamp);
+        ritual.endTimestamp = ritual.initTimestamp + duration;
 
         address previous = address(0);
         for(uint256 i=0; i < length; i++){
@@ -140,9 +189,11 @@ contract Coordinator is Ownable {
             newParticipant.provider = current;
             previous = current;
         }
+
+        processRitualPayment(id, providers, duration);
 
         // TODO: Include cohort fingerprint in StartRitual event?
-        emit StartRitual(id, msg.sender, providers);
+        emit StartRitual(id, ritual.authority, providers);
         return id;
     }
 
@@ -151,6 +202,8 @@ contract Coordinator is Ownable {
     }
 
     function postTranscript(uint32 ritualId, bytes calldata transcript) external {
+        uint256 initialGasLeft = gasleft();
+
         Ritual storage ritual = rituals[ritualId];
         require(
             getRitualState(ritual) == RitualState.AWAITING_TRANSCRIPTS,
@@ -181,6 +234,7 @@ contract Coordinator is Ownable {
         if (ritual.totalTranscripts == ritual.dkgSize){
             emit StartAggregationRound(ritualId);
         }
+        processReimbursement(initialGasLeft);
     }
 
     function postAggregation(
@@ -189,6 +243,8 @@ contract Coordinator is Ownable {
         BLS12381.G1Point calldata publicKey,
         bytes calldata decryptionRequestStaticKey
     ) external {
+        uint256 initialGasLeft = gasleft();
+
         Ritual storage ritual = rituals[ritualId];
         require(
             getRitualState(ritual) == RitualState.AWAITING_AGGREGATIONS,
@@ -209,13 +265,18 @@ contract Coordinator is Ownable {
 
         require(
             participant.decryptionRequestStaticKey.length == 0,
-            "Node already provided request encrypting key"
+            "Node already provided decryption request static key"
+        );
+
+        require(
+            decryptionRequestStaticKey.length == 42,
+            "Invalid length for decryption request static key"
         );
 
         // nodes commit to their aggregation result
         bytes32 aggregatedTranscriptDigest = keccak256(aggregatedTranscript);
         participant.aggregated = true;
-        participant.decryptionRequestStaticKey = decryptionRequestStaticKey;  // TODO validation?
+        participant.decryptionRequestStaticKey = decryptionRequestStaticKey;
         emit AggregationPosted(ritualId, provider, aggregatedTranscriptDigest);
 
         if (ritual.aggregatedTranscript.length == 0) {
@@ -228,22 +289,23 @@ contract Coordinator is Ownable {
             ritual.aggregationMismatch = true;
             emit EndRitual({
                 ritualId: ritualId,
-                initiator: ritual.initiator,
                 successful: false
             });
-            // TODO: Consider freeing ritual storage
-            return;
         }
 
-        ritual.totalAggregations++;
-        if (ritual.totalAggregations == ritual.dkgSize){
-            emit EndRitual({
-                ritualId: ritualId,
-                initiator: ritual.initiator,
-                successful: true
-            });
-            // TODO: Consider including public key in event
+        if(!ritual.aggregationMismatch){
+            ritual.totalAggregations++;
+            if (ritual.totalAggregations == ritual.dkgSize){
+                processPendingFee(ritualId);
+                emit EndRitual({
+                    ritualId: ritualId,
+                    successful: true
+                });
+                // TODO: Consider including public key in event
+            }
         }
+
+        processReimbursement(initialGasLeft);
     }
 
     function getParticipantFromProvider(
@@ -267,4 +329,55 @@ contract Coordinator is Ownable {
     ) external view returns (Participant memory) {
         return getParticipantFromProvider(rituals[ritualID], provider);
     }
+
+    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.safeTransferFrom(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
+            uint256 gasUsed = initialGasLeft - gasleft();
+            try reimbursementPool.refund(gasUsed, msg.sender) {
+                return;
+            } catch {
+                return;
+            }
+        }
+    }
 }

contracts/contracts/coordination/FlateRateFeeModel.sol

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+pragma solidity ^0.8.0;
+
+import "./IFeeModel.sol";
+import "../../threshold/IAccessControlApplication.sol";
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
+/**
+* @title FlatRateFeeModel
+* @notice FlateRateFeeModel
+*/
+contract FlatRateFeeModel is IFeeModel {
+
+    IERC20 public immutable currency;
+    uint256 public immutable feeRatePerSecond;
+    IAccessControlApplication public immutable stakes;
+
+    constructor(IERC20 _currency, uint256 _feeRatePerSecond, address _stakes){
+        currency = _currency;
+        feeRatePerSecond = _feeRatePerSecond;
+        stakes = IAccessControlApplication(_stakes);
+    }
+
+    function getRitualInitiationCost(
+        address[] calldata providers,
+        uint32 duration
+    ) external view returns(uint256){
+        uint256 size = providers.length;
+        require(duration > 0, "Invalid ritual duration");
+        require(size > 0, "Invalid ritual size");
+        return feeRatePerSecond * size * duration;
+    }
+}

contracts/contracts/coordination/IFeeModel.sol

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "../../threshold/IAccessControlApplication.sol";
+
+/**
+* @title IFeeModel
+* @notice IFeeModel
+*/
+interface IFeeModel {
+    function currency() external view returns(IERC20);
+    function stakes() external view returns(IAccessControlApplication);
+    function getRitualInitiationCost(address[] calldata providers, uint32 duration) external view returns(uint256);
+}