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 override {
         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/libraries/SafeERC20.sol

@@ -15,6 +15,12 @@ import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 /// To use this library you can add a `using SafeERC20 for IERC20;` statement to your contract,
 /// which allows you to call the safe operations as `token.safeTransfer(...)`, etc.
 library SafeERC20 {
+    /// @notice Emits when safeTransfer fails.
+    /// @param _token Token to transfer.
+    /// @param _to Recipient address.
+    /// @param _value Amount transferred.
+    event SafeTransferFailed(IERC20 _token, address _to, uint256 _value);
+
     /// @notice Increases the allowance granted to `spender` by the caller.
     /// @param _token Token to transfer.
     /// @param _spender The address which will spend the funds.
@@ -31,7 +37,9 @@ library SafeERC20 {
     /// @return Whether transfer succeeded or not.
     function safeTransfer(IERC20 _token, address _to, uint256 _value) internal returns (bool) {
         (bool success, bytes memory data) = address(_token).call(abi.encodeCall(IERC20.transfer, (_to, _value)));
-        return (success && (data.length == 0 || abi.decode(data, (bool))));
+        bool ok = success && (data.length == 0 || abi.decode(data, (bool)));
+        if (!ok) emit SafeTransferFailed(_token, _to, _value);
+        return ok;
     }
 
     /// @notice Calls transferFrom() without reverting.

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();
+}