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 8 commits into
base: dev
Choose a base branch
from
Open

fix(SortitionModule): total stake update #2199

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 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
a016740
feat(KC): add event for failed transfer
unknownunknown1 Dec 12, 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
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