Forge Testing Leveling

Dr Dimaz Wijaya | Estimated read time: 8 min Image credits: Photo by GrumpyBeere on Pixabay

A few tips and tricks for improving smart contract testing with Foundry.

Forge Testing Leveling

Testing is a powerful tool in a smart contract security's essential toolkit. It is usually done in a smart contract development suite, such as Forge (part of the Foundry suite). In this article, we discuss ways to enhance our Forge testing with two goals in mind. First, to achieve greater coverage through fuzzing, and second, to reuse our test codes for invariant testing.

We assume that readers are already familiar with smart contract testing and fuzzing using the Forge framework, either for decentralised application development or as part of a security assessment.

Part 1 — Coverage Increase Through Fuzzing

Improving test coverage often requires significant effort. One reason is that the target function needs a complex setup to reach all possible routes (both happy and sad paths) in the code. Often, testers write a test function for each path by copy-pasting identical code and modifying the necessary parts for the targeted routes. While this method seems efficient, it has a drawback: if the copied code contains bugs, fixing them requires going through all copies of the code.

One way to alleviate this problem is through fuzzing. Consider the following toy Vault contract as a target for testing.

contract Vault {
    ...
    function deposit() public payable {
        ...
        if (msg.value == 0) {
            revert;
        } else {
            emit DepositReceived();
        }
        ...
    }
    ...
}

To test Vault.deposit(), we can fuzz msg.value and evaluate each branch based on the randomized inputs. Below is an example of how one test function covers both the happy path and the sad path, and then evaluates a value upon a successful call.


    event DepositReceived();

    function test_deposit(uint256 msg_value) public {
        deal(address(this), msg_value);

        uint256 vaultBalanceBefore = vault.balance;

        if (msg_value == 0) {
            vm.expectRevert();
        } else {
            vm.expectEmit();
            emit DepositReceived();
        }

        vault.deposit{msg.value: msg_value}();

        // Validation
        if (msg_value > 0) {
            assertEq(vault.balance, vaultBalanceBefore + msg_value);
        }

    }

