feat(DK): add permission for flushing

kleros · Jul 31, 2024 · 748c9b0 · 748c9b0
1 parent 87bc6c0
commit 748c9b0
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 9 deletions.
diff --git a/contracts/src/arbitration/KlerosCoreBase.sol b/contracts/src/arbitration/KlerosCoreBase.sol
@@ -471,19 +471,19 @@ abstract contract KlerosCoreBase is IArbitratorV2 {
         _setStake(msg.sender, _courtID, _newStake, false, OnError.Revert);
     }
 
-    /// @dev Sets the stake of a specified account in a court, typically to apply a delayed stake or unstake inactive jurors.
+    /// @dev Sets the stake of a specified account in a court, typically to apply a delayed stake or unstake inactive/ineligible jurors.
     /// @param _account The account whose stake is being set.
     /// @param _courtID The ID of the court.
     /// @param _newStake The new stake.
     /// @param _alreadyTransferred Whether the PNKs have already been transferred to the contract.
-    function setStakeBySortitionModule(
+    function setStakeBySortitionModuleOrDK(
         address _account,
         uint96 _courtID,
         uint256 _newStake,
         bool _alreadyTransferred
     ) external {
-        // TODO: generalize this function so it can be used by DK too. Rename to setStakeWhitelist or something. So it can be done for someone not just for sender.
-        if (msg.sender != address(sortitionModule)) revert SortitionModuleOnly();
+        if (msg.sender != address(sortitionModule) && disputeKitIDs[IDisputeKit(msg.sender)] == 0)
+            revert SortitionModuleOrDKOnly();
         _setStake(_account, _courtID, _newStake, _alreadyTransferred, OnError.Return);
     }
 
@@ -1148,7 +1148,7 @@ abstract contract KlerosCoreBase is IArbitratorV2 {
     error GovernorOnly();
     error GuardianOrGovernorOnly();
     error DisputeKitOnly();
-    error SortitionModuleOnly();
+    error SortitionModuleOrDKOnly();
     error UnsuccessfulCall();
     error InvalidDisputKitParent();
     error DepthLevelMax();

diff --git a/contracts/src/arbitration/SortitionModuleBase.sol b/contracts/src/arbitration/SortitionModuleBase.sol
@@ -206,7 +206,7 @@ abstract contract SortitionModuleBase is ISortitionModule {
             DelayedStake storage delayedStake = delayedStakes[i];
             // Delayed stake could've been manually removed already. In this case simply move on to the next item.
             if (delayedStake.account != address(0)) {
-                core.setStakeBySortitionModule(
+                core.setStakeBySortitionModuleOrDK(
                     delayedStake.account,
                     delayedStake.courtID,
                     delayedStake.stake,
@@ -433,7 +433,7 @@ abstract contract SortitionModuleBase is ISortitionModule {
     function setJurorInactive(address _account) external override onlyByCore {
         uint96[] memory courtIDs = getJurorCourtIDs(_account);
         for (uint256 j = courtIDs.length; j > 0; j--) {
-            core.setStakeBySortitionModule(_account, courtIDs[j - 1], 0, false);
+            core.setStakeBySortitionModuleOrDK(_account, courtIDs[j - 1], 0, false);
         }
     }
 

diff --git a/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol b/contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol
@@ -475,8 +475,7 @@ contract DisputeKitSybilResistant is IDisputeKit, Initializable, UUPSProxiable {
         for (uint256 i = 0; i < _courts.length; i++) {
             uint96 courtID = _courts[i];
             require(core.isSupported(courtID, disputeKitID), "DK not supported by court");
-            // TODO: generalize the function name and its permission.
-            core.setStakeBySortitionModule(_juror, courtID, 0, false);
+            core.setStakeBySortitionModuleOrDK(_juror, courtID, 0, false);
         }
     }