feat: fuzzing of staking invariants for ProtocolStaking contract with Foundry

[Seth-Schmidt]

Summary

Invariant fuzzing for the ProtocolStaking contract using Foundry. Uses a handler-based setup to exercise staking, rewards, and unstaking flows under randomized sequences.

Structure

• Handler (ProtocolStakingHandler.sol): Wraps protocol actions, bounds inputs, tracks ghost state, and exposes equivalence scenarios.
• Invariant test (ProtocolStakingInvariantTest.t.sol): Defines invariants and runs them after each handler call.

How to run

cd contracts/staking
npm install
forge install foundry-rs/forge-std --no-git
forge test

Files Modified

• contracts/staking/test/foundry/ProtocolStakingInvariantTest.t.sol
• contracts/staking/test/foundry/handlers/ProtocolStakingHandler.sol
• contracts/staking/test/foundry/harness/ProtocolStakingHarness.sol
• contracts/staking/test/foundry/TESTING.md (setup and invariants)
• contracts/staking/foundry.toml
• contracts/staking/package.json

closes: https://github.com/zama-ai/protocol-apps-internal/issues/101

[jatZama]

I don't think you need this whole list of selectors and a call to targetSelector here. Foundry by default will fuzz all available public/external functions from the targetContract.

[Seth-Schmidt]

Fixed here: 18e7bc9

[jatZama]

On the other hand, since Foundry will use random msg.sender addresses in invariant testing, maybe it would make sense to use targetSenders cheatcode here, and to pass the 5 actors to it (ie the eligible account) + few other random addresses, this is in order to reduce bias. Because my main concern here is that most calls would be done with current setup with uneligible accounts, thus reducing the coverage of the fuzzing campaign. Wdyt?

[jatZama]

Actually just noticed that you are using vm.prank calls in all functions of the handler contract to handle this correctly.
But even despite this trick, I think you are restricting a bit too much the fuzzer. Technically, the ProtocolStaking allows uneligible account to call stake() function without reverting. IMO it would be best to also allow some uneligible account to call it, to improve coverage (while still biasing for eligible accounts). There are few methods to do this, I will let you decide what is best approach.

[Seth-Schmidt]

The caller should always be a known actor through the combination of bounding the actor index actorIndex = bound(actorIndex, 0, actors.length - 1); and a subsequent prank(account) before each staking action in the handler (or the manager if it is an admin action). However, I do think it makes sense to be explicit and use targetSenders for traces and the removal of ambiguity.

[Seth-Schmidt]

Actually just noticed that you are using vm.prank calls in all functions of the handler contract to handle this correctly. But even despite this trick, I think you are restricting a bit too much the fuzzer. Technically, the ProtocolStaking allows uneligible account to call stake() function without reverting. IMO it would be best to also allow some uneligible account to call it, to improve coverage (while still biasing for eligible accounts). There are few methods to do this, I will let you decide what is best approach.

Got it. You can ignore my previous reply, the comment had not refreshed yet. I think we came to the same conclusion on this.

[Seth-Schmidt]

1bea63a

The actors have now been added as targetSender's

To enforce ineligibility, I made an isOutgroup flag that prevents ~20% of actors from ever becoming eligible, however they can still call the rest of the functions.

[jatZama]

I am a bit confused, I think you don't need to use targetSender here, since you are using vm.prank in all handler's methods?

[Seth-Schmidt]

I have removed the vm.prank for account level actions in the handler, but the manager prank for admin actions remain. The targetSender is preferable in my opinion because of the reduction in pranks and not having to fuzz the actor index for each action.

[jatZama]

I thought from your previous comment #50 (comment) that using vm.randomUint in the setup would be problematic, so I am a bit confused seeing you are using it here... Was it not an issue ultimately?

[melanciani]

is this assert needed ? sounds more like a unit test logic (like we already have)

I'm asking because I don't see any other assets in other functions so it looks a bit inconsistent

[Seth-Schmidt]

This assert is necessary to check that any fuzzed state does not allow for dust in the earned calculation after claiming. I agree it should be moved outside of the handler action itself to a modifier to be inline with the other transitional invariants.

[melanciani]

for clarity I would maybe rename this function into addEligibleMsgSender or similar

same for removeEligibleAccount -> removeEligibleMsgSender

alternatively, you could add an address account input to these functions

[Seth-Schmidt]

There initially was an account_index input for the fuzzer to use for account choosing, but foundry has a feature of targetSender that allows the passing of addresses for the fuzzer to randomly use. It's these addresses that will be used to call all of the handler functions by the fuzzing engine:

protocol-apps/contracts/staking/test/foundry/ProtocolStakingInvariantTest.t.sol

Lines 89 to 94 in ed53d1e

	handler = new ProtocolStakingHandler(protocolStaking, zama, manager, actorsList);
	targetContract(address(handler));
	for (uint256 i = 0; i < actorCount; i++) {
	targetSender(actorsList[i]);
	}

I think I will add some documentation around this, because it's not exactly intuitive, especially for testing.