It is also easy to see that with this structure, we can still manually inspect each path by overriding the value of msg_value.

    function test_deposit(uint256 msg_value) public {
        msg_value = 0; // Additional line for manual inspection msg.value == 0
        deal(address(this), msg_value);
        ...

We now have a streamlined test function capable of evaluating multiple paths. However, the complexity of the code can increase significantly if a single test encompasses too many paths. Testers need to be mindful of this and strive for an optimal balance between simplicity and complexity. A strategy for managing a large function is to divide it into distinct sections:

  • Verifications or checks and auxiliary elements (e.g. formatters);
  • Principal logic;
  • Calls to internal functions.

It is possible to test each section independently. For testing purposes, internal functions can be made accessible via helper contracts.

Part 2 — Test Contract As a Handler for Invariant Testing

In this second part, we use WETH9 as an example. WETH9 (Wrapped Ether) is a simple yet popular contract to wrap a native asset Ether into an ERC20-compliant token.

Invariant testing is a powerful method to rigorously check for counter-examples in the code by running different paths or sequences of function calls. Invariant testing mostly consists of two parts: the invariant statement and the handler contract.

An invariant statement defines a condition of the system that should never be violated. For example, in WETH9 we should ensure that the totalSupply() is always equal to the Ether balance held by the contract. With this, we can construct an invariant statement as follows.

    ...
    WETH9 public weth;
    ...
    function invariant_eth_equality() public {
            assertEq(address(weth).balance, weth.totalSupply());
    }

Invariant statements are evaluated on every call path executed during testing. Writing an invariant statement is not trivial and requires a deep understanding of the protocol's mechanics. We will not discuss invariant statement selection and construction in this article; instead, we will focus on the second part: the handler contract.

A handler contract restricts the paths chosen by the test executor. It wraps the targeted functions by providing the necessary environment for the functions to execute successfully. For example, consider the deposit() function from the WETH9 contract below.


    function deposit() public payable {
        balanceOf[msg.sender] += msg.value;
        emit Deposit(msg.sender, msg.value);
    }

Function deposit() will be meaningful in a test case if msg.value is not zero, because only then the storage changes. Therefore, we will have the following function in our handler contract.


    function deposit(uint256 depositAmount) public payable {
        deal(address(this), depositAmount);
        weth.deposit{value: depositAmount}();
    }

If we already have a test contract for WETH9 (let's call it WETH9Test), we can use it as a handler for our invariant testing that targets WETH9. In this case, we do not need to write another handler contract. However, some adjustments are necessary.

A normal test in Forge is stateless, meaning the state resets with every function call. This also applies to a standard fuzzing test, where the state resets with every input variation. Invariant testing, on the other hand, is stateful. The state resets only when the call sequence and the invariant statements are executed and evaluated. Therefore, we need a way to differentiate between stateless and stateful tests, among other adjustments.

Counter as A Flag

Since stateful testing keeps state data throughout the test, we can use a storage-based counter to differentiate between stateful and stateless tests. There are many ways to implement this, but let's modify the famous Counter contract from Remix.

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

contract Counter {
    uint256 public number;

    modifier functionCalled() {
        increment();
        _;
    }

    constructor() {
        number = 0;
    }

    function increment() public {
        number++;
    }

    function isStateless() public view returns (bool) {
        if (number == 1) {
            return true;
        } else {
            return false;
        }
    }
}

We use the Counter contract as a parent for our test contract. This gives us access to the modifier functionCalled() and the view function isStateless().


contract WETH9Test is Counter {
...
}

It is easy to understand that the modifier functionCalled() increments the variable number. The increment operation can also be done through the function increment(). However, the modifier allows minimal changes to our existing test code structure. The modifier is attached to a test function so that we can later identify whether the test is in stateless or stateful mode by calling the Counter.isStateless() function.

    function test_deposit(uint256 depositAmount) public functionCalled {
        if (isStateless()) {
            deal(address(this), depositAmount);
        }
        weth.deposit{value: address(this).balance}();
    }

In the snippet above, the command deal is executed only during stateless testing, not stateful testing. This means that in stateful testing, the deposited amount depends on the contract's balance at that point.

It's worth noting that any runs in the invariant testing may execute the isStateless() block. This occurs when test_deposit() is the first or even the only function called in the session. However, this has no impact on the result of invariant testing, as test_deposit() will pass anyway.

Another interesting behavior to consider when writing a reusable test function is that invariant testing (stateful testing), by default, continues its execution even if some calls revert (where fail_on_revert = false, a setting which can be overridden in forge.toml). On the other hand, we do not want to revert in standard testing (stateless testing) for trivial matters. With that in mind, the following snippet works for both stateless and stateful testing, regardless of the user's balance.

    function test_deposit(address user) public functionCalled {
        if (isStateless()) {
            deal(user, 1 ether);
        }
        vm.startPrank(user);
        weth.deposit{value: 1 ether}();
        vm.stopPrank();
    }

Limiting Variation

The success of a fuzzing campaign (including stateless and stateful testing methods) often depends on the precision of the search scope. For example, in an attack scenario involving the function test_deposit(address), we do not want any random address as user to be used in the calls. Instead, we can focus on a few addresses, where one of them potentially behaves like an attacker.

The following snippet provides sample helper functions that behave similarly to a bound in Forge, but for addresses. It limits the options of the fuzzed address user to return only either ALICE or BOB.

contract WETH9Test is Counter {

    address public ALICE = makeAddr("alice");
    address public BOB = makeAddr("bob");

    address[] users = [ALICE, BOB];

    function _helper_address_to_uint256(address addr) internal view returns (uint256) {
        uint256 i = uint256(uint160(addr));
        return i;
    }

    function _helper_index_to_user(uint256 index) internal view returns (address) {
        index = bound(index, 0, users.length - 1);
        return users[index];
    }

    function helper_address_to_user(address addr) public view returns (address) {
        // Check if address is already a user
        for (uint256 i = 0; i < users.length; i++) {
            if (users[i] == addr) {
                return addr;
            }
        }
        uint256 idx = _helper_address_to_uint256(addr);
        return _helper_index_to_user(idx);
    }
}

Then we can simply use them in the modified test_deposit() below.

    function test_deposit(address user) public functionCalled {
        if (isStateless()) {
            deal(user, 1 ether);
        } else {
            user = helper_address_to_user(user);
        }
        vm.startPrank(user);
        weth.deposit{value: 1 ether}();
        vm.stopPrank();
    }

Setting Up Targets

Now that we have all the ingredients, we need to set the test functions as the targets. The snippet below assumes there are two test functions: test_deposit() and test_withdraw().

Similar to a standard fuzzing test, function setup() is called when a new sequence of calls starts.

// SPDX-License-Identifier: MIT

pragma solidity ^0.8.20;

import {Test, console} from "forge-std/Test.sol";
import {WETH9Test as Handler} from "./WETH9Test.t.sol";

contract WETH9TestInvariant is Test, Counter {

    Handler public handler;

    function setUp() public {
        handler = new Handler();
        handler.setUp();
        handler.overrideStateless();

        // Set selectors
        bytes4[] memory selectors = new bytes4[](2);
        selectors[0] = handler.test_deposit.selector;
        selectors[1] = handler.test_withdraw.selector;

        targetSelector(FuzzSelector({
            addr: address(handler),
            selectors: selectors
        }));

        // set targetContract
        targetContract(address(handler));
    }
}

There are several things worth noting in the snippet above. First, we call handler.setUp(), which is WETH9Test.setUp() on our invariant test contract because it is not automatically called during stateful testing. Second, there is a call to handler.overrideStateless(). This prepares our stateful/stateless flag by setting the value of Counter.number() to one. The next calls to the modifier functionCalled() set the value of Counter.number() to be more than one, indicating stateful testing (Counter.isStateless() == false).

The function overrideStateless() in our handler contract is simply:

    function overrideStateless() public functionCalled {}

Conclusion

Enhancing Forge testing through fuzzing and invariant testing provides a robust framework for smart contract security. Fuzzing helps in increasing test coverage by generating diverse input scenarios, while invariant testing ensures the integrity of the system through rigorous checks. By combining these methods and utilizing reusable test contracts, developers can create more comprehensive and efficient testing procedures. This approach not only identifies potential vulnerabilities but also ensures that the smart contracts behave as expected under various conditions, ultimately contributing to more secure and reliable blockchain applications.