Skip to content

fix(SortitionModule): total stake update #2199

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
unknownunknown1 wants to merge 7 commits into
base: dev
Choose a base branch
from
Open

fix(SortitionModule): total stake update #2199

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
1006d82
fix(SortitionModule): total stake update
unknownunknown1 Nov 29, 2025
a2b5d76
feat(SortitionModule): add certora PoC to Foundry
unknownunknown1 Dec 2, 2025
719ea3e
fix(KC): malicious arbitrable revert
unknownunknown1 Dec 6, 2025
5799a3b
fix(SortitionModule): sync totalStaked with stakedPnk
unknownunknown1 Dec 6, 2025
cb492fd
fix(KC): check safeTransfer for withdrawLeftover
unknownunknown1 Dec 8, 2025
cae435e
fix(KC): remove try catch
unknownunknown1 Dec 8, 2025
8c256aa
test: adding back the reverting arbitrable test, small nitpicks
jaybuidl Dec 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions contracts/src/arbitration/KlerosCore.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1363,6 +1363,7 @@ contract KlerosCore is IArbitratorV2, Initializable, UUPSProxiable {
}
}
sortitionModule.setStake(_account, _courtID, pnkDeposit, pnkWithdrawal, _newStake);
sortitionModule.updateTotalStake(pnkDeposit, pnkWithdrawal);

return true;
}
Expand Down
12 changes: 10 additions & 2 deletions contracts/src/arbitration/SortitionModule.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -311,19 +311,27 @@ 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);
}

/// @inheritdoc ISortitionModule
function updateTotalStake(uint256 _pnkDeposit, uint256 _pnkWithdrawal) external override onlyByCore {
// Note that we don't update totalStake in setStake() function because it doesn't always change total (e.g. during rewards/penalties).
if (_pnkDeposit > 0) {
totalStaked += _pnkDeposit;
} else {
totalStaked -= _pnkWithdrawal;
}
}

/// @inheritdoc ISortitionModule
function setStake(
address _account,
Expand Down
5 changes: 5 additions & 0 deletions contracts/src/arbitration/interfaces/ISortitionModule.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,11 @@ interface ISortitionModule {
bool _noDelay
) external returns (uint256 pnkDeposit, uint256 pnkWithdrawal, StakingResult stakingResult);

/// @notice Updates the total amount staked in all courts.
/// @param _pnkDeposit The amount of PNK that increases total stake.
/// @param _pnkWithdrawal The amount of PNK that decreases total stake.
function updateTotalStake(uint256 _pnkDeposit, uint256 _pnkWithdrawal) external;

/// @notice Update the state of the stakes, called by KC at the end of setStake flow.
///
/// @dev `O(n + p * log_k(j))` where
Expand Down
5 changes: 5 additions & 0 deletions contracts/src/arbitration/university/SortitionModuleUniversity.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,11 @@ contract SortitionModuleUniversity is ISortitionModuleUniversity, UUPSProxiable,
// NOP
}

/// @inheritdoc ISortitionModule
function updateTotalStake(uint256 _pnkDeposit, uint256 _pnkWithdrawal) external override onlyByCore {
// NOP
}

/// @inheritdoc ISortitionModule
function createDisputeHook(uint256 /*_disputeID*/, uint256 /*_roundID*/) external override onlyByCore {
disputesWithoutJurors++;
Expand Down
104 changes: 104 additions & 0 deletions contracts/test/foundry/KlerosCore_Execution.t.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ 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";

/// @title KlerosCore_ExecutionTest
/// @dev Tests for KlerosCore execution, rewards, and ruling finalization
Expand Down Expand Up @@ -749,6 +750,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;
Expand Down Expand Up @@ -847,4 +894,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
}
}
26 changes: 26 additions & 0 deletions contracts/test/foundry/KlerosCore_Staking.t.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
}
Comment on lines +144 to +168
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# First, search for totalStaked references in the test file
rg -n "totalStaked" contracts/test/foundry/KlerosCore_Staking.t.sol -B2 -A2