The idea for leaving removeEligibleAccount is to keep the 1 to 1 relationship between handler and ProtocolStaking functions. It is more of an override/wrapper than a standalone function. I have no problem with adjusting the naming, but just want to be clear on the initial idea.

[melanciani]

best practice with unmet conditions is to return instead of revert here ? just asking to be sure some scenari are not silently failing because of this !

[Seth-Schmidt]

There seems to be two schools of thought around handling reverts: either reverts are encouraged or avoided. Too many reverts can inhibit the building of a fuzzed state. I've leaned more towards avoiding, but the only action that is gated by an early return is unstake to avoid transferFrom issues. Thinking about it now, we can probably expect and assert the revert instead to be more explicit with the expected behavior.

[melanciani]

I think these are great but imo we should be a bit more rigorous here, by specifically detailing what are the different variables in each invariant

for ex in Staked funds solvency:, it is not 100% clear what is totalStaked and released (a sum ? for a single account ? a function to call ?) !

[melanciani]

why say "strictly non-decreasing" instead of "strictly increasing" ?

[melanciani]

i think it should be strictly increasing timestamps .

Also do you think the and cumulative amounts. should be an invariant of itself ?

and in our invariant page we also had this one which I think should be checked as well, right ? :

The _released[account] is always inferior or equal to the latest unstake request in _unstakeRequest[account].latest() , ie awaitingRelease view function should never revert.

[Seth-Schmidt]

Also do you think the and cumulative amounts. should be an invariant of itself ?

Yes maybe there can be some additional clarity around the notion invariant: "The _unstakeRequests checkpoints must have increasing timestamps and “cumulative amounts”, ie only increasing for same timestamps keys."

I took the literal interpretation of only increasing for same timestamps keys:

protocol-apps/contracts/staking/test/foundry/handlers/ProtocolStakingHandler.sol

Lines 143 to 148 in ed53d1e

	if (hadCheckpoint) {
	assertGe(postKey, preKey, "unstake request keys must be non-decreasing");
	if (postKey == preKey) {
	assertGe(postValue, preValue, "unstake request values must be non-decreasing for same key");
	}
	}

but I will double check how checkpoints are handled, intuitively the amounts should always be increasing.

and in our invariant page we also had this one which I think should be checked as well, right ?

This one gets handled inherently by the awaitingRelease(account) call:

protocol-apps/contracts/staking/test/foundry/handlers/ProtocolStakingHandler.sol

Lines 126 to 129 in ed53d1e

	// inherent check that awaiting release does not revert
	// _released[account] is always inferior or equal to the latest unstake request in _unstakeRequest[account].latest()
	uint256 postAwaitingRelease = protocolStaking.awaitingRelease(account);
	assertGe(postAwaitingRelease, preAwaitingRelease, "awaitingRelease must not decrease except after release");

which has that check directly in the function:

protocol-apps/contracts/staking/contracts/ProtocolStaking.sol

Line 318 in ed53d1e

return $._unstakeRequests[account].latest() - $._released[account];

I can update FUZZING.md to be more explicit about this.

[melanciani]

we should also add the same claimRewards no ?

[melanciani]

weight is also approx no ? due to rounding errors

[melanciani]

same

[melanciani]

nit: maybe best to keep same invariant order vs what you've described in the md file !

[melanciani]

is using N=3 enough here ? shouldn't we try with like 5 or more to be sure we actually check N-1 and not something close like N/2 or 1+1 ?

[melanciani]

are we not missing equivalence invariant tests ? I can't find where stakeEquivalenceScenarioand unstakeEquivalenceScenario` functions are called

[Seth-Schmidt]

The scenario tests were initially structured as invariants:

protocol-apps/contracts/staking/test/foundry/ProtocolStakingInvariantTest.t.sol

Lines 252 to 266 in da83054

	function invariant_StakeEquivalence() public {
	if (!handler.ghost_lastCallWasStakeEquivalenceScenario()) return;
	assertEq(handler.ghost_sharesDouble(), handler.ghost_sharesSingle(), "stake equivalence: shares");
	// TODO: Weight is not expected to be strictly equal, might want to try to break the equivalence invariant
	// have not found a counter example for now
	assertEq(handler.ghost_weightDouble(), handler.ghost_weightSingle(), "stake equivalence: weight");
	assertApproxEqAbs(
	handler.ghost_earnedDouble(),
	handler.ghost_earnedSingle(),
	handler.EQUIVALENCE_EARNED_TOLERANCE(),
	"stake equivalence: earned"
	);
	handler.clearEquivalenceScenarioFlags();
	}

The fuzzer will call the scenario functions as a part of the fuzzing campaign at random points in the fuzzed state. The handler functions must be atomic to leverage the vm.snapshot cheat code that allows comparing the different terminal states against the same starting state. So, these "invariants" are not true invariants and are only checked when the function is called by the fuzzer.

I removed the invariant tests for these because it causes the engine to check the same ghost states after each function call when those variables will only change when the fuzzer calls one of the handler functions, so we can check directly in the function itself.

[melanciani]

what about equivalence invariants ?

[melanciani]

overall looks really great, but I think we are missing a few docs and tests !