Merge pull request #422 from cygnusv/handover-gas

Simplify management of participants public key history

Author: cygnusv

contracts/contracts/coordination/Coordinator.sol

@@ -129,7 +129,6 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
         BLS12381.G2Point publicKey;
     }
 
-    bytes32 public constant TREASURY_ROLE = keccak256("TREASURY_ROLE");
     bytes32 public constant FEE_MODEL_MANAGER_ROLE = keccak256("FEE_MODEL_MANAGER_ROLE");
     bytes32 public constant HANDOVER_SUPERVISOR_ROLE = keccak256("HANDOVER_SUPERVISOR_ROLE");
 
@@ -140,7 +139,7 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
     uint32 public immutable handoverTimeout;
 
     Ritual[] private ritualsStub; // former rituals, "internal" for testing only
-    uint32 public dkgTimeoutStub;
+    uint32 private timeoutStub; // former (dkg) timeout
     uint16 public maxDkgSize;
     bool private stub1; // former isInitiationPublic
 
@@ -287,6 +286,7 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
     function setProviderPublicKey(BLS12381.G2Point calldata publicKey) external {
         uint32 lastRitualId = uint32(numberOfRituals);
         address stakingProvider = application.operatorToStakingProvider(msg.sender);
+        require(!isProviderKeySet(msg.sender), "Provider has already set public key");
         require(stakingProvider != address(0), "Operator has no bond with staking provider");
 
         ParticipantKey memory newRecord = ParticipantKey(lastRitualId, publicKey);
@@ -300,25 +300,18 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
 
     function getProviderPublicKey(
         address provider,
-        uint256 ritualId
+        uint256
     ) external view returns (BLS12381.G2Point memory) {
         ParticipantKey[] storage participantHistory = participantKeysHistory[provider];
-
-        for (uint256 i = participantHistory.length; i > 0; i--) {
-            if (participantHistory[i - 1].lastRitualId <= ritualId) {
-                return participantHistory[i - 1].publicKey;
-            }
-        }
-
-        revert("No keys found prior to the provided ritual");
+        return participantHistory[participantHistory.length - 1].publicKey;
     }
 
-    /**
-     * @dev This method is deprecated. Use `isProviderKeySet` instead.
-     */
-    function isProviderPublicKeySet(address) external view returns (bool) {
-        revert("Deprecated method. Upgrade your node to latest version");
-    }
+    // /**
+    //  * @dev This method is deprecated. Use `isProviderKeySet` instead.
+    //  */
+    // function isProviderPublicKeySet(address) external view returns (bool) {
+    //     revert("Deprecated method. Upgrade your node to latest version");
+    // }
 
     function isProviderKeySet(address provider) public view returns (bool) {
         ParticipantKey[] storage participantHistory = participantKeysHistory[provider];
@@ -409,10 +402,6 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
         return id;
     }
 
-    function cohortFingerprint(address[] calldata nodes) public pure returns (bytes32) {
-        return keccak256(abi.encode(nodes));
-    }
-
     function expectedTranscriptSize(
         uint16 dkgSize,
         uint16 threshold
@@ -591,6 +580,7 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
         bytes calldata transcript,
         bytes calldata decryptionRequestStaticKey
     ) external {
+        uint256 initialGasLeft = gasleft();
         require(isRitualActive(ritualId), "Ritual is not active");
         require(transcript.length > 0, "Parameters can't be empty");
         require(
@@ -609,9 +599,11 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
         handover.transcript = transcript;
         handover.decryptionRequestStaticKey = decryptionRequestStaticKey;
         emit HandoverTranscriptPosted(ritualId, departingParticipant, provider);
+        processReimbursement(initialGasLeft);
     }
 
     function postBlindedShare(uint32 ritualId, bytes calldata blindedShare) external {
+        uint256 initialGasLeft = gasleft();
         require(isRitualActive(ritualId), "Ritual is not active");
 
         address provider = application.operatorToStakingProvider(msg.sender);
@@ -624,6 +616,7 @@ contract Coordinator is Initializable, AccessControlDefaultAdminRulesUpgradeable
 
         handover.blindedShare = blindedShare;
         emit BlindedSharePosted(ritualId, provider);
+        processReimbursement(initialGasLeft);
     }
 
     function cancelHandover(

scripts/simulate_coordinator_upgrade.py

@@ -1,7 +1,7 @@
 # Usage:
 #  > ape run --interactive simulate_coordinator_upgrade --network polygon:mainnet-fork:foundry
 
-from ape import accounts, chain, project
+from ape import accounts, chain, networks, project
 
 from deployment.constants import CONSTRUCTOR_PARAMS_DIR
 from deployment.params import Deployer

tests/test_coordinator.py

@@ -44,10 +44,10 @@ def deployer(accounts):
 
 
 @pytest.fixture(scope="module")
-def treasury(accounts):
-    treasury_index = MAX_DKG_SIZE + 3
-    assert len(accounts) >= treasury_index
-    return accounts[treasury_index]
+def fee_manager(accounts):
+    fee_manager_index = MAX_DKG_SIZE + 3
+    assert len(accounts) >= fee_manager_index
+    return accounts[fee_manager_index]
 
 
 @pytest.fixture()
@@ -87,13 +87,12 @@ def coordinator(project, deployer, application, oz_dependency):
 
 
 @pytest.fixture()
-def fee_model(project, deployer, coordinator, erc20, treasury):
+def fee_model(project, deployer, coordinator, erc20, fee_manager):
     contract = project.FlatRateFeeModel.deploy(
         coordinator.address, erc20.address, FEE_RATE, sender=deployer
     )
-    coordinator.grantRole(coordinator.FEE_MODEL_MANAGER_ROLE(), treasury, sender=deployer)
-    coordinator.approveFeeModel(contract.address, sender=treasury)
-    coordinator.grantRole(coordinator.TREASURY_ROLE(), treasury, sender=deployer)
+    coordinator.grantRole(coordinator.FEE_MODEL_MANAGER_ROLE(), fee_manager, sender=deployer)
+    coordinator.approveFeeModel(contract.address, sender=fee_manager)
     return contract
 
 
@@ -178,7 +177,7 @@ def initiate_ritual(coordinator, fee_model, erc20, authority, nodes, allow_logic
 
 
 def test_initiate_ritual(
-    coordinator, nodes, initiator, erc20, fee_model, deployer, treasury, global_allow_list
+    coordinator, nodes, initiator, erc20, fee_model, deployer, fee_manager, global_allow_list
 ):
     authority, tx = initiate_ritual(
         coordinator=coordinator,
@@ -222,7 +221,7 @@ def test_initiate_ritual(
     assert fee_model.pendingFees(ritualID) == fee
 
     with ape.reverts():
-        fee_model.withdrawTokens(1, sender=treasury)
+        fee_model.withdrawTokens(1, sender=fee_manager)
 
     with ape.reverts("Can't withdraw pending fees"):
         fee_model.withdrawTokens(1, sender=deployer)
@@ -479,7 +478,7 @@ def test_get_participant(nodes, coordinator, initiator, erc20, fee_model, global
 
 
 def test_post_aggregation(
-    coordinator, nodes, initiator, erc20, fee_model, treasury, deployer, global_allow_list
+    coordinator, nodes, initiator, erc20, fee_model, fee_manager, deployer, global_allow_list
 ):
     initiate_ritual(
         coordinator=coordinator,
@@ -527,7 +526,7 @@ def test_post_aggregation(
     assert retrieved_public_key == dkg_public_key
     assert coordinator.getRitualIdFromPublicKey(dkg_public_key) == ritualID
 
-    fee_model.processPendingFee(ritualID, sender=treasury)
+    fee_model.processPendingFee(ritualID, sender=fee_manager)
     fee = fee_model.getRitualCost(len(nodes), DURATION)
     assert erc20.balanceOf(fee_model) == fee
     assert fee_model.totalPendingFees() == 0
@@ -539,7 +538,7 @@ def test_post_aggregation(
 
 
 def test_post_aggregation_fails(
-    coordinator, nodes, initiator, erc20, fee_model, treasury, deployer, global_allow_list
+    coordinator, nodes, initiator, erc20, fee_model, fee_manager, deployer, global_allow_list
 ):
     initiator_balance_before_payment = erc20.balanceOf(initiator)
 
@@ -593,7 +592,7 @@ def test_post_aggregation_fails(
     initiator_balance_before_refund = erc20.balanceOf(initiator)
     assert initiator_balance_before_refund == initiator_balance_before_payment - fee
 
-    fee_model.processPendingFee(ritualID, sender=treasury)
+    fee_model.processPendingFee(ritualID, sender=fee_manager)
 
     initiator_balance_after_refund = erc20.balanceOf(initiator)
     fee_model_balance_after_refund = erc20.balanceOf(fee_model)