Repository: kleros/kleros-v2

Length of output: 7231

🏁 Script executed:

#!/bin/bash
# Search for tests that combine delayed stakes with totalStaked checks
rg -n "delayedStake.*totalStaked|totalStaked.*delayedStake" contracts/test/foundry/ -i

Repository: kleros/kleros-v2

Length of output: 42

🏁 Script executed:

#!/bin/bash
# Search for all test functions that mention delayed stake
rg -n "function test.*[Dd]elayed" contracts/test/foundry/KlerosCore_Staking.t.sol

Repository: kleros/kleros-v2

Length of output: 182

🏁 Script executed:

#!/bin/bash
# Check the entire context around delayed stake tests to see if totalStaked is tested
rg -n "test_delayedStake" contracts/test/foundry/KlerosCore_Staking.t.sol -A 50 | head -100

Repository: kleros/kleros-v2

Length of output: 42

🏁 Script executed:

#!/bin/bash
# Read the test_executeDelayedStakes function
sed -n '324,450p' contracts/test/foundry/KlerosCore_Staking.t.sol

Repository: kleros/kleros-v2

Length of output: 6960

🏁 Script executed:

#!/bin/bash
# Find the SortitionModule contract file
fd -t f "SortitionModule\.sol" contracts/

Repository: kleros/kleros-v2

Length of output: 163

🏁 Script executed:

#!/bin/bash
# Search for the totalStaked function/variable definition
rg -n "totalStaked" contracts/src/ -B2 -A5 | head -100

Repository: kleros/kleros-v2

Length of output: 7695

🏁 Script executed:

#!/bin/bash
# Search for executeDelayedStakes implementation
rg -n "function executeDelayedStakes" contracts/src/ -A 30

Repository: kleros/kleros-v2

Length of output: 9796

🏁 Script executed:

#!/bin/bash
# Search for setStakeBySortitionModule in KlerosCore
rg -n "function setStakeBySortitionModule" contracts/src/ -A 15

Repository: kleros/kleros-v2

Length of output: 3466

🏁 Script executed:

#!/bin/bash
# Confirm _setStake calls the totalStaked update
rg -n "function _setStake" contracts/src/arbitration/KlerosCore.sol -A 40 | head -60

Repository: kleros/kleros-v2

Length of output: 2059

Test coverage for delayed stakes affecting totalStaked should be added.

The verification confirms the review comment's concern is valid. While the current test_setStake_totalStaked() correctly validates immediate stakes, there is a test coverage gap:

  • Delayed stakes do NOT immediately update sortitionModule.totalStaked() (they return early in _setStake() at line 1351)
  • When executeDelayedStakes() is called during the Staking phase, it invokes setStakeBySortitionModule() which calls _setStake() with _noDelay=true, which then calls updateTotalStake() to update the global total
  • The existing test_executeDelayedStakes() validates delayed stake execution but never checks that sortitionModule.totalStaked() is correctly updated after execution

A test case should verify that executing delayed stakes properly increments the global totalStaked metric.

🤖 Prompt for AI Agents 
In contracts/test/foundry/KlerosCore_Staking.t.sol around lines 144 to 168, add
a test that verifies delayed stakes do not immediately affect
sortitionModule.totalStaked and that calling executeDelayedStakes during the
Staking phase updates totalStaked appropriately: create delayed stakes (use
callers that will trigger the delayed path), assert
sortitionModule.totalStaked() remains unchanged, advance the contract to the
Staking phase (or ensure executeDelayedStakes is callable), call
core.executeDelayedStakes(), then assert sortitionModule.totalStaked() equals
the expected new total after delayed stakes are applied.


function test_setStake_maxStakePathCheck() public {
uint256[] memory supportedDK = new uint256[](1);
supportedDK[0] = DISPUTE_KIT_CLASSIC;
Expand Down
Loading