fix(SortitionModule): total stake update

contracts/src/arbitration/KlerosCore.sol

@@ -627,7 +627,7 @@ contract KlerosCore is IArbitratorV2, Initializable, UUPSProxiable {
     function transferBySortitionModule(address _account, uint256 _amount) external {
         if (msg.sender != address(sortitionModule)) revert SortitionModuleOnly();
         // Note eligibility is checked in SortitionModule.
-        pinakion.safeTransfer(_account, _amount);
+        if (!pinakion.safeTransfer(_account, _amount)) revert TransferFailed();
     }
 
     /// @inheritdoc IArbitratorV2

contracts/src/arbitration/SortitionModule.sol

@@ -311,15 +311,13 @@ contract SortitionModule is ISortitionModule, Initializable, UUPSProxiable {
         // Current phase is Staking: set stakes.
         if (stakeIncrease) {
             pnkDeposit = stakeChange;
-            totalStaked += stakeChange;
         } else {
             pnkWithdrawal = stakeChange;
             uint256 possibleWithdrawal = juror.stakedPnk > juror.lockedPnk ? juror.stakedPnk - juror.lockedPnk : 0;
             if (pnkWithdrawal > possibleWithdrawal) {
                 // Ensure locked tokens remain in the contract. They can only be released during Execution.
                 pnkWithdrawal = possibleWithdrawal;
             }
-            totalStaked -= pnkWithdrawal;
         }
         return (pnkDeposit, pnkWithdrawal, StakingResult.Successful);
     }
@@ -391,8 +389,10 @@ contract SortitionModule is ISortitionModule, Initializable, UUPSProxiable {
             }
             // Increase juror's balance by deposited amount.
             juror.stakedPnk += _pnkDeposit;
+            totalStaked += _pnkDeposit;
         } else {
             juror.stakedPnk -= _pnkWithdrawal;
+            totalStaked -= _pnkWithdrawal;
             if (_newStake == 0) {
                 // Cleanup
                 for (uint256 i = juror.courtIDs.length; i > 0; i--) {
@@ -460,6 +460,7 @@ contract SortitionModule is ISortitionModule, Initializable, UUPSProxiable {
         uint256 amount = getJurorLeftoverPNK(_account);
         if (amount == 0) revert NotEligibleForWithdrawal();
         jurors[_account].stakedPnk = 0;
+        totalStaked -= amount;
         core.transferBySortitionModule(_account, amount);
         emit LeftoverPNKWithdrawn(_account, amount);
     }

contracts/src/arbitration/arbitrables/ArbitrableExample.sol

@@ -149,7 +149,7 @@ contract ArbitrableExample is IArbitrableV2 {
     }
 
     /// @inheritdoc IArbitrableV2
-    function rule(uint256 _arbitratorDisputeID, uint256 _ruling) external override {
+    function rule(uint256 _arbitratorDisputeID, uint256 _ruling) external virtual {
         uint256 localDisputeID = externalIDtoLocalID[_arbitratorDisputeID];
         DisputeStruct storage dispute = disputes[localDisputeID];
         if (msg.sender != address(arbitrator)) revert ArbitratorOnly();

contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol

@@ -475,8 +475,8 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
         address payable _beneficiary,
         uint256 _choice
     ) external returns (uint256 amount) {
-        (, , , bool isRuled, ) = core.disputes(_coreDisputeID);
-        if (!isRuled) revert DisputeNotResolved();
+        (, , KlerosCore.Period period, , ) = core.disputes(_coreDisputeID);
+        if (period != KlerosCore.Period.execution) revert DisputeNotResolved();
         if (core.paused()) revert CoreIsPaused();
         if (!coreDisputeIDToActive[_coreDisputeID].dispute) revert DisputeUnknownInThisDisputeKit();
 

contracts/src/test/MaliciousArbitrableMock.sol

@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.24;
+
+import {IArbitratorV2, IDisputeTemplateRegistry, IERC20, ArbitrableExample} from "../arbitration/arbitrables/ArbitrableExample.sol";
+
+/// @title MaliciousArbitrableMock
+/// A mock contract to check intentional rule() revert.
+contract MaliciousArbitrableMock is ArbitrableExample {
+    bool public doRevert;
+
+    function changeBehaviour(bool _doRevert) external {
+        doRevert = _doRevert;
+    }
+
+    constructor(
+        IArbitratorV2 _arbitrator,
+        string memory _templateData,
+        string memory _templateDataMappings,
+        bytes memory _arbitratorExtraData,
+        IDisputeTemplateRegistry _templateRegistry,
+        IERC20 _weth
+    )
+        ArbitrableExample(
+            _arbitrator,
+            _templateData,
+            _templateDataMappings,
+            _arbitratorExtraData,
+            _templateRegistry,
+            _weth
+        )
+    {
+        doRevert = true;
+    }
+
+    function rule(uint256 _arbitratorDisputeID, uint256 _ruling) external override {
+        if (doRevert) revert RuleReverted();
+
+        uint256 localDisputeID = externalIDtoLocalID[_arbitratorDisputeID];
+        DisputeStruct storage dispute = disputes[localDisputeID];
+        if (msg.sender != address(arbitrator)) revert ArbitratorOnly();
+        if (_ruling > dispute.numberOfRulingOptions) revert RulingOutOfBounds();
+        if (dispute.isRuled) revert DisputeAlreadyRuled();
+
+        dispute.isRuled = true;
+        dispute.ruling = _ruling;
+
+        emit Ruling(IArbitratorV2(msg.sender), _arbitratorDisputeID, dispute.ruling);
+    }
+
+    error RuleReverted();
+}

contracts/test/foundry/KlerosCore_Execution.t.sol

@@ -8,6 +8,8 @@ import {DisputeKitClassicBase} from "../../src/arbitration/dispute-kits/DisputeK
 import {IArbitratorV2, IArbitrableV2} from "../../src/arbitration/KlerosCore.sol";
 import {IERC20} from "../../src/libraries/SafeERC20.sol";
 import "../../src/libraries/Constants.sol";
+import {console} from "forge-std/console.sol";
+import {MaliciousArbitrableMock} from "../../src/test/MaliciousArbitrableMock.sol";
 
 /// @title KlerosCore_ExecutionTest
 /// @dev Tests for KlerosCore execution, rewards, and ruling finalization
@@ -96,6 +98,9 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
             "Wrong penalty coherence 2 vote ID"
         );
 
+        assertEq(pinakion.balanceOf(address(core)), 22000, "Wrong token balance of the core");
+        assertEq(sortitionModule.totalStaked(), 22000, "Total staked should be equal to the balance in this test");
+
         vm.expectEmit(true, true, true, true);
         emit SortitionModule.StakeLocked(staker1, 1000, true);
         vm.expectEmit(true, true, true, true);
@@ -141,9 +146,11 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
         assertEq(staker1.balance, 0, "Wrong balance of the staker1");
         assertEq(staker2.balance, 0.09 ether, "Wrong balance of the staker2");
 
-        assertEq(pinakion.balanceOf(address(core)), 22000, "Wrong token balance of the core"); // Was 21500. 1000 was transferred to staker2
+        assertEq(pinakion.balanceOf(address(core)), 22000, "Token balance of the core shouldn't change after rewards");
+        assertEq(sortitionModule.totalStaked(), 22000, "Total staked shouldn't change after rewards");
+
         assertEq(pinakion.balanceOf(staker1), 999999999999998000, "Wrong token balance of staker1");
-        assertEq(pinakion.balanceOf(staker2), 999999999999980000, "Wrong token balance of staker2"); // 20k stake and 1k added as a reward, thus -19k from the default
+        assertEq(pinakion.balanceOf(staker2), 999999999999980000, "Wrong token balance of staker2");
     }
 
     function test_execute_NoCoherence() public {
@@ -477,13 +484,28 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
         vm.prank(owner);
         core.transferBySortitionModule(staker1, 1000);
 
+        assertEq(pinakion.balanceOf(address(core)), 1000, "Wrong token balance of the core");
+        assertEq(sortitionModule.totalStaked(), 1000, "Wrong totalStaked before withdrawal");
+
+        vm.prank(address(core));
+        pinakion.transfer(staker2, 1000); // Manually send the balance to trigger the revert
+
+        vm.expectRevert(KlerosCore.TransferFailed.selector);
+        sortitionModule.withdrawLeftoverPNK(staker1);
+
+        vm.prank(staker2);
+        pinakion.transfer(address(core), 1000); // Transfer the tokens back to execute withdrawal
+
         vm.expectEmit(true, true, true, true);
         emit SortitionModule.LeftoverPNKWithdrawn(staker1, 1000);
         sortitionModule.withdrawLeftoverPNK(staker1);
 
         (totalStaked, , , ) = sortitionModule.getJurorBalance(staker1, GENERAL_COURT);
         assertEq(totalStaked, 0, "Should be unstaked fully");
 
+        assertEq(pinakion.balanceOf(address(core)), 0, "Wrong token balance of the core");
+        assertEq(sortitionModule.totalStaked(), 0, "Wrong totalStaked after withdrawal");
+
         // Check that everything is withdrawn now
         assertEq(pinakion.balanceOf(address(core)), 0, "Core balance should be empty");
         assertEq(pinakion.balanceOf(staker1), 1 ether, "All PNK should be withdrawn");
@@ -718,12 +740,13 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
         disputeKit.fundAppeal{value: 0.41 ether}(disputeID, 2); // Underpay a bit to not create an appeal and withdraw the funded sum fully
 
         vm.warp(block.timestamp + timesPerPeriod[3]);
-        core.passPeriod(disputeID); // Execution
 
         vm.expectRevert(DisputeKitClassicBase.DisputeNotResolved.selector);
         disputeKit.withdrawFeesAndRewards(disputeID, payable(staker1), 1);
 
-        core.executeRuling(disputeID);
+        core.passPeriod(disputeID); // Execution
+        // executeRuling() should be irrelevant for withdrawals in case malicious arbitrable reverts rule()
+        //core.executeRuling(disputeID);
 
         vm.prank(owner);
         core.pause();
@@ -749,6 +772,52 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
         assertEq(address(disputeKit).balance, 0, "Wrong balance of the DK");
     }
 
+    function test_inflatedTotalStaked_whenDelayedStakeExecute_whenJurorHasNoFunds() public {
+        // pre conditions
+        // 1. there is a dispute in drawing phase
+        // 2. juror call setStake with an amount greater than his PNK balance
+        // 3. draw jurors, move to voting phase and execute voting
+        // 4. move sortition to staking phase
+        uint256 disputeID = 0;
+        uint256 amountToStake = 20000;
+        _stakePnk_createDispute_moveToDrawingPhase(disputeID, staker1, amountToStake);
+
+        KlerosCore.Round memory round = core.getRoundInfo(disputeID, 0);
+        uint256 pnkAtStakePerJuror = round.pnkAtStakePerJuror;
+        _stakeBalanceForJuror(staker1, type(uint256).max);
+        _drawJurors_advancePeriodToVoting(disputeID);
+        _vote_execute(disputeID, staker1);
+        sortitionModule.passPhase(); // set it to staking phase
+        _assertJurorBalance(
+            disputeID,
+            staker1,
+            amountToStake,
+            pnkAtStakePerJuror * DEFAULT_NB_OF_JURORS,
+            amountToStake,
+            1
+        );
+
+        console.log("totalStaked before: %e", sortitionModule.totalStaked());
+
+        // execution: execute delayed stake
+        sortitionModule.executeDelayedStakes(1);
+
+        // post condition: inflated totalStaked
+        console.log("totalStaked after: %e", sortitionModule.totalStaked());
+        _assertJurorBalance(
+            disputeID,
+            staker1,
+            amountToStake,
+            pnkAtStakePerJuror * DEFAULT_NB_OF_JURORS,
+            amountToStake,
+            1
+        );
+
+        // new juror tries to stake but totalStaked already reached type(uint256).max
+        // it reverts with "arithmetic underflow or overflow (0x11)"
+        _stakeBalanceForJuror(staker2, 20000);
+    }
+
     function testFuzz_executeIterations(uint256 iterations) public {
         uint256 disputeID = 0;
         uint256 roundID = 0;
@@ -847,4 +916,61 @@ contract KlerosCore_ExecutionTest is KlerosCore_TestBase {
         assertEq(totalLocked, (pnkAtStake * nbJurors) - unlockedTokens, "Wrong amount locked");
         assertEq(stakedInCourt, 2000, "Wrong amount staked in court");
     }
+
+    ///////// Internal //////////
+
+    function _assertJurorBalance(
+        uint256 disputeID,
+        address juror,
+        uint256 totalStakedPnk,
+        uint256 totalLocked,
+        uint256 stakedInCourt,
+        uint256 nbCourts
+    ) internal {
+        (uint256 totalStakedPnk, uint256 totalLocked, uint256 stakedInCourt, uint256 nbCourts) = sortitionModule
+            .getJurorBalance(juror, GENERAL_COURT);
+        assertEq(totalStakedPnk, totalStakedPnk, "Wrong totalStakedPnk"); // jurors total staked a.k.a juror.stakedPnk
+        assertEq(totalLocked, totalLocked, "Wrong totalLocked");
+        assertEq(stakedInCourt, stakedInCourt, "Wrong stakedInCourt"); // juror staked in court a.k.a _stakeOf
+        assertEq(nbCourts, nbCourts, "Wrong nbCourts");
+    }
+
+    function _stakeBalanceForJuror(address juror, uint256 amount) internal {
+        console.log("actual juror PNK balance before staking: %e", pinakion.balanceOf(juror));
+        vm.prank(juror);
+        core.setStake(GENERAL_COURT, amount);
+    }
+
+    function _stakePnk_createDispute_moveToDrawingPhase(uint256 disputeID, address juror, uint256 amount) internal {
+        vm.prank(juror);
+        core.setStake(GENERAL_COURT, amount);
+        vm.prank(disputer);
+        arbitrable.createDispute{value: feeForJuror * DEFAULT_NB_OF_JURORS}("Action");
+        vm.warp(block.timestamp + minStakingTime);
+        sortitionModule.passPhase(); // Generating
+        vm.warp(block.timestamp + rngLookahead);
+        sortitionModule.passPhase(); // Drawing phase
+
+        assertEq(sortitionModule.totalStaked(), amount, "!totalStaked");
+    }
+
+    function _drawJurors_advancePeriodToVoting(uint256 disputeID) internal {
+        core.draw(disputeID, DEFAULT_NB_OF_JURORS);
+        vm.warp(block.timestamp + timesPerPeriod[0]);
+        core.passPeriod(disputeID); // Vote
+    }
+
+    function _vote_execute(uint256 disputeID, address juror) internal {
+        uint256[] memory voteIDs = new uint256[](3);
+        voteIDs[0] = 0;
+        voteIDs[1] = 1;
+        voteIDs[2] = 2;
+
+        vm.prank(juror);
+        disputeKit.castVote(disputeID, voteIDs, 2, 0, "XYZ");
+        core.passPeriod(disputeID); // Appeal
+
+        vm.warp(block.timestamp + timesPerPeriod[3]);
+        core.passPeriod(disputeID); // Execution
+    }
 }

contracts/test/foundry/KlerosCore_Staking.t.sol

@@ -141,6 +141,32 @@ contract KlerosCore_StakingTest is KlerosCore_TestBase {
         );
     }
 
+    function test_setStake_totalStaked() public {
+        // Increase
+        vm.prank(staker1);
+        core.setStake(GENERAL_COURT, 4000);
+        vm.prank(staker1);
+        core.setStake(GENERAL_COURT, 5001);
+        vm.prank(staker2);
+        core.setStake(GENERAL_COURT, 1000);
+        vm.prank(staker2);
+        core.setStake(GENERAL_COURT, 1500);
+
+        assertEq(sortitionModule.totalStaked(), 6501, "Wrong totalStaked");
+
+        // Decrease
+        vm.prank(staker1);
+        core.setStake(GENERAL_COURT, 3000);
+        vm.prank(staker1);
+        core.setStake(GENERAL_COURT, 2500);
+        vm.prank(staker2);
+        core.setStake(GENERAL_COURT, 1400);
+        vm.prank(staker2);
+        core.setStake(GENERAL_COURT, 1200);
+
+        assertEq(sortitionModule.totalStaked(), 3700, "Wrong totalStaked");
+    }
+
     function test_setStake_maxStakePathCheck() public {
         uint256[] memory supportedDK = new uint256[](1);
         supportedDK[0] = DISPUTE_KIT_CLASSIC;