From 748c9b0a551bae9a395b0f4531b3e8da61ebee8f Mon Sep 17 00:00:00 2001
From: unknownunknown1 <danil@kleros.io>
Date: Wed, 31 Jul 2024 18:51:15 +1000
Subject: [PATCH] feat(DK): add permission for flushing

---
 contracts/src/arbitration/KlerosCoreBase.sol           | 10 +++++-----
 contracts/src/arbitration/SortitionModuleBase.sol      |  4 ++--
 .../dispute-kits/DisputeKitSybilResistant.sol          |  3 +--
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/contracts/src/arbitration/KlerosCoreBase.sol b/contracts/src/arbitration/KlerosCoreBase.sol
index c1cce81d6..ce1e60d03 100644
--- 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
index 24e32e6ea..8872f90dc 100644
--- 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
index f20986c4d..d61d9fc41 100644
--- 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);
         }
     }