The Tincho Exercise

While taking a look at Cyfrin Updraft, more on that in a later post, potentially, instructor Patrick Collins highly suggests to take onto a small challenge: a fuzz test exercise proposed by Tincho.
You can find the repo of the challenge here.
While Tincho already suggests his own solution right underneath the exercise, I thought to write this to go through my thought process developing the test, as it turned out to be structurally different, and equally effective, as Tincho's.

The Challenge

The original code to test and the test file already provided are reported here.

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

contract Registry {
    error PaymentNotEnough(uint256 expected, uint256 actual);

    uint256 public constant PRICE = 1 ether;

    mapping(address account => bool registered) private registry;

    function register() external payable {
        if(msg.value < PRICE) {
            revert PaymentNotEnough(PRICE, msg.value);
        }

        registry[msg.sender] = true;
    }

    function isRegistered(address account) external view returns (bool) {
        return registry[account];
    }
}

Regarding this code, there is not much to say. If you pay more than enough, you will be "registered". Otherwise, the transaction reverts. Function isRegistered allows to check if a user has registered or not. Pretty easy and straightforward.

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.13;

import {Test, console2} from "forge-std/Test.sol";
import {Registry} from "../src/Registry.sol";

contract RegistryTest is Test {
    Registry registry;
    address alice;

    function setUp() public {
        alice = makeAddr("alice");

        registry = new Registry();
    }

    function test_register() public {
        uint256 amountToPay = registry.PRICE();

        vm.deal(alice, amountToPay);
        vm.startPrank(alice);

        uint256 aliceBalanceBefore = address(alice).balance;

        registry.register{value: amountToPay}();

        uint256 aliceBalanceAfter = address(alice).balance;

        assertTrue(registry.isRegistered(alice), "Did not register user");
        assertEq(address(registry).balance, registry.PRICE(), "Unexpected registry balance");
        assertEq(aliceBalanceAfter, aliceBalanceBefore - registry.PRICE(), "Unexpected user balance");
    }

    /** Code your fuzz test here */
}

The test suite checks for the basic scenario. If Alice sends enough (exactly enough, wink wink), will the code behave as expected?
In order to measure the baseline of expected behavior, the test suite will:

• check that Alice is now registered;
• verify that the balance of the registry is now exactly the amount necessary to register, as only Alice has registered so far;
• assert that Alice has the balance she had before executing the transaction, minus the amount she sent to register to the Registry contract.

Checking for these three properties, the developers of the protocol defined what we can identify as the invariants of this project. A violation of any of these might represent a valid bug and a definite issue to raise during an audit.

Writing The Fuzz Test

We would like to test for more scenarios, not just the one presented above. So what we do is we fuzz test the amount to send, to make sure the code always respects the three invariants defined above, as we understood that they represent what "expected behavior" is for the developers of the protocol.
In order to run a fuzz test, we must use the keyword fuzz_ when defining our new test function. Additionally, we have to define the parameter that will be fuzzed. In our case, it is going to be the amount sent.
We can declare the new function as:

    function test_fuzz_registerFunction(uint256 amount) public {

The initial steps of dealing some ETH to Alice and starting the prank can be copied from above.
It is clear to see that we have two possible scenarios running this fuzz test:

• the random value picked can be either more or equal than PRICE, meaning the smart contract should not revert;
• the random value is less than PRICE, meaning we should expect a revert.

That is why our test function will look like this, reflecting this duality:

if(amount >= registry.PRICE()){
            registry.register{value: amount}();
            assertTrue(registry.isRegistered(alice), "Did not register user");
            assertEq(address(registry).balance, registry.PRICE(), "Unexpected registry balance");
            assertEq(address(alice).balance, amount - registry.PRICE());
        }
        else {
            vm.expectRevert(abi.encodeWithSelector(PaymentNotEnough.selector,registry.PRICE(), amount));
            registry.register{value: amount}();
        }

Precisely as discussed, we coded our fuzzing test to:

• either perform a correct transaction, and expect it to follow all the invariants;
• or to attempt a transaction with too little ETH, so the smart contract should revert and emit a custom error.

For clarity, the whole function is:

function test_fuzz_registerFunction(uint256 amount) public {
        vm.deal(alice, amount);
        vm.startPrank(alice);

        if(amount >= registry.PRICE()){
            registry.register{value: amount}();
            assertTrue(registry.isRegistered(alice), "Did not register user");
            assertEq(address(registry).balance, registry.PRICE(), "Unexpected registry balance");
            assertEq(address(alice).balance, amount - registry.PRICE());
        }
        else {
            vm.expectRevert(abi.encodeWithSelector(PaymentNotEnough.selector,registry.PRICE(), amount));
            registry.register{value: amount}();
        }
    }

The Result

The above function will identify a scenario where one of the three invariants will break. No spoilers, so I won't point out which one (or did I already?), but you can always copy paste this function and run forge test -vv to see for yourself.

The Bonus

Finding bugs is great, but how do we fix this?
I decided to propose a potential solution, following the CEI pattern (sorry Patrick, I know you do not love this). The change is small, but fundamental.
Basically, after the sender has been registered (apply the effects before the interactions), we can check if the user sent too much ETH and send back the unnecessary amount. If the transfer back to the user fails, we can also decide to revert on the whole transaction.
The code I added, and that successfully passes the tests, is this:

if(msg.value > PRICE) {
        (bool success, ) = msg.sender.call{value: msg.value-PRICE}("");
        if(!success) revert();
    }

If you find bugs in it, please let me know. Merry hacking ;